《游戏服务端编程实践》17.2 服务端脚本热更（Lua/WASM）

// file: ch17/hotswap/hotswap.go
// tier: L1
// verify: go test ./ch17/hotswap/

// Package hotswap 实现版本化 handler 的原子切换与引用计数卸载。
package hotswap

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

var (
	ErrEmptyVersion     = errors.New("hotswap: empty version")
	ErrEmptyScript      = errors.New("hotswap: empty script")
	ErrDuplicateVersion = errors.New("hotswap: duplicate version")
	ErrUnknownVersion   = errors.New("hotswap: unknown version")
	ErrNoActiveVersion  = errors.New("hotswap: no active version")
	ErrNoLoader         = errors.New("hotswap: loader is nil")
	ErrLoaderPanic      = errors.New("hotswap: loader panic")
	ErrScriptPanic      = errors.New("hotswap: script panic")
	ErrTimeout          = errors.New("hotswap: script timeout")
)

// Request 是传给脚本的最小请求单元。Payload 由调用方保证只读，
// 脚本不得修改它，否则会影响同一请求的后续处理。
type Request struct {
	Path    string
	Payload []byte
}

// Response 是脚本返回的结果。Version 字段由 handler 填写，方便
// 在灰度期间确认某次响应来自哪个版本。
type Response struct {
	Version string
	Body    []byte
}

// Handler 是一个已加载的脚本版本。同一个 handler 会被多个请求
// 并发调用，脚本自身必须保证并发安全。
type Handler func(context.Context, Request) (Response, error)

// Loader 把脚本文本编译为 Handler。编译失败时不返回 Handler。
// Loader 内部 panic 会被捕获并转成 ErrLoaderPanic。
type Loader func(script []byte) (Handler, error)

// Version 是一个已加载但未必激活的脚本版本。
type Version struct {
	ID       string
	Script   []byte
	handler  Handler
	loadedAt time.Time

	refs       atomic.Int64
	retired    atomic.Bool
	unloadOnce sync.Once
	unloaded   chan struct{}
}

// Retired 表示该版本已不再接受新请求。
func (v *Version) Retired() bool {
	return v.retired.Load()
}

// Unloaded 在版本正式卸载后关闭。retire 时若仍有在途请求，
// 该 channel 会等到最后一个请求释放引用才关闭。
func (v *Version) Unloaded() <-chan struct{} {
	return v.unloaded
}

func (v *Version) finalize() {
	v.unloadOnce.Do(func() {
		close(v.unloaded)
	})
}

// Option 调整 Registry 的行为。
type Option func(*Registry)

// WithTimeout 为每次脚本调用设置超时。超时是协作式的：handler 必须
// 观察 ctx.Done() 才能及时返回；不观察的 handler 仍会继续占用 goroutine。
func WithTimeout(d time.Duration) Option {
	return func(r *Registry) {
		if d > 0 {
			r.timeout = d
		}
	}
}

// Registry 保存已加载的版本，并以原子方式切换当前版本。
type Registry struct {
	mu       sync.RWMutex
	loader   Loader
	timeout  time.Duration
	current  *Version
	versions map[string]*Version
}

// New 创建注册表。loader 负责把脚本编译成 handler，不能为 nil。
func New(loader Loader, opts ...Option) *Registry {
	r := &Registry{
		loader:   loader,
		versions: make(map[string]*Version),
	}
	for _, opt := range opts {
		opt(r)
	}
	return r
}

// Load 编译脚本并登记版本，但不会让它成为当前版本。
func (r *Registry) Load(id string, script []byte) (*Version, error) {
	if id == "" {
		return nil, ErrEmptyVersion
	}
	if len(script) == 0 {
		return nil, ErrEmptyScript
	}
	if r.loader == nil {
		return nil, ErrNoLoader
	}
	handler, err := safeLoad(r.loader, script)
	if err != nil {
		return nil, err
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, ok := r.versions[id]; ok {
		return nil, fmt.Errorf("%w: %s", ErrDuplicateVersion, id)
	}
	owned := append([]byte(nil), script...)
	v := &Version{
		ID:       id,
		Script:   owned,
		handler:  handler,
		loadedAt: time.Now(),
		unloaded: make(chan struct{}),
	}
	r.versions[id] = v
	return v, nil
}

// Activate 把已加载且未退役的版本设为当前版本。切换只改一个指针，
// 已经在执行的请求继续持有旧版本引用。
func (r *Registry) Activate(id string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	v, ok := r.versions[id]
	if !ok {
		return fmt.Errorf("%w: %s", ErrUnknownVersion, id)
	}
	if v.retired.Load() {
		return fmt.Errorf("%w: retired %s", ErrUnknownVersion, id)
	}
	r.current = v
	return nil
}

// Invoke 获取当前版本的一个引用并执行 handler。引用在请求结束时
// 释放；如果版本已退役且引用归零，则在这里完成卸载。
func (r *Registry) Invoke(ctx context.Context, req Request) (Response, error) {
	r.mu.RLock()
	v := r.current
	if v == nil {
		r.mu.RUnlock()
		return Response{}, ErrNoActiveVersion
	}
	v.refs.Add(1)
	r.mu.RUnlock()
	defer r.release(v)

	if r.timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, r.timeout)
		defer cancel()
	}
	resp, err := call(ctx, v, req)
	if err != nil {
		if errors.Is(err, context.DeadlineExceeded) {
			return Response{}, fmt.Errorf("%w: %s", ErrTimeout, v.ID)
		}
		return Response{}, err
	}
	if r.timeout > 0 && ctx.Err() != nil {
		return Response{}, fmt.Errorf("%w: %s", ErrTimeout, v.ID)
	}
	return resp, nil
}

// Retire 停止版本接收新请求。若没有在途请求，立即卸载；否则等待
// 最后一个请求释放引用。重复调用是幂等的。
func (r *Registry) Retire(id string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	v, ok := r.versions[id]
	if !ok {
		return fmt.Errorf("%w: %s", ErrUnknownVersion, id)
	}
	if v.retired.Swap(true) {
		return nil
	}
	if r.current == v {
		r.current = nil
	}
	if v.refs.Load() == 0 {
		v.finalize()
	}
	return nil
}

// Current 返回当前激活版本的 ID。
func (r *Registry) Current() (string, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	if r.current == nil {
		return "", false
	}
	return r.current.ID, true
}

// InFlight 返回指定版本的在途请求数。
func (r *Registry) InFlight(id string) int64 {
	r.mu.RLock()
	defer r.mu.RUnlock()
	v, ok := r.versions[id]
	if !ok {
		return 0
	}
	return v.refs.Load()
}

func (r *Registry) release(v *Version) {
	if v.refs.Add(-1) != 0 {
		return
	}
	r.mu.Lock()
	if v.retired.Load() {
		v.finalize()
	}
	r.mu.Unlock()
}

func call(ctx context.Context, v *Version, req Request) (resp Response, err error) {
	defer func() {
		if p := recover(); p != nil {
			err = fmt.Errorf("%w: %s: %v", ErrScriptPanic, v.ID, p)
		}
	}()
	return v.handler(ctx, req)
}

func safeLoad(loader Loader, script []byte) (handler Handler, err error) {
	defer func() {
		if p := recover(); p != nil {
			handler = nil
			err = fmt.Errorf("%w: %v", ErrLoaderPanic, p)
		}
	}()
	return loader(script)
}

cd .local/checks/gs-practice-lab
gofmt -l ch17/hotswap
go vet ./ch17/hotswap/
go test ./ch17/hotswap/ -count=1