Go Fail Fast Wait Group

failfast.go

package failfast

import (
	"context"
	"sync"
)

// FailFastWG allow synchronization between multiple parallel goroutines.
// If one fails, the common context is cancelled and the others should return.
type WaitGroup struct {
	wg  *sync.WaitGroup
	err error
	ctx context.Context

	cancelCtx context.CancelFunc
	once      sync.Once
}

func NewWaitGroup(ctx context.Context) *WaitGroup {
	ctx, cancelCtx := context.WithCancel(ctx)
	return &WaitGroup{
		wg:        &sync.WaitGroup{},
		ctx:       ctx,
		cancelCtx: cancelCtx,
		once:      sync.Once{},
	}
}

func (ff *WaitGroup) Add(n int) {
	ff.wg.Add(n)
}

func (ff *WaitGroup) Done() {
	ff.wg.Done()
}

// Fail records the error and cancells the context. Can be called multiple times but only the first error is recorded.
func (ff *WaitGroup) Fail(err error) {
	ff.once.Do(func() {
		ff.err = err
		ff.cancelCtx()
	})
}

// Context Returns the common context.
func (ff *WaitGroup) Context() context.Context {
	return ff.ctx
}

// Wait for all the go routines to finish. returns an error if any of the go routines fail.
func (ff *WaitGroup) Wait() error {
	ff.wg.Wait()
	ff.cancelCtx()
	return ff.err
}