background

Some times when using a sync.WaitGroup, I want to bound the whole action with a timeout.

I thought this had been called out in #39863 or some other issue, but I could not find anything.

description

We should be able to add a context aware wait method:

package sync

func (wg *WaitGroup) WaitContext(ctx context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	go func() { wg.Wait(); cancel() }()
	<-ctx.Done()
}

costs

It may be weird to have the callers of sync.WaitGroup.Done untethered, since sync.WaitGroup.WaitContext can return early, but if you want to add a watchdog thread it can still use sync.WaitGroup.Wait since there are no changes to that implementation. That being said, it does add complexity.

alternatives

Since this method somewhat consumes the context.Context callers may prefer that context.Context.Err be returned. Unfortunately, this mixes the intentions, since the same error is available to the caller, but is not a painful implementation:

func (wg *WaitGroup) WaitContext(ctx context.Context) error {
	ctx1, cancel := context.WithCancel(ctx)
	go func() { wg.Wait(); cancel() }()
	<-ctx1.Done()
	return ctx.Err()
}

If we instead wanted to embed the context.Context into the struct directly, so that sync.WaitGroup.Wait used it implicitly, we could either add a public field, or a private field with a setter. I assume we would want the latter since that is what was done for http.Request. Continuing on, we may want to expose the error that the embedded context.Context has, requiring a func (wg *WaitGroup) Err() error` method. Unfortunately, all of this is pretty complex, and would require reworking several extant method implementations, and did not seem worth the cost.