✨ Implement waiterr.WaitForError

waiterr.go

@@ -7,16 +7,38 @@ import (
 
 // WaitErr wraps a sync.WaitGroup with error handling.
 type WaitErr struct {
-	wg   sync.WaitGroup
+	wg            sync.WaitGroup
-	errs []error
+	errs          []error
-	mut  sync.RWMutex
+	mut           sync.RWMutex
+	firstErr      error
+	firstErrOnce  sync.Once
+	errCh         chan error // Buffered channel of size 1
+	initErrChOnce sync.Once
 }
 
 // Go runs f in its own goroutine. When f returns, its error is stored, and returned
-// with we.Wait()
+// with we.Wait().
 func (we *WaitErr) Go(f func() error) {
+	we.initErrChOnce.Do(func() {
+		we.errCh = make(chan error, 1)
+	})
 	wrap := func() {
 		err := f()
+
+		if err != nil {
+			we.firstErrOnce.Do(func() {
+				we.mut.Lock() // Acquire lock before writing to firstErr
+				we.firstErr = err
+				we.mut.Unlock() // Release lock after writing
+
+				// Non-blocking send to errCh
+				select {
+				case we.errCh <- err:
+				default:
+				}
+			})
+		}
+
 		we.mut.Lock()
 		defer we.mut.Unlock()
 		we.errs = append(we.errs, err)
@@ -27,10 +49,35 @@ func (we *WaitErr) Go(f func() error) {
 // WaitForError waits for the first error to be returned by one of our go routines, and immediately returns
 // with that error. If all functions return successfully, a nil is returned.
 func (we *WaitErr) WaitForError() error {
-	// Implement this
+	if we.errCh == nil {
-	// If we already have an error, return it immediately without waiting
+		panic("WaitForError called before Go, errCh is nil")
-	// If no error has yet been returned, wait for the very first error and return it
+	}
-	return nil
+	// Check if an error has already been set
+	we.mut.RLock()
+	if we.firstErr != nil {
+		err := we.firstErr
+		we.mut.RUnlock()
+		return err
+	}
+	we.mut.RUnlock()
+
+	// Create a channel to signal when all goroutines are done
+	done := make(chan struct{})
+	go func() {
+		we.wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case err := <-we.errCh:
+		return err
+	case <-done:
+		// All goroutines finished, and no error was sent to errCh
+		// Re-check firstErr in case it was set just before 'done' was closed
+		we.mut.RLock()
+		defer we.mut.RUnlock()
+		return we.firstErr // This will be nil if no error occurred
+	}
 }
 
 // Wait for all current goroutines to finish. Return an error that combines all errors returned