Support configurable ShouldPing func when acquiring conns

Author: christiankiely

pgxpool/pool.go

@@ -87,6 +87,7 @@ type Pool struct {
 	prepareConn           func(context.Context, *pgx.Conn) (bool, error)
 	afterRelease          func(*pgx.Conn) bool
 	beforeClose           func(*pgx.Conn)
+	shouldPing            func(context.Context, ShouldPingParams) bool
 	minConns              int32
 	minIdleConns          int32
 	maxConns              int32
@@ -104,6 +105,12 @@ type Pool struct {
 	closeChan chan struct{}
 }
 
+// ShouldPingParams are the parameters passed to ShouldPing.
+type ShouldPingParams struct {
+	Conn         *pgx.Conn
+	IdleDuration time.Duration
+}
+
 // Config is the configuration struct for creating a pool. It must be created by [ParseConfig] and then it can be
 // modified.
 type Config struct {
@@ -143,6 +150,10 @@ type Config struct {
 	// BeforeClose is called right before a connection is closed and removed from the pool.
 	BeforeClose func(*pgx.Conn)
 
+	// ShouldPing is called after a connection is acquired from the pool. If it returns true, the connection is pinged to check for liveness.
+	// If this func is not set, the default behavior is to ping connections that have been idle for at least 1 second.
+	ShouldPing func(context.Context, ShouldPingParams) bool
+
 	// MaxConnLifetime is the duration since creation after which a connection will be automatically closed.
 	MaxConnLifetime time.Duration
 
@@ -238,6 +249,14 @@ func NewWithConfig(ctx context.Context, config *Config) (*Pool, error) {
 		p.releaseTracer = t
 	}
 
+	if config.ShouldPing != nil {
+		p.shouldPing = config.ShouldPing
+	} else {
+		p.shouldPing = func(ctx context.Context, params ShouldPingParams) bool {
+			return params.IdleDuration > time.Second
+		}
+	}
+
 	var err error
 	p.p, err = puddle.NewPool(
 		&puddle.Config[*connResource]{
@@ -578,7 +597,8 @@ func (p *Pool) Acquire(ctx context.Context) (c *Conn, err error) {
 
 		cr := res.Value()
 
-		if res.IdleDuration() > time.Second {
+		shouldPingParams := ShouldPingParams{Conn: cr.conn, IdleDuration: res.IdleDuration()}
+		if p.shouldPing(ctx, shouldPingParams) {
 			err := cr.conn.Ping(ctx)
 			if err != nil {
 				res.Destroy()

pgxpool/pool_test.go

@@ -204,6 +204,48 @@ func TestPoolAcquireChecksIdleConns(t *testing.T) {
 	require.NotContains(t, pids, cPID)
 }
 
+func TestPoolAcquireChecksIdleConnsWithShouldPing(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	controllerConn, err := pgx.Connect(ctx, os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	defer controllerConn.Close(ctx)
+	pgxtest.SkipCockroachDB(t, controllerConn, "Server does not support pg_terminate_backend() (https://github.com/cockroachdb/cockroach/issues/35897)")
+
+	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+
+	// Replace the default ShouldPing func
+	var shouldPingLastCalledWith *pgxpool.ShouldPingParams
+	config.ShouldPing = func(ctx context.Context, params pgxpool.ShouldPingParams) bool {
+		shouldPingLastCalledWith = &params
+		return false
+	}
+
+	pool, err := pgxpool.NewWithConfig(ctx, config)
+	require.NoError(t, err)
+	defer pool.Close()
+
+	c, err := pool.Acquire(ctx)
+	require.NoError(t, err)
+	c.Release()
+
+	time.Sleep(time.Millisecond * 200)
+
+	c, err = pool.Acquire(ctx)
+	require.NoError(t, err)
+	conn := c.Conn()
+
+	require.NotNil(t, shouldPingLastCalledWith)
+	assert.Equal(t, conn, shouldPingLastCalledWith.Conn)
+	assert.InDelta(t, time.Millisecond*200, shouldPingLastCalledWith.IdleDuration, float64(time.Millisecond*100))
+
+	c.Release()
+}
+
 func TestPoolAcquireFunc(t *testing.T) {
 	t.Parallel()