errgroup: use WithCancelCause to cancel context

Fixes golang/go#59355

Change-Id: Ib6a88e7e5fefe7b0d5672035af16d109aabcbf1e
Reviewed-on: https://go-review.googlesource.com/c/sync/+/481255
TryBot-Result: Gopher Robot <[email protected]>
Run-TryBot: Bryan Mills <[email protected]>
Reviewed-by: Bryan Mills <[email protected]>
Run-TryBot: Ian Lance Taylor <[email protected]>
Reviewed-by: Michael Knyszek <[email protected]>
Auto-Submit: Bryan Mills <[email protected]>

errgroup/errgroup.go

@@ -20,7 +20,7 @@ type token struct{}
 // A zero Group is valid, has no limit on the number of active goroutines,
 // and does not cancel on error.
 type Group struct {
- cancel func()
+ cancel func(error)
 
  wg sync.WaitGroup
 
@@ -43,7 +43,7 @@ func (g *Group) done() {
 // returns a non-nil error or the first time Wait returns, whichever occurs
 // first.
 func WithContext(ctx context.Context) (*Group, context.Context) {
- ctx, cancel := context.WithCancel(ctx)
+ ctx, cancel := withCancelCause(ctx)
  return &Group{cancel: cancel}, ctx
 }
 
@@ -52,7 +52,7 @@ func WithContext(ctx context.Context) (*Group, context.Context) {
 func (g *Group) Wait() error {
  g.wg.Wait()
  if g.cancel != nil {
- g.cancel()
+ g.cancel(g.err)
  }
  return g.err
 }
@@ -76,7 +76,7 @@ func (g *Group) Go(f func() error) {
  g.errOnce.Do(func() {
  g.err = err
  if g.cancel != nil {
- g.cancel()
+ g.cancel(g.err)
  }
  })
  }
@@ -105,7 +105,7 @@ func (g *Group) TryGo(f func() error) bool {
  g.errOnce.Do(func() {
  g.err = err
  if g.cancel != nil {
- g.cancel()
+ g.cancel(g.err)
  }
  })
  }

errgroup/go120.go

@@ -0,0 +1,14 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build go1.20
+// +build go1.20
+
+package errgroup
+
+import "context"
+
+func withCancelCause(parent context.Context) (context.Context, func(error)) {
+ return context.WithCancelCause(parent)
+}

errgroup/go120_test.go

@@ -0,0 +1,55 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build go1.20
+// +build go1.20
+
+package errgroup_test
+
+import (
+ "context"
+ "errors"
+ "testing"
+
+ "golang.org/x/sync/errgroup"
+)
+
+func TestCancelCause(t *testing.T) {
+ errDoom := errors.New("group_test: doomed")
+
+ cases := []struct {
+ errs []error
+ want error
+ }{
+ {want: nil},
+ {errs: []error{nil}, want: nil},
+ {errs: []error{errDoom}, want: errDoom},
+ {errs: []error{errDoom, nil}, want: errDoom},
+ }
+
+ for _, tc := range cases {
+ g, ctx := errgroup.WithContext(context.Background())
+
+ for _, err := range tc.errs {
+ err := err
+ g.TryGo(func() error { return err })
+ }
+
+ if err := g.Wait(); err != tc.want {
+ t.Errorf("after %T.TryGo(func() error { return err }) for err in %v\n"+
+ "g.Wait() = %v; want %v",
+ g, tc.errs, err, tc.want)
+ }
+
+ if tc.want == nil {
+ tc.want = context.Canceled
+ }
+
+ if err := context.Cause(ctx); err != tc.want {
+ t.Errorf("after %T.TryGo(func() error { return err }) for err in %v\n"+
+ "context.Cause(ctx) = %v; tc.want %v",
+ g, tc.errs, err, tc.want)
+ }
+ }
+}

errgroup/pre_go120.go

@@ -0,0 +1,15 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build !go1.20
+// +build !go1.20
+
+package errgroup
+
+import "context"
+
+func withCancelCause(parent context.Context) (context.Context, func(error)) {
+ ctx, cancel := context.WithCancel(parent)
+ return ctx, func(error) { cancel() }
+}