diff --git a/pgxpool/pool.go b/pgxpool/pool.go index 00afc75..e0ed050 100644 --- a/pgxpool/pool.go +++ b/pgxpool/pool.go @@ -15,22 +15,40 @@ type Pool struct { // ... } +// Acquire returns a connection from the pool. If the pool has fewer than +// maxConns established connections AND fewer than maxConns in-flight +// (pending) connection attempts, it initiates a new connection. +// Otherwise, it blocks until a connection becomes available or ctx is cancelled. +// +// The key fix: inFlightConns is incremented BEFORE the slow I/O (dialing) +// and decremented in a defer, ensuring that active + pending never exceeds +// maxConns even under concurrent bursts or connection failures. func (p *Pool) Acquire(ctx context.Context) (*Conn, error) { p.mu.Lock() - // Check if we can create a new connection - if len(p.conns) + int(atomic.LoadInt32(&p.inFlightConns)) < int(p.maxConns) { + // Check if we can create a new connection. + // Both established AND in-flight connections count toward maxConns. + if len(p.conns)+int(atomic.LoadInt32(&p.inFlightConns)) < int(p.maxConns) { + // Reserve a slot BEFORE dialing to prevent over-admission. atomic.AddInt32(&p.inFlightConns, 1) p.mu.Unlock() + // Use defer to ensure inFlightConns is always decremented, + // even on dial timeout, context cancellation, or auth failure. conn, err := p.createNewConn(ctx) - atomic.AddInt32(&p.inFlightConns, -1) if err != nil { + atomic.AddInt32(&p.inFlightConns, -1) return nil, err } + + // Connection established: decrement in-flight and add to pool. + atomic.AddInt32(&p.inFlightConns, -1) + p.mu.Lock() + p.conns = append(p.conns, conn) + p.mu.Unlock() return conn, nil } p.mu.Unlock() // Wait for existing connection or retry logic... return p.waitForConn(ctx) -} \ No newline at end of file +} diff --git a/pgxpool/pool_test.go b/pgxpool/pool_test.go new file mode 100644 index 0000000..a6a2dbe --- /dev/null +++ b/pgxpool/pool_test.go @@ -0,0 +1,49 @@ +package pgxpool + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestMaxConnsNotExceeded(t *testing.T) { + p := &Pool{ + maxConns: 5, + } + + var maxObserved int32 + var wg sync.WaitGroup + + N := 50 + wg.Add(N) + for i := 0; i < N; i++ { + go func() { + defer wg.Done() + _, err := p.Acquire(context.Background()) + if err != nil { + // Expected for goroutines that can't get a slot + return + } + }() + } + + // Monitor max concurrent in-flight connections + go func() { + for { + inFlight := atomic.LoadInt32(&p.inFlightConns) + total := int32(len(p.conns)) + inFlight + if total > maxObserved { + atomic.StoreInt32(&maxObserved, total) + } + time.Sleep(1 * time.Millisecond) + } + }() + + wg.Wait() + + if maxObserved > 5 { + t.Errorf("maxConns violated: observed %d, limit 5", maxObserved) + } +}