database/sql: fix flake in TestContextCancelDuringRawBytesScan

If the cancellation takes effect between Next and Scan,
then Scan returns context.Canceled, but the test wasn't
allowing this case.

The old flake reproduced easily with:

	go test -c
	stress ./sql.test -test.count=100 -test.run=TestContextCancelDuringRawBytesScan

The new test modes exercise that path directly instead of needing stress.

The new check for context.Canceled fixes the new test mode "top".

Fixes #60445.

Change-Id: I3870039a0fbe0a43c3e261b43b175ef83f818765
Reviewed-on: https://go-review.googlesource.com/c/go/+/502876
Reviewed-by: Ian Lance Taylor <iant@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Auto-Submit: Russ Cox <rsc@golang.org>
Run-TryBot: Russ Cox <rsc@golang.org>

1 files changed, 34 insertions, 3 deletions

diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 4f2a2d83ef..718056c351 100644
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -4385,8 +4385,16 @@ func TestRowsScanProperlyWrapsErrors(t *testing.T) {
 	}
 }
 
-// From go.dev/issue/60304
 func TestContextCancelDuringRawBytesScan(t *testing.T) {
+	for _, mode := range []string{"nocancel", "top", "bottom", "go"} {
+		t.Run(mode, func(t *testing.T) {
+			testContextCancelDuringRawBytesScan(t, mode)
+		})
+	}
+}
+
+// From go.dev/issue/60304
+func testContextCancelDuringRawBytesScan(t *testing.T, mode string) {
 	db := newTestDB(t, "people")
 	defer closeDB(t, db)
 
@@ -4394,6 +4402,8 @@ func TestContextCancelDuringRawBytesScan(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	// cancel used to call close asynchronously.
+	// This test checks that it waits so as not to interfere with RawBytes.
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
@@ -4404,9 +4414,22 @@ func TestContextCancelDuringRawBytesScan(t *testing.T) {
 	numRows := 0
 	var sink byte
 	for r.Next() {
+		if mode == "top" && numRows == 2 {
+			// cancel between Next and Scan is observed by Scan as err = context.Canceled.
+			// The sleep here is only to make it more likely that the cancel will be observed.
+			// If not, the test should still pass, like in "go" mode.
+			cancel()
+			time.Sleep(100 * time.Millisecond)
+		}
 		numRows++
 		var s RawBytes
 		err = r.Scan(&s)
+		if numRows == 3 && err == context.Canceled {
+			if r.closemuScanHold {
+				t.Errorf("expected closemu NOT to be held")
+			}
+			break
+		}
 		if !r.closemuScanHold {
 			t.Errorf("expected closemu to be held")
 		}
@@ -4414,8 +4437,16 @@ func TestContextCancelDuringRawBytesScan(t *testing.T) {
 			t.Fatal(err)
 		}
 		t.Logf("read %q", s)
-		if numRows == 2 {
-			cancel() // invalidate the context, which used to call close asynchronously
+		if mode == "bottom" && numRows == 2 {
+			// cancel before Next should be observed by Next, exiting the loop.
+			// The sleep here is only to make it more likely that the cancel will be observed.
+			// If not, the test should still pass, like in "go" mode.
+			cancel()
+			time.Sleep(100 * time.Millisecond)
+		}
+		if mode == "go" && numRows == 2 {
+			// cancel at any future time, to catch other cases
+			go cancel()
 		}
 		for _, b := range s { // some operation reading from the raw memory
 			sink += b