context: don't return a non-nil from Err before Done is closed

The Context.Err documentation states that it returns nil if the
context's done channel is not closed. Fix a race condition introduced
by CL 653795 where Err could return a non-nil error slightly before
the Done channel is closed.

No impact on Err performance when returning nil.

Slows down Err when returning non-nil by about 3x,
but that's still almost 2x faster than before CL 653795
and the performance of this path is less important.
(A tight loop checking Err for doneness will be terminated
by the first Err call to return a non-nil result.)

    goos: darwin
    goarch: arm64
    pkg: context
    cpu: Apple M4 Pro
                   │ /tmp/bench.0 │            /tmp/bench.1             │
                   │    sec/op    │   sec/op     vs base                │
    ErrOK-14          1.806n ± 1%   1.774n ± 0%    -1.77% (p=0.000 n=8)
    ErrCanceled-14    1.821n ± 1%   7.525n ± 3%  +313.23% (p=0.000 n=8)
    geomean           1.813n        3.654n       +101.47%

Fixes golang#75533

Change-Id: Iea22781a199ace7e7f70cf65168c36e090cd2e2a
Reviewed-on: https://go-review.googlesource.com/c/go/+/705235
TryBot-Bypass: Damien Neil <[email protected]>
Reviewed-by: Nicholas Husin <[email protected]>
Reviewed-by: Nicholas Husin <[email protected]>
Auto-Submit: Damien Neil <[email protected]>

Author: neild

src/context/context.go

@@ -463,6 +463,8 @@ func (c *cancelCtx) Done() <-chan struct{} {
 func (c *cancelCtx) Err() error {
 	// An atomic load is ~5x faster than a mutex, which can matter in tight loops.
 	if err := c.err.Load(); err != nil {
+		// Ensure the done channel has been closed before returning a non-nil error.
+		<-c.Done()
 		return err.(error)
 	}
 	return nil

src/context/x_test.go

@@ -1177,3 +1177,23 @@ func (c *customContext) Err() error {
 func (c *customContext) Value(key any) any {
 	return c.parent.Value(key)
 }
+
+// Issue #75533.
+func TestContextErrDoneRace(t *testing.T) {
+	// 4 iterations reliably reproduced #75533.
+	for range 10 {
+		ctx, cancel := WithCancel(Background())
+		donec := ctx.Done()
+		go cancel()
+		for ctx.Err() == nil {
+			if runtime.GOARCH == "wasm" {
+				runtime.Gosched() // need to explicitly yield
+			}
+		}
+		select {
+		case <-donec:
+		default:
+			t.Fatalf("ctx.Err is non-nil, but ctx.Done is not closed")
+		}
+	}
+}