《游戏服务端编程实践》1.3.2 运行期事件与资源管理

// file: ch01/runtime/timer.go
// tier: L1
// verify: go test ./ch01/runtime/
package runtime

import (
	"container/heap"
	"errors"
	"sync"
	"time"
)

var (
	ErrClosed      = errors.New("runtime: timer closed")
	ErrBadInterval = errors.New("runtime: interval must be positive")
	ErrNilCallback = errors.New("runtime: nil callback")
)

// TimerID 是定时任务的唯一标识。
type TimerID uint64

// Task 是定时器堆中的一个任务。
type Task struct {
	ID       TimerID
	At       time.Time
	Interval time.Duration
	Fn       func()
	index    int
	seq      uint64
}

type taskHeap []*Task

func (h taskHeap) Len() int { return len(h) }
func (h taskHeap) Less(i, j int) bool {
	if h[i].At.Equal(h[j].At) {
		return h[i].seq < h[j].seq
	}
	return h[i].At.Before(h[j].At)
}
func (h taskHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
	h[i].index = i
	h[j].index = j
}
func (h *taskHeap) Push(x any) {
	task := x.(*Task)
	task.index = len(*h)
	*h = append(*h, task)
}
func (h *taskHeap) Pop() any {
	old := *h
	n := len(old)
	task := old[n-1]
	old[n-1] = nil
	task.index = -1
	*h = old[:n-1]
	return task
}

// Timer 负责到期判定和回调调度，回调本身必须保持轻量。
type Timer struct {
	mu      sync.Mutex
	tasks   taskHeap
	byID    map[TimerID]*Task
	nextID  TimerID
	seq     uint64
	closed  bool
	running bool
	wake    chan struct{}
	stop    chan struct{}
	done    chan struct{}
	now     func() time.Time
}

// New 创建一个使用系统时钟的定时器。
func New() *Timer { return NewWithClock(time.Now) }

// NewWithClock 允许测试注入逻辑时钟，配合 Advance 使用。
func NewWithClock(now func() time.Time) *Timer {
	t := &Timer{
		byID: make(map[TimerID]*Task),
		wake: make(chan struct{}, 1),
		stop: make(chan struct{}),
		done: make(chan struct{}),
		now:  now,
	}
	heap.Init(&t.tasks)
	return t
}

// After 注册一个延迟执行的一次性任务。
func (t *Timer) After(d time.Duration, fn func()) (TimerID, error) {
	if d < 0 {
		d = 0
	}
	return t.add(t.now().Add(d), 0, fn)
}

// Every 注册周期任务；从本次处理的 now 重新计时，避免补跑风暴。
func (t *Timer) Every(interval time.Duration, fn func()) (TimerID, error) {
	if interval <= 0 {
		return 0, ErrBadInterval
	}
	return t.add(t.now().Add(interval), interval, fn)
}

// At 注册指定绝对时间执行的一次性任务，主要供测试使用。
func (t *Timer) At(at time.Time, fn func()) (TimerID, error) {
	return t.add(at, 0, fn)
}

func (t *Timer) add(at time.Time, interval time.Duration, fn func()) (TimerID, error) {
	if fn == nil {
		return 0, ErrNilCallback
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.closed {
		return 0, ErrClosed
	}
	t.nextID++
	task := &Task{ID: t.nextID, At: at, Interval: interval, Fn: fn, seq: t.seq}
	t.seq++
	heap.Push(&t.tasks, task)
	t.byID[task.ID] = task
	t.signal()
	return task.ID, nil
}

// Cancel 取消尚未触发的任务。
func (t *Timer) Cancel(id TimerID) bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	task, ok := t.byID[id]
	if !ok {
		return false
	}
	heap.Remove(&t.tasks, task.index)
	delete(t.byID, id)
	return true
}

// Advance 执行所有 At <= now 的任务，返回触发数量。
func (t *Timer) Advance(now time.Time) int {
	t.mu.Lock()
	var due []*Task
	for len(t.tasks) > 0 && !t.tasks[0].At.After(now) {
		task := heap.Pop(&t.tasks).(*Task)
		due = append(due, task)
		if task.Interval > 0 {
			task.At = now.Add(task.Interval)
			task.seq = t.seq
			t.seq++
			heap.Push(&t.tasks, task)
		} else {
			delete(t.byID, task.ID)
		}
	}
	t.mu.Unlock()
	for _, task := range due {
		task.Fn()
	}
	return len(due)
}

// Start 启动后台循环。
func (t *Timer) Start() {
	t.mu.Lock()
	if t.running || t.closed {
		t.mu.Unlock()
		return
	}
	t.running = true
	t.mu.Unlock()
	go t.loop()
}

func (t *Timer) loop() {
	defer close(t.done)
	for {
		t.mu.Lock()
		if t.closed {
			t.mu.Unlock()
			return
		}
		if len(t.tasks) == 0 {
			t.mu.Unlock()
			select {
			case <-t.wake:
			case <-t.stop:
				return
			}
			continue
		}
		wait := time.Until(t.tasks[0].At)
		t.mu.Unlock()
		if wait <= 0 {
			t.Advance(time.Now())
			continue
		}
		timer := time.NewTimer(wait)
		select {
		case <-timer.C:
		case <-t.wake:
			if !timer.Stop() {
				<-timer.C
			}
		case <-t.stop:
			if !timer.Stop() {
				<-timer.C
			}
			return
		}
	}
}

func (t *Timer) signal() {
	select {
	case t.wake <- struct{}{}:
	default:
	}
}

// Close 停止后台循环并丢弃未触发任务，可重复调用。
func (t *Timer) Close() {
	t.mu.Lock()
	if t.closed {
		t.mu.Unlock()
		return
	}
	t.closed = true
	t.tasks = nil
	t.byID = make(map[TimerID]*Task)
	running := t.running
	t.mu.Unlock()
	close(t.stop)
	if running {
		<-t.done
	}
}

cd .local/checks/gs-practice-lab
gofmt -l ch01/runtime
go vet ./ch01/runtime/
go test ./ch01/runtime/