《后端架构设计·高性能与高可用》5.1 连接池、线程池与协程池

// file: internal/platform/pool/bound.go
// tier: L1
// verify: go test ./internal/platform/pool/
package pool

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

var (
	ErrQueueFull     = errors.New("pool: queue full")
	ErrWaitTimeout   = errors.New("pool: wait timeout")
	ErrPoolExhausted = errors.New("pool: connection wait timeout")
	ErrLeaseReleased = errors.New("pool: lease already released")
)

type Policy int

const (
	Block Policy = iota // wait for room or caller cancellation
	RejectOnFull
	RejectOnTimeout
)

type Config struct {
	Workers, QueueSize int
	WaitTimeout        time.Duration
	Policy             Policy
}

type Snapshot struct {
	Submitted, Completed, Rejected, WaitTimeouts int64
	InFlight, Queued                             int64
}

// Pool keeps queue depth and in-flight count as separate limits: a full queue
// means work has not started, high in-flight means the downstream is the
// bottleneck. Neither number alone explains saturation.
type Pool struct {
	cfg                                         Config
	tasks                                       chan func()
	wg                                          sync.WaitGroup
	closed                                      atomic.Bool
	submitted, completed, rejected, waitTimeout atomic.Int64
	inFlight                                    atomic.Int64
}

func New(cfg Config) *Pool {
	if cfg.Workers < 1 {
		cfg.Workers = 1
	}
	p := &Pool{cfg: cfg, tasks: make(chan func(), cfg.QueueSize)}
	p.wg.Add(cfg.Workers)
	for i := 0; i < cfg.Workers; i++ {
		go p.worker()
	}
	return p
}

func (p *Pool) worker() {
	defer p.wg.Done()
	for task := range p.tasks {
		p.inFlight.Add(1)
		task()
		p.inFlight.Add(-1)
		p.completed.Add(1)
	}
}

// Submit never invents capacity by spawning a goroutine per task: the bounded
// queue is the only buffer, and the policy decides what a full queue means.
func (p *Pool) Submit(ctx context.Context, task func()) error {
	p.submitted.Add(1)
	select {
	case p.tasks <- task:
		return nil
	default:
	}
	if p.cfg.Policy == RejectOnFull {
		p.rejected.Add(1)
		return ErrQueueFull
	}
	var timeout <-chan time.Time
	if p.cfg.Policy == RejectOnTimeout {
		timer := time.NewTimer(p.cfg.WaitTimeout)
		defer timer.Stop()
		timeout = timer.C
	}
	select {
	case p.tasks <- task:
		return nil
	case <-timeout:
		p.rejected.Add(1)
		p.waitTimeout.Add(1)
		return ErrWaitTimeout
	case <-ctx.Done():
		p.rejected.Add(1)
		return ctx.Err()
	}
}

func (p *Pool) Stats() Snapshot {
	return Snapshot{
		Submitted: p.submitted.Load(), Completed: p.completed.Load(),
		Rejected: p.rejected.Load(), WaitTimeouts: p.waitTimeout.Load(),
		InFlight: p.inFlight.Load(), Queued: int64(len(p.tasks)),
	}
}

// Close drains accepted tasks. Callers must stop submitting first.
func (p *Pool) Close() {
	if p.closed.CompareAndSwap(false, true) {
		close(p.tasks)
	}
	p.wg.Wait()
}

type Conn struct{ ID int }

type Lease struct {
	Conn     *Conn
	pool     *ConnPool
	acquired time.Time
	released atomic.Bool
}

// ConnPool separates maxOpen (hard cap on real connections) from maxIdle
// (connections kept warm). The lease is the accounting unit: every Acquire must
// be paired with Release, and leaked leases stay visible.
type ConnPool struct {
	mu       sync.Mutex
	maxOpen  int
	maxIdle  int
	nextID   int
	open     int
	idle     []*Conn
	leases   map[int]*Lease
	notify   chan struct{}
	timedOut atomic.Int64
}

func NewConnPool(maxOpen, maxIdle int) *ConnPool {
	if maxIdle > maxOpen {
		maxIdle = maxOpen
	}
	return &ConnPool{maxOpen: maxOpen, maxIdle: maxIdle,
		leases: make(map[int]*Lease), notify: make(chan struct{}, maxOpen)}
}

// Acquire waits on one notification channel instead of one channel per waiter.
// The caller's deadline, not a pool-level knob, bounds the wait.
func (p *ConnPool) Acquire(ctx context.Context) (*Lease, error) {
	for {
		p.mu.Lock()
		if n := len(p.idle); n > 0 {
			c := p.idle[n-1]
			p.idle = p.idle[:n-1]
			l := p.leaseLocked(c)
			p.mu.Unlock()
			return l, nil
		}
		if p.open < p.maxOpen {
			c := &Conn{ID: p.nextID}
			p.nextID++
			p.open++
			l := p.leaseLocked(c)
			p.mu.Unlock()
			return l, nil
		}
		p.mu.Unlock()
		select {
		case <-ctx.Done():
			p.timedOut.Add(1)
			signal(p.notify)
			return nil, ErrPoolExhausted
		case <-p.notify:
		}
	}
}

