pool.go

package pool

import (
	"context"
	"errors"
	"net"
	"sync"
	"time"
)

var ErrPoolClosed = errors.New("pool is closed")

type connEntry struct {
	conn   net.Conn
	idleAt time.Time
}

type Pool struct {
	mu          sync.Mutex
	maxSize     int
	idleTimeout time.Duration
	factory     func() (net.Conn, error)
	HealthCheck func(net.Conn) bool

	idle        []connEntry
	inFlight    int
	closed      bool
	notify      chan struct{}
	cleanupStop chan struct{}
}

func NewPool(maxSize int, idleTimeout time.Duration, factory func() (net.Conn, error)) *Pool {
	p := &Pool{
		maxSize:     maxSize,
		idleTimeout: idleTimeout,
		factory:     factory,
		HealthCheck: func(net.Conn) bool { return true },
		notify:      make(chan struct{}),
		cleanupStop: make(chan struct{}),
	}
	if idleTimeout > 0 {
		go p.cleanupLoop()
	}
	return p
}

func (p *Pool) Lock(fn func()) {
	p.mu.Lock()
	defer p.mu.Unlock()
	fn()
}

func (p *Pool) cleanupLoop() {
	ticker := time.NewTicker(p.idleTimeout / 2)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			p.removeExpired()
		case <-p.cleanupStop:
			return
		}
	}
}

func (p *Pool) removeExpired() {
	now := time.Now()
	p.mu.Lock()
	var toClose []net.Conn
	alive := p.idle[:0]
	for _, e := range p.idle {
		if now.Sub(e.idleAt) > p.idleTimeout {
			toClose = append(toClose, e.conn)
		} else {
			alive = append(alive, e)
		}
	}
	p.idle = alive
	p.mu.Unlock()
	for _, c := range toClose {
		c.Close()
	}
}

func (p *Pool) broadcast() { close(p.notify); p.notify = make(chan struct{}) }

// acquireState is the result of a single lock-protected decision pass.
type acquireState int

const (
	acquireReuse   acquireState = iota // got an idle candidate, validate outside lock
	acquireCreate                      // room to create a new connection
	acquireWait                        // at cap, must wait for a release
)

func (p *Pool) tryAcquire() (acquireState, net.Conn, chan struct{}) {
	// All lock-protected logic in one place, caller handles I/O outside.
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.closed {
		return acquireWait, nil, nil // caller checks closed first
	}

	// Pop from idle — skip expired entries (time-only check, no I/O).
	for len(p.idle) > 0 {
		e := p.idle[len(p.idle)-1]
		p.idle = p.idle[:len(p.idle)-1]
		if p.idleTimeout > 0 && time.Since(e.idleAt) > p.idleTimeout {
			e.conn.Close() // cheap for expired conns
			continue
		}
		p.inFlight++
		return acquireReuse, e.conn, nil
	}

	if p.inFlight < p.maxSize {
		p.inFlight++
		return acquireCreate, nil, nil
	}

	return acquireWait, nil, p.notify
}

func (p *Pool) Acquire(ctx context.Context) (net.Conn, error) {
	for {
		state, conn, ch := p.tryAcquire()

		switch state {
		case acquireReuse:
			if p.HealthCheck(conn) {
				return &pooledConn{pool: p, Conn: conn}, nil
			}
			conn.Close()
			p.mu.Lock()
			p.inFlight--
			p.broadcast()
			p.mu.Unlock()

		case acquireCreate:
			c, err := p.factory()
			if err != nil {
				p.mu.Lock()
				p.inFlight--
				p.broadcast()
				p.mu.Unlock()
				return nil, err
			}
			return &pooledConn{pool: p, Conn: c}, nil

		case acquireWait:
			if ch == nil { // pool is closed
				return nil, ErrPoolClosed
			}
			select {
			case <-ch:
			case <-ctx.Done():
				return nil, ctx.Err()
			}
		}
	}
}

func (p *Pool) release(conn net.Conn) {
	healthy := p.HealthCheck(conn) // I/O outside lock

	p.mu.Lock()
	p.inFlight--
	shouldClose := p.closed || !healthy
	if !shouldClose {
		p.idle = append(p.idle, connEntry{conn, time.Now()})
	}
	p.broadcast()
	p.mu.Unlock()

	if shouldClose {
		conn.Close() // outside lock
	}
}

func (p *Pool) Close() error {
	return p.CloseWithContext(context.Background())
}

