《MMO 服务器开发》2.1 注册与健康检查

t0  scene-01 启动  → 注册：建 10s 租约，写键 nodes/scene-01
t1  正常心跳       → 每 ~3s 续租一次，键始终存在
t2  实例过载       → 自报 degraded：键还在，但权重调低、停止派发新流量
t3  进程被 kill -9 → 心跳断，租约到期，etcd 自动删键 → 其它节点看到 offline
t4  进程恢复       → 重新注册，键重现，回到 healthy

// file: internal/discovery/registry.go
// tier: L2
// verify: go test ./internal/discovery/ -run TestEtcd -v
package discovery

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

	clientv3 "go.etcd.io/etcd/client/v3"
)

// NodeID 是集群节点标识，形如 "scene-01"、"gateway-03"。
type NodeID string

// HealthState 是节点的健康三态。注意三态来源不同：healthy / degraded
// 由节点自报，offline 由租约过期推导——节点崩溃后没法再自报，只能靠
// 租约自动过期把它判死。
type HealthState uint8

const (
	StateHealthy  HealthState = iota + 1 // 正常，承接全部流量
	StateDegraded                        // 存活但受限，应停止派发新流量
	StateOffline                         // 租约过期，键被 etcd 自动删除
)

// String 便于日志与断言输出。
func (s HealthState) String() string {
	switch s {
	case StateHealthy:
		return "healthy"
	case StateDegraded:
		return "degraded"
	case StateOffline:
		return "offline"
	default:
		return "unknown"
	}
}

// NodeRecord 是写进 etcd 的注册记录。State 只会是 healthy 或 degraded：
// 节点还活着才写得了记录，offline 是记录消失后才推导出来的。
type NodeRecord struct {
	ID     NodeID      `json:"id"`
	Addr   string      `json:"addr"`
	Weight int         `json:"weight"`
	State  HealthState `json:"state"`
}

// Member 是注册表的读侧视图：记录 + etcd 分配的租约 ID。
type Member struct {
	NodeRecord
	LeaseID int64
}

// Registry 把一批节点注册到 etcd，用 TTL 租约 + 心跳续约维护存活。
type Registry struct {
	cli      *clientv3.Client
	prefix   string
	ttl      int64
	mu       sync.Mutex
	sessions map[NodeID]*session
}

type session struct {
	lease  clientv3.LeaseID
	cancel context.CancelFunc
}

// NewRegistry 创建挂在 cli 上的注册表，键前缀 prefix（须以 ':' 结尾），租约 TTL 为 ttl。
func NewRegistry(cli *clientv3.Client, prefix string, ttl time.Duration) *Registry {
	return &Registry{cli: cli, prefix: prefix, ttl: int64(ttl.Seconds()),
		sessions: make(map[NodeID]*session)}
}

func (r *Registry) key(id NodeID) string { return r.prefix + "nodes/" + string(id) }

// Register 注册一个节点：建租约、写键、起心跳续约。重复注册同一 ID 报错。
func (r *Registry) Register(ctx context.Context, rec NodeRecord) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, ok := r.sessions[rec.ID]; ok {
		return fmt.Errorf("discovery: node %s already registered", rec.ID)
	}
	if rec.State == 0 {
		rec.State = StateHealthy
	}
	grant, err := r.cli.Grant(ctx, r.ttl)
	if err != nil {
		return fmt.Errorf("discovery: grant lease: %w", err)
	}
	val, _ := json.Marshal(rec)
	if _, err := r.cli.Put(ctx, r.key(rec.ID), string(val), clientv3.WithLease(grant.ID)); err != nil {
		_, _ = r.cli.Revoke(context.Background(), grant.ID)
		return fmt.Errorf("discovery: put: %w", err)
	}
	// 心跳续约：KeepAlive 返回一个通道，etcd 每隔 TTL/3 自动续一次。
	// 通道一旦关闭（网络断、etcd 重启、租约被撤），节点就该认为自己已掉线。
	hbCtx, cancel := context.WithCancel(context.Background())
	kaCh, err := r.cli.KeepAlive(hbCtx, grant.ID)
	if err != nil {
		cancel()
		_, _ = r.cli.Revoke(context.Background(), grant.ID)
		return fmt.Errorf("discovery: keepalive: %w", err)
	}
	r.sessions[rec.ID] = &session{lease: grant.ID, cancel: cancel}
	go func() {
		for range kaCh {
			// 续约成功什么都不用做；通道关闭即代表租约已失效。
		}
	}()
	return nil
}

