Function TestConnectContextCancelHandlerRace

pgconn/pgconn_test.go:4896–4948

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 CancelRequestConte

(t *testing.T)

//
// 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()