func (p *ConnPool) leaseLocked(c *Conn) *Lease {
	l := &Lease{Conn: c, pool: p, acquired: time.Now()}
	p.leases[c.ID] = l
	return l
}

func (l *Lease) Release() error {
	if l.released.Swap(true) {
		return ErrLeaseReleased
	}
	p := l.pool
	p.mu.Lock()
	delete(p.leases, l.Conn.ID)
	if len(p.idle) < p.maxIdle {
		p.idle = append(p.idle, l.Conn)
	} else {
		p.open--
	}
	p.mu.Unlock()
	signal(p.notify)
	return nil
}

func signal(ch chan struct{}) {
	select {
	case ch <- struct{}{}:
	default:
	}
}

// Leaked feeds a lease-age alarm. It must never reap a lease still in use.
func (p *ConnPool) Leaked(age time.Duration) []int {
	p.mu.Lock()
	defer p.mu.Unlock()
	now := time.Now()
	var ids []int
	for id, l := range p.leases {
		if now.Sub(l.acquired) >= age {
			ids = append(ids, id)
		}
	}
	return ids
}

func (p *ConnPool) Stats() (open, idle, inUse int, timedOut int64) {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.open, len(p.idle), len(p.leases), p.timedOut.Load()
}

// file: internal/platform/pool/pool_test.go
// tier: L1
// verify: go test ./internal/platform/pool/ -count=1 -v
package pool

import (
	"context"
	"errors"
	"testing"
	"time"
)

func TestRejectAndTimeout(t *testing.T) {
	ctx := context.Background()
	release, started := make(chan struct{}), make(chan struct{}, 2)
	p := New(Config{Workers: 2, QueueSize: 1, Policy: RejectOnFull})
	for i := 0; i < 2; i++ {
		if err := p.Submit(ctx, func() { started <- struct{}{}; <-release }); err != nil {
			t.Fatal(err)
		}
		<-started
	}
	if err := p.Submit(ctx, func() {}); err != nil {
		t.Fatal(err)
	}
	if err := p.Submit(ctx, func() {}); !errors.Is(err, ErrQueueFull) {
		t.Fatal(err)
	}
	close(release)
	p.Close()
	p = New(Config{Workers: 1, QueueSize: 0, WaitTimeout: 30 * time.Millisecond, Policy: RejectOnTimeout})
	block, entered := make(chan struct{}), make(chan struct{})
	p.Submit(ctx, func() { close(entered); <-block })
	<-entered
	begin := time.Now()
	err := p.Submit(ctx, func() {})
	waited := time.Since(begin)
	if !errors.Is(err, ErrWaitTimeout) || waited < 25*time.Millisecond {
		t.Fatalf("err=%v waited=%s", err, waited)
	}
	close(block)
	p.Close()
	t.Logf("queue_full_rejected=1 wait_timeout_waited=%s", waited.Round(time.Millisecond))
}

func TestConnPoolLease(t *testing.T) {
	ctx := context.Background()
	p := NewConnPool(2, 1)
	l1, _ := p.Acquire(ctx)
	l2, _ := p.Acquire(ctx)
	timed, cancel := context.WithTimeout(ctx, 30*time.Millisecond)
	defer cancel()
	if _, err := p.Acquire(timed); !errors.Is(err, ErrPoolExhausted) {
		t.Fatal(err)
	}
	_ = l1.Release()
	if open, idle, inUse, timedOut := p.Stats(); open != 2 || idle != 1 || inUse != 1 || timedOut != 1 {
		t.Fatalf("open=%d idle=%d inUse=%d timedOut=%d", open, idle, inUse, timedOut)
	}
	time.Sleep(10 * time.Millisecond)
	if leaked := p.Leaked(5 * time.Millisecond); len(leaked) != 1 || leaked[0] != l2.Conn.ID {
		t.Fatalf("leaked = %v", leaked)
	}
	_ = l2.Release()
	if open, idle, inUse, _ := p.Stats(); open != 1 || idle != 1 || inUse != 0 {
		t.Fatalf("open=%d idle=%d inUse=%d", open, idle, inUse)
	}
	t.Logf("timed_out=1 leaked=[%d] final_open=1 final_idle=1", l2.Conn.ID)
}

$ go test ./internal/platform/pool/ -run TestPoolSizeSweep -count=1 -v
workers=1  elapsed=2.85s  throughput=421/s  avg_wait=231.278ms rejected=0
workers=4  elapsed=710ms  throughput=1690/s avg_wait=57.698ms  rejected=0
workers=16 elapsed=176ms  throughput=6825/s avg_wait=14.357ms  rejected=0
workers=64 elapsed=49ms   throughput=24729/s avg_wait=3.833ms  rejected=0