Fix race condition if exec's context is canceled just after completion

all_test.go

@@ -2204,3 +2204,86 @@ func TestBeginMode(t *testing.T) {
 		}
 	}
 }
+
+// https://gitlab.com/cznic/sqlite/-/issues/94
+func TestCancelRace(t *testing.T) {
+	tempDir, err := ioutil.TempDir("", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		os.RemoveAll(tempDir)
+	}()
+
+	db, err := sql.Open("sqlite", filepath.Join(tempDir, "testcancelrace.sqlite"))
+	if err != nil {
+		t.Fatalf("Failed to open database: %v", err)
+	}
+	defer db.Close()
+
+	tests := []struct {
+		name string
+		f    func(context.Context, *sql.DB) error
+	}{
+		{
+			"db.ExecContext",
+			func(ctx context.Context, d *sql.DB) error {
+				_, err := db.ExecContext(ctx, "select 1")
+				return err
+			},
+		},
+		{
+			"db.QueryContext",
+			func(ctx context.Context, d *sql.DB) error {
+				_, err := db.QueryContext(ctx, "select 1")
+				return err
+			},
+		},
+		{
+			"tx.ExecContext",
+			func(ctx context.Context, d *sql.DB) error {
+				tx, err := db.BeginTx(ctx, &sql.TxOptions{})
+				if err != nil {
+					return err
+				}
+				defer tx.Rollback()
+				if _, err := tx.ExecContext(ctx, "select 1"); err != nil {
+					return err
+				}
+				return tx.Rollback()
+			},
+		},
+		{
+			"tx.QueryContext",
+			func(ctx context.Context, d *sql.DB) error {
+				tx, err := db.BeginTx(ctx, &sql.TxOptions{})
+				if err != nil {
+					return err
+				}
+				defer tx.Rollback()
+				if _, err := tx.QueryContext(ctx, "select 1"); err != nil {
+					return err
+				}
+				return tx.Rollback()
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// this is a race condition, so it's not guaranteed to fail on any given run,
+			// but with a moderate number of iterations it will eventually catch it
+			iterations := 100
+			for i := 0; i < iterations; i++ {
+				// none of these iterations should ever fail, because we never cancel their
+				// context until after they complete
+				ctx, cancel := context.WithCancel(context.Background())
+				if err := tt.f(ctx, db); err != nil {
+					t.Fatalf("Failed to run test query on iteration %d: %v", i, err)
+				}
+				cancel()
+			}
+		})
+	}
+}