test: add delay proxy for ping test

Author: amirsalarsafaei

pgxpool/helper_test.go

@@ -0,0 +1,146 @@
+package pgxpool_test
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+// delayProxy is a TCP proxy that introduces a configurable delay on reads from the database.
+type delayProxy struct {
+	listener      net.Listener
+	targetAddr    string
+	targetNetwork string
+	proxyConfig   *pgxpool.Config
+	readDelay     time.Duration
+	wg            sync.WaitGroup
+	closed        atomic.Bool
+}
+
+func newDelayProxy(t *testing.T, targetConnString string, readDelay time.Duration) *delayProxy {
+	t.Helper()
+
+	config, err := pgxpool.ParseConfig(targetConnString)
+	require.NoError(t, err)
+
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+
+	// Determine target network and address
+	var targetNetwork, targetAddr string
+	if strings.HasPrefix(config.ConnConfig.Host, "/") {
+		targetNetwork = "unix"
+		targetAddr = filepath.Join(config.ConnConfig.Host, fmt.Sprintf(".s.PGSQL.%d", config.ConnConfig.Port))
+	} else {
+		targetNetwork = "tcp"
+		targetAddr = net.JoinHostPort(config.ConnConfig.Host, strconv.Itoa(int(config.ConnConfig.Port)))
+	}
+
+	// Create proxy config with listener address
+	proxyAddr := listener.Addr().(*net.TCPAddr)
+	config.ConnConfig.Host = proxyAddr.IP.String()
+	config.ConnConfig.Port = uint16(proxyAddr.Port)
+
+	p := &delayProxy{
+		listener:      listener,
+		targetAddr:    targetAddr,
+		targetNetwork: targetNetwork,
+		proxyConfig:   config,
+		readDelay:     readDelay,
+	}
+
+	p.wg.Add(1)
+	go p.run()
+
+	t.Cleanup(func() { _ = p.Close() })
+
+	return p
+}
+
+func (p *delayProxy) run() {
+	defer p.wg.Done()
+
+	for {
+		clientConn, err := p.listener.Accept()
+		if err != nil {
+			if p.closed.Load() {
+				return
+			}
+			continue
+		}
+
+		p.wg.Add(1)
+		go p.handleConn(clientConn)
+	}
+}
+
+func (p *delayProxy) handleConn(clientConn net.Conn) {
+	defer p.wg.Done()
+	defer clientConn.Close()
+
+	serverConn, err := net.Dial(p.targetNetwork, p.targetAddr)
+	if err != nil {
+		return
+	}
+	defer serverConn.Close()
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	// Client -> Server (no delay)
+	go func() {
+		defer wg.Done()
+		_, _ = io.Copy(serverConn, clientConn)
+		if tcpConn, ok := serverConn.(*net.TCPConn); ok {
+			_ = tcpConn.CloseWrite()
+		}
+	}()
+
+	// Server -> Client (with delay)
+	go func() {
+		defer wg.Done()
+		delayedReader := &delayedReader{r: serverConn, delay: p.readDelay}
+		_, _ = io.Copy(clientConn, delayedReader)
+		if tcpConn, ok := clientConn.(*net.TCPConn); ok {
+			_ = tcpConn.CloseWrite()
+		}
+	}()
+
+	wg.Wait()
+}
+
+type delayedReader struct {
+	r     io.Reader
+	delay time.Duration
+}
+
+func (dr *delayedReader) Read(b []byte) (int, error) {
+	if dr.delay > 0 {
+		time.Sleep(dr.delay)
+	}
+	return dr.r.Read(b)
+}
+
+func (p *delayProxy) Config() *pgxpool.Config {
+	return p.proxyConfig.Copy()
+}
+
+func (p *delayProxy) Close() error {
+	if !p.closed.CompareAndSwap(false, true) {
+		return nil
+	}
+
+	err := p.listener.Close()
+	p.wg.Wait()
+	return err
+}

pgxpool/pool_test.go

@@ -1288,21 +1288,20 @@ func TestPoolAcquirePingTimeout(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
 	defer cancel()
 
-	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
-	require.NoError(t, err)
+	delayProxy := newDelayProxy(t, os.Getenv("PGX_TEST_DATABASE"), 500*time.Millisecond)
 
-	// Set a very short ping timeout to force timeout during ping and destruction of the connection
-	config.PingTimeout = 1 * time.Nanosecond
+	config := delayProxy.Config()
+	config.PingTimeout = 200 * 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