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

Most existing answers suggest leaking goroutines. The idiomatic way to assign a timeout to WaitGroup.Wait is to use underlying sync/atomic package primitives. I took code from @icza answer and rewrote it using the atomic package, and added context cancelation as that's an idiomatic way to notify of a timeout.

package main

import (
    "context"
    "fmt"
    "sync/atomic"
    "time"
)

func main() {
    var submitCount int32
    // run this instead of wg.Add(1)
    atomic.AddInt32(&submitCount, 1)

    // run this instead of wg.Done()
    // atomic.AddInt32(&submitCount, -1)

    timeout := time.Second
    ctx, cancel := context.WithTimeout(context.Background(), timeout)
    defer cancel()
    fmt.Printf("Wait for waitgroup (up to %s)\n", timeout)

    waitWithCtx(ctx, &submitCount)

    fmt.Println("Free at last")
}

// waitWithCtx returns when passed counter drops to zero
// or when context is cancelled
func waitWithCtx(ctx context.Context, counter *int32) {
    ticker := time.NewTicker(10 * time.Millisecond)
    for {
        select {
        case <-ctx.Done():
            return
        case <-ticker.C:
            if atomic.LoadInt32(counter) == 0 {
                return
            }
        }
    }
}

Same code in Go Playground