proposal: sync: add WaitContext to WaitGroup

[carnott-snap]

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.

[ianlancetaylor]

CC @bcmills

[bcmills]

In most cases when you have the possibility of a timeout, you want to cancel the work rather than just the action of waiting on the work, and you want to distinguish between “work completed” (err != nil) and “work timed out” (err == nil). For those cases, errgroup.WithContext is usually a better fit.

For the rare cases where you really do want to wait on a best-effort basis but also allow the work to continue, it's easy enough to wrap yourself. You don't even need to burn a full context.Context and goroutine per call; you can use a sync.Map to deduplicate: https://play.golang.org/p/MaefhJWdU2u

That said, I don't see any reason to put this in the sync package. It's not entirely trivial, but can be implemented as an independent library.

[rogpeppe]

@bcmills one issue with your example above is that it can potentially leave a goroutine hanging around indefinitely if the context is cancelled and the waitgroup takes a long time to complete. It's also seem like quite a relatively heavyweight operation compared to the very lightweight Wait call. I wonder if it might be possible to do better with runtime support.