func (p *Pool) CloseWithContext(ctx context.Context) error {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return nil
	}
	p.closed = true
	toClose := make([]net.Conn, len(p.idle))
	for i, e := range p.idle {
		toClose[i] = e.conn
	}
	p.idle = nil
	p.broadcast()
	p.mu.Unlock()

	close(p.cleanupStop)
	for _, c := range toClose {
		c.Close()
	}

	for {
		p.mu.Lock()
		if p.inFlight == 0 {
			p.mu.Unlock()
			return nil
		}
		ch := p.notify
		p.mu.Unlock()
		select {
		case <-ch:
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

type pooledConn struct {
	net.Conn
	pool *Pool
	once sync.Once
}

func (pc *pooledConn) Close() error {
	pc.once.Do(func() { pc.pool.release(pc.Conn) })
	return nil
}

pool_test.go

package pool

import (
	"context"
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

// mockConn is a simple net.Conn stub for testing.
type mockConn struct {
	closed int32 // atomic
}

func (m *mockConn) Read(b []byte) (n int, err error)  { return 0, net.ErrClosed }
func (m *mockConn) Write(b []byte) (n int, err error) { return 0, net.ErrClosed }
func (m *mockConn) Close() error                       { atomic.StoreInt32(&m.closed, 1); return nil }
func (m *mockConn) LocalAddr() net.Addr                { return nil }
func (m *mockConn) RemoteAddr() net.Addr               { return nil }
func (m *mockConn) SetDeadline(t time.Time) error      { return nil }
func (m *mockConn) SetReadDeadline(t time.Time) error  { return nil }
func (m *mockConn) SetWriteDeadline(t time.Time) error { return nil }

func (m *mockConn) isClosed() bool { return atomic.LoadInt32(&m.closed) == 1 }

func newFactory() (func() (net.Conn, error), *atomic.Int32) {
	var count atomic.Int32
	factory := func() (net.Conn, error) {
		count.Add(1)
		return &mockConn{}, nil
	}
	return factory, &count
}

// --- Tests ---

func TestBasicAcquireRelease(t *testing.T) {
	factory, count := newFactory()
	p := NewPool(3, 0, factory)

	c1, err := p.Acquire(context.Background())
	if err != nil {
		t.Fatalf("Acquire: %v", err)
	}
	if count.Load() != 1 {
		t.Fatalf("expected 1 connection created, got %d", count.Load())
	}

	// Return to pool.
	c1.Close()

	// Should reuse the same connection.
	c2, err := p.Acquire(context.Background())
	if err != nil {
		t.Fatalf("Acquire: %v", err)
	}
	if count.Load() != 1 {
		t.Fatalf("expected reuse (1 conn created), got %d", count.Load())
	}
	c2.Close()
	p.Close()
}

func TestMaxSizeBlocking(t *testing.T) {
	factory, _ := newFactory()
	p := NewPool(2, 0, factory)

	c1, _ := p.Acquire(context.Background())
	c2, _ := p.Acquire(context.Background())

	// Third acquire should block.
	done := make(chan struct{})
	go func() {
		c3, err := p.Acquire(context.Background())
		if err != nil {
			t.Errorf("blocked Acquire returned error: %v", err)
		} else {
			c3.Close()
		}
		close(done)
	}()

	// Give the goroutine time to block.
	time.Sleep(100 * time.Millisecond)

	// Release one connection.
	c1.Close()

	// The goroutine should unblock.
	select {
	case <-done:
		// good
	case <-time.After(time.Second):
		t.Fatal("Acquire should have unblocked after Release")
	}

	c2.Close()
	p.Close()
}

func TestContextCancellation(t *testing.T) {
	factory, _ := newFactory()
	p := NewPool(1, 0, factory)

	c1, _ := p.Acquire(context.Background())

	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()

	_, err := p.Acquire(ctx)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("expected context.DeadlineExceeded, got %v", err)
	}

	c1.Close()
	p.Close()
}

func TestIdleTimeout(t *testing.T) {
	factory, count := newFactory()
	p := NewPool(2, 100*time.Millisecond, factory)

	c1, _ := p.Acquire(context.Background())
	c1.Close() // returns to idle

	// Wait for idle timeout.
	time.Sleep(200 * time.Millisecond)

	// Should create a new connection since the old one expired.
	c2, _ := p.Acquire(context.Background())
	if count.Load() != 2 {
		t.Fatalf("expected 2 connections (expired one replaced), got %d", count.Load())
	}
	c2.Close()
	p.Close()
}

func TestCloseRejectsNewAcquires(t *testing.T) {
	factory, _ := newFactory()
	p := NewPool(2, 0, factory)

	p.Close()

	_, err := p.Acquire(context.Background())
	if !errors.Is(err, ErrPoolClosed) {
		t.Fatalf("expected ErrPoolClosed, got %v", err)
	}
}

func TestCloseWaitsForInFlight(t *testing.T) {
	factory, _ := newFactory()
	p := NewPool(2, 0, factory)

	c1, _ := p.Acquire(context.Background())
	c2, _ := p.Acquire(context.Background())

	closeDone := make(chan struct{})
	go func() {
		p.Close()
		close(closeDone)
	}()

	// Close should not finish yet.
	select {
	case <-closeDone:
		t.Fatal("Close should wait for in-flight connections")
	case <-time.After(100 * time.Millisecond):
		// expected
	}

	c1.Close()
	c2.Close()

	select {
	case <-closeDone:
		// good
	case <-time.After(time.Second):
		t.Fatal("Close should have finished after all connections returned")
	}
}

func TestDoubleClose(t *testing.T) {
	factory, _ := newFactory()
	p := NewPool(2, 0, factory)
	p.Close()
	p.Close() // should not panic
}

func TestConcurrentAcquireRelease(t *testing.T) {
	factory, _ := newFactory()
	p := NewPool(10, 0, factory)

	var wg sync.WaitGroup
	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			c, err := p.Acquire(context.Background())
			if err != nil {
				t.Errorf("Acquire: %v", err)
				return
			}
			// Simulate some work.
			time.Sleep(time.Millisecond)
			c.Close()
		}()
	}
	wg.Wait()
	p.Close()
}