// SetState 更新节点的自报健康态（healthy <-> degraded）。offline 不能自报。
func (r *Registry) SetState(ctx context.Context, id NodeID, state HealthState) error {
	if state == StateOffline {
		return fmt.Errorf("discovery: offline is derived, not self-reported")
	}
	r.mu.Lock()
	s, ok := r.sessions[id]
	r.mu.Unlock()
	if !ok {
		return fmt.Errorf("discovery: node %s not registered", id)
	}
	resp, err := r.cli.Get(ctx, r.key(id))
	if err != nil || len(resp.Kvs) == 0 {
		return fmt.Errorf("discovery: node %s key vanished", id)
	}
	var rec NodeRecord
	_ = json.Unmarshal(resp.Kvs[0].Value, &rec)
	rec.State = state
	val, _ := json.Marshal(rec)
	_, err = r.cli.Put(ctx, r.key(id), string(val), clientv3.WithLease(s.lease))
	return err
}

// Deregister 主动下线：撤销租约，键立即消失。
func (r *Registry) Deregister(ctx context.Context, id NodeID) error {
	r.mu.Lock()
	s, ok := r.sessions[id]
	if ok {
		delete(r.sessions, id)
	}
	r.mu.Unlock()
	if !ok {
		return fmt.Errorf("discovery: node %s not registered", id)
	}
	s.cancel()
	_, err := r.cli.Revoke(ctx, s.lease)
	return err
}

// Suspend 只停心跳、不撤租约，模拟「进程还在但心跳断了」的崩溃：
// 租约到期后键被自动删除，节点被监视器判 offline。
func (r *Registry) Suspend(id NodeID) {
	r.mu.Lock()
	s, ok := r.sessions[id]
	if ok {
		delete(r.sessions, id)
	}
	r.mu.Unlock()
	if ok {
		s.cancel()
	}
}

// List 拉取当前注册表里的全部节点（一次强一致读）。
func (r *Registry) List(ctx context.Context) ([]Member, error) {
	resp, err := r.cli.Get(ctx, r.prefix+"nodes/", clientv3.WithPrefix())
	if err != nil {
		return nil, err
	}
	out := make([]Member, 0, len(resp.Kvs))
	for _, kv := range resp.Kvs {
		var rec NodeRecord
		if err := json.Unmarshal(kv.Value, &rec); err != nil {
			return nil, err
		}
		out = append(out, Member{NodeRecord: rec, LeaseID: kv.Lease})
	}
	return out, nil
}

// Close 取消所有本地心跳协程（不撤租约，等 TTL 自然过期）。
func (r *Registry) Close() {
	r.mu.Lock()
	defer r.mu.Unlock()
	for id, s := range r.sessions {
		s.cancel()
		delete(r.sessions, id)
	}
}

// nodeIDFromKey 从 "prefix/nodes/<id>" 里抠出 <id>。
func nodeIDFromKey(prefix, key string) NodeID {
	return NodeID(strings.TrimPrefix(key, prefix+"nodes/"))
}

// snippet: HealthMonitor.Run 的事件分派（节选自 internal/discovery/watch.go）
switch ev.Type {
case clientv3.EventTypePut:
	var rec NodeRecord
	_ = json.Unmarshal(ev.Kv.Value, &rec)
	mem := Member{NodeRecord: rec, LeaseID: ev.Kv.Lease}
	m.mu.Lock()
	_, existed := m.state[id]
	m.state[id] = mem
	m.mu.Unlock()
	if existed {
		e = Event{Type: EventUpdated, Member: mem} // healthy <-> degraded
	} else {
		e = Event{Type: EventAdded, Member: mem}
	}
case clientv3.EventTypeDelete:
	m.mu.Lock()
	mem, ok := m.state[id]
	if !ok {
		mem = Member{NodeRecord: NodeRecord{ID: id}}
	}
	// 键消失 = 租约过期或主动撤销，一律判 offline。
	// 保留在视图里标 offline（而不是删掉），三态才看得到「谁掉线了」。
	mem.State = StateOffline
	m.state[id] = mem
	m.mu.Unlock()
	e = Event{Type: EventRemoved, Member: mem}
}

=== RUN   TestEtcdRegisterHealthAndOffline
--- PASS: TestEtcdRegisterHealthAndOffline (0.06s)
=== RUN   TestEtcdLeaseExpiry
--- PASS: TestEtcdLeaseExpiry (2.49s)
PASS
ok  	aegis/internal/discovery	3.023s