✨ Add WithContext function for cancellation

waiterr.go

@@ -1,6 +1,7 @@
 package waiterr
 
 import (
+	"context"
 	"errors"
 	"sync"
 )
@@ -8,10 +9,20 @@ import (
 func New() WaitErr {
 	var we waitErr
 	we.errCh = make(chan error, 1)
+	we.cancel = func(error) {}
 
 	return &we
 }
 
+func WithContext(ctx context.Context) (WaitErr, context.Context) {
+	var we waitErr
+	we.errCh = make(chan error, 1)
+	cCtx, canc := context.WithCancelCause(ctx)
+	we.cancel = canc
+
+	return &we, cCtx
+}
+
 // WaitErr provides a way to run multiple goroutines and wait for their completion,
 // collecting any errors they return.
 type WaitErr interface {
@@ -37,6 +48,7 @@ type waitErr struct {
 	firstErr     error
 	firstErrOnce sync.Once
 	errCh        chan error // Buffered channel of size 1
+	cancel       context.CancelCauseFunc
 }
 
 // Go runs f in its own goroutine. When f returns, its error is stored, and returned
@@ -49,6 +61,7 @@ func (we *waitErr) Go(f func() error) {
 			we.firstErrOnce.Do(func() {
 				we.mut.Lock() // Acquire lock before writing to firstErr
 				we.firstErr = err
+				we.cancel(err)
 				we.mut.Unlock() // Release lock after writing
 
 				// Non-blocking send to errCh
@@ -103,7 +116,9 @@ func (we *waitErr) Wait() error {
 	we.wg.Wait()
 	we.mut.RLock()
 	defer we.mut.RUnlock()
-	return errors.Join(we.errs...)
+	ret := errors.Join(we.errs...)
+	we.cancel(ret)
+	return ret
 }
 
 // Unwrap returns all non-nil errors returned by our functions.