From 09c8b65e5b81cbb156730744dbdeb448de962778 Mon Sep 17 00:00:00 2001
From: Chris Lu <chrislusf@users.noreply.github.com>
Date: Wed, 21 Oct 2026 14:42:30 -0700
Subject: [PATCH] ec distribution: register nodes implied by shard locations

---
 .../erasure_coding/distribution/analysis.go   | 26 ++++++++++++++-----
 .../distribution/distribution_test.go         | 20 ++++++++++++++
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/weed/storage/erasure_coding/distribution/analysis.go b/weed/storage/erasure_coding/distribution/analysis.go
index 5f60efda56..d71d77fc3d 100644
--- a/weed/storage/erasure_coding/distribution/analysis.go
+++ b/weed/storage/erasure_coding/distribution/analysis.go
@@ -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)
 	}
 }
diff --git a/weed/storage/erasure_coding/distribution/distribution_test.go b/weed/storage/erasure_coding/distribution/distribution_test.go
index 97e053d192..ef73b8d814 100644
--- a/weed/storage/erasure_coding/distribution/distribution_test.go
+++ b/weed/storage/erasure_coding/distribution/distribution_test.go
@@ -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))