streamingccl: deflake TestPartitionedStreamReplicationClient

pkg/ccl/streamingccl/streamclient/BUILD.bazel

@@ -59,11 +59,13 @@ go_test(
         "//pkg/security/securitytest",
         "//pkg/server",
         "//pkg/sql/catalog/desctestutils",
+        "//pkg/sql/pgwire/pgcode",
         "//pkg/streaming",
         "//pkg/testutils",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/skip",
         "//pkg/testutils/testcluster",
+        "//pkg/util/cancelchecker",
         "//pkg/util/ctxgroup",
         "//pkg/util/hlc",
         "//pkg/util/leaktest",
@@ -72,6 +74,8 @@ go_test(
         "//pkg/util/randutil",
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
+        "@com_github_cockroachdb_errors//:errors",
+        "@com_github_lib_pq//:pq",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/ccl/streamingccl/streamclient/partitioned_stream_client_test.go

@@ -11,6 +11,7 @@ package streamclient
 import (
 	"context"
 	"fmt"
+	"strings"
 	"testing"
 	"time"
 
@@ -23,16 +24,20 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/desctestutils"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/streaming"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
+	"github.com/cockroachdb/cockroach/pkg/util/cancelchecker"
 	"github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	"github.com/cockroachdb/errors"
+	"github.com/lib/pq"
 	"github.com/stretchr/testify/require"
 )
 
@@ -189,7 +194,12 @@ INSERT INTO d.t2 VALUES (2);
 
 	// Test if Subscribe can react to cancellation signal.
 	cancelFn()
-	require.ErrorIs(t, cg.Wait(), context.Canceled)
+
+	// When the context is cancelled, lib/pq sends a query cancellation message to
+	// the server. Occasionally, we see the error from this cancellation before
+	// the subscribe function sees our local context cancellation.
+	err = cg.Wait()
+	require.True(t, errors.Is(err, context.Canceled) || isQueryCanceledError(err))
 
 	// Testing client.Complete()
 	err = client.Complete(ctx, streaming.StreamID(999))
@@ -205,3 +215,12 @@ SET CLUSTER SETTING stream_replication.stream_liveness_track_frequency = '200ms'
 	h.SysDB.CheckQueryResultsRetry(t,
 		fmt.Sprintf("SELECT status FROM [SHOW JOBS] WHERE job_id = %d", streamID), [][]string{{"succeeded"}})
 }
+
+// isQueryCanceledError returns true if the error appears to be a query cancelled error.
+func isQueryCanceledError(err error) bool {
+	var pqErr pq.Error
+	if ok := errors.As(err, &pqErr); ok {
+		return pqErr.Code == pq.ErrorCode(pgcode.QueryCanceled.String())
+	}
+	return strings.Contains(err.Error(), cancelchecker.QueryCanceledError.Error())
+}