Short-circuit new goroutines if an error already occured

This commit is contained in:
Dan Jones 2025-09-07 23:21:44 -05:00
commit 8640832a27
2 changed files with 55 additions and 7 deletions

View file

@ -13,6 +13,7 @@ import (
"context"
"fmt"
"sync"
"sync/atomic"
)
type token struct{}
@ -30,7 +31,7 @@ type Group struct {
sem chan token
errOnce sync.Once
err error
err atomic.Value
}
func (g *Group) done() {
@ -50,14 +51,22 @@ func WithContext(ctx context.Context) (*Group, context.Context) {
return &Group{cancel: cancel}, ctx
}
func (g *Group) error() error {
v := g.err.Load()
if v == nil {
return nil
}
return v.(error)
}
// Wait blocks until all function calls from the Go method have returned, then
// returns the first non-nil error (if any) from them.
func (g *Group) Wait() error {
g.wg.Wait()
if g.cancel != nil {
g.cancel(g.err)
g.cancel(g.error())
}
return g.err
return g.error()
}
// Go calls the given function in a new goroutine.
@ -70,6 +79,9 @@ func (g *Group) Wait() error {
// cancel the associated Context, if any. The error will be returned
// by Wait.
func (g *Group) Go(f func() error) {
if g.error() != nil {
return
}
if g.sem != nil {
g.sem <- token{}
}
@ -92,9 +104,9 @@ func (g *Group) Go(f func() error) {
if err := f(); err != nil {
g.errOnce.Do(func() {
g.err = err
g.err.Store(err)
if g.cancel != nil {
g.cancel(g.err)
g.cancel(err)
}
})
}
@ -106,6 +118,9 @@ func (g *Group) Go(f func() error) {
//
// The return value reports whether the goroutine was started.
func (g *Group) TryGo(f func() error) bool {
if g.error() != nil {
return false
}
if g.sem != nil {
select {
case g.sem <- token{}:
@ -121,9 +136,9 @@ func (g *Group) TryGo(f func() error) bool {
if err := f(); err != nil {
g.errOnce.Do(func() {
g.err = err
g.err.Store(err)
if g.cancel != nil {
g.cancel(g.err)
g.cancel(err)
}
})
}