pgxpool: Acquire destroys a healthy idle connection when the caller's context is cancelled during the liveness ping #2680

shouldPingParams := ShouldPingParams{Conn: cr.conn, IdleDuration: res.IdleDuration()}
if p.shouldPing(ctx, shouldPingParams) {
    err := func() error {
        pingCtx := ctx
        ...
        return cr.conn.Ping(pingCtx)
    }()
    if err != nil {
        res.Destroy()
        continue
    }
}

// Repro: pgxpool destroys a healthy idle connection when the caller's context is
// cancelled while Acquire is running its liveness ping.
//
// usage: go run . -dsn 'postgres://...' [-delay 20ms] [-cancel-after 5ms] [-rounds 5] [-detached]
package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/url"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// delayProxy forwards TCP to target, delaying every server->client chunk by d,
// so a ping takes at least d (stands in for a network round trip).
func delayProxy(target string, d time.Duration) string {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		log.Fatal(err)
	}
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go func() {
				s, err := net.Dial("tcp", target)
				if err != nil {
					c.Close()
					return
				}
				go func() { io.Copy(s, c); s.Close() }()
				buf := make([]byte, 32<<10)
				for {
					n, err := s.Read(buf)
					if n > 0 {
						time.Sleep(d)
						if _, werr := c.Write(buf[:n]); werr != nil {
							break
						}
					}
					if err != nil {
						break
					}
				}
				c.Close()
			}()
		}
	}()
	return ln.Addr().String()
}

func main() {
	dsn := flag.String("dsn", "", "postgres connection string")
	delay := flag.Duration("delay", 20*time.Millisecond, "added server->client latency")
	cancelAfter := flag.Duration("cancel-after", 5*time.Millisecond, "cancel the acquirer's context after this long")
	rounds := flag.Int("rounds", 5, "rounds")
	detached := flag.Bool("detached", false, "use a ShouldPing hook that pings under context.WithoutCancel (proposed behaviour)")
	flag.Parse()

	u, err := url.Parse(*dsn)
	if err != nil {
		log.Fatal(err)
	}
	u.Host = delayProxy(u.Host, *delay)
	cfg, err := pgxpool.ParseConfig(u.String())
	if err != nil {
		log.Fatal(err)
	}
	cfg.MaxConns, cfg.MinConns = 1, 0
	if *detached {
		cfg.ShouldPing = func(ctx context.Context, p pgxpool.ShouldPingParams) bool {
			if p.IdleDuration <= time.Second { // pgx's default rule
				return false
			}
			pctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 2*time.Second)
			defer cancel()
			return p.Conn.Ping(pctx) != nil // only a failed ping lets pgx ping again and destroy
		}
		// If the caller left during the detached ping, hand the healthy connection back
		// (ok=true + error => pgx releases it) and return the caller's context error.
		cfg.PrepareConn = func(ctx context.Context, _ *pgx.Conn) (bool, error) { return true, ctx.Err() }
	}
	ctx := context.Background()
	pool, err := pgxpool.NewWithConfig(ctx, cfg)
	if err != nil {
		log.Fatal(err)
	}
	defer pool.Close()

	backendPID := func() uint32 {
		c, err := pool.Acquire(ctx)
		if err != nil {
			log.Fatal(err)
		}
		defer c.Release()
		return c.Conn().PgConn().PID()
	}
	fmt.Printf("pgx config: delay=%s cancel-after=%s detached=%v\n", *delay, *cancelAfter, *detached)
	pid := backendPID()
	for i := 1; i <= *rounds; i++ {
		time.Sleep(1100 * time.Millisecond) // idle > 1s, so the next Acquire pings
		cctx, cancel := context.WithCancel(ctx)
		time.AfterFunc(*cancelAfter, cancel)
		c, aerr := pool.Acquire(cctx)
		if c != nil {
			c.Release()
		}
		cancel()
		if aerr == nil || !errors.Is(aerr, context.Canceled) {
			fmt.Printf("round %d: acquire did not fail with context.Canceled: %v\n", i, aerr)
		}
		next := backendPID()
		st := pool.Stat()
		fmt.Printf("round %d: cancelled acquire -> %v | backend pid %d -> %d (same connection: %v) | NewConnsCount=%d MaxLifetimeDestroy=%d MaxIdleDestroy=%d\n",
			i, aerr, pid, next, pid == next, st.NewConnsCount(), st.MaxLifetimeDestroyCount(), st.MaxIdleDestroyCount())
		pid = next
	}
}