database/sql: prevent internal context error from being returned from Rows.Err()

CL 497675 modified Rows such that context errors are propagated through
Rows.Err(). This caused an issue where calling Close meant that an
internal cancellation error would (eventually) be returned from Err:

1. A caller makes a query using a cancellable context.
2. initContextClose sees that either the query context or the
   transaction context can be canceled, so will need to spawn a
   goroutine to capture their errors.
3. initContextClose derives a context from the query context via
   WithCancel and sets rs.cancel.
4. When a user calls Close, rs.cancel is called. awaitDone's ctx is
   cancelled, which is good, since we don't want it to hang forever.
5. This internal cancellation (after CL 497675) has its error saved on
   contextDone.
6. Later, calling Err will return the error in contextDone if present.

This leads to a race condition depending on how quickly Err is called
after Close.

The docs for Close and Err state that calling Close should have no
affect on the return result for Err. So, a potential fix is to ensure
that awaitDone does not save the error when the cancellation comes from
a Close via rs.cancel.

This CL does that, using a new context not derived from the query
context, whose error is ignored as the query context's error used to be
before the original bugfix.

The included test fails before the CL, and passes afterward.

Fixes #60932

Change-Id: I2bf4c549efd83d62b86e298c9c45ebd06a3ad89a
Reviewed-on: https://go-review.googlesource.com/c/go/+/505397
Auto-Submit: Russ Cox <rsc@golang.org>
Run-TryBot: Russ Cox <rsc@golang.org>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Russ Cox <rsc@golang.org>
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

2 files changed, 36 insertions, 6 deletions

diff --git a/src/database/sql/sql.go b/src/database/sql/sql.go
index a77d63dc5e..0764c7d17a 100644
--- a/src/database/sql/sql.go
+++ b/src/database/sql/sql.go
@@ -2944,15 +2944,17 @@ func (rs *Rows) initContextClose(ctx, txctx context.Context) {
 	if bypassRowsAwaitDone {
 		return
 	}
-	ctx, rs.cancel = context.WithCancel(ctx)
-	go rs.awaitDone(ctx, txctx)
+	closectx, cancel := context.WithCancel(ctx)
+	rs.cancel = cancel
+	go rs.awaitDone(ctx, txctx, closectx)
 }
 
-// awaitDone blocks until either ctx or txctx is canceled. The ctx is provided
-// from the query context and is canceled when the query Rows is closed.
+// awaitDone blocks until ctx, txctx, or closectx is canceled. 
+// The ctx is provided from the query context.
 // If the query was issued in a transaction, the transaction's context
-// is also provided in txctx to ensure Rows is closed if the Tx is closed.
-func (rs *Rows) awaitDone(ctx, txctx context.Context) {
+// is also provided in txctx, to ensure Rows is closed if the Tx is closed.
+// The closectx is closed by an explicit call to rs.Close.
+func (rs *Rows) awaitDone(ctx, txctx, closectx context.Context) {
 	var txctxDone <-chan struct{}
 	if txctx != nil {
 		txctxDone = txctx.Done()
@@ -2964,6 +2966,9 @@ func (rs *Rows) awaitDone(ctx, txctx context.Context) {
 	case <-txctxDone:
 		err := txctx.Err()
 		rs.contextDone.Store(&err)
+	case <-closectx.Done():
+		// rs.cancel was called via Close(); don't store this into contextDone
+		// to ensure Err() is unaffected.
 	}
 	rs.close(ctx.Err())
 }
diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 718056c351..e6a5cd912a 100644
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -4493,6 +4493,31 @@ func TestContextCancelBetweenNextAndErr(t *testing.T) {
 	}
 }
 
+func TestNilErrorAfterClose(t *testing.T) {
+	db := newTestDB(t, "people")
+	defer closeDB(t, db)
+
+	// This WithCancel is important; Rows contains an optimization to avoid
+	// spawning a goroutine when the query/transaction context cannot be
+	// canceled, but this test tests a bug which is caused by said goroutine.
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	r, err := db.QueryContext(ctx, "SELECT|people|name|")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := r.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	time.Sleep(10 * time.Millisecond) // increase odds of seeing failure
+	if err := r.Err(); err != nil {
+		t.Fatal(err)
+	}
+}
+
 // badConn implements a bad driver.Conn, for TestBadDriver.
 // The Exec method panics.
 type badConn struct{}