《游戏服务端编程实践》11.3 服务注册与发现机制

// file: ch11/registry/registry.go
// tier: L1
// verify: go test ./ch11/registry/

package registry

import (
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"
)

var (
	ErrInvalidTTL   = errors.New("registry: ttl must be positive")
	ErrEmptyService = errors.New("registry: empty service")
	ErrEmptyID      = errors.New("registry: empty instance id")
	ErrEmptyAddr    = errors.New("registry: empty address")
	ErrNotFound     = errors.New("registry: instance not found")
)

type EventKind string

const (
	EventRegister   EventKind = "register"
	EventRenew      EventKind = "renew"
	EventDeregister EventKind = "deregister"
	EventExpire     EventKind = "expire"
)

// Instance 是一次注册的完整描述；ExpiresAt 由注册中心写入，调用方传值会被覆盖。
type Instance struct {
	Service   string
	ID        string
	Addr      string
	Metadata  map[string]string
	ExpiresAt time.Time
}

type Event struct {
	Kind     EventKind
	Instance Instance
	At       time.Time
}

// Watcher 的通道满时事件会被丢弃；客户端发现缺口后应重新 List 全量对齐。
type Watcher struct {
	id        uint64
	service   string
	ch        chan Event
	reg       *Registry
	closeOnce sync.Once
}

func (w *Watcher) C() <-chan Event {
	return w.ch
}

func (w *Watcher) Close() {
	w.closeOnce.Do(func() { w.reg.unwatch(w) })
}

// Registry 是带 TTL 的进程内注册表，读写与过期清理共用一把锁。
type Registry struct {
	mu          sync.Mutex
	ttl         time.Duration
	now         func() time.Time
	instances   map[string]Instance
	watchers    map[uint64]*Watcher
	nextWatcher uint64
}

func New(ttl time.Duration) (*Registry, error) {
	return NewWithClock(ttl, time.Now)
}

// NewWithClock 允许测试注入时钟，TTL 行为不需要真实等待。
func NewWithClock(ttl time.Duration, now func() time.Time) (*Registry, error) {
	if ttl <= 0 {
		return nil, ErrInvalidTTL
	}
	if now == nil {
		now = time.Now
	}
	return &Registry{
		ttl:       ttl,
		now:       now,
		instances: make(map[string]Instance),
		watchers:  make(map[uint64]*Watcher),
	}, nil
}

func key(service, id string) string {
	return service + "\x00" + id
}

// Register 写入新实例；同 ID 重复注册视为覆盖，并重新计算过期时间。
func (r *Registry) Register(inst Instance) (Instance, error) {
	service := strings.TrimSpace(inst.Service)
	id := strings.TrimSpace(inst.ID)
	addr := strings.TrimSpace(inst.Addr)
	switch {
	case service == "":
		return Instance{}, ErrEmptyService
	case id == "":
		return Instance{}, ErrEmptyID
	case addr == "":
		return Instance{}, ErrEmptyAddr
	}
	inst.Service, inst.ID, inst.Addr = service, id, addr
	inst.Metadata = cloneMetadata(inst.Metadata)

	r.mu.Lock()
	defer r.mu.Unlock()
	now := r.now()
	inst.ExpiresAt = now.Add(r.ttl)
	r.instances[key(service, id)] = inst
	r.notifyLocked(Event{Kind: EventRegister, Instance: cloneInstance(inst), At: now})
	return cloneInstance(inst), nil
}

// Renew 只续约仍然存活的实例；已经过期的实例不能靠续约复活。
func (r *Registry) Renew(service, id string) (Instance, error) {
	service = strings.TrimSpace(service)
	id = strings.TrimSpace(id)
	r.mu.Lock()
	defer r.mu.Unlock()
	k := key(service, id)
	inst, ok := r.instances[k]
	if !ok {
		return Instance{}, fmt.Errorf("%w: %s/%s", ErrNotFound, service, id)
	}
	now := r.now()
	if !inst.ExpiresAt.After(now) {
		delete(r.instances, k)
		r.notifyLocked(Event{Kind: EventExpire, Instance: cloneInstance(inst), At: now})
		return Instance{}, fmt.Errorf("%w: %s/%s expired", ErrNotFound, service, id)
	}
	inst.ExpiresAt = now.Add(r.ttl)
	r.instances[k] = inst
	r.notifyLocked(Event{Kind: EventRenew, Instance: cloneInstance(inst), At: now})
	return cloneInstance(inst), nil
}

func (r *Registry) Deregister(service, id string) error {
	service = strings.TrimSpace(service)
	id = strings.TrimSpace(id)
	r.mu.Lock()
	defer r.mu.Unlock()
	k := key(service, id)
	inst, ok := r.instances[k]
	if !ok {
		return fmt.Errorf("%w: %s/%s", ErrNotFound, service, id)
	}
	delete(r.instances, k)
	r.notifyLocked(Event{Kind: EventDeregister, Instance: cloneInstance(inst), At: r.now()})
	return nil
}

// List 返回存活实例快照，并在同一次加锁中清理所有已过期实例。
func (r *Registry) List(service string) []Instance {
	service = strings.TrimSpace(service)
	r.mu.Lock()
	defer r.mu.Unlock()
	now := r.now()
	out := make([]Instance, 0, len(r.instances))
	for k, inst := range r.instances {
		if !inst.ExpiresAt.After(now) {
			delete(r.instances, k)
			r.notifyLocked(Event{Kind: EventExpire, Instance: cloneInstance(inst), At: now})
			continue
		}
		if service == "" || inst.Service == service {
			out = append(out, cloneInstance(inst))
		}
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].Service != out[j].Service {
			return out[i].Service < out[j].Service
		}
		return out[i].ID < out[j].ID
	})
	return out
}

// Watch 订阅后续事件；初始全量仍由调用方先 List 一次获得。
func (r *Registry) Watch(service string, buffer int) (*Watcher, error) {
	service = strings.TrimSpace(service)
	if buffer < 0 {
		buffer = 0
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.nextWatcher++
	w := &Watcher{id: r.nextWatcher, service: service, ch: make(chan Event, buffer), reg: r}
	r.watchers[w.id] = w
	return w, nil
}

func (r *Registry) unwatch(w *Watcher) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, ok := r.watchers[w.id]; !ok {
		return
	}
	delete(r.watchers, w.id)
	close(w.ch)
}

// Sweep 主动清理过期实例并返回数量；后台 Run 会周期性调用它。
func (r *Registry) Sweep() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	now := r.now()
	removed := 0
	for k, inst := range r.instances {
		if inst.ExpiresAt.After(now) {
			continue
		}
		delete(r.instances, k)
		removed++
		r.notifyLocked(Event{Kind: EventExpire, Instance: cloneInstance(inst), At: now})
	}
	return removed
}

func (r *Registry) notifyLocked(ev Event) {
	for _, w := range r.watchers {
		if w.service != "" && ev.Instance.Service != w.service {
			continue
		}
		select {
		case w.ch <- ev:
		default:
		}
	}
}

func cloneInstance(in Instance) Instance {
	out := in
	out.Metadata = cloneMetadata(in.Metadata)
	return out
}

func cloneMetadata(in map[string]string) map[string]string {
	if in == nil {
		return nil
	}
	out := make(map[string]string, len(in))
	for k, v := range in {
		out[k] = v
	}
	return out
}

cd .local/checks/gs-practice-lab && gofmt -l ch11/registry && go vet ./ch11/registry/ && go test ./ch11/registry/ -count=1

ok  	gspractice/ch11/registry	0.890s