Fix data race when context cancelled during connect

The application-supplied BuildContextWatcherHandler was armed during the
auth handshake. If the context was cancelled mid-connect, handlers that
read *PgConn fields (notably CancelRequestContextWatcherHandler, which
reads pgConn.pid and pgConn.secretKey via CancelRequest) would race with
the connect goroutine's writes to those fields when handling
BackendKeyData. CancelRequest semantics also don't make sense before the
connection is established.

Use the deadline-only handler during connect and swap in the
application-supplied handler at ReadyForQuery, before any queries
ValidateConnect might run.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: jackc

pgconn/pgconn.go

@@ -362,7 +362,13 @@ func connectOne(ctx context.Context, config *Config, connectConfig *connectOneCo
 		}
 	}
 
-	pgConn.contextWatcher = ctxwatch.NewContextWatcher(config.BuildContextWatcherHandler(pgConn))
+	// Use a deadline-only watcher during connect. The application-supplied
+	// BuildContextWatcherHandler may read *PgConn fields (e.g.
+	// CancelRequestContextWatcherHandler reads pgConn.pid and
+	// pgConn.secretKey), which would race with this function's writes to
+	// those fields when handling BackendKeyData. The application handler is
+	// installed below, after the connection reaches connStatusIdle.
+	pgConn.contextWatcher = ctxwatch.NewContextWatcher(&DeadlineContextWatcherHandler{Conn: pgConn.conn})
 	pgConn.contextWatcher.Watch(ctx)
 	defer pgConn.contextWatcher.Unwatch()
 
@@ -454,14 +460,13 @@ func connectOne(ctx context.Context, config *Config, connectConfig *connectOneCo
 			}
 		case *pgproto3.ReadyForQuery:
 			pgConn.status = connStatusIdle
-			if config.ValidateConnect != nil {
-				// ValidateConnect may execute commands that cause the context to be watched again. Unwatch first to avoid
-				// the watch already in progress panic. This is that last thing done by this method so there is no need to
-				// restart the watch after ValidateConnect returns.
-				//
-				// See https://github.com/jackc/pgconn/issues/40.
-				pgConn.contextWatcher.Unwatch()
+			// The connect-phase deadline-only watcher is no longer needed; replace
+			// it with the application-supplied watcher so subsequent operations
+			// (including any queries run by ValidateConnect) use it.
+			pgConn.contextWatcher.Unwatch()
+			pgConn.contextWatcher = ctxwatch.NewContextWatcher(config.BuildContextWatcherHandler(pgConn))
 
+			if config.ValidateConnect != nil {
 				err := config.ValidateConnect(ctx, pgConn)
 				if err != nil {
 					if _, ok := err.(*NotPreferredError); ignoreNotPreferredErr && ok {

pgconn/pgconn_test.go

@@ -4841,6 +4841,91 @@ func TestCancelRequestContextWatcherHandler(t *testing.T) {
 	}
 }
 
+// raceProbeHandler implements ctxwatch.Handler. On HandleCancel it synchronously reads pgConn.PID() and
+// pgConn.SecretKey() — the same fields that CancelRequestContextWatcherHandler's spawned goroutine reads via
+// CancelRequest. Reading synchronously inside HandleCancel guarantees the read happens before HandleUnwatchAfterCancel
+// can cancel anything, eliminating the timing-window flakiness of the real handler and making this a reliable
+// race-detector reproducer for the underlying issue: during connect, *PgConn-reading handlers race with the auth loop's
+// writes to pgConn.pid and pgConn.secretKey when BackendKeyData arrives.
+type raceProbeHandler struct {
+	conn *pgconn.PgConn
+	pid  uint32
+	key  []byte
+}
+
+func (h *raceProbeHandler) HandleCancel(context.Context) {
+	h.pid = h.conn.PID()
+	h.key = h.conn.SecretKey()
+}
+
+func (h *raceProbeHandler) HandleUnwatchAfterCancel() {}
+
+// TestConnectContextCancelHandlerRace reproduces the data race that occurred when the application-supplied
+// BuildContextWatcherHandler was armed during the connect handshake. If the context was cancelled while connectOne was
+// still running, a handler that reads *PgConn fields (such as CancelRequestContextWatcherHandler, which calls
+// CancelRequest reading pgConn.pid and pgConn.secretKey) would race with the auth loop's writes to those fields when
+// handling BackendKeyData.
+//
+// The fix is to use the deadline-only handler during connect and only swap in the application-supplied handler after
+// the connection is established.
+//
+// Run with -race to detect the race.
+//
+// Race originally reported in https://github.com/jackc/pgx/pull/2534.
+func TestConnectContextCancelHandlerRace(t *testing.T) {
+	t.Parallel()
+
+	for i := range 10 {
+		t.Run(fmt.Sprintf("Stress %d", i), func(t *testing.T) {
+			t.Parallel()
+
+			script := &pgmock.Script{
+				Steps: []pgmock.Step{
+					pgmock.ExpectAnyMessage(&pgproto3.StartupMessage{ProtocolVersion: pgproto3.ProtocolVersion30, Parameters: map[string]string{}}),
+					pgmock.SendMessage(&pgproto3.AuthenticationOk{}),
+					pgmock.SendMessage(&pgproto3.BackendKeyData{ProcessID: 12345, SecretKey: []byte{1, 2, 3, 4}}),
+					pgmockWaitStep(200 * time.Millisecond),
+					pgmock.SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}),
+				},
+			}
+
+			ln, err := net.Listen("tcp", "127.0.0.1:")
+			require.NoError(t, err)
+			defer ln.Close()
+
+			go func() {
+				conn, err := ln.Accept()
+				if err != nil {
+					return
+				}
+				defer conn.Close()
+				_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
+				_ = script.Run(pgproto3.NewBackend(conn, conn))
+			}()
+
+			host, port, _ := strings.Cut(ln.Addr().String(), ":")
+			connStr := fmt.Sprintf("sslmode=disable host=%s port=%s user=test database=test", host, port)
+			config, err := pgconn.ParseConfig(connStr)
+			require.NoError(t, err)
+			config.BuildContextWatcherHandler = func(conn *pgconn.PgConn) ctxwatch.Handler {
+				return &raceProbeHandler{conn: conn}
+			}
+			config.ConnectTimeout = 5 * time.Second
+
+			ctx, cancel := context.WithCancel(context.Background())
+			// Cancel during the 200ms wait between BackendKeyData and ReadyForQuery,
+			// after the auth loop has written pid/secretKey.
+			time.AfterFunc(50*time.Millisecond, cancel)
+			defer cancel()
+
+			pgConn, err := pgconn.ConnectConfig(ctx, config)
+			if err == nil {
+				closeConn(t, pgConn)
+			}
+		})
+	}
+}
+
 func TestConnectProtocolVersion32(t *testing.T) {
 	t.Parallel()