《游戏服务端架构·规模化、全球化与成本》4.2 一致性哈希与迁移

// file: internal/consistenthash/ring.go
// tier: L1
// verify: go test ./internal/consistenthash/
package consistenthash

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

var (
	ErrNoNodes   = errors.New("consistenthash: ring needs at least one node")
	ErrVnodes    = errors.New("consistenthash: virtual nodes must be positive")
	ErrEmptyNode = errors.New("consistenthash: empty node name")
	ErrDuplicate = errors.New("consistenthash: duplicate node")
	ErrUnknown   = errors.New("consistenthash: node is not on the ring")
)

// token is one virtual node position. The node name is kept beside the hash so
// a collision is still deterministic: ordering falls back to the name.
type token struct {
	hash uint64
	node string
}

// Ring maps keys to owners. vnodes is fixed at construction because changing
// it changes every token, which is a full migration even if the node list is
// unchanged.
type Ring struct {
	vnodes int
	nodes  map[string]bool
	tokens []token
}

func New(vnodes int, nodes ...string) (*Ring, error) {
	if vnodes <= 0 {
		return nil, ErrVnodes
	}
	r := &Ring{vnodes: vnodes, nodes: make(map[string]bool)}
	for _, node := range nodes {
		if err := r.Add(node); err != nil {
			return nil, err
		}
	}
	if len(r.nodes) == 0 {
		return nil, ErrNoNodes
	}
	return r, nil
}

func hashString(s string) uint64 {
	h := fnv.New64a()
	_, _ = h.Write([]byte(s))
	return h.Sum64()
}

// Add inserts a node and all of its virtual nodes. It does not move data; the
// caller must run a plan and copy the pixels, then switch routing.
func (r *Ring) Add(node string) error {
	if node == "" {
		return ErrEmptyNode
	}
	if r.nodes[node] {
		return fmt.Errorf("%w: %q", ErrDuplicate, node)
	}
	r.nodes[node] = true
	for i := 0; i < r.vnodes; i++ {
		r.tokens = append(r.tokens, token{
			hash: hashString(node + "#" + strconv.Itoa(i)),
			node: node,
		})
	}
	r.sortTokens()
	return nil
}

// Remove refuses to empty the ring. A caller that really wants to stop serving
// should remove the last node from routing, not from the safety check.
func (r *Ring) Remove(node string) error {
	if !r.nodes[node] {
		return fmt.Errorf("%w: %q", ErrUnknown, node)
	}
	if len(r.nodes) == 1 {
		return ErrNoNodes
	}
	delete(r.nodes, node)
	kept := r.tokens[:0]
	for _, tok := range r.tokens {
		if tok.node != node {
			kept = append(kept, tok)
		}
	}
	r.tokens = kept
	return nil
}

func (r *Ring) sortTokens() {
	sort.Slice(r.tokens, func(i, j int) bool {
		if r.tokens[i].hash != r.tokens[j].hash {
			return r.tokens[i].hash < r.tokens[j].hash
		}
		return r.tokens[i].node < r.tokens[j].node
	})
}

// Nodes returns the node names in sorted order.
func (r *Ring) Nodes() []string {
	out := make([]string, 0, len(r.nodes))
	for node := range r.nodes {
		out = append(out, node)
	}
	sort.Strings(out)
	return out
}

// Owner finds the first token at or after the key hash, wrapping at the end.
func (r *Ring) Owner(key string) (string, bool) {
	if len(r.tokens) == 0 {
		return "", false
	}
	h := hashString(key)
	i := sort.Search(len(r.tokens), func(i int) bool { return r.tokens[i].hash >= h })
	if i == len(r.tokens) {
		i = 0
	}
	return r.tokens[i].node, true
}

// Move is one sampled key whose owner changed.
type Move struct {
	Key  string
	From string
	To   string
}

// Plan is a sample-based estimate, not a copy list. Keys outside the sample
// are not proven unchanged; the migration worker still needs a durable scan.
type Plan struct {
	Moves []Move
	Stay  int
}

func (p Plan) MovedTo(node string) int {
	n := 0
	for _, m := range p.Moves {
		if m.To == node {
			n++
		}
	}
	return n
}

func (p Plan) Total() int { return len(p.Moves) + p.Stay }

// BuildPlan compares two rings over a sorted key sample. Sorting makes the
// output stable across runs, which matters when two operators diff plans.
func BuildPlan(old, next *Ring, keys []string) Plan {
	sorted := append([]string(nil), keys...)
	sort.Strings(sorted)
	var plan Plan
	for _, key := range sorted {
		from, ok := old.Owner(key)
		if !ok {
			continue
		}
		to, ok := next.Owner(key)
		if !ok {
			continue
		}
		if from == to {
			plan.Stay++
			continue
		}
		plan.Moves = append(plan.Moves, Move{Key: key, From: from, To: to})
	}
	return plan
}