《游戏服务端编程实践》1.4.2 分布式与微服务

// file: ch01/services/registry.go
// tier: L1
// verify: go test ./ch01/services/
package services

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

var (
	// ErrInvalidInstance 表示注册信息缺少服务名、实例 ID 或地址。
	ErrInvalidInstance = errors.New("services: invalid instance")
	// ErrNotFound 表示服务或实例不存在。
	ErrNotFound = errors.New("services: instance not found")
	// ErrNoHealthy 表示没有可用的健康实例。
	ErrNoHealthy = errors.New("services: no healthy instance")
)

// Instance 是服务注册表中的一个实例。
type Instance struct {
	ID       string
	Service  string
	Addr     string
	Weight   int
	Healthy  bool
	LastSeen time.Time

	current int
}

// Registry 是一个进程内服务注册表。它适合单机实验和单元测试，
// 生产环境应把同样的语义放到 etcd、Consul 或 Nacos 中。
type Registry struct {
	mu       sync.RWMutex
	services map[string]map[string]*Instance
	now      func() time.Time
}

// NewRegistry 创建一个使用系统时钟的注册表。
func NewRegistry() *Registry {
	return NewRegistryWithClock(time.Now)
}

// NewRegistryWithClock 允许注入逻辑时钟，便于测试心跳过期。
func NewRegistryWithClock(now func() time.Time) *Registry {
	return &Registry{
		services: make(map[string]map[string]*Instance),
		now:      now,
	}
}

// Register 注册或覆盖一个实例。重复注册同一 ID 视为幂等更新，
// 并把健康状态重置为 true。
func (r *Registry) Register(instance Instance) error {
	if instance.ID == "" || instance.Service == "" || instance.Addr == "" {
		return ErrInvalidInstance
	}
	if instance.Weight <= 0 {
		instance.Weight = 1
	}
	if instance.LastSeen.IsZero() {
		instance.LastSeen = r.now()
	}
	instance.Healthy = true
	instance.current = 0

	r.mu.Lock()
	defer r.mu.Unlock()
	byID := r.services[instance.Service]
	if byID == nil {
		byID = make(map[string]*Instance)
		r.services[instance.Service] = byID
	}
	copied := instance
	byID[instance.ID] = &copied
	return nil
}

// Unregister 摘除实例，返回实例是否存在。
func (r *Registry) Unregister(service, id string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	byID := r.services[service]
	if byID == nil {
		return false
	}
	if _, ok := byID[id]; !ok {
		return false
	}
	delete(byID, id)
	if len(byID) == 0 {
		delete(r.services, service)
	}
	return true
}

// Heartbeat 刷新实例的存活时间并恢复健康状态。
func (r *Registry) Heartbeat(service, id string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	instance, ok := r.lookup(service, id)
	if !ok {
		return ErrNotFound
	}
	instance.Healthy = true
	instance.LastSeen = r.now()
	return nil
}

// Report 记录一次健康检查结果。probe 返回错误时实例被标记为不健康，
// 但不会被立即删除，便于观察和恢复。
func (r *Registry) Report(service, id string, probeErr error) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	instance, ok := r.lookup(service, id)
	if !ok {
		return ErrNotFound
	}
	instance.Healthy = probeErr == nil
	instance.LastSeen = r.now()
	return nil
}

// Check 执行一次带超时的健康探测并把结果写入注册表。
func (r *Registry) Check(ctx context.Context, service, id string, timeout time.Duration, probe func(context.Context) error) error {
	if probe == nil {
		return errors.New("services: nil probe")
	}
	checkCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	probeErr := probe(checkCtx)
	if err := r.Report(service, id, probeErr); err != nil {
		return err
	}
	return probeErr
}

// Sweep 把超过 ttl 没有心跳的实例标记为不健康，返回本次状态翻转数量。
func (r *Registry) Sweep(ttl time.Duration) int {
	r.mu.Lock()
	defer r.mu.Unlock()
	now := r.now()
	changed := 0
	for _, byID := range r.services {
		for _, instance := range byID {
			if instance.Healthy && now.Sub(instance.LastSeen) > ttl {
				instance.Healthy = false
				changed++
			}
		}
	}
	return changed
}

// Healthy 返回某服务当前健康的实例，按 ID 排序。
func (r *Registry) Healthy(service string) []Instance {
	r.mu.RLock()
	defer r.mu.RUnlock()
	byID := r.services[service]
	out := make([]Instance, 0, len(byID))
	for _, instance := range byID {
		if instance.Healthy {
			out = append(out, *instance)
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
	return out
}

// Pick 用平滑加权轮询选择一个健康实例。exclude 用于重试：
// 调用方把已经失败的实例 ID 放进去，下一次自然会选到别的实例。
func (r *Registry) Pick(service string, exclude map[string]struct{}) (Instance, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	byID := r.services[service]
	var selected *Instance
	total := 0
	for _, instance := range byID {
		if !instance.Healthy {
			continue
		}
		if _, skip := exclude[instance.ID]; skip {
			continue
		}
		instance.current += instance.Weight
		total += instance.Weight
		if selected == nil || instance.current > selected.current ||
			(instance.current == selected.current && instance.ID < selected.ID) {
			selected = instance
		}
	}
	if selected == nil {
		return Instance{}, ErrNoHealthy
	}
	selected.current -= total
	return *selected, nil
}

func (r *Registry) lookup(service, id string) (*Instance, bool) {
	byID := r.services[service]
	if byID == nil {
		return nil, false
	}
	instance, ok := byID[id]
	return instance, ok
}

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