《游戏服务端编程实践》13.2 无锁结构与内存复用

// file: ch13/ring/ring.go
// tier: L1
// verify: go test ./ch13/ring/

package ring

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

var (
	// ErrFull 表示 Drop 策略下缓冲区已满。
	ErrFull = errors.New("ring: full")
	// ErrEmpty 表示 Drop 策略下缓冲区为空。
	ErrEmpty = errors.New("ring: empty")
	// ErrClosed 表示缓冲区已关闭。
	ErrClosed = errors.New("ring: closed")
	// ErrNil 表示向缓冲区投递了 nil Frame。
	ErrNil = errors.New("ring: nil frame")
)

// Policy 决定缓冲区满或空时的行为。
type Policy uint8

const (
	// Drop 在满时立即返回 ErrFull，在空时立即返回 ErrEmpty。
	// 适合可以丢弃旧帧或中间帧的场景，例如位置同步。
	Drop Policy = iota
	// Block 在满/空时等待，直到有空间/数据或缓冲区关闭。
	// 适合不能丢消息的控制流，但要求调用方接受阻塞。
	Block
)

// Frame 是通过环形缓冲区传递的消息。
//
// Data 使用池化切片，PutFrame 时只截断长度、保留容量，
// 下一次 GetFrame 拿到的对象因此既干净又可复用。
type Frame struct {
	Seq  uint64
	Kind uint16
	Data []byte
}

// Reset 把 Frame 恢复到可复用的零值状态，同时保留 Data 的容量。
func (f *Frame) Reset() {
	f.Seq = 0
	f.Kind = 0
	f.Data = f.Data[:0]
}

var framePool = sync.Pool{
	New: func() any { return &Frame{} },
}

// GetFrame 从池中取一个已重置的 Frame。
func GetFrame() *Frame {
	return framePool.Get().(*Frame)
}

// PutFrame 重置并归还 Frame。同一个 Frame 不能重复归还。
func PutFrame(f *Frame) {
	if f == nil {
		return
	}
	f.Reset()
	framePool.Put(f)
}

// Ring 是单生产者单消费者（SPSC）的定长环形缓冲区。
//
// 容量会被向上取整到 2 的幂，索引用位掩码代替取模。
// 生产者只推进 head，消费者只推进 tail，进度通过原子变量
// 发布，因此正常路径没有共享锁。Block 策略下用条件变量
// 睡眠等待，唤醒后必须在锁内重新检查条件，避免丢失唤醒。
type Ring struct {
	buf    []*Frame
	mask   uint64
	size   uint64
	policy Policy

	head atomic.Uint64 // 生产者写入位置
	tail atomic.Uint64 // 消费者读取位置

	closed atomic.Bool

	mu       sync.Mutex
	notFull  *sync.Cond
	notEmpty *sync.Cond
}

// New 创建容量为 capacity 的环形缓冲区。
//
// capacity 小于 2 时按 2 处理，其余向上取整到 2 的幂。
func New(capacity int, policy Policy) *Ring {
	if capacity < 2 {
		capacity = 2
	}
	size := uint64(1)
	for size < uint64(capacity) {
		size <<= 1
	}
	r := &Ring{
		buf:    make([]*Frame, size),
		mask:   size - 1,
		size:   size,
		policy: policy,
	}
	r.notFull = sync.NewCond(&r.mu)
	r.notEmpty = sync.NewCond(&r.mu)
	return r
}

// Push 由唯一的生产者调用。
//
// Drop 策略在满时返回 ErrFull；Block 策略等待消费者腾出位置。
// 关闭后 Push 返回 ErrClosed。Close 必须发生在生产者停止之后，
// 否则最后一条正在写入的帧可能落在关闭之后。
func (r *Ring) Push(f *Frame) error {
	if f == nil {
		return ErrNil
	}
	for {
		if r.closed.Load() {
			return ErrClosed
		}
		head := r.head.Load()
		if head-r.tail.Load() < r.size {
			r.buf[head&r.mask] = f
			r.head.Store(head + 1)
			if r.policy == Block {
				r.mu.Lock()
				r.notEmpty.Broadcast()
				r.mu.Unlock()
			}
			return nil
		}
		if r.policy == Drop {
			return ErrFull
		}

		r.mu.Lock()
		if r.closed.Load() {
			r.mu.Unlock()
			return ErrClosed
		}
		// 拿到锁后重新检查，防止消费者在上面的检查与 Wait 之间
		// 完成出队并广播，造成丢失唤醒。
		if r.head.Load()-r.tail.Load() >= r.size {
			r.notFull.Wait()
		}
		r.mu.Unlock()
	}
}

// Pop 由唯一的消费者调用。
//
// Drop 策略在空时返回 ErrEmpty。Block 策略等待生产者写入；
// 缓冲区关闭且已排空时返回 ErrClosed。关闭后仍可继续 Pop
// 出已经入队的数据，直到排空为止。
func (r *Ring) Pop() (*Frame, error) {
	for {
		tail := r.tail.Load()
		if r.head.Load() != tail {
			f := r.buf[tail&r.mask]
			r.buf[tail&r.mask] = nil
			r.tail.Store(tail + 1)
			if r.policy == Block {
				r.mu.Lock()
				r.notFull.Broadcast()
				r.mu.Unlock()
			}
			return f, nil
		}
		if r.closed.Load() {
			return nil, ErrClosed
		}
		if r.policy == Drop {
			return nil, ErrEmpty
		}

		r.mu.Lock()
		if r.closed.Load() && r.head.Load() == r.tail.Load() {
			r.mu.Unlock()
			return nil, ErrClosed
		}
		if r.head.Load() == r.tail.Load() {
			r.notEmpty.Wait()
		}
		r.mu.Unlock()
	}
}

// Close 关闭缓冲区并唤醒所有等待者。
//
// 已经入队的数据不会被丢弃：Block 策略下消费者可以继续 Pop
// 到排空为止，之后返回 ErrClosed。
func (r *Ring) Close() {
	if r.closed.Swap(true) {
		return
	}
	r.mu.Lock()
	r.notFull.Broadcast()
	r.notEmpty.Broadcast()
	r.mu.Unlock()
}

// Len 返回当前已入队但未出队的帧数。
func (r *Ring) Len() int {
	return int(r.head.Load() - r.tail.Load())
}

// Cap 返回向上取整后的实际容量，一定是 2 的幂。
func (r *Ring) Cap() int {
	return int(r.size)
}

// Policy 返回创建时选择的满/空处理策略。
func (r *Ring) Policy() Policy {
	return r.policy
}

// Closed 报告缓冲区是否已经关闭。
func (r *Ring) Closed() bool {
	return r.closed.Load()
}

cd .local/checks/gs-practice-lab
gofmt -l ch13/ring
go vet ./ch13/ring/
go test ./ch13/ring/ -count=1
go test -race ./ch13/ring/ -count=1