context: don't return the wrong error when Cause races cancellation

Check to see if a context is canceled at all
before checking for the cancellaion cause.
If we can't find a cause, use the original error.

Avoids a data race where we look for a cause,
find none (because the context is not canceled),
the context is canceled,
and we then return ctx.Err() (even though there is now a cause).

Fixes #73390

Change-Id: I97f44aef25c6b02871d987970abfb4c215c5c80e
Reviewed-on: https://go-review.googlesource.com/c/go/+/679835
Reviewed-by: Sean Liao <sean@liao.dev>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Cherry Mui <cherryyz@google.com>

1 files changed, 9 insertions, 9 deletions

diff --git a/src/context/context.go b/src/context/context.go
index 232fc55d87..d00ac67e38 100644
--- a/src/context/context.go
+++ b/src/context/context.go
@@ -286,6 +286,10 @@ func withCancel(parent Context) *cancelCtx {
 // Otherwise Cause(c) returns the same value as c.Err().
 // Cause returns nil if c has not been canceled yet.
 func Cause(c Context) error {
+	err := c.Err()
+	if err == nil {
+		return nil
+	}
 	if cc, ok := c.Value(&cancelCtxKey).(*cancelCtx); ok {
 		cc.mu.Lock()
 		cause := cc.cause
@@ -293,16 +297,12 @@ func Cause(c Context) error {
 		if cause != nil {
 			return cause
 		}
-		// Either this context is not canceled,
-		// or it is canceled and the cancellation happened in a
-		// custom context implementation rather than a *cancelCtx.
+		// The parent cancelCtx doesn't have a cause,
+		// so c must have been canceled in some custom context implementation.
 	}
-	// There is no cancelCtxKey value with a cause, so we know that c is
-	// not a descendant of some canceled Context created by WithCancelCause.
-	// Therefore, there is no specific cause to return.
-	// If this is not one of the standard Context types,
-	// it might still have an error even though it won't have a cause.
-	return c.Err()
+	// We don't have a cause to return from a parent cancelCtx,
+	// so return the context's error.
+	return err
 }
 
 // AfterFunc arranges to call f in its own goroutine after ctx is canceled.