From 9bde556742f775fb6b0d8553be1ead78867fec11 Mon Sep 17 00:00:00 2001
From: Rafi Shamim <rafi@cockroachlabs.com>
Date: Wed, 7 Jun 2023 16:10:13 -0400
Subject: [PATCH 1/2] sql/tests: fix data race in
 TestErrorDuringExtendedProtocolCommit

Release note: None
---
 pkg/sql/tests/autocommit_extended_protocol_test.go | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pkg/sql/tests/autocommit_extended_protocol_test.go b/pkg/sql/tests/autocommit_extended_protocol_test.go
index 066aab7a04ff..68e0bd80e634 100644
--- a/pkg/sql/tests/autocommit_extended_protocol_test.go
+++ b/pkg/sql/tests/autocommit_extended_protocol_test.go
@@ -16,6 +16,7 @@ import (
 	"errors"
 	"net/url"
 	"strings"
+	"sync/atomic"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
@@ -27,7 +28,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
-	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
 	"github.com/jackc/pgx/v4"
 	"github.com/stretchr/testify/require"
 )
@@ -158,14 +158,14 @@ func TestErrorDuringExtendedProtocolCommit(t *testing.T) {
 	var db *gosql.DB
 
 	var shouldErrorOnAutoCommit syncutil.AtomicBool
-	var traceID tracingpb.TraceID
+	var traceID atomic.Uint64
 	params, _ := CreateTestServerParams()
 	params.Knobs.SQLExecutor = &sql.ExecutorTestingKnobs{
 		DisableAutoCommitDuringExec: true,
 		BeforeExecute: func(ctx context.Context, stmt string, descriptors *descs.Collection) {
 			if strings.Contains(stmt, "SELECT 'cat'") {
 				shouldErrorOnAutoCommit.Set(true)
-				traceID = tracing.SpanFromContext(ctx).TraceID()
+				traceID.Store(uint64(tracing.SpanFromContext(ctx).TraceID()))
 			}
 		},
 		BeforeAutoCommit: func(ctx context.Context, stmt string) error {
@@ -174,7 +174,7 @@ func TestErrorDuringExtendedProtocolCommit(t *testing.T) {
 				// saw when executing our test query. This is so we know that this
 				// autocommit corresponds to our test qyery rather than an internal
 				// query.
-				if traceID == tracing.SpanFromContext(ctx).TraceID() {
+				if traceID.Load() == uint64(tracing.SpanFromContext(ctx).TraceID()) {
 					shouldErrorOnAutoCommit.Set(false)
 					return errors.New("injected error")
 				}

From 36724e21f7097456cb2d140e158469b1d27f2e08 Mon Sep 17 00:00:00 2001
From: Rafi Shamim <rafi@cockroachlabs.com>
Date: Wed, 7 Jun 2023 16:20:06 -0400
Subject: [PATCH 2/2] sql: deflake TestCancelQueriesRace

It can flake due to a goroutine that is out of our control. Now, we
ignore that goroutine.

Release note: None
---
 pkg/sql/show_test.go          | 6 ++++--
 pkg/sql/tests/BUILD.bazel     | 1 -
 pkg/util/leaktest/leaktest.go | 2 ++
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/pkg/sql/show_test.go b/pkg/sql/show_test.go
index 242c9189cdc9..ab6ffe22f933 100644
--- a/pkg/sql/show_test.go
+++ b/pkg/sql/show_test.go
@@ -1202,12 +1202,14 @@ func TestCancelQueriesRace(t *testing.T) {
 		_, _ = sqlDB.ExecContext(ctx, `SELECT pg_sleep(10)`)
 		close(waiter)
 	}()
-	_, _ = sqlDB.ExecContext(ctx, `CANCEL QUERIES (
+	_, err := sqlDB.ExecContext(ctx, `CANCEL QUERIES (
 		SELECT query_id FROM [SHOW QUERIES] WHERE query LIKE 'SELECT pg_sleep%'
 	)`)
-	_, _ = sqlDB.ExecContext(ctx, `CANCEL QUERIES (
+	require.NoError(t, err)
+	_, err = sqlDB.ExecContext(ctx, `CANCEL QUERIES (
 		SELECT query_id FROM [SHOW QUERIES] WHERE query LIKE 'SELECT pg_sleep%'
 	)`)
+	require.NoError(t, err)
 
 	cancel()
 	<-waiter
diff --git a/pkg/sql/tests/BUILD.bazel b/pkg/sql/tests/BUILD.bazel
index 572c3bed0ea7..8a7a7a610ad2 100644
--- a/pkg/sql/tests/BUILD.bazel
+++ b/pkg/sql/tests/BUILD.bazel
@@ -122,7 +122,6 @@ go_test(
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "//pkg/util/tracing",
-        "//pkg/util/tracing/tracingpb",
         "//pkg/util/uuid",
         "@com_github_cockroachdb_cockroach_go_v2//crdb",
         "@com_github_cockroachdb_cockroach_go_v2//crdb/crdbpgx",
diff --git a/pkg/util/leaktest/leaktest.go b/pkg/util/leaktest/leaktest.go
index b013b54b4d6f..fcdfd237c41a 100644
--- a/pkg/util/leaktest/leaktest.go
+++ b/pkg/util/leaktest/leaktest.go
@@ -62,6 +62,8 @@ func interestingGoroutines() map[int64]string {
 			strings.Contains(stack, "github.com/jackc/pgconn.(*PgConn).asyncClose.func1") ||
 			// Ignore pgconn which creates a goroutine to watch context cancellation.
 			strings.Contains(stack, "github.com/jackc/pgconn/internal/ctxwatch.(*ContextWatcher).Watch.func1") ||
+			// Ignore pq goroutine that watches for context cancellation.
+			strings.Contains(stack, "github.com/lib/pq.(*conn).watchCancel") ||
 			// Seems to be gccgo specific.
 			(runtime.Compiler == "gccgo" && strings.Contains(stack, "testing.T.Parallel")) ||
 			// Ignore intentionally long-running logging goroutines that live for the