sqlproxyccl: pass in ctx when transferring connections

Previously, the span used when transferring connections was reused from
the forwarder. However, transferring connections has async components,
so it's possible that the forwarder calls `Finish` on the span before
the connection transfer is done using the span. This leads to a panic:
`panic: use of Span after Finish`. To address this, this patch passes a
context in when transferring a connection, instead of using the
forwarder's context.

Fixes: #153569
Release note: None

Author: davidwding

pkg/ccl/sqlproxyccl/balancer/balancer.go

@@ -371,7 +371,7 @@ func (b *Balancer) processQueue(ctx context.Context) {
 
 			// Each request is retried up to maxTransferAttempts.
 			for i := 0; i < maxTransferAttempts && ctx.Err() == nil; i++ {
-				err := req.conn.TransferConnection()
+				err := req.conn.TransferConnection(ctx)
 				if err == nil || errors.Is(err, context.Canceled) {
 					break
 				}

pkg/ccl/sqlproxyccl/balancer/conn.go

@@ -19,7 +19,7 @@ type ConnectionHandle interface {
 	// TransferConnection performs a connection migration on the connection
 	// handle. Invoking this blocks until the connection migration process has
 	// been completed.
-	TransferConnection() error
+	TransferConnection(context.Context) error
 
 	// IsIdle returns true if the connection is idle, and false otherwise.
 	IsIdle() bool

pkg/ccl/sqlproxyccl/balancer/conn_test.go

@@ -44,7 +44,7 @@ func (h *testConnHandle) Close() {
 }
 
 // TransferConnection implements the ConnectionHandle interface.
-func (h *testConnHandle) TransferConnection() error {
+func (h *testConnHandle) TransferConnection(ctx context.Context) error {
 	h.mu.Lock()
 	defer h.mu.Unlock()
 	h.mu.onTransferConnectionCount++

pkg/ccl/sqlproxyccl/conn_migration.go

@@ -127,7 +127,7 @@ var errTransferCannotStart = errors.New("transfer cannot be started")
 // error).
 //
 // TransferConnection implements the balancer.ConnectionHandle interface.
-func (f *forwarder) TransferConnection() (retErr error) {
+func (f *forwarder) TransferConnection(ctx context.Context) (retErr error) {
 	// A previous non-recoverable transfer would have closed the forwarder, so
 	// return right away.
 	if f.ctx.Err() != nil {
@@ -146,18 +146,18 @@ func (f *forwarder) TransferConnection() (retErr error) {
 	// whenever the context expires. We have to close the forwarder because
 	// the transfer may be blocked on I/O, and the only way for now is to close
 	// the connections. This then allow TransferConnection to return and cleanup.
-	ctx, cancel := newTransferContext(f.ctx)
+	transferCtx, cancel := newTransferContext(ctx)
 	defer cancel()
 
 	// Use a separate handler for timeouts. This is the only way to handle
 	// blocked I/Os as described above.
 	go func() {
-		<-ctx.Done()
+		<-transferCtx.Done()
 		// This Close call here in addition to the one in the defer callback
 		// below is on purpose. This would help unblock situations where we're
 		// blocked on sending/reading messages from connections that couldn't
 		// be handled with context.Context.
-		if !ctx.isRecoverable() {
+		if !transferCtx.isRecoverable() {
 			f.Close()
 		}
 	}()
@@ -179,7 +179,7 @@ func (f *forwarder) TransferConnection() (retErr error) {
 
 		// When TransferConnection returns, it's either the forwarder has been
 		// closed, or the procesors have been resumed.
-		if !ctx.isRecoverable() {
+		if !transferCtx.isRecoverable() {
 			log.Dev.Infof(logCtx, "transfer failed: connection closed, latency=%v, err=%v", latencyDur, retErr)
 			f.metrics.ConnMigrationErrorFatalCount.Inc(1)
 			f.Close()
@@ -201,16 +201,16 @@ func (f *forwarder) TransferConnection() (retErr error) {
 
 	// Suspend both processors before starting the transfer.
 	request, response := f.getProcessors()
-	if err := request.suspend(ctx); err != nil {
+	if err := request.suspend(transferCtx); err != nil {
 		return errors.Wrap(err, "suspending request processor")
 	}
-	if err := response.suspend(ctx); err != nil {
+	if err := response.suspend(transferCtx); err != nil {
 		return errors.Wrap(err, "suspending response processor")
 	}
 
 	// Transfer the connection.
 	clientConn, serverConn := f.getConns()
-	newServerConn, err := transferConnection(ctx, f, f.connector, f.metrics, clientConn, serverConn)
+	newServerConn, err := transferConnection(transferCtx, f, f.connector, f.metrics, clientConn, serverConn)
 	if err != nil {
 		return errors.Wrap(err, "transferring connection")
 	}

pkg/ccl/sqlproxyccl/proxy_handler_test.go

@@ -2018,12 +2018,12 @@ func TestConnectionMigration(t *testing.T) {
 		require.Equal(t, totalAttempts, count)
 	}
 
-	transferConnWithRetries := func(t *testing.T, f *forwarder) error {
+	transferConnWithRetries := func(t *testing.T, ctx context.Context, f *forwarder) error {
 		t.Helper()
 
 		var nonRetriableErrSeen bool
 		err := testutils.SucceedsSoonError(func() error {
-			err := f.TransferConnection()
+			err := f.TransferConnection(ctx)
 			if err == nil {
 				return nil
 			}
@@ -2091,7 +2091,7 @@ func TestConnectionMigration(t *testing.T) {
 			require.NoError(t, err)
 
 			// Show that we get alternating SQL pods when we transfer.
-			require.NoError(t, transferConnWithRetries(t, f))
+			require.NoError(t, transferConnWithRetries(t, tCtx, f))
 			require.Equal(t, int64(1), f.metrics.ConnMigrationSuccessCount.Count())
 			require.Equal(t, tenant2.SQLAddr(), queryAddr(tCtx, t, db))
 
@@ -2102,7 +2102,7 @@ func TestConnectionMigration(t *testing.T) {
 			_, err = db.Exec("SET application_name = 'bar'")
 			require.NoError(t, err)
 
-			require.NoError(t, transferConnWithRetries(t, f))
+			require.NoError(t, transferConnWithRetries(t, tCtx, f))
 			require.Equal(t, int64(2), f.metrics.ConnMigrationSuccessCount.Count())
 			require.Equal(t, tenant1.SQLAddr(), queryAddr(tCtx, t, db))
 
@@ -2120,14 +2120,14 @@ func TestConnectionMigration(t *testing.T) {
 			go func() {
 				defer wg.Done()
 				for subCtx.Err() == nil {
-					_ = f.TransferConnection()
+					_ = f.TransferConnection(tCtx)
 					time.Sleep(100 * time.Millisecond)
 				}
 			}()
 
 			// This loop will run approximately 5 seconds.
 			var tenant1Addr, tenant2Addr int
-			for i := 0; i < 100; i++ {
+			for range 100 {
 				addr := queryAddr(tCtx, t, db)
 				if addr == tenant1.SQLAddr() {
 					tenant1Addr++
@@ -2167,7 +2167,7 @@ func TestConnectionMigration(t *testing.T) {
 			err = crdb.ExecuteTx(tCtx, db, nil /* txopts */, func(tx *gosql.Tx) error {
 				// Run multiple times to ensure that connection isn't closed.
 				for i := 0; i < 5; {
-					err := f.TransferConnection()
+					err := f.TransferConnection(tCtx)
 					if err == nil {
 						return errors.New("no error")
 					}
@@ -2201,7 +2201,7 @@ func TestConnectionMigration(t *testing.T) {
 			require.Equal(t, int64(0), f.metrics.ConnMigrationErrorFatalCount.Count())
 
 			// Once the transaction is closed, transfers should work.
-			require.NoError(t, transferConnWithRetries(t, f))
+			require.NoError(t, transferConnWithRetries(t, tCtx, f))
 			require.NotEqual(t, initAddr, queryAddr(tCtx, t, db))
 			require.Nil(t, f.ctx.Err())
 			require.Equal(t, initSuccessCount+1, f.metrics.ConnMigrationSuccessCount.Count())
@@ -2223,7 +2223,7 @@ func TestConnectionMigration(t *testing.T) {
 			lookupAddrDelayDuration = 10 * time.Second
 			defer testutils.TestingHook(&defaultTransferTimeout, 3*time.Second)()
 
-			err := f.TransferConnection()
+			err := f.TransferConnection(tCtx)
 			require.Error(t, err)
 			require.Regexp(t, "injected delays", err.Error())
 			require.Equal(t, initAddr, queryAddr(tCtx, t, db))
@@ -2319,7 +2319,7 @@ func TestConnectionMigration(t *testing.T) {
 		time.Sleep(2 * time.Second)
 		// This should be an error because the transfer timed out. Connection
 		// should automatically be closed.
-		require.Error(t, f.TransferConnection())
+		require.Error(t, f.TransferConnection(tCtx))
 
 		select {
 		case <-time.After(10 * time.Second):