pgxpool: don't destroy a healthy connection on caller cancel during ping

pgxpool/pool.go

@@ -655,11 +655,25 @@ func (p *Pool) Acquire(ctx context.Context) (c *Conn, err error) {
 
 		shouldPingParams := ShouldPingParams{Conn: cr.conn, IdleDuration: res.IdleDuration()}
 		if p.shouldPing(ctx, shouldPingParams) {
+			// If the caller already cancelled before we got here there is no
+			// point pinging: the acquire will have to fail with that error
+			// anyway, and running the ping with a cancelled context would
+			// trip the connection's context watcher and destroy a connection
+			// the pool had no reason to lose. Release and return early.
+			if err := ctx.Err(); err != nil {
+				res.Release()
+				return nil, err
+			}
+
+			// Detach the ping from the caller's context so a cancellation
+			// mid-ping cannot break the socket (the pgconn context watcher
+			// would send a cancel request and tear down the connection).
+			// pingTimeout still bounds how long we wait.
 			err := func() error {
-				pingCtx := ctx
+				pingCtx := context.WithoutCancel(ctx)
 				if p.pingTimeout > 0 {
 					var cancel context.CancelFunc
-					pingCtx, cancel = context.WithTimeout(ctx, p.pingTimeout)
+					pingCtx, cancel = context.WithTimeout(pingCtx, p.pingTimeout)
 					defer cancel()
 				}
 				return cr.conn.Ping(pingCtx)
@@ -668,6 +682,14 @@ func (p *Pool) Acquire(ctx context.Context) (c *Conn, err error) {
 				res.Destroy()
 				continue
 			}
+
+			// The ping succeeded, but the caller may have cancelled while
+			// it was in flight. The connection is still healthy, so put it
+			// back and surface the caller's error without destroying it.
+			if err := ctx.Err(); err != nil {
+				res.Release()
+				return nil, err
+			}
 		}
 
 		if p.prepareConn != nil {

pgxpool/pool_test.go

@@ -385,6 +385,59 @@ func TestPoolAcquireChecksIdleConnsWithShouldPing(t *testing.T) {
 	c.Release()
 }
 
+
+// Regression test: https://github.com/jackc/pgx/issues/2680
+// If the caller cancels its context during the liveness ping, the pool must
+// not destroy a healthy connection. Previously the ping used the caller's
+// context, so cancellation tripped the connection's context watcher, the
+// ping failed, and res.Destroy() was called even though the connection was
+// fine.
+func TestPoolAcquirePingDoesNotDestroyOnCallerCancel(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	config.MaxConns = 1
+	// Force the ping on every acquire so the race is deterministic.
+	config.ShouldPing = func(context.Context, pgxpool.ShouldPingParams) bool { return true }
+
+	pool, err := pgxpool.NewWithConfig(ctx, config)
+	require.NoError(t, err)
+	defer pool.Close()
+
+	// Prime the pool so a backend PID exists to compare against.
+	c, err := pool.Acquire(ctx)
+	require.NoError(t, err)
+	pidBefore := c.Conn().PgConn().PID()
+	c.Release()
+
+	before := pool.Stat()
+
+	// Cancel the acquire context immediately so the caller is already done
+	// by the time the pool starts (or finishes) the ping.
+	cancelCtx, cancelAcquire := context.WithCancel(ctx)
+	cancelAcquire()
+	_, err = pool.Acquire(cancelCtx)
+	require.ErrorIs(t, err, context.Canceled, "caller should still see context.Canceled")
+
+	after := pool.Stat()
+
+	// The pool must not have thrown away a connection just because the caller
+	// cancelled. If NewConnsCount went up or the backend PID changed, the
+	// bug has regressed.
+	assert.Equal(t, before.NewConnsCount(), after.NewConnsCount(),
+		"pool created a new connection; the previous one was destroyed on caller cancel")
+
+	c, err = pool.Acquire(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, pidBefore, c.Conn().PgConn().PID(),
+		"backend PID changed; the pool destroyed a healthy connection")
+	c.Release()
+}
+
 // https://github.com/jackc/pgx/issues/2379
 func TestPoolAcquireWithMaxConnsEqualsMaxInt32(t *testing.T) {
 	t.Parallel()