ec distribution: register nodes implied by shard locations

This commit is contained in:
Chris Lu
2026-10-21 14:42:30 -07:00
parent f04701b542
commit 09c8b65e5b
2 changed files with 40 additions and 6 deletions
@@ -26,6 +26,13 @@ func (a *TopologyAnalysis) AddNode(node *TopologyNode) {
if !contains(a.DCToRacks[node.DataCenter], node.Rack) {
a.DCToRacks[node.DataCenter] = append(a.DCToRacks[node.DataCenter], node.Rack)
}
seen := node.ShardIDs[:0]
for _, shardID := range node.ShardIDs {
if !contains(seen, shardID) {
seen = append(seen, shardID)
}
}
node.ShardIDs = seen
for _, shardID := range node.ShardIDs {
a.AddShardLocation(ShardLocation{
ShardID: shardID,
@@ -38,19 +45,26 @@ func (a *TopologyAnalysis) AddNode(node *TopologyNode) {
// AddShardLocation records one shard's placement in the topology. Repeated
// calls for a shard already known to its node are ignored, so callers may
// either set ShardIDs on the TopologyNode or register each shard here.
// either set ShardIDs on the TopologyNode or register each shard here. A
// location naming an unregistered node registers that node so the rebalancer
// can see the shard; for a registered node the node's own dc/rack win.
func (a *TopologyAnalysis) AddShardLocation(loc ShardLocation) {
node, ok := a.AllNodes[loc.NodeID]
if !ok {
node = &TopologyNode{NodeID: loc.NodeID, DataCenter: loc.DataCenter, Rack: loc.Rack}
a.AddNode(node)
}
if contains(a.NodeToShards[loc.NodeID], loc.ShardID) {
return
}
a.ShardsByDC[loc.DataCenter]++
a.ShardsByRack[rackKey(loc.DataCenter, loc.Rack)]++
a.ShardsByDC[node.DataCenter]++
a.ShardsByRack[rackKey(node.DataCenter, node.Rack)]++
a.ShardsByNode[loc.NodeID]++
a.DCToShards[loc.DataCenter] = append(a.DCToShards[loc.DataCenter], loc.ShardID)
a.RackToShards[rackKey(loc.DataCenter, loc.Rack)] = append(a.RackToShards[rackKey(loc.DataCenter, loc.Rack)], loc.ShardID)
a.DCToShards[node.DataCenter] = append(a.DCToShards[node.DataCenter], loc.ShardID)
a.RackToShards[rackKey(node.DataCenter, node.Rack)] = append(a.RackToShards[rackKey(node.DataCenter, node.Rack)], loc.ShardID)
a.NodeToShards[loc.NodeID] = append(a.NodeToShards[loc.NodeID], loc.ShardID)
a.TotalShards++
if node, ok := a.AllNodes[loc.NodeID]; ok && !contains(node.ShardIDs, loc.ShardID) {
if !contains(node.ShardIDs, loc.ShardID) {
node.ShardIDs = append(node.ShardIDs, loc.ShardID)
}
}
@@ -341,6 +341,26 @@ func TestPlanRebalanceNoDuplicateShardInDomain(t *testing.T) {
assert.Equal(t, []int{3}, plan.Analysis.AllNodes["b1"].ShardIDs)
}
// Shard locations may be registered without registering the node first; the
// rebalancer must still see the shard, counted under the node's own dc/rack.
func TestPlanRebalanceShardForUnregisteredNode(t *testing.T) {
a := NewTopologyAnalysis()
a.AddShardLocation(ShardLocation{ShardID: 0, NodeID: "n1", DataCenter: "dc1", Rack: "r1"})
a.AddShardLocation(ShardLocation{ShardID: 1, NodeID: "n1", DataCenter: "dc1", Rack: "r1"})
a.AddNode(makeNode("n2", "dc1", "r1", 10))
a.Finalize()
assert.Equal(t, 1, a.TotalDCs)
assert.Equal(t, 2, a.ShardsByNode["n1"])
assert.Equal(t, []int{0, 1}, a.AllNodes["n1"].ShardIDs)
plan, err := NewRebalancer(DefaultECConfig(), ReplicationConfig{MinDataCenters: 1, MinRacksPerDC: 1, MinNodesPerRack: 2}).PlanRebalance(a)
require.NoError(t, err)
assert.Equal(t, 1, plan.TotalMoves)
assert.Equal(t, 1, plan.Analysis.ShardsByNode["n1"])
assert.Equal(t, 1, plan.Analysis.ShardsByNode["n2"])
}
func TestPlanRebalanceNoDuplicateShardOnNode(t *testing.T) {
a := NewTopologyAnalysis()
a.AddNode(makeNode("n1", "dc1", "r1", 10, 0, 1, 2, 3))