《游戏服务端编程实践》10.3.3 负载均衡与消息分片

// file: ch10/hub/shard.go
// tier: L1
// verify: go test ./ch10/hub/

package hub

import (
	"errors"
	"fmt"
	"hash/fnv"
	"sort"
	"sync"
)

var (
	ErrNoInstance        = errors.New("shard: no healthy instance")
	ErrUnknownInstance   = errors.New("shard: unknown instance")
	ErrDuplicateInstance = errors.New("shard: duplicate instance")
	ErrBadReplicas       = errors.New("shard: replicas must be positive")
)

type ringEntry struct {
	hash uint32
	id   string
}

// ShardRouter 用一致性哈希把频道固定到聊天服实例；热点频道可以单独指定。
type ShardRouter struct {
	mu       sync.RWMutex
	replicas int
	entries  []ringEntry
	healthy  map[string]bool
	hot      map[string]string
}

func NewShardRouter(replicas int, ids ...string) (*ShardRouter, error) {
	if replicas <= 0 {
		return nil, ErrBadReplicas
	}
	r := &ShardRouter{
		replicas: replicas,
		healthy:  make(map[string]bool),
		hot:      make(map[string]string),
	}
	for _, id := range ids {
		if err := r.AddInstance(id); err != nil {
			return nil, err
		}
	}
	return r, nil
}

// AddInstance 为实例生成 replicas 个虚拟节点；重复 ID 会被拒绝。
func (r *ShardRouter) AddInstance(id string) error {
	if id == "" {
		return fmt.Errorf("%w: empty id", ErrUnknownInstance)
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, ok := r.healthy[id]; ok {
		return ErrDuplicateInstance
	}
	r.healthy[id] = true
	for i := 0; i < r.replicas; i++ {
		r.entries = append(r.entries, ringEntry{
			hash: hashKey(fmt.Sprintf("%s#%d", id, i)),
			id:   id,
		})
	}
	sort.Slice(r.entries, func(i, j int) bool {
		if r.entries[i].hash == r.entries[j].hash {
			return r.entries[i].id < r.entries[j].id
		}
		return r.entries[i].hash < r.entries[j].hash
	})
	return nil
}

// RemoveInstance 摘除实例，并返回因此失去专属实例的热点频道。
func (r *ShardRouter) RemoveInstance(id string) []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, ok := r.healthy[id]; !ok {
		return nil
	}
	delete(r.healthy, id)
	kept := r.entries[:0]
	for _, entry := range r.entries {
		if entry.id != id {
			kept = append(kept, entry)
		}
	}
	r.entries = kept
	var lost []string
	for channel, target := range r.hot {
		if target == id {
			delete(r.hot, channel)
			lost = append(lost, channel)
		}
	}
	sort.Strings(lost)
	return lost
}

// SetHealthy 只影响选择，不删除虚拟节点；恢复后频道归属立刻复原。
func (r *ShardRouter) SetHealthy(id string, healthy bool) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, ok := r.healthy[id]; !ok {
		return ErrUnknownInstance
	}
	r.healthy[id] = healthy
	return nil
}

// AssignHot 把热点频道固定到指定实例。
func (r *ShardRouter) AssignHot(channel, instanceID string) error {
	if err := ValidateChannel(channel); err != nil {
		return err
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, ok := r.healthy[instanceID]; !ok {
		return ErrUnknownInstance
	}
	r.hot[channel] = instanceID
	return nil
}

// Pick 返回频道归属实例；热点固定优先，其次沿哈希环找到第一个健康实例。
func (r *ShardRouter) Pick(channel string) (string, error) {
	if err := ValidateChannel(channel); err != nil {
		return "", err
	}
	r.mu.RLock()
	defer r.mu.RUnlock()
	if len(r.entries) == 0 {
		return "", ErrNoInstance
	}
	if target, ok := r.hot[channel]; ok && r.healthy[target] {
		return target, nil
	}
	start := sort.Search(len(r.entries), func(i int) bool {
		return r.entries[i].hash >= hashKey(channel)
	})
	for i := 0; i < len(r.entries); i++ {
		entry := r.entries[(start+i)%len(r.entries)]
		if r.healthy[entry.id] {
			return entry.id, nil
		}
	}
	return "", ErrNoInstance
}

// Instances 返回全部实例 ID（含不健康实例），按名称排序。
func (r *ShardRouter) Instances() []string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]string, 0, len(r.healthy))
	for id := range r.healthy {
		out = append(out, id)
	}
	sort.Strings(out)
	return out
}

// HotChannels 返回热点频道映射的副本。
func (r *ShardRouter) HotChannels() map[string]string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make(map[string]string, len(r.hot))
	for channel, id := range r.hot {
		out[channel] = id
	}
	return out
}

func hashKey(s string) uint32 {
	h := fnv.New32a()
	_, _ = h.Write([]byte(s))
	return h.Sum32()
}

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

ok  	gspractice/ch10/hub	0.275s