《MMO 服务器开发》9.3 分片与再均衡

// file: internal/store/shard.go
// tier: L1
// verify: go test ./internal/store/ -run TestConsistentHashRing -v

package store

import (
	"hash/crc32"
	"sort"
	"strconv"
)

// Ring is a consistent-hash ring. Each shard is placed at vnodes points, and a
// key maps to the first shard at or after hash(key) walking the ring. Adding a
// shard only steals the keys that fall in the new points; with enough virtual
// nodes that is roughly (new-old)/new of all keys, not all of them.
type Ring struct {
	points []ringPoint
	shards int
}

type ringPoint struct {
	hash  uint32
	shard int
}

// NewRing builds a ring over `shards` shards, each with `vnodes` virtual
// points. More virtual nodes smooth the key distribution across shards.
func NewRing(shards, vnodes int) *Ring {
	r := &Ring{shards: shards}
	for s := 0; s < shards; s++ {
		for v := 0; v < vnodes; v++ {
			key := "shard-" + strconv.Itoa(s) + "#" + strconv.Itoa(v)
			r.points = append(r.points, ringPoint{
				hash:  crc32.ChecksumIEEE([]byte(key)),
				shard: s,
			})
		}
	}
	sort.Slice(r.points, func(i, j int) bool {
		if r.points[i].hash != r.points[j].hash {
			return r.points[i].hash < r.points[j].hash
		}
		return r.points[i].shard < r.points[j].shard
	})
	return r
}

// ShardOf returns the shard a key belongs to. It is deterministic: the same key
// always lands on the same shard for a given ring.
func (r *Ring) ShardOf(key string) int {
	h := crc32.ChecksumIEEE([]byte(key))
	i := sort.Search(len(r.points), func(i int) bool { return r.points[i].hash >= h })
	if i == len(r.points) {
		i = 0
	}
	return r.points[i].shard
}

// Shards reports how many shards the ring covers.
func (r *Ring) Shards() int { return r.shards }

// RangeShard assigns a whole server (range) to a shard by id. Range sharding
// keeps a server's players together but cannot split one hot server.
func RangeShard(serverID uint16, shards int) int {
	return int(serverID) % shards
}

// snippet: internal/store/shard_pg.go 中 Rebalance 的核心循环
	for s := 0; s < oldRing.Shards(); s++ {
		ids, _ := PlayerIDs(ctx, pool, s)
		for _, id := range ids {
			dst := newRing.ShardOf(strconv.FormatUint(id, 10))
			if dst == s {
				continue // 落点没变，原地不动
			}
			copyPlayerRow(ctx, pool, s, dst, id)          // 1. 双写进新分片
			same, _ := samePlayerRow(ctx, pool, s, dst, id)
			if !same {
				return moved, verified, fmt.Errorf("verify failed for player %d", id) // 2. 校验
			}
			verified++
			pool.Exec(ctx, "DELETE FROM "+shardTable(s)+" WHERE player_id=$1", id) // 3. 切读后删旧
			moved++
		}
	}

// file: internal/store/shard_pg_test.go
// tier: L2
// verify: go test ./internal/store/ -run TestRebalanceFrom4To8 -v

package store

import (
	"context"
	"fmt"
	"strconv"
	"testing"
)

// TestRebalanceFrom4To8 grows a 4-shard cluster to 8 shards and asserts the two
// things a rebalance must never get wrong: no data is lost, and after the move
// every player is reachable through the new ring.
func TestRebalanceFrom4To8(t *testing.T) {
	pool := testPool(t)
	defer pool.Close()
	resetSchema(t, pool)
	ctx := context.Background()

	const n = 1000
	oldRing := NewRing(4, 128)
	newRing := NewRing(8, 128)

	if err := CreateShardTables(ctx, pool, 8); err != nil {
		t.Fatal(err)
	}

	// Seed n players into the shards the OLD ring assigns them to.
	for i := 0; i < n; i++ {
		id := uint64(100000 + i)
		p := &Player{ID: id, Account: fmt.Sprintf("acc-%d", id), Name: "p",
			Level: 1, Gold: int64(id), Version: 1}
		s := oldRing.ShardOf(strconv.FormatUint(id, 10))
		if err := InsertPlayerInto(ctx, pool, s, p); err != nil {
			t.Fatal(err)
		}
	}
	before, err := TotalCount(ctx, pool, 4) // only shards 0..3 hold data
	if err != nil {
		t.Fatal(err)
	}
	if before != n {
		t.Fatalf("before rebalance total=%d want %d", before, n)
	}

	moved, verified, err := Rebalance(ctx, pool, oldRing, newRing)
	if err != nil {
		t.Fatalf("rebalance: %v", err)
	}
	if moved == 0 {
		t.Fatal("rebalance moved nothing; the test is not exercising the migration")
	}

	// No loss: the total across all 8 shards is unchanged. A failed cut would
	// show up here as a total greater than n (a duplicated row).
	after, err := TotalCount(ctx, pool, 8)
	if err != nil {
		t.Fatal(err)
	}
	if after != n {
		t.Fatalf("after rebalance total=%d want %d (data lost or duplicated)", after, n)
	}

	// Routing correct: every player now lives in exactly the shard the NEW ring
	// says it should.
	for i := 0; i < n; i++ {
		id := uint64(100000 + i)
		want := newRing.ShardOf(strconv.FormatUint(id, 10))
		ok, err := PlayerRowExists(ctx, pool, want, id)
		if err != nil {
			t.Fatal(err)
		}
		if !ok {
			t.Fatalf("player %d not found in its new shard %d", id, want)
		}
	}

	// And it is gone from any old shard it no longer belongs to.
	for i := 0; i < n; i++ {
		id := uint64(100000 + i)
		key := strconv.FormatUint(id, 10)
		oldShard := oldRing.ShardOf(key)
		newShard := newRing.ShardOf(key)
		if oldShard == newShard {
			continue
		}
		ok, err := PlayerRowExists(ctx, pool, oldShard, id)
		if err != nil {
			t.Fatal(err)
		}
		if ok {
			t.Fatalf("player %d still lingers in old shard %d", id, oldShard)
		}
	}

	t.Logf("rebalance ok: moved=%d verified=%d, total=%d across 8 shards", moved, verified, after)
}