func TestReleaseUnhealthyConnection(t *testing.T) {
	factory, count := newFactory()
	p := NewPool(2, 0, factory)
	p.HealthCheck = func(c net.Conn) bool {
		if mc, ok := c.(*mockConn); ok {
			return !mc.isClosed()
		}
		return true
	}

	c1, _ := p.Acquire(context.Background())
	raw := c1.(*pooledConn).Conn

	// Close the underlying connection to make it "unhealthy".
	raw.(*mockConn).Close()

	// Releasing should discard it (not return to idle).
	c1.Close()

	// Next acquire should create a new connection.
	c2, _ := p.Acquire(context.Background())
	if count.Load() != 2 {
		t.Fatalf("expected 2 connections (unhealthy one discarded), got %d", count.Load())
	}
	c2.Close()
	p.Close()
}

func TestReleaseWhenPoolAtCapacity(t *testing.T) {
	factory, _ := newFactory()
	p := NewPool(2, 0, factory)

	c1, _ := p.Acquire(context.Background())
	c2, _ := p.Acquire(context.Background())

	// Release both — pool should hold them in idle.
	c1.Close()
	c2.Close()
	p.Close()
}

func TestFactoryError(t *testing.T) {
	errFactory := func() (net.Conn, error) {
		return nil, errors.New("factory error")
	}
	p := NewPool(2, 0, errFactory)

	_, err := p.Acquire(context.Background())
	if err == nil || err.Error() != "factory error" {
		t.Fatalf("expected factory error, got %v", err)
	}
	p.Close()
}

// --- Concurrent stress tests ---

func TestConcurrentHighContention(t *testing.T) {
	factory, created := newFactory()
	pool := NewPool(5, 0, factory)

	const workers = 200
	const rounds = 50

	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < rounds; j++ {
				c, err := pool.Acquire(context.Background())
				if err != nil {
					t.Errorf("Acquire: %v", err)
					return
				}
				c.Close()
			}
		}()
	}
	wg.Wait()

	if n := created.Load(); n > int32(5) {
		t.Errorf("created %d connections, expected at most 5", n)
	}
	pool.Close()
}

func TestConcurrentWithContextCancel(t *testing.T) {
	factory, _ := newFactory()
	pool := NewPool(3, 0, factory)

	// Hold all 3 connections.
	conns := make([]net.Conn, 3)
	for i := range conns {
		conns[i], _ = pool.Acquire(context.Background())
	}

	var (
		wg        sync.WaitGroup
		cancelled atomic.Int32
		succeeded atomic.Int32
	)

	// 20 goroutines with short timeout — will all cancel.
	for i := 0; i < 20; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
			defer cancel()
			c, err := pool.Acquire(ctx)
			if err != nil {
				cancelled.Add(1)
				return
			}
			succeeded.Add(1)
			c.Close()
		}()
	}

	// 5 goroutines with long timeout — will succeed after release.
	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
			defer cancel()
			c, err := pool.Acquire(ctx)
			if err != nil {
				cancelled.Add(1)
				return
			}
			succeeded.Add(1)
			c.Close()
		}()
	}

	// Wait for short-timeout goroutines to cancel, then release.
	time.Sleep(200 * time.Millisecond)
	for _, c := range conns {
		c.Close()
	}

	wg.Wait()

	if cancelled.Load() == 0 {
		t.Error("expected some cancellations")
	}
	if succeeded.Load() == 0 {
		t.Error("expected some successes after release")
	}
	t.Logf("cancelled=%d, succeeded=%d", cancelled.Load(), succeeded.Load())
	pool.Close()
}

