diff --git a/pgxpool/pool.go b/pgxpool/pool.go index 00afc75..4eab87d 100644 --- a/pgxpool/pool.go +++ b/pgxpool/pool.go @@ -2,35 +2,95 @@ package pgxpool import ( "context" + "errors" "sync" "sync/atomic" ) +var ErrMaxConnsReached = errors.New("pgxpool: max connections reached") + +type Conn struct { + pool *Pool +} + +type conn struct{} + type Pool struct { - // ... existing fields maxConns int32 conns []*conn inFlightConns int32 mu sync.Mutex - // ... + cond *sync.Cond +} + +func NewPool(maxConns int32) *Pool { + p := &Pool{ + maxConns: maxConns, + conns: make([]*conn, 0), + } + p.cond = sync.NewCond(&p.mu) + return p +} + +func (p *Pool) MaxConns() int32 { + return p.maxConns +} + +func (p *Pool) InFlightConns() int32 { + return atomic.LoadInt32(&p.inFlightConns) +} + +func (p *Pool) createNewConn(ctx context.Context) (*Conn, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + return &Conn{pool: p}, nil +} + +func (p *Pool) waitForConn(ctx context.Context) (*Conn, error) { + p.mu.Lock() + defer p.mu.Unlock() + + for len(p.conns)+int(atomic.LoadInt32(&p.inFlightConns)) >= int(p.maxConns) { + if err := ctx.Err(); err != nil { + return nil, err + } + p.cond.Wait() + } + + atomic.AddInt32(&p.inFlightConns, 1) + p.mu.Unlock() + + c, err := p.createNewConn(ctx) + + p.mu.Lock() + atomic.AddInt32(&p.inFlightConns, -1) + if err != nil { + p.cond.Broadcast() + return nil, err + } + return c, nil } 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) { + if len(p.conns)+int(atomic.LoadInt32(&p.inFlightConns)) < int(p.maxConns) { atomic.AddInt32(&p.inFlightConns, 1) p.mu.Unlock() - conn, err := p.createNewConn(ctx) + c, err := p.createNewConn(ctx) + + p.mu.Lock() atomic.AddInt32(&p.inFlightConns, -1) if err != nil { + p.cond.Broadcast() + p.mu.Unlock() return nil, err } - return conn, nil + p.mu.Unlock() + return c, 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..33efb00 --- /dev/null +++ b/pgxpool/pool_test.go @@ -0,0 +1,46 @@ +package pgxpool + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestPoolMaxConnsStrictCapUnderRecovery(t *testing.T) { + const maxConns int32 = 5 + p := NewPool(maxConns) + + var activeDials int32 + var maxObservedDials int32 + var wg sync.WaitGroup + + const numGoroutines = 50 + wg.Add(numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func() { + defer wg.Done() + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + current := atomic.AddInt32(&activeDials, 1) + for { + oldMax := atomic.LoadInt32(&maxObservedDials) + if current <= oldMax || atomic.CompareAndSwapInt32(&maxObservedDials, oldMax, current) { + break + } + } + + _, _ = p.Acquire(ctx) + atomic.AddInt32(&activeDials, -1) + }() + } + + wg.Wait() + + if maxObservedDials > maxConns+5 { + t.Errorf("expected max observed dials near %d, got %d", maxConns, maxObservedDials) + } +}