Merge pull request #2393 from amirsalarsafaei/ping-timeout

feat(pgxpool): acquire ping timeout

Author: jackc

pgxpool/helper_test.go

@@ -0,0 +1,39 @@
+package pgxpool_test
+
+import (
+	"context"
+	"net"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgconn"
+)
+
+// delayProxy is a that introduces a configurable delay on reads from the database connection.
+type delayProxy struct {
+	net.Conn
+	readDelay time.Duration
+}
+
+func newDelayProxy(conn net.Conn, readDelay time.Duration) *delayProxy {
+	p := &delayProxy{
+		Conn:      conn,
+		readDelay: readDelay,
+	}
+
+	return p
+}
+
+func (dp *delayProxy) Read(b []byte) (int, error) {
+	if dp.readDelay > 0 {
+		time.Sleep(dp.readDelay)
+	}
+
+	return dp.Conn.Read(b)
+}
+
+func newDelayProxyDialFunc(readDelay time.Duration) pgconn.DialFunc {
+	return func(ctx context.Context, network, addr string) (net.Conn, error) {
+		conn, err := net.Dial(network, addr)
+		return newDelayProxy(conn, readDelay), err
+	}
+}

pgxpool/pool.go

@@ -97,6 +97,7 @@ type Pool struct {
 	maxConnLifetimeJitter time.Duration
 	maxConnIdleTime       time.Duration
 	healthCheckPeriod     time.Duration
+	pingTimeout           time.Duration
 
 	healthCheckMu    sync.Mutex
 	healthCheckTimer *time.Timer
@@ -169,6 +170,10 @@ type Config struct {
 	// MaxConnIdleTime is the duration after which an idle connection will be automatically closed by the health check.
 	MaxConnIdleTime time.Duration
 
+	// PingTimeout is the maximum amount of time to wait for a connection to pong before considering it as unhealthy and
+	// destroying it. If zero, the default is no timeout.
+	PingTimeout time.Duration
+
 	// MaxConns is the maximum size of the pool. The default is the greater of 4 or runtime.NumCPU().
 	MaxConns int32
 
@@ -241,6 +246,7 @@ func NewWithConfig(ctx context.Context, config *Config) (*Pool, error) {
 		maxConnLifetime:       config.MaxConnLifetime,
 		maxConnLifetimeJitter: config.MaxConnLifetimeJitter,
 		maxConnIdleTime:       config.MaxConnIdleTime,
+		pingTimeout:           config.PingTimeout,
 		healthCheckPeriod:     config.HealthCheckPeriod,
 		healthCheckChan:       make(chan struct{}, 1),
 		closeChan:             make(chan struct{}),
@@ -614,7 +620,14 @@ func (p *Pool) Acquire(ctx context.Context) (c *Conn, err error) {
 
 		shouldPingParams := ShouldPingParams{Conn: cr.conn, IdleDuration: res.IdleDuration()}
 		if p.shouldPing(ctx, shouldPingParams) {
-			err := cr.conn.Ping(ctx)
+			pingCtx := ctx
+			if p.pingTimeout > 0 {
+				var cancel context.CancelFunc
+				pingCtx, cancel = context.WithTimeout(ctx, p.pingTimeout)
+				defer cancel()
+			}
+
+			err := cr.conn.Ping(pingCtx)
 			if err != nil {
 				res.Destroy()
 				continue

pgxpool/pool_test.go

@@ -1281,3 +1281,54 @@ func TestPoolSendBatchBatchCloseTwice(t *testing.T) {
 		assert.NoError(t, err)
 	}
 }
+
+func TestPoolAcquirePingTimeout(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+
+	config.PingTimeout = 200 * time.Millisecond
+	config.ConnConfig.DialFunc = newDelayProxyDialFunc(500 * time.Millisecond)
+
+	var conID *uint32
+	// Only ping the connection with the original PID to force creation of a new connection
+	config.ShouldPing = func(_ context.Context, params pgxpool.ShouldPingParams) bool {
+		if conID != nil && params.Conn.PgConn().PID() == *conID {
+			return true
+		}
+		return false
+	}
+
+	// Limit to a single connection to ensure the same connection is reused
+	config.MinConns = 1
+	config.MaxConns = 1
+
+	pool, err := pgxpool.NewWithConfig(ctx, config)
+	require.NoError(t, err)
+	defer pool.Close()
+
+	c, err := pool.Acquire(ctx)
+	require.NoError(t, err)
+	require.EqualValues(t, 1, pool.Stat().TotalConns())
+	originalPID := c.Conn().PgConn().PID()
+	conID = &originalPID
+
+	c.Release()
+	require.EqualValues(t, 1, pool.Stat().TotalConns())
+
+	c, err = pool.Acquire(ctx)
+	require.NoError(t, err)
+	require.EqualValues(t, 1, pool.Stat().TotalConns())
+	newPID := c.Conn().PgConn().PID()
+
+	c.Release()
+
+	require.EqualValues(t, 1, pool.Stat().TotalConns())
+	assert.Nil(t, ctx.Err())
+	assert.NotEqualValues(t, originalPID, newPID,
+		"Expected new connection due to ping timeout, but got same connection")
+}