func TestConcurrentCloseWhileInFlight(t *testing.T) {
	factory, _ := newFactory()
	pool := NewPool(10, 0, factory)

	const workers = 50
	var (
		wg        sync.WaitGroup
		inFlight  sync.WaitGroup
		released  atomic.Int32
		acqErrors atomic.Int32
	)

	// Phase 1: all workers acquire a connection.
	inFlight.Add(workers)
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			c, err := pool.Acquire(context.Background())
			if err != nil {
				acqErrors.Add(1)
				inFlight.Done()
				return
			}
			inFlight.Done()
			// Hold until signal.
			time.Sleep(time.Duration(50+((i%5)*20)) * time.Millisecond)
			c.Close()
			released.Add(1)
		}()
	}

	// Wait until all workers have either acquired or failed.
	inFlight.Wait()

	// Phase 2: close pool while connections are still held.
	closeDone := make(chan struct{})
	go func() {
		pool.Close()
		close(closeDone)
	}()

	// Close should block until all connections are returned.
	select {
	case <-closeDone:
		t.Fatal("Close should wait for in-flight connections")
	case <-time.After(50 * time.Millisecond):
		// expected
	}

	// Wait for Close to finish.
	select {
	case <-closeDone:
		// good
	case <-time.After(5 * time.Second):
		t.Fatal("Close did not finish in time")
	}
	wg.Wait()

	// After Close, new acquires must fail.
	_, err := pool.Acquire(context.Background())
	if !errors.Is(err, ErrPoolClosed) {
		t.Fatalf("expected ErrPoolClosed after Close, got %v", err)
	}
}

func TestConcurrentIdleTimeout(t *testing.T) {
	factory, created := newFactory()
	pool := NewPool(5, 50*time.Millisecond, factory)

	const workers = 100
	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			c, err := pool.Acquire(context.Background())
			if err != nil {
				t.Errorf("Acquire: %v", err)
				return
			}
			time.Sleep(time.Millisecond)
			c.Close()
			// Random delay to let some connections expire.
			time.Sleep(time.Duration((i%10)*10) * time.Millisecond)
		}()
	}
	wg.Wait()

	total := created.Load()
	if total == 0 {
		t.Error("expected at least 1 connection created")
	}
	t.Logf("created %d connections (pool size 5, idle timeout 50ms)", total)
	pool.Close()
}

func TestConcurrentMixedOps(t *testing.T) {
	factory, created := newFactory()
	pool := NewPool(4, 200*time.Millisecond, factory)
	pool.HealthCheck = func(c net.Conn) bool {
		if mc, ok := c.(*mockConn); ok {
			return !mc.isClosed()
		}
		return true
	}

	const iterations = 500
	var wg sync.WaitGroup

	// Workers that acquire, maybe kill the conn, then release.
	for i := 0; i < iterations; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
			defer cancel()

			c, err := pool.Acquire(ctx)
			if err != nil {
				return // timeout or pool closed
			}

			// Simulate work.
			time.Sleep(time.Microsecond * time.Duration(i%100))

			// 10% chance: deliberately kill the underlying connection.
			if i%10 == 0 {
				if pc, ok := c.(*pooledConn); ok {
					if mc, ok := pc.Conn.(*mockConn); ok {
						mc.Close()
					}
				}
			}
			c.Close()
		}(i)
	}

	wg.Wait()
	t.Logf("created %d connections for %d iterations", created.Load(), iterations)
	pool.Close()
}

func TestConcurrentDoubleClose(t *testing.T) {
	factory, _ := newFactory()
	pool := NewPool(5, 0, factory)

	// Acquire some connections.
	for i := 0; i < 3; i++ {
		c, _ := pool.Acquire(context.Background())
		go func() {
			time.Sleep(50 * time.Millisecond)
			c.Close()
		}()
	}

	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			pool.Close() // concurrent Close should be safe
		}()
	}
	wg.Wait()
}

func TestConcurrentReleaseAndClose(t *testing.T) {
	factory, _ := newFactory()
	pool := NewPool(10, 0, factory)

	conns := make([]net.Conn, 10)
	for i := range conns {
		conns[i], _ = pool.Acquire(context.Background())
	}

	var wg sync.WaitGroup
	// Release all concurrently.
	for i, c := range conns {
		wg.Add(1)
		go func(c net.Conn) {
			defer wg.Done()
			time.Sleep(time.Duration((i%5)*10) * time.Millisecond)
			c.Close()
		}(c)
	}

	// Concurrently close the pool.
	go func() {
		time.Sleep(20 * time.Millisecond)
		pool.Close()
	}()

	wg.Wait()
}