diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..fd358aa --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module github.com/laurentketterle-hub/pgx-1 + +go 1.26.5 diff --git a/main.go b/main.go index 49f4dee..8072d70 100644 --- a/main.go +++ b/main.go @@ -1,7 +1,17 @@ package main -import "fmt" +import ( + "fmt" + + "github.com/laurentketterle-hub/pgx-1/pgxpool" +) func main() { - fmt.Println("Hello, Bounty Hunter!") + // Create a pool with MaxConns=5 and a simple dialer for demonstration. + pool := pgxpool.NewPool(5, pgxpool.SimpleDialFunc()) + defer pool.Close() + + stat := pool.Stat() + fmt.Printf("Pool: MaxConns=%d Active=%d Idle=%d InFlight=%d\n", + stat.MaxConns, stat.ActiveConns, stat.IdleConns, stat.InFlightConns) } diff --git a/pgxpool/pool.go b/pgxpool/pool.go index 00afc75..7a07650 100644 --- a/pgxpool/pool.go +++ b/pgxpool/pool.go @@ -2,35 +2,244 @@ package pgxpool import ( "context" + "errors" + "fmt" "sync" "sync/atomic" ) +// Conn represents a connection acquired from the pool. +// Release() must be called to return the connection to the pool. +type Conn struct { + pool *Pool + conn *conn +} + +// Release returns the connection to the pool. +func (c *Conn) Release() { + if c.pool == nil || c.conn == nil { + return + } + c.pool.release(c.conn) + c.pool = nil + c.conn = nil +} + +// conn is an internal tracked connection. +type conn struct { + id int32 +} + +// DialFunc is the signature for a function that creates a new connection. +type DialFunc func(ctx context.Context) (*conn, error) + +// Pool is a concurrency-safe bounded connection pool. +// +// It tracks both active (acquired) connections and pending (in-flight) +// connection attempts to prevent connection storms. The invariant +// +// activeCount + inFlightConns + len(idleConns) <= maxConns +// +// is maintained under the pool mutex, ensuring MaxConns is never +// breached even during database recovery scenarios. type Pool struct { - // ... existing fields maxConns int32 - conns []*conn - inFlightConns int32 + idleConns []*conn // idle connections ready for reuse + activeCount int32 // currently acquired (checked-out) connections + inFlightConns int32 // connection attempts in progress (dial/handshake) mu sync.Mutex - // ... + cond *sync.Cond // signaled when a connection is released or a dial fails + closed bool + dialFunc DialFunc + nextID int32 +} + +// NewPool creates a new Pool that limits total connections to maxConns. +// dialFunc is called to establish new physical connections. +func NewPool(maxConns int32, dialFunc DialFunc) *Pool { + p := &Pool{ + maxConns: maxConns, + dialFunc: dialFunc, + idleConns: make([]*conn, 0, maxConns), + } + p.cond = sync.NewCond(&p.mu) + return p } +// Acquire obtains a connection from the pool, blocking until one is +// available or ctx is cancelled. +// +// Connection storm prevention: +// - Before dialing, Acquire atomically checks that +// activeCount + inFlightConns < maxConns and increments +// inFlightConns — all under the pool mutex. +// - If the limit is reached, the caller blocks on a condition +// variable until a connection is released or a pending dial +// completes (success or failure). +// - On dial failure, inFlightConns is decremented and waiters +// are woken so another goroutine can attempt a connection. +// - On dial success, the connection is added to the pool's +// accounting before inFlightConns is decremented, preserving +// the invariant. 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) { - atomic.AddInt32(&p.inFlightConns, 1) - p.mu.Unlock() - - conn, err := p.createNewConn(ctx) - atomic.AddInt32(&p.inFlightConns, -1) - if err != nil { - return nil, err + defer p.mu.Unlock() + + for { + if p.closed { + return nil, errors.New("pgxpool: pool is closed") + } + + // Check context before blocking + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + + // 1. Try to return an idle connection immediately. + if len(p.idleConns) > 0 { + c := p.idleConns[len(p.idleConns)-1] + p.idleConns = p.idleConns[:len(p.idleConns)-1] + p.activeCount++ + return &Conn{pool: p, conn: c}, nil } - return conn, nil + + // 2. Check if we're allowed to start a new dial. + // activeCount + inFlightConns < maxConns means there is + // at least one free slot (neither idle nor pending). + totalInUse := p.activeCount + atomic.LoadInt32(&p.inFlightConns) + if totalInUse < p.maxConns { + // Reserve a pending slot before releasing the lock. + atomic.AddInt32(&p.inFlightConns, 1) + p.mu.Unlock() + + // Perform the slow I/O outside the lock. + c, dialErr := p.dialFunc(ctx) + + // Reacquire the lock to atomically update pool state. + p.mu.Lock() + + if dialErr != nil { + // Dial failed: release the reserved slot and wake + // waiters so another caller can try. + atomic.AddInt32(&p.inFlightConns, -1) + p.cond.Broadcast() + return nil, fmt.Errorf("pgxpool: dial failed: %w", dialErr) + } + + // Dial succeeded: add to active set and release the + // pending reservation — both under the lock so the + // invariant + // activeCount + inFlightConns + len(idle) <= maxConns + // is never violated. + p.activeCount++ + atomic.AddInt32(&p.inFlightConns, -1) + + // Wake any waiters — another slot may have opened up. + p.cond.Broadcast() + return &Conn{pool: p, conn: c}, nil + } + + // 3. Pool is fully saturated (active + inFlight >= maxConns). + // Wait for a connection release or a dial completion, but + // respect context cancellation by waking when ctx is done. + done := p.waitWithContext(ctx) + if done { + // Context was cancelled while waiting; the deferred unlock + // will run, and the loop will catch ctx.Done() on the next + // iteration. + continue + } + // p.cond.Wait() was woken by a signal; loop again to + // re-check conditions. } - p.mu.Unlock() +} - // Wait for existing connection or retry logic... - return p.waitForConn(ctx) -} \ No newline at end of file +// waitWithContext waits on the pool's condition variable, but also +// respects ctx cancellation. It returns true if the context was +// cancelled (caller should re-check and return the error), false +// if the cond was signaled normally. +// +// Must be called with p.mu held. +func (p *Pool) waitWithContext(ctx context.Context) (cancelled bool) { + // Use a channel to detect context cancellation. + // We can't select on cond.Wait, so we spawn a goroutine + // that broadcasts the cond when ctx is done. + done := make(chan struct{}) + defer close(done) + + go func() { + select { + case <-ctx.Done(): + // Context cancelled: wake up the cond waiter. + p.mu.Lock() + p.cond.Broadcast() + p.mu.Unlock() + case <-done: + // Normal exit: context wasn't cancelled before + // cond.Wait returned. + } + }() + + p.cond.Wait() + + // Check if we woke up because of context cancellation. + select { + case <-ctx.Done(): + return true + default: + return false + } +} + +// release returns a connection to the pool's idle set and wakes a waiter. +func (p *Pool) release(c *conn) { + p.mu.Lock() + defer p.mu.Unlock() + + p.activeCount-- + p.idleConns = append(p.idleConns, c) + p.cond.Signal() +} + +// Close shuts down the pool. Subsequent Acquire calls return an error. +func (p *Pool) Close() { + p.mu.Lock() + defer p.mu.Unlock() + + if p.closed { + return + } + p.closed = true + p.idleConns = nil + p.cond.Broadcast() +} + +// Stat returns a snapshot of the pool's current state. +type Stat struct { + MaxConns int32 + ActiveConns int32 + IdleConns int32 + InFlightConns int32 +} + +// Stat returns current pool statistics. +func (p *Pool) Stat() Stat { + p.mu.Lock() + defer p.mu.Unlock() + + return Stat{ + MaxConns: p.maxConns, + ActiveConns: p.activeCount, + IdleConns: int32(len(p.idleConns)), + InFlightConns: atomic.LoadInt32(&p.inFlightConns), + } +} + +// Len returns the number of idle connections in the pool. +func (p *Pool) Len() int { + p.mu.Lock() + defer p.mu.Unlock() + return len(p.idleConns) +} diff --git a/pgxpool/pool_test.go b/pgxpool/pool_test.go new file mode 100644 index 0000000..713b52f --- /dev/null +++ b/pgxpool/pool_test.go @@ -0,0 +1,294 @@ +package pgxpool + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" +) + +// countDialer creates connections with an incrementing ID and tracks how +// many concurrent dials are in progress. +type countDialer struct { + mu sync.Mutex + maxConcurrent int32 // highest number of concurrent dials observed + current int32 // dials currently in progress + id int32 + // slow indicates dials should take some time (simulates network latency) + slow time.Duration +} + +func (d *countDialer) dial(ctx context.Context) (*conn, error) { + cur := atomic.AddInt32(&d.current, 1) + defer atomic.AddInt32(&d.current, -1) + + // Track the maximum concurrent dials ever observed + for { + max := atomic.LoadInt32(&d.maxConcurrent) + if cur <= max || atomic.CompareAndSwapInt32(&d.maxConcurrent, max, cur) { + break + } + } + + if d.slow > 0 { + select { + case <-time.After(d.slow): + case <-ctx.Done(): + return nil, ctx.Err() + } + } + + id := atomic.AddInt32(&d.id, 1) + return &conn{id: id}, nil +} + +func (d *countDialer) maxConcurrentObserved() int32 { + return atomic.LoadInt32(&d.maxConcurrent) +} + +func TestAcquireRespectsMaxConns(t *testing.T) { + ctx := context.Background() + maxConns := int32(5) + dialer := &countDialer{slow: 50 * time.Millisecond} + + pool := NewPool(maxConns, dialer.dial) + defer pool.Close() + + var wg sync.WaitGroup + numGoroutines := 30 + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + conn, err := pool.Acquire(ctx) + if err != nil { + t.Errorf("Acquire failed: %v", err) + return + } + // Hold the connection briefly + time.Sleep(20 * time.Millisecond) + conn.Release() + }() + } + + wg.Wait() + + maxObserved := dialer.maxConcurrentObserved() + if maxObserved > maxConns { + t.Errorf("max concurrent dials = %d, want <= %d (MaxConns violation!)", maxObserved, maxConns) + } else { + t.Logf("max concurrent dials = %d (within MaxConns=%d ✓)", maxObserved, maxConns) + } + + // All connections should be idle after releases + time.Sleep(100 * time.Millisecond) + if l := pool.Len(); l != int(maxConns) { + t.Logf("idle connections = %d (expected %d)", l, maxConns) + } +} + +func TestAcquireBlocksWhenAtCapacity(t *testing.T) { + ctx := context.Background() + maxConns := int32(3) + dialer := &countDialer{slow: 100 * time.Millisecond} + + pool := NewPool(maxConns, dialer.dial) + defer pool.Close() + + // Acquire all connections + conns := make([]*Conn, 0, maxConns) + for i := int32(0); i < maxConns; i++ { + c, err := pool.Acquire(ctx) + if err != nil { + t.Fatalf("Acquire %d failed: %v", i, err) + } + conns = append(conns, c) + } + + // Verify pool is saturated + stat := pool.Stat() + if stat.ActiveConns != maxConns { + t.Errorf("ActiveConns = %d, want %d", stat.ActiveConns, maxConns) + } + + // Try to acquire another - should block + acquired := make(chan struct{}) + go func() { + c, err := pool.Acquire(ctx) + if err != nil { + t.Logf("blocked Acquire returned: %v", err) + } else { + c.Release() + } + close(acquired) + }() + + // Should NOT acquire immediately + select { + case <-acquired: + t.Error("Acquire should have blocked but returned immediately") + case <-time.After(200 * time.Millisecond): + t.Log("Acquire correctly blocked at capacity ✓") + } + + // Release one connection - should unblock the waiter + conns[0].Release() + + select { + case <-acquired: + t.Log("Acquire unblocked after release ✓") + case <-time.After(2 * time.Second): + t.Error("Acquire did not unblock after release") + } + + // Release remaining + for i := 1; i < len(conns); i++ { + conns[i].Release() + } +} + +func TestPendingCounterDecrementedOnFailure(t *testing.T) { + ctx := context.Background() + maxConns := int32(2) + + // Dialer that fails for the first 3 attempts + var failCount int32 = 3 + failingDialer := func(ctx context.Context) (*conn, error) { + if atomic.AddInt32(&failCount, -1) >= 0 { + return nil, fmt.Errorf("simulated dial failure") + } + return &conn{id: 99}, nil + } + + pool := NewPool(maxConns, failingDialer) + defer pool.Close() + + // First few acquires should fail but not deadlock the pool + var successCount int32 + var wg sync.WaitGroup + + for i := 0; i < 5; i++ { + wg.Add(1) + go func() { + defer wg.Done() + conn, err := pool.Acquire(ctx) + if err == nil { + atomic.AddInt32(&successCount, 1) + time.Sleep(10 * time.Millisecond) + conn.Release() + } + }() + } + + wg.Wait() + + // After failures, the pool should still be functional + if s := atomic.LoadInt32(&successCount); s < 1 { + t.Errorf("expected at least 1 successful acquire, got %d", s) + } + + stat := pool.Stat() + t.Logf("Final stat: active=%d idle=%d inFlight=%d max=%d", + stat.ActiveConns, stat.IdleConns, stat.InFlightConns, stat.MaxConns) +} + +func TestContextCancellation(t *testing.T) { + maxConns := int32(1) + slowDialer := &countDialer{slow: 500 * time.Millisecond} + + pool := NewPool(maxConns, slowDialer.dial) + defer pool.Close() + + // Saturate the pool + c1, err := pool.Acquire(context.Background()) + if err != nil { + t.Fatalf("initial acquire failed: %v", err) + } + + // Try to acquire with a short deadline + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + _, err = pool.Acquire(ctx) + if err == nil { + t.Error("expected context.DeadlineExceeded but got nil") + } else if !errors.Is(err, context.DeadlineExceeded) { + t.Logf("got error: %v (expected DeadlineExceeded)", err) + } + + // Pool should still work after cancelled acquire + c1.Release() + c2, err := pool.Acquire(context.Background()) + if err != nil { + t.Fatalf("acquire after cancellation failed: %v", err) + } + c2.Release() +} + +func TestConcurrentReleaseAndAcquire(t *testing.T) { + ctx := context.Background() + maxConns := int32(5) + dialer := &countDialer{slow: 30 * time.Millisecond} + + pool := NewPool(maxConns, dialer.dial) + defer pool.Close() + + var wg sync.WaitGroup + iterations := 100 + goroutines := 10 + + var totalAcquires int32 + + for g := 0; g < goroutines; g++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < iterations; i++ { + c, err := pool.Acquire(ctx) + if err != nil { + t.Errorf("Acquire failed: %v", err) + return + } + atomic.AddInt32(&totalAcquires, 1) + time.Sleep(time.Millisecond) + c.Release() + } + }() + } + + wg.Wait() + + total := atomic.LoadInt32(&totalAcquires) + if total != int32(goroutines*iterations) { + t.Errorf("total acquires = %d, want %d", total, goroutines*iterations) + } + + maxObserved := dialer.maxConcurrentObserved() + t.Logf("max concurrent dials = %d (MaxConns=%d)", maxObserved, maxConns) + if maxObserved > maxConns { + t.Errorf("max concurrent dials %d exceeded MaxConns %d", maxObserved, maxConns) + } +} + +func TestPoolClose(t *testing.T) { + pool := NewPool(1, func(ctx context.Context) (*conn, error) { + return &conn{id: 1}, nil + }) + + c, err := pool.Acquire(context.Background()) + if err != nil { + t.Fatalf("acquire failed: %v", err) + } + c.Release() + + pool.Close() + + _, err = pool.Acquire(context.Background()) + if err == nil { + t.Error("expected error after Close()") + } +}