《游戏服务端编程实践》5.4.3 压测与连接池管理

// file: ch05/backend/pool.go
// tier: L1
// verify: go test ./ch05/backend/

package backend

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

var (
	ErrPoolClosed     = errors.New("backend: pool closed")
	ErrAcquireTimeout = errors.New("backend: acquire timeout")
)

// PoolConfig 的零值不可用：Dial 必填，MaxOpen 至少为 1。
type PoolConfig struct {
	Dial           func(context.Context) (net.Conn, error)
	MaxOpen        int
	MaxIdle        int
	AcquireTimeout time.Duration
	IdleTimeout    time.Duration
	Check          func(net.Conn) error
}

type Pool struct {
	cfg    PoolConfig
	mu     sync.Mutex
	idle   []idleConn
	open   int
	closed bool
	wait   chan struct{}
}

type idleConn struct {
	conn net.Conn
	at   time.Time
}

// Conn 是借出的连接。确认损坏时先 MarkBroken 再 Put，池会关闭它而不是复用。
type Conn struct {
	net.Conn
	pool     *Pool
	mu       sync.Mutex
	broken   bool
	returned bool
}

func NewPool(cfg PoolConfig) *Pool {
	if cfg.MaxOpen <= 0 {
		cfg.MaxOpen = 1
	}
	if cfg.MaxIdle <= 0 || cfg.MaxIdle > cfg.MaxOpen {
		cfg.MaxIdle = cfg.MaxOpen
	}
	return &Pool{cfg: cfg, wait: make(chan struct{})}
}

func (p *Pool) Get(ctx context.Context) (*Conn, error) {
	if p.cfg.AcquireTimeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, p.cfg.AcquireTimeout)
		defer cancel()
	}
	for {
		p.mu.Lock()
		if p.closed {
			p.mu.Unlock()
			return nil, ErrPoolClosed
		}
		if len(p.idle) > 0 {
			ic := p.idle[len(p.idle)-1]
			p.idle = p.idle[:len(p.idle)-1]
			p.mu.Unlock()
			if p.cfg.IdleTimeout > 0 && time.Since(ic.at) > p.cfg.IdleTimeout {
				p.discard(ic.conn)
				continue
			}
			if p.cfg.Check != nil && p.cfg.Check(ic.conn) != nil {
				p.discard(ic.conn)
				continue
			}
			return &Conn{Conn: ic.conn, pool: p}, nil
		}
		if p.open < p.cfg.MaxOpen {
			p.open++
			p.mu.Unlock()
			conn, err := p.cfg.Dial(ctx)
			if err != nil {
				p.mu.Lock()
				p.open--
				p.notifyLocked()
				p.mu.Unlock()
				return nil, err
			}
			return &Conn{Conn: conn, pool: p}, nil
		}
		wait := p.wait
		p.mu.Unlock()
		select {
		case <-wait:
		case <-ctx.Done():
			return nil, ErrAcquireTimeout
		}
	}
}

func (p *Pool) discard(conn net.Conn) {
	_ = conn.Close()
	p.mu.Lock()
	p.open--
	p.notifyLocked()
	p.mu.Unlock()
}

func (p *Pool) Put(c *Conn) {
	if c == nil || c.Conn == nil {
		return
	}
	c.mu.Lock()
	if c.returned {
		c.mu.Unlock()
		return
	}
	c.returned = true
	broken := c.broken
	c.mu.Unlock()

	p.mu.Lock()
	if broken || p.closed || len(p.idle) >= p.cfg.MaxIdle {
		p.open--
		p.notifyLocked()
		p.mu.Unlock()
		_ = c.Conn.Close()
		return
	}
	p.idle = append(p.idle, idleConn{conn: c.Conn, at: time.Now()})
	p.notifyLocked()
	p.mu.Unlock()
}

func (c *Conn) MarkBroken() {
	c.mu.Lock()
	c.broken = true
	c.mu.Unlock()
}

func (p *Pool) notifyLocked() {
	close(p.wait)
	p.wait = make(chan struct{})
}

func (p *Pool) Close() error {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return nil
	}
	p.closed = true
	idle := p.idle
	p.idle = nil
	p.open -= len(idle)
	p.notifyLocked()
	p.mu.Unlock()
	for _, ic := range idle {
		_ = ic.conn.Close()
	}
	return nil
}

// file: ch05/backend/pool_test.go
// tier: L1
// verify: go test ./ch05/backend/

package backend

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

func pipeDialer() (func(context.Context) (net.Conn, error), *atomic.Int32) {
	var calls atomic.Int32
	dial := func(context.Context) (net.Conn, error) {
		calls.Add(1)
		client, server := net.Pipe()
		go func() {
			_, _ = io.Copy(io.Discard, server)
			_ = server.Close()
		}()
		return client, nil
	}
	return dial, &calls
}

func TestPoolReuseAndDiscard(t *testing.T) {
	dial, calls := pipeDialer()
	p := NewPool(PoolConfig{Dial: dial, MaxOpen: 2, MaxIdle: 2})
	defer p.Close()
	c1, err := p.Get(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	p.Put(c1)
	c2, err := p.Get(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if c2.Conn != c1.Conn || calls.Load() != 1 {
		t.Fatalf("connection not reused: calls=%d", calls.Load())
	}
	c2.MarkBroken()
	p.Put(c2)
	c3, err := p.Get(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if c3.Conn == c2.Conn || calls.Load() != 2 {
		t.Fatalf("broken connection reused: calls=%d", calls.Load())
	}
	p.Put(c3)
	if err := p.Close(); err != nil {
		t.Fatal(err)
	}
	if _, err := p.Get(context.Background()); !errors.Is(err, ErrPoolClosed) {
		t.Fatalf("get after close: %v", err)
	}
}

func TestPoolAcquireTimeout(t *testing.T) {
	dial, _ := pipeDialer()
	p := NewPool(PoolConfig{Dial: dial, MaxOpen: 1, MaxIdle: 1, AcquireTimeout: 30 * time.Millisecond})
	defer p.Close()
	c1, err := p.Get(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	start := time.Now()
	if _, err := p.Get(context.Background()); !errors.Is(err, ErrAcquireTimeout) || time.Since(start) < 20*time.Millisecond {
		t.Fatalf("err = %v, elapsed = %v", err, time.Since(start))
	}
	p.Put(c1)
	c2, err := p.Get(context.Background())
	if err != nil {
		t.Fatalf("after put: %v", err)
	}
	p.Put(c2)
}

func BenchmarkPoolGetPut(b *testing.B) {
	p := NewPool(PoolConfig{Dial: func(context.Context) (net.Conn, error) {
		client, server := net.Pipe()
		go func() {
			_, _ = io.Copy(io.Discard, server)
			_ = server.Close()
		}()
		return client, nil
	}, MaxOpen: 8, MaxIdle: 8})
	defer p.Close()
	b.ReportAllocs()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			c, err := p.Get(context.Background())
			if err != nil {
				b.Fatal(err)
			}
			p.Put(c)
		}
	})
}

go test ./ch05/backend/ -run '^$' -bench 'BenchmarkPoolGetPut' -benchmem -count=3

BenchmarkPoolGetPut-10  2801174  416.8 ns/op  160 B/op  2 allocs/op
BenchmarkPoolGetPut-10  2852833  415.8 ns/op  160 B/op  2 allocs/op
BenchmarkPoolGetPut-10  2895294  416.9 ns/op  160 B/op  2 allocs/op