ec distribution: plan rebalances on a copy of the topology

recordMove mutated the caller's TopologyAnalysis while planning, so an
interrupted or partial execution left the input describing moves that
never happened, and replanning on it returned nothing. PlanRebalance
now simulates on a cloned analysis; plan.Analysis reports the layout
the moves would produce and the caller's input is left untouched.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Chris Lu
2026-10-21 14:21:38 -07:00
co-authored by Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
parent 0d13730600
commit 8163febdaf
2 changed files with 41 additions and 25 deletions
@@ -185,6 +185,11 @@ func TestPlanRebalanceAcrossDCs(t *testing.T) {
assert.Equal(t, 7, plan.Analysis.ShardsByDC["dcA"])
assert.Equal(t, 7, plan.Analysis.ShardsByDC["dcB"])
// the given analysis still describes the un-moved layout
assert.Equal(t, 14, a.ShardsByDC["dcA"])
assert.Equal(t, 0, a.ShardsByDC["dcB"])
assert.Equal(t, []int{0, 1, 2, 3}, a.AllNodes["a1"].ShardIDs)
// parity shards (10-13) move before data shards
assert.True(t, plan.Moves[0].ShardID >= 10)
assert.Equal(t, "b1", plan.Moves[0].DestNode.NodeID)
@@ -226,6 +231,7 @@ func TestPlanRebalanceWithinRack(t *testing.T) {
assert.Equal(t, 2, plan.Analysis.ShardsByNode["n1"])
assert.Equal(t, 2, plan.Analysis.ShardsByNode["n2"])
assert.Equal(t, []int{3, 2}, plan.Analysis.AllNodes["n2"].ShardIDs)
assert.Equal(t, []int{0, 1, 2, 3}, a.AllNodes["n1"].ShardIDs)
}
func TestPlanRebalanceRespectsFreeSlots(t *testing.T) {
@@ -241,6 +247,7 @@ func TestPlanRebalanceRespectsFreeSlots(t *testing.T) {
assert.Equal(t, 3, plan.Analysis.ShardsByNode["n1"])
assert.Equal(t, 1, plan.Analysis.ShardsByNode["n2"])
assert.Equal(t, 0, plan.Analysis.AllNodes["n2"].FreeSlots)
assert.Equal(t, 1, a.AllNodes["n2"].FreeSlots)
}
func TestPlanRebalanceNoopWhenBalanced(t *testing.T) {
@@ -17,7 +17,7 @@ type ShardMove struct {
type RebalancePlan struct {
Moves []ShardMove
Distribution *ECDistribution
Analysis *TopologyAnalysis
Analysis *TopologyAnalysis // simulated layout once every move applies; move endpoints reference its nodes
// Statistics
TotalMoves int
@@ -41,14 +41,23 @@ func NewRebalancer(ec ECConfig, rep ReplicationConfig) *Rebalancer {
// the balanced distribution: first across data centers, then racks within each
// data center, then nodes within each rack. Moves stop when source and
// destination both reach their even share, or when destination nodes have no
// free slots.
// free slots. The moves are planned on a copy of the topology, so the given
// analysis still describes the cluster as it is; plan.Analysis shows the
// layout the moves would produce.
func (r *Rebalancer) PlanRebalance(analysis *TopologyAnalysis) (*RebalancePlan, error) {
if analysis == nil {
return nil, fmt.Errorf("cannot plan rebalancing without a topology analysis")
}
sim := NewTopologyAnalysis()
for _, id := range sortedKeys(analysis.AllNodes) {
node := *analysis.AllNodes[id]
node.ShardIDs = append([]int(nil), node.ShardIDs...)
sim.AddNode(&node)
}
sim.Finalize()
plan := &RebalancePlan{
Distribution: CalculateDistribution(r.ecConfig, r.repConfig),
Analysis: analysis,
Analysis: sim,
}
recordMove := func(shardID int, src, dst *TopologyNode, level, domain string) {
@@ -56,18 +65,18 @@ func (r *Rebalancer) PlanRebalance(analysis *TopologyAnalysis) (*RebalancePlan,
dst.ShardIDs = append(dst.ShardIDs, shardID)
src.FreeSlots++
dst.FreeSlots--
analysis.ShardsByDC[src.DataCenter]--
analysis.ShardsByDC[dst.DataCenter]++
analysis.ShardsByRack[rackKey(src.DataCenter, src.Rack)]--
analysis.ShardsByRack[rackKey(dst.DataCenter, dst.Rack)]++
analysis.ShardsByNode[src.NodeID]--
analysis.ShardsByNode[dst.NodeID]++
analysis.DCToShards[src.DataCenter] = removeShard(analysis.DCToShards[src.DataCenter], shardID)
analysis.DCToShards[dst.DataCenter] = append(analysis.DCToShards[dst.DataCenter], shardID)
analysis.RackToShards[rackKey(src.DataCenter, src.Rack)] = removeShard(analysis.RackToShards[rackKey(src.DataCenter, src.Rack)], shardID)
analysis.RackToShards[rackKey(dst.DataCenter, dst.Rack)] = append(analysis.RackToShards[rackKey(dst.DataCenter, dst.Rack)], shardID)
analysis.NodeToShards[src.NodeID] = removeShard(analysis.NodeToShards[src.NodeID], shardID)
analysis.NodeToShards[dst.NodeID] = append(analysis.NodeToShards[dst.NodeID], shardID)
sim.ShardsByDC[src.DataCenter]--
sim.ShardsByDC[dst.DataCenter]++
sim.ShardsByRack[rackKey(src.DataCenter, src.Rack)]--
sim.ShardsByRack[rackKey(dst.DataCenter, dst.Rack)]++
sim.ShardsByNode[src.NodeID]--
sim.ShardsByNode[dst.NodeID]++
sim.DCToShards[src.DataCenter] = removeShard(sim.DCToShards[src.DataCenter], shardID)
sim.DCToShards[dst.DataCenter] = append(sim.DCToShards[dst.DataCenter], shardID)
sim.RackToShards[rackKey(src.DataCenter, src.Rack)] = removeShard(sim.RackToShards[rackKey(src.DataCenter, src.Rack)], shardID)
sim.RackToShards[rackKey(dst.DataCenter, dst.Rack)] = append(sim.RackToShards[rackKey(dst.DataCenter, dst.Rack)], shardID)
sim.NodeToShards[src.NodeID] = removeShard(sim.NodeToShards[src.NodeID], shardID)
sim.NodeToShards[dst.NodeID] = append(sim.NodeToShards[dst.NodeID], shardID)
plan.Moves = append(plan.Moves, ShardMove{
ShardID: shardID,
SourceNode: src,
@@ -149,9 +158,9 @@ func (r *Rebalancer) PlanRebalance(analysis *TopologyAnalysis) (*RebalancePlan,
}
}
dcList := sortedKeys(analysis.DCToRacks)
dcList := sortedKeys(sim.DCToRacks)
dcRackKeys := func(dc string) []string {
racks := analysis.DCToRacks[dc]
racks := sim.DCToRacks[dc]
keys := make([]string, 0, len(racks))
for _, rack := range racks {
keys = append(keys, rackKey(dc, rack))
@@ -160,23 +169,23 @@ func (r *Rebalancer) PlanRebalance(analysis *TopologyAnalysis) (*RebalancePlan,
}
totalShards := 0
for _, dc := range dcList {
totalShards += shardCount(dcRackKeys(dc), analysis.RackToNodes)
totalShards += shardCount(dcRackKeys(dc), sim.RackToNodes)
}
dcCount := func(dc string) int { return shardCount(dcRackKeys(dc), analysis.RackToNodes) }
dcCount := func(dc string) int { return shardCount(dcRackKeys(dc), sim.RackToNodes) }
balance("dc", dcList, dcCount, ceilDivide(totalShards, len(dcList)),
func(dc string) []*TopologyNode { return nodesOfRacks(dcRackKeys(dc), analysis.RackToNodes) },
func(dc string) []*TopologyNode { return nodesOfRacks(dcRackKeys(dc), sim.RackToNodes) },
func(string) []string { return dcList })
for _, dc := range dcList {
racks := analysis.DCToRacks[dc]
rackCount := func(rack string) int { return shardCount([]string{rackKey(dc, rack)}, analysis.RackToNodes) }
racks := sim.DCToRacks[dc]
rackCount := func(rack string) int { return shardCount([]string{rackKey(dc, rack)}, sim.RackToNodes) }
balance("rack", racks, rackCount, ceilDivide(dcCount(dc), len(racks)),
func(rack string) []*TopologyNode { return analysis.RackToNodes[rackKey(dc, rack)] },
func(rack string) []*TopologyNode { return sim.RackToNodes[rackKey(dc, rack)] },
func(string) []string { return racks })
}
for _, rack := range sortedKeys(analysis.RackToNodes) {
nodes := analysis.RackToNodes[rack]
for _, rack := range sortedKeys(sim.RackToNodes) {
nodes := sim.RackToNodes[rack]
nodeCount := func(nodeID string) int {
for _, n := range nodes {
if n.NodeID == nodeID {