Merge #84394

84394: streamingccl: switch PartitionedStreamClient to use pgx library r=gh-casper a=gh-casper

    We have been switching to pgx from lib/pq across teams in CRL.
    This PR continues this journey and makes PartitionedStreamClient 
     to use pgx and also in a thread-safe way.

    To avoid overhead of creating connection, APIs except Subscribe
    share the same connection guarded by a mutex, and Subcribe owns
    a private connection which makes sense as two Subscribe calls
    don't want to share a connection.

    Fixes data race we have been seen in the past: 
    Closes: #83694

    Also unskip stream_replication_e2e_test under race.

    Release note: None

Co-authored-by: Casper <[email protected]>

Author: craig[bot]

pkg/ccl/streamingccl/streamclient/BUILD.bazel

@@ -35,6 +35,7 @@ go_library(
         "//pkg/util/timeutil",
         "//pkg/util/tracing",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_jackc_pgx_v4//:pgx",
     ],
 )
 
@@ -67,7 +68,6 @@ go_test(
         "//pkg/streaming",
         "//pkg/testutils",
         "//pkg/testutils/serverutils",
-        "//pkg/testutils/skip",
         "//pkg/testutils/testcluster",
         "//pkg/util/cancelchecker",
         "//pkg/util/ctxgroup",

pkg/ccl/streamingccl/streamclient/client.go

@@ -78,7 +78,7 @@ type Client interface {
 	) (Subscription, error)
 
 	// Close releases all the resources used by this client.
-	Close() error
+	Close(ctx context.Context) error
 
 	// Complete completes a replication stream consumption.
 	Complete(ctx context.Context, streamID streaming.StreamID) error
@@ -123,7 +123,9 @@ type Subscription interface {
 }
 
 // NewStreamClient creates a new stream client based on the stream address.
-func NewStreamClient(streamAddress streamingccl.StreamAddress) (Client, error) {
+func NewStreamClient(
+	ctx context.Context, streamAddress streamingccl.StreamAddress,
+) (Client, error) {
 	var streamClient Client
 	streamURL, err := streamAddress.URL()
 	if err != nil {
@@ -134,7 +136,7 @@ func NewStreamClient(streamAddress streamingccl.StreamAddress) (Client, error) {
 	case "postgres", "postgresql":
 		// The canonical PostgreSQL URL scheme is "postgresql", however our
 		// own client commands also accept "postgres".
-		return newPartitionedStreamClient(streamURL)
+		return newPartitionedStreamClient(ctx, streamURL)
 	case RandomGenScheme:
 		streamClient, err = newRandomStreamClient(streamURL)
 		if err != nil {

pkg/ccl/streamingccl/streamclient/client_test.go

@@ -50,7 +50,7 @@ func (sc testStreamClient) Heartbeat(
 }
 
 // Close implements the Client interface.
-func (sc testStreamClient) Close() error {
+func (sc testStreamClient) Close(ctx context.Context) error {
 	return nil
 }
 
@@ -109,11 +109,12 @@ func (t testStreamSubscription) Err() error {
 // client could be used.
 func ExampleClient() {
 	client := testStreamClient{}
+	ctx := context.Background()
 	defer func() {
-		_ = client.Close()
+		_ = client.Close(ctx)
 	}()
 
-	id, err := client.Create(context.Background(), roachpb.MakeTenantID(1))
+	id, err := client.Create(ctx, roachpb.MakeTenantID(1))
 	if err != nil {
 		panic(err)
 	}
@@ -125,7 +126,7 @@ func ExampleClient() {
 
 	done := make(chan struct{})
 
-	grp := ctxgroup.WithContext(context.Background())
+	grp := ctxgroup.WithContext(ctx)
 	grp.GoCtx(func(ctx context.Context) error {
 		ticker := time.NewTicker(time.Second * 30)
 		for {
@@ -150,14 +151,14 @@ func ExampleClient() {
 		ts := ingested.ts
 		ingested.Unlock()
 
-		topology, err := client.Plan(context.Background(), id)
+		topology, err := client.Plan(ctx, id)
 		if err != nil {
 			panic(err)
 		}
 
 		for _, partition := range topology {
 			// TODO(dt): use Subscribe helper and partition.SrcAddr
-			sub, err := client.Subscribe(context.Background(), id, partition.SubscriptionToken, ts)
+			sub, err := client.Subscribe(ctx, id, partition.SubscriptionToken, ts)
 			if err != nil {
 				panic(err)
 			}

pkg/ccl/streamingccl/streamclient/partitioned_stream_client.go

@@ -10,7 +10,6 @@ package streamclient
 
 import (
 	"context"
-	gosql "database/sql"
 	"net"
 	"net/url"
 
@@ -19,35 +18,45 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/streaming"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/errors"
+	"github.com/jackc/pgx/v4"
 )
 
 type partitionedStreamClient struct {
-	srcDB          *gosql.DB // DB handle to the source cluster
 	urlPlaceholder url.URL
+	pgxConfig      *pgx.ConnConfig
 
 	mu struct {
 		syncutil.Mutex
 
 		closed              bool
 		activeSubscriptions map[*partitionedStreamSubscription]struct{}
+		srcConn             *pgx.Conn // pgx connection to the source cluster
 	}
 }
 
-func newPartitionedStreamClient(remote *url.URL) (*partitionedStreamClient, error) {
-	db, err := gosql.Open("postgres", remote.String())
+func newPartitionedStreamClient(
+	ctx context.Context, remote *url.URL,
+) (*partitionedStreamClient, error) {
+	config, err := pgx.ParseConfig(remote.String())
+	if err != nil {
+		return nil, err
+	}
+	conn, err := pgx.ConnectConfig(ctx, config)
 	if err != nil {
 		return nil, err
 	}
 
 	client := &partitionedStreamClient{
-		srcDB:          db,
 		urlPlaceholder: *remote,
+		pgxConfig:      config,
 	}
 	client.mu.activeSubscriptions = make(map[*partitionedStreamSubscription]struct{})
+	client.mu.srcConn = conn
 	return client, nil
 }
 
@@ -60,19 +69,16 @@ func (p *partitionedStreamClient) Create(
 	ctx, sp := tracing.ChildSpan(ctx, "streamclient.Client.Create")
 	defer sp.Finish()
 
-	streamID := streaming.InvalidStreamID
-
-	conn, err := p.srcDB.Conn(ctx)
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	var streamID streaming.StreamID
+	row := p.mu.srcConn.QueryRow(ctx, `SELECT crdb_internal.start_replication_stream($1)`, tenantID.ToUint64())
+	err := row.Scan(&streamID)
 	if err != nil {
-		return streamID, err
-	}
-
-	row := conn.QueryRowContext(ctx, `SELECT crdb_internal.start_replication_stream($1)`, tenantID.ToUint64())
-	if row.Err() != nil {
-		return streamID, errors.Wrapf(row.Err(), "error creating replication stream for tenant %s", tenantID.String())
+		return streaming.InvalidStreamID,
+			errors.Wrapf(err, "error creating replication stream for tenant %s", tenantID.String())
 	}
 
-	err = row.Scan(&streamID)
 	return streamID, err
 }
 
@@ -83,21 +89,14 @@ func (p *partitionedStreamClient) Heartbeat(
 	ctx, sp := tracing.ChildSpan(ctx, "streamclient.Client.Heartbeat")
 	defer sp.Finish()
 
-	conn, err := p.srcDB.Conn(ctx)
-	if err != nil {
-		return streampb.StreamReplicationStatus{}, err
-	}
-
-	row := conn.QueryRowContext(ctx,
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	row := p.mu.srcConn.QueryRow(ctx,
 		`SELECT crdb_internal.replication_stream_progress($1, $2)`, streamID, consumed.String())
-	if row.Err() != nil {
-		return streampb.StreamReplicationStatus{},
-			errors.Wrapf(row.Err(), "error sending heartbeat to replication stream %d", streamID)
-	}
-
 	var rawStatus []byte
 	if err := row.Scan(&rawStatus); err != nil {
-		return streampb.StreamReplicationStatus{}, err
+		return streampb.StreamReplicationStatus{},
+			errors.Wrapf(err, "error sending heartbeat to replication stream %d", streamID)
 	}
 	var status streampb.StreamReplicationStatus
 	if err := protoutil.Unmarshal(rawStatus, &status); err != nil {
@@ -121,23 +120,19 @@ func (p *partitionedStreamClient) postgresURL(servingAddr string) (url.URL, erro
 func (p *partitionedStreamClient) Plan(
 	ctx context.Context, streamID streaming.StreamID,
 ) (Topology, error) {
-	conn, err := p.srcDB.Conn(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	row := conn.QueryRowContext(ctx, `SELECT crdb_internal.replication_stream_spec($1)`, streamID)
-	if row.Err() != nil {
-		return nil, errors.Wrapf(row.Err(), "error planning replication stream %d", streamID)
-	}
-
-	var rawSpec []byte
-	if err := row.Scan(&rawSpec); err != nil {
-		return nil, err
-	}
 	var spec streampb.ReplicationStreamSpec
-	if err := protoutil.Unmarshal(rawSpec, &spec); err != nil {
-		return nil, err
+	{
+		p.mu.Lock()
+		defer p.mu.Unlock()
+
+		row := p.mu.srcConn.QueryRow(ctx, `SELECT crdb_internal.replication_stream_spec($1)`, streamID)
+		var rawSpec []byte
+		if err := row.Scan(&rawSpec); err != nil {
+			return nil, errors.Wrapf(err, "error planning replication stream %d", streamID)
+		}
+		if err := protoutil.Unmarshal(rawSpec, &spec); err != nil {
+			return nil, err
+		}
 	}
 
 	topology := Topology{}
@@ -163,7 +158,7 @@ func (p *partitionedStreamClient) Plan(
 }
 
 // Close implements Client interface.
-func (p *partitionedStreamClient) Close() error {
+func (p *partitionedStreamClient) Close(ctx context.Context) error {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
@@ -176,7 +171,7 @@ func (p *partitionedStreamClient) Close() error {
 		close(sub.closeChan)
 		delete(p.mu.activeSubscriptions, sub)
 	}
-	return p.srcDB.Close()
+	return p.mu.srcConn.Close(ctx)
 }
 
 // Subscribe implements Client interface.
@@ -198,11 +193,11 @@ func (p *partitionedStreamClient) Subscribe(
 	}
 
 	res := &partitionedStreamSubscription{
-		eventsChan: make(chan streamingccl.Event),
-		db:         p.srcDB,
-		specBytes:  specBytes,
-		streamID:   stream,
-		closeChan:  make(chan struct{}),
+		eventsChan:    make(chan streamingccl.Event),
+		srcConnConfig: p.pgxConfig,
+		specBytes:     specBytes,
+		streamID:      stream,
+		closeChan:     make(chan struct{}),
 	}
 	p.mu.Lock()
 	defer p.mu.Unlock()
@@ -215,21 +210,19 @@ func (p *partitionedStreamClient) Complete(ctx context.Context, streamID streami
 	ctx, sp := tracing.ChildSpan(ctx, "streamclient.Client.Complete")
 	defer sp.Finish()
 
-	conn, err := p.srcDB.Conn(ctx)
-	if err != nil {
-		return err
-	}
-	row := conn.QueryRowContext(ctx, `SELECT crdb_internal.complete_replication_stream($1)`, streamID)
-	if row.Err() != nil {
-		return errors.Wrapf(row.Err(), "error completing replication stream %d", streamID)
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	row := p.mu.srcConn.QueryRow(ctx, `SELECT crdb_internal.complete_replication_stream($1)`, streamID)
+	if err := row.Scan(&streamID); err != nil {
+		return errors.Wrapf(err, "error completing replication stream %d", streamID)
 	}
 	return nil
 }
 
 type partitionedStreamSubscription struct {
-	err        error
-	db         *gosql.DB
-	eventsChan chan streamingccl.Event
+	err           error
+	srcConnConfig *pgx.ConnConfig
+	eventsChan    chan streamingccl.Event
 	// Channel to send signal to close the subscription.
 	closeChan chan struct{}
 
@@ -274,16 +267,22 @@ func (p *partitionedStreamSubscription) Subscribe(ctx context.Context) error {
 	defer sp.Finish()
 
 	defer close(p.eventsChan)
-	conn, err := p.db.Conn(ctx)
+	// Each subscription has its own pgx connection.
+	srcConn, err := pgx.ConnectConfig(ctx, p.srcConnConfig)
 	if err != nil {
 		return err
 	}
+	defer func() {
+		if err != nil {
+			log.Warningf(ctx, "error when closing subscription connection: %v", err)
+		}
+	}()
 
-	_, err = conn.ExecContext(ctx, `SET avoid_buffering = true`)
+	_, err = srcConn.Exec(ctx, `SET avoid_buffering = true`)
 	if err != nil {
 		return err
 	}
-	rows, err := conn.QueryContext(ctx, `SELECT * FROM crdb_internal.stream_partition($1, $2)`,
+	rows, err := srcConn.Query(ctx, `SELECT * FROM crdb_internal.stream_partition($1, $2)`,
 		p.streamID, p.specBytes)
 	if err != nil {
 		return err

pkg/ccl/streamingccl/streamclient/partitioned_stream_client_test.go

@@ -28,7 +28,6 @@ import (
 	"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"
@@ -58,7 +57,6 @@ func (f *subscriptionFeedSource) Close(ctx context.Context) {}
 
 func TestPartitionedStreamReplicationClient(t *testing.T) {
 	defer leaktest.AfterTest(t)()
-	skip.UnderRaceWithIssue(t, 83694)
 	defer log.Scope(t).Close(t)
 
 	h, cleanup := streamingtest.NewReplicationHelper(t,
@@ -90,9 +88,9 @@ INSERT INTO d.t1 (i) VALUES (42);
 INSERT INTO d.t2 VALUES (2);
 `)
 
-	client, err := newPartitionedStreamClient(&h.PGUrl)
+	client, err := newPartitionedStreamClient(ctx, &h.PGUrl)
 	defer func() {
-		require.NoError(t, client.Close())
+		require.NoError(t, client.Close(ctx))
 	}()
 	require.NoError(t, err)
 	expectStreamState := func(streamID streaming.StreamID, status jobs.Status) {
@@ -166,9 +164,9 @@ INSERT INTO d.t2 VALUES (2);
 	url, err := streamingccl.StreamAddress(top[0].SrcAddr).URL()
 	require.NoError(t, err)
 	// Create a new stream client with the given partition address.
-	subClient, err := newPartitionedStreamClient(url)
+	subClient, err := newPartitionedStreamClient(ctx, url)
 	defer func() {
-		require.NoError(t, subClient.Close())
+		require.NoError(t, subClient.Close(ctx))
 	}()
 	require.NoError(t, err)
 	sub, err := subClient.Subscribe(ctx, streamID, encodeSpec("t1"), hlc.Timestamp{})

pkg/ccl/streamingccl/streamclient/random_stream_client.go

@@ -339,7 +339,7 @@ func (m *randomStreamClient) getDescriptorAndNamespaceKVForTableID(
 }
 
 // Close implements the Client interface.
-func (m *randomStreamClient) Close() error {
+func (m *randomStreamClient) Close(ctx context.Context) error {
 	return nil
 }