taskgroup: add an OnError method to Group

taskgroup.go

@@ -22,8 +22,7 @@ type Task func() error
 // non-nil error reported by any task (and not otherwise filtered) is returned
 // from the Wait method.
 type Group struct {
-	wg      sync.WaitGroup // counter for active goroutines
-	onError ErrorFunc      // called each time a task returns non-nil
+	wg sync.WaitGroup // counter for active goroutines
 
 	// active is nonzero when the group is "active", meaning there has been at
 	// least one call to Go since the group was created or the last Wait.
@@ -32,8 +31,9 @@ type Group struct {
 	// path reads active and only acquires μ if it discovers setup is needed.
 	active atomic.Uint32
 
-	μ   sync.Mutex // guards err
-	err error      // error returned from Wait
+	μ       sync.Mutex // guards err
+	err     error      // error returned from Wait
+	onError ErrorFunc  // called each time a task returns non-nil
 }
 
 // activate resets the state of the group and marks it as active.  This is
@@ -47,13 +47,25 @@ func (g *Group) activate() {
 	}
 }
 
-// New constructs a new empty group.  If ef != nil, it is called for each error
-// reported by a task running in the group.  The value returned by ef replaces
-// the task's error. If ef == nil, errors are not filtered.
+// New constructs a new empty group with the specified error filter.
+// See [Group.OnError] for a description of how errors are filtered.
+// If ef == nil, no filtering is performed.
+func New(ef ErrorFunc) *Group { return new(Group).OnError(ef) }
+
+// OnError sets the error filter for g. If ef == nil, the error filter is
+// removed and errors are no longer filtered. Otherwise, each non-nil error
+// reported by a task running in g is passed to ef, and the value it returns
+// replaces the task's error.
 //
-// Calls to ef are issued by a single goroutine, so it is safe for ef to
-// manipulate local data structures without additional locking.
-func New(ef ErrorFunc) *Group { return &Group{onError: ef} }
+// Calls to ef are synchronized so that it is safe for ef to manipulate local
+// data structures without additional locking. It is safe to call OnError while
+// tasks are active in g.
+func (g *Group) OnError(ef ErrorFunc) *Group {
+	g.μ.Lock()
+	defer g.μ.Unlock()
+	g.onError = ef
+	return g
+}
 
 // Go runs task in a new goroutine in g.
 func (g *Group) Go(task Task) {

taskgroup_test.go

@@ -30,7 +30,7 @@ func TestBasic(t *testing.T) {
 	t.Logf("Group value is %d bytes", reflect.TypeOf((*taskgroup.Group)(nil)).Elem().Size())
 
 	// Verify that the group works at all.
-	g := taskgroup.New(nil)
+	var g taskgroup.Group
 	g.Go(busyWork(25, nil))
 	if err := g.Wait(); err != nil {
 		t.Errorf("Unexpected task error: %v", err)
@@ -44,7 +44,7 @@ func TestBasic(t *testing.T) {
 	}
 
 	t.Run("Zero", func(t *testing.T) {
-		var g taskgroup.Group
+		g := taskgroup.New(nil)
 		g.Go(busyWork(30, nil))
 		if err := g.Wait(); err != nil {
 			t.Errorf("Unexpected task error: %v", err)
@@ -62,11 +62,18 @@ func TestErrorPropagation(t *testing.T) {
 	defer leaktest.Check(t)()
 
 	var errBogus = errors.New("bogus")
+
 	var g taskgroup.Group
 	g.Go(func() error { return errBogus })
 	if err := g.Wait(); err != errBogus {
 		t.Errorf("Wait: got error %v, wanted %v", err, errBogus)
 	}
+
+	g.OnError(func(error) error { return nil }) // discard
+	g.Go(func() error { return errBogus })
+	if err := g.Wait(); err != nil {
+		t.Errorf("Wait: got error %v, wanted nil", err)
+	}
 }
 
 func TestCancellation(t *testing.T) {
@@ -154,7 +161,7 @@ func TestCapacity(t *testing.T) {
 func TestRegression(t *testing.T) {
 	t.Run("WaitRace", func(t *testing.T) {
 		ready := make(chan struct{})
-		g := taskgroup.New(nil)
+		var g taskgroup.Group
 		g.Go(func() error {
 			<-ready
 			return nil
@@ -174,7 +181,7 @@ func TestRegression(t *testing.T) {
 				t.Errorf("Unexpected panic: %v", x)
 			}
 		}()
-		g := taskgroup.New(nil)
+		var g taskgroup.Group
 		g.Wait()
 	})
 }
@@ -208,7 +215,7 @@ func TestSingleTask(t *testing.T) {
 			return <-release
 		})
 
-		g := taskgroup.New(nil)
+		var g taskgroup.Group
 		g.Run(func() {
 			if err := s.Wait(); err != sentinel {
 				t.Errorf("Background Wait: got %v, want %v", err, sentinel)
@@ -231,7 +238,7 @@ func TestWaitMoreTasks(t *testing.T) {
 		results++
 	})
 
-	g := taskgroup.New(nil)
+	var g taskgroup.Group
 
 	// Test that if a task spawns more tasks on its own recognizance, waiting
 	// correctly waits for all of them provided we do not let the group go empty
@@ -283,7 +290,7 @@ func TestCollector(t *testing.T) {
 	c := taskgroup.Collect(func(v int) { sum += v })
 
 	vs := rand.Perm(15)
-	g := taskgroup.New(nil)
+	var g taskgroup.Group
 
 	for i, v := range vs {
 		v := v