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>

diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 4f2a2d83ef5cdd854d47ff1dd4f2eaffb966eea5..718056c3512606d47352b4ac8d96f953f2e59005 100644 (file)
--- 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