Timeout for WaitGroup.Wait()

Question

What is an idiomatic way to assign a timeout to WaitGroup.Wait() ?

The reason I want to do this, is to safeguard my \'scheduler\' from potentially awaiting an errant \'w

Answer 1

Another solution without leaking wg.Wait() routine: just use (well-suported and widely-used) golang.org/x/sync/semaphore:

• Instead of sync.WaitGroup{} use sem.NewWeighted(N) (you have to know N in advance)
• Instead of wg.Add(1) use err := sem.Acquire(ctx, 1)
• Instead of defer wg.Done() use defer sem.Release(1)
• Instead of wg.Wait() you can use sem.Acquire(ctx, N) with context with timeout.
• Watch out, this is only equivalent to sync.WaitGroup in this specific use-case (when you only call Add(1) and Release(1) N times). Read the documentation carefully.

Example:

package main

import (
    "context"
    "log"
    "time"

    "golang.org/x/sync/semaphore"
)

func worker(n int) {
    time.Sleep(time.Duration(n) * time.Second)
    log.Printf("Worker %v finished", n)
}

func main() {

    const N = 5
    sem := semaphore.NewWeighted(N)

    for i := 0; i < N; i++ {

        err := sem.Acquire(context.Background(), 1)
        if err != nil {
            log.Fatal("sem.Acquire err", err)
        }
        go func(n int) {
            defer sem.Release(1)
            worker(n)
        }(i)
    }

    ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
    defer cancel()

    err := sem.Acquire(ctx, N)
    if err != nil {
        log.Println("sem.Acquire err:", err)
        return
    }

    log.Println("sem.Acquire ok")
}

Which results in:

2009/11/10 23:00:00 Worker 0 finished
2009/11/10 23:00:01 Worker 1 finished
2009/11/10 23:00:02 Worker 2 finished
2009/11/10 23:00:02 sem.Acquire err: context deadline exceeded