Timeout for WaitGroup.Wait()

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 'worker' for ever. This leads to some philosophical questions (i.e. how can the system reliably continue once it has errant workers?), but I think that's out of scope for this question.

I have an answer which I'll provide. Now that I've written it down, it doesn't seem so bad but it still feels more convoluted than it ought to. I'd like to know if there's something available which is simpler, more idiomatic, or even an alternative approach which doesn't use WaitGroups.

Mostly your solution you posted below is as good as it can get. Couple of tips to improve it:

• Alternatively you may close the channel to signal completion instead of sending a value on it, a receive operation on a closed channel can always proceed immediately.
• And it's better to use defer statement to signal completion, it is executed even if a function terminates abruptly.
• Also if there is only one "job" to wait for, you can completely omit the WaitGroup and just send a value or close the channel when job is complete (the same channel you use in your select statement).
• Specifying 1 second duration is as simple as: timeout := time.Second. Specifying 2 seconds for example is: timeout := 2 * time.Second. You don't need the conversion, time.Second is already of type time.Duration, multiplying it with an untyped constant like 2 will also yield a value of type time.Duration.

I would also create a helper / utility function wrapping this functionality. Note that WaitGroup must be passed as a pointer else the copy will not get "notified" of the WaitGroup.Done() calls. Something like:

// waitTimeout waits for the waitgroup for the specified max timeout.
// Returns true if waiting timed out.
func waitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool {
	c := make(chan struct{})
	go func() {
		defer close(c)
		wg.Wait()
	}()
	select {
   	case <-c:
		return false // completed normally
	case <-time.After(timeout):
		return true // timed out
	}
}

Using it:

if waitTimeout(&wg, time.Second) {
	fmt.Println("Timed out waiting for wait group")
} else {
	fmt.Println("Wait group finished")
}

Try it on the Go Playground.

I did it like this: http://play.golang.org/p/eWv0fRlLEC

go func() {
	wg.Wait()
	c <- struct{}{}
}()
timeout := time.Duration(1) * time.Second
fmt.Printf("Wait for waitgroup (up to %s)\n", timeout)
select {
case <-c:
	fmt.Printf("Wait group finished\n")
case <-time.After(timeout):
	fmt.Printf("Timed out waiting for wait group\n")
}
fmt.Printf("Free at last\n")

It works fine, but is it the best way to do it?

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

Throwing in a solution which does not leak a goroutine, or rely on polling (sleeps):

import "atomic"

type WaitGroup struct {
	count int32
	done chan struct{}
}

func NewWaitGroup() *WaitGroup {
	return &WaitGroup{
		done: make(chan struct{}),
	}
}

func (wg *WaitGroup) Add(i int32) {
	select {
	case <-wg.done:
		panic("use of an already closed WaitGroup")
	default:
	}
	atomic.AddInt32(&wg.count, i)
}

func (wg *WaitGroup) Done() {
	i := atomic.AddInt32(&wg.count, -1)
	if i == 0 {
		close(wg.done)
	}
	if i < 0 {
		panic("too many Done() calls")
	}
}

func (wg *WaitGroup) C() <-chan struct{} {
	return wg.done
}

Usage:

wg := NewWaitGroup()
wg.Add(1)
go func() {
  // do stuff
  wg.Done()
}

select {
case <-wg.C():
  fmt.Printf("Completed!\n")
case <-time.After(time.Second):
  fmt.Printf("Timed out!\n")
}

This is a bad idea. Do not abandon goroutines, doing so may introduce races, resource leaks and unexpected conditions, ultimately impacting the stability of your application.

Instead use timeouts throughout your code consistently in order to make sure no goroutine is blocked forever or takes too long to run.

The idiomatic way for achieving that is via context.WithTimeout():

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

// Now perform any I/O using the given ctx:
go func() {
  err = example.Connect(ctx)
  if err != nil { /* handle err and exit goroutine */ }
  . . .
}()

Now you can safely use WaitGroup.Wait(), knowing it will always finish in a timely manner.

The following will not introduce any leaking goroutines

func callingFunc() {
	...
	wg := new(sync.WaitGroup)
	for _, msg := range msgs {
		wg.Add(1)
		go wrapperParallelCall(ctx, params, wg)
	}

	wg.Wait()
}

func wrapperParallelCall(ctx, params, wg) {
	ctx, cancel := context.WithTimeout(ctx, time.Second)
	defer wg.Done()
	defer cancel()

	originalSequenceCall(ctx, params)
}

func originalSequenceCall(ctx, params) {...}

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

I wrote a library that encapsulates the concurrency logic https://github.com/shomali11/parallelizer which you can also pass a timeout.

Here is an example without a timeout:

func main() {
	group := parallelizer.DefaultGroup()

    group.Add(func() {
	    for char := 'a'; char < 'a'+3; char++ {
		    fmt.Printf("%c ", char)
	    }
    })

	group.Add(func() {
    	for number := 1; number < 4; number++ {
	    	fmt.Printf("%d ", number)
	    }
    })

    err := group.Run()

	fmt.Println()
    fmt.Println("Done")
    fmt.Printf("Error: %v", err)
}

Output:

a 1 b 2 c 3 
Done
Error: <nil>

Here is an example with a timeout:

func main() {
    options := &parallelizer.Options{Timeout: time.Second}
	group := parallelizer.NewGroup(options)

    group.Add(func() {
	    time.Sleep(time.Minute)

		for char := 'a'; char < 'a'+3; char++ {
    		fmt.Printf("%c ", char)
    	}
    })

	group.Add(func() {
    	time.Sleep(time.Minute)

	    for number := 1; number < 4; number++ {
		    fmt.Printf("%d ", number)
	    }
    })

    err := group.Run()

	fmt.Println()
    fmt.Println("Done")
    fmt.Printf("Error: %v", err)
}

Output:

Done
Error: timeout

This is not an actual answer to this question but was the (much simpler) solution to my little problem when I had this question.

My 'workers' were doing http.Get() requests so I just set the timeout on the http client.

urls := []string{"http://1.jpg", "http://2.jpg"}
wg := &sync.WaitGroup{}
for _, url := range urls {
	wg.Add(1)
	go func(url string) {
		client := http.Client{
			Timeout: time.Duration(3 * time.Second), // only want very fast responses
		}
		resp, err := client.Get(url)
		//... check for errors
		//... do something with the image when there are no errors
		//...

		wg.Done()
	}(url)
	
}
wg.Wait()

We had the same need for one of our systems. by passing a context to goroutines and closing that context when we are facing timeout, we would prevent goroutine leaks.

func main() {
	ctx := context.Background()
	ctxWithCancel, cancelFunc := context.WithCancel(ctx)
	var wg sync.WaitGroup
	Provide(ctxWithCancel, 5, &wg)
	Provide(ctxWithCancel, 5, &wg)
	c := make(chan struct{})
	go func() {
		wg.Wait()
		c <- struct{}{}
		fmt.Println("closed")
	}()

	select {
	case <-c:
	case <-time.After(20 * time.Millisecond):
		cancelFunc()
		fmt.Println("timeout")
	}
}

func Work(ctx context.Context, to int) {
	for i := 0; i < to; i++ {
		select {
		case <-ctx.Done():
			return
		default:
			fmt.Println(i)
			time.Sleep(10 * time.Millisecond)
		}
	}
}

func Provide(ctx context.Context, to int, wg *sync.WaitGroup) {
	wg.Add(1)
	go func() {
		Work(ctx, to)
		wg.Done()
	}()
}