context: support non-standard Context in Cause

If Cause is called on a non-standard Context, call ctx.Err.

Fixes #62582

Change-Id: Iac4ed93203eb5529f8839eb479b6ee2ee5ff6cbc
Reviewed-on: https://go-review.googlesource.com/c/go/+/527277
Reviewed-by: Bryan Mills <bcmills@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Auto-Submit: Ian Lance Taylor <iant@google.com>
Reviewed-by: Ian Lance Taylor <iant@google.com>

diff --git a/src/context/context.go b/src/context/context.go
index ee66b43c8500c7b0db6b79d645844560231625ab..80e1787576364967f99dfd207d739361b130c593 100644 (file)
--- a/src/context/context.go
+++ b/src/context/context.go
@@ -286,7 +286,12 @@ func Cause(c Context) error {
                defer cc.mu.Unlock()
                return cc.cause
        }
-       return nil
+       // There is no cancelCtxKey value, so we know that c is
+       // not a descendant of some 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()
 }
 
 // AfterFunc arranges to call f in its own goroutine after ctx is done

diff --git a/src/context/x_test.go b/src/context/x_test.go
index 57fe60b4ee060bd5660feb77dd7f1928b1971815..e006e53470f2e7f7ed18e87809842d3513b66ada 100644 (file)
--- a/src/context/x_test.go
+++ b/src/context/x_test.go
@@ -867,6 +867,126 @@ func TestCustomContextPropagation(t *testing.T) {
        }
 }
 
+// customCauseContext is a custom Context used to test context.Cause.
+type customCauseContext struct {
+       mu   sync.Mutex
+       done chan struct{}
+       err  error
+
+       cancelChild CancelFunc
+}
+
+func (ccc *customCauseContext) Deadline() (deadline time.Time, ok bool) {
+       return
+}
+
+func (ccc *customCauseContext) Done() <-chan struct{} {
+       ccc.mu.Lock()
+       defer ccc.mu.Unlock()
+       return ccc.done
+}
+
+func (ccc *customCauseContext) Err() error {
+       ccc.mu.Lock()
+       defer ccc.mu.Unlock()
+       return ccc.err
+}
+
+func (ccc *customCauseContext) Value(key any) any {
+       return nil
+}
+
+func (ccc *customCauseContext) cancel() {
+       ccc.mu.Lock()
+       ccc.err = Canceled
+       close(ccc.done)
+       cancelChild := ccc.cancelChild
+       ccc.mu.Unlock()
+
+       if cancelChild != nil {
+               cancelChild()
+       }
+}
+
+func (ccc *customCauseContext) setCancelChild(cancelChild CancelFunc) {
+       ccc.cancelChild = cancelChild
+}
+
+func TestCustomContextCause(t *testing.T) {
+       // Test if we cancel a custom context, Err and Cause return Canceled.
+       ccc := &customCauseContext{
+               done: make(chan struct{}),
+       }
+       ccc.cancel()
+       if got := ccc.Err(); got != Canceled {
+               t.Errorf("ccc.Err() = %v, want %v", got, Canceled)
+       }
+       if got := Cause(ccc); got != Canceled {
+               t.Errorf("Cause(ccc) = %v, want %v", got, Canceled)
+       }
+
+       // Test that if we pass a custom context to WithCancelCause,
+       // and then cancel that child context with a cause,
+       // that the cause of the child canceled context is correct
+       // but that the parent custom context is not canceled.
+       ccc = &customCauseContext{
+               done: make(chan struct{}),
+       }
+       ctx, causeFunc := WithCancelCause(ccc)
+       cause := errors.New("TestCustomContextCause")
+       causeFunc(cause)
+       if got := ctx.Err(); got != Canceled {
+               t.Errorf("after CancelCauseFunc ctx.Err() = %v, want %v", got, Canceled)
+       }
+       if got := Cause(ctx); got != cause {
+               t.Errorf("after CancelCauseFunc Cause(ctx) = %v, want %v", got, cause)
+       }
+       if got := ccc.Err(); got != nil {
+               t.Errorf("after CancelCauseFunc ccc.Err() = %v, want %v", got, nil)
+       }
+       if got := Cause(ccc); got != nil {
+               t.Errorf("after CancelCauseFunc Cause(ccc) = %v, want %v", got, nil)
+       }
+
+       // Test that if we now cancel the parent custom context,
+       // the cause of the child canceled context is still correct,
+       // and the parent custom context is canceled without a cause.
+       ccc.cancel()
+       if got := ctx.Err(); got != Canceled {
+               t.Errorf("after CancelCauseFunc ctx.Err() = %v, want %v", got, Canceled)
+       }
+       if got := Cause(ctx); got != cause {
+               t.Errorf("after CancelCauseFunc Cause(ctx) = %v, want %v", got, cause)
+       }
+       if got := ccc.Err(); got != Canceled {
+               t.Errorf("after CancelCauseFunc ccc.Err() = %v, want %v", got, Canceled)
+       }
+       if got := Cause(ccc); got != Canceled {
+               t.Errorf("after CancelCauseFunc Cause(ccc) = %v, want %v", got, Canceled)
+       }
+
+       // Test that if we associate a custom context with a child,
+       // then canceling the custom context cancels the child.
+       ccc = &customCauseContext{
+               done: make(chan struct{}),
+       }
+       ctx, cancelFunc := WithCancel(ccc)
+       ccc.setCancelChild(cancelFunc)
+       ccc.cancel()
+       if got := ctx.Err(); got != Canceled {
+               t.Errorf("after CancelCauseFunc ctx.Err() = %v, want %v", got, Canceled)
+       }
+       if got := Cause(ctx); got != Canceled {
+               t.Errorf("after CancelCauseFunc Cause(ctx) = %v, want %v", got, Canceled)
+       }
+       if got := ccc.Err(); got != Canceled {
+               t.Errorf("after CancelCauseFunc ccc.Err() = %v, want %v", got, Canceled)
+       }
+       if got := Cause(ccc); got != Canceled {
+               t.Errorf("after CancelCauseFunc Cause(ccc) = %v, want %v", got, Canceled)
+       }
+}
+
 func TestAfterFuncCalledAfterCancel(t *testing.T) {
        ctx, cancel := WithCancel(Background())
        donec := make(chan struct{})