Preserve in-flight shards in shardMap across /_cat/shards relocation window

sean- · sean- · commit 73325d42f32a · 2026-05-22T04:22:34.000-07:00
/_cat/shards parsing skips RELOCATING and INITIALIZING rows (see
discovery.go:getShardPlacement), so a shard mid-relocation is absent
from placement.ShardToNodes for the duration of the relocation window.
The previous unconditional store would clobber a complete cached map
with this partial response, leaving Shards[N] == nil for the in-flight
shard and disabling shard-exact routing for it.

Carry forward entries the new response is missing: a shard in
RELOCATING state still serves reads from the source until the
destination flips to STARTED, so the prior entry remains a valid
routing target for the relocation window. New data wins where present;
entries beyond NumberOfPrimaryShards are dropped so a shrunk index
doesn't retain unreachable high-numbered shards.

Use CompareAndSwap to bound the visible Load+merge+Store window: if a
concurrent writer (e.g. an external DiscoverNodes racing the
discoveryRefreshLoop's cat-only refresh) replaced the map, we drop our
update rather than clobber theirs - both racers derived from
near-simultaneous /_cat/shards calls so retrying buys nothing.

Signed-off-by: Sean Chittenden &lt;sean.chittenden@crowdstrike.com&gt;
diff --git a/opensearchtransport/index_routing_cache.go b/opensearchtransport/index_routing_cache.go
@@ -7,6 +7,7 @@
 package opensearchtransport
 
 import (
+	"maps"
 	"math"
 	"sync"
 	"sync/atomic"
@@ -310,22 +311,7 @@ func (c *indexSlotCache) updateFromDiscovery(shardPlacement map[string]*indexSha
 				nodes := placement.Nodes
 				slot.shardNodeCount.Store(int32(len(nodes))) //nolint:gosec // Node count bounded by cluster size.
 				slot.shardNodeNames.Store(&nodes)
-
-				// Update per-shard-number map for murmur3 routing.
-				// Only store a new map when placement data is complete.
-				// When ShardToNodes is empty (e.g., transient /_cat/shards
-				// response during shard relocation), preserve the existing
-				// map rather than niling it -- a stale shard map still
-				// routes correctly; a nil one disables shard-exact routing
-				// entirely until the next successful discovery cycle.
-				if len(placement.ShardToNodes) > 0 {
-					sm := &indexShardMap{
-						NumberOfPrimaryShards: placement.NumberOfPrimaryShards,
-						RoutingNumShards:      placement.RoutingNumShards,
-						Shards:                placement.ShardToNodes,
-					}
-					slot.shardMap.Store(sm)
-				}
+				slot.mergeShardMap(placement)
 			} else {
 				// Index not in shard data -- may have been deleted or
 				// the /_cat/shards response was truncated/stale. Clear
@@ -379,6 +365,59 @@ func (c *indexSlotCache) updateFromDiscovery(shardPlacement map[string]*indexSha
 	}
 }
 
+// mergeShardMap updates the cached per-shard placement map from a fresh
+// /_cat/shards-derived placement.
+//
+// /_cat/shards rows for RELOCATING (source) and INITIALIZING (destination)
+// shards are skipped during parsing, so a shard in flight between two nodes
+// is absent from placement.ShardToNodes for the duration of the relocation.
+// Replacing the cached map with a partial response would leave Shards[N] == nil
+// for the in-flight shard and disable shard-exact routing for it.
+//
+// Carry forward entries the new response is missing: a shard in RELOCATING
+// state still serves reads from the source node until the destination flips
+// to STARTED, so the prior entry remains a valid routing target for the
+// relocation window. New data wins where present; entries beyond the current
+// primary count are dropped so a shrunk index doesn't retain unreachable
+// high-numbered shards.
+//
+// Empty placements are ignored to preserve the previous map across transient
+// /_cat/shards failures (a stale map still routes correctly; a nil one
+// disables shard-exact routing entirely).
+//
+// The Load+CAS pattern bounds the visible window for the merge: if another
+// writer (e.g. an external DiscoverNodes racing the discoveryRefreshLoop's
+// cat-only refresh) replaced the map between this Load and the CAS, the swap
+// fails and we drop our update. Both racers derived their placements from
+// near-simultaneous /_cat/shards calls, so the winner's data is at least as
+// fresh as ours; retrying would just produce equivalent work.
+func (slot *indexSlot) mergeShardMap(placement *indexShardPlacement) {
+	if len(placement.ShardToNodes) == 0 {
+		return
+	}
+
+	old := slot.shardMap.Load()
+	merged := placement.ShardToNodes
+	if old != nil && len(old.Shards) > len(merged) {
+		merged = make(map[int]*shardNodes, len(old.Shards))
+		maps.Copy(merged, old.Shards)
+		maps.Copy(merged, placement.ShardToNodes)
+		if placement.NumberOfPrimaryShards > 0 {
+			for k := range merged {
+				if k >= placement.NumberOfPrimaryShards {
+					delete(merged, k)
+				}
+			}
+		}
+	}
+
+	slot.shardMap.CompareAndSwap(old, &indexShardMap{
+		NumberOfPrimaryShards: placement.NumberOfPrimaryShards,
+		RoutingNumShards:      placement.RoutingNumShards,
+		Shards:                merged,
+	})
+}
+
 // shardNodeNameSet returns the set of node names hosting shards for an index,
 // or nil if unknown. Used by rendezvousTopK for shard-aware partitioning.
 func (slot *indexSlot) shardNodeNameSet() map[string]struct{} {
diff --git a/opensearchtransport/index_routing_cache_internal_test.go b/opensearchtransport/index_routing_cache_internal_test.go
@@ -295,6 +295,155 @@ func TestIndexSlotCacheUpdateFromDiscovery(t *testing.T) {
 	})
 }
 
+// TestIndexSlotCacheShardMapMerge covers the merge semantics of updateFromDiscovery
+// for the per-shard placement map: a new partial /_cat/shards response (e.g. during
+// a relocation window where source/destination rows are skipped) must not clobber
+// shards from a previously-complete cached map, while a genuine shrink in primary
+// count must drop unreachable old entries.
+func TestIndexSlotCacheShardMapMerge(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name      string
+		prior     *indexShardMap // nil means no prior map
+		incoming  *indexShardPlacement
+		wantCount int
+		wantPrims int
+		wantShard map[int]*shardNodes // expected entries; nil means "must not be present"
+	}{
+		{
+			name:  "no prior map: store as-is",
+			prior: nil,
+			incoming: &indexShardPlacement{
+				NumberOfPrimaryShards: 2,
+				RoutingNumShards:      256,
+				ShardToNodes: map[int]*shardNodes{
+					0: {Primary: "node0"},
+					1: {Primary: "node1"},
+				},
+			},
+			wantCount: 2,
+			wantPrims: 2,
+			wantShard: map[int]*shardNodes{
+				0: {Primary: "node0"},
+				1: {Primary: "node1"},
+			},
+		},
+		{
+			name: "complete new placement overwrites prior",
+			prior: &indexShardMap{
+				NumberOfPrimaryShards: 2,
+				RoutingNumShards:      256,
+				Shards: map[int]*shardNodes{
+					0: {Primary: "old0"},
+					1: {Primary: "old1"},
+				},
+			},
+			incoming: &indexShardPlacement{
+				NumberOfPrimaryShards: 2,
+				RoutingNumShards:      256,
+				ShardToNodes: map[int]*shardNodes{
+					0: {Primary: "new0"},
+					1: {Primary: "new1"},
+				},
+			},
+			wantCount: 2,
+			wantPrims: 2,
+			wantShard: map[int]*shardNodes{
+				0: {Primary: "new0"},
+				1: {Primary: "new1"},
+			},
+		},
+		{
+			name: "relocation window: missing shard preserved from prior",
+			prior: &indexShardMap{
+				NumberOfPrimaryShards: 3,
+				RoutingNumShards:      384,
+				Shards: map[int]*shardNodes{
+					0: {Primary: "node0", Replicas: []string{"node1"}},
+					1: {Primary: "node1", Replicas: []string{"node2"}},
+					2: {Primary: "node2", Replicas: []string{"node0"}},
+				},
+			},
+			incoming: &indexShardPlacement{
+				NumberOfPrimaryShards: 3,
+				RoutingNumShards:      384,
+				ShardToNodes: map[int]*shardNodes{
+					// Shard 1 absent: source is RELOCATING, destination is INITIALIZING,
+					// both filtered out by the /_cat/shards parser.
+					0: {Primary: "node0", Replicas: []string{"node3"}}, // replica moved
+					2: {Primary: "node2", Replicas: []string{"node0"}},
+				},
+			},
+			wantCount: 3,
+			wantPrims: 3,
+			wantShard: map[int]*shardNodes{
+				0: {Primary: "node0", Replicas: []string{"node3"}}, // new wins
+				1: {Primary: "node1", Replicas: []string{"node2"}}, // carried forward
+				2: {Primary: "node2", Replicas: []string{"node0"}}, // new wins
+			},
+		},
+		{
+			name: "shrunk index drops shards beyond new primary count",
+			prior: &indexShardMap{
+				NumberOfPrimaryShards: 4,
+				RoutingNumShards:      512,
+				Shards: map[int]*shardNodes{
+					0: {Primary: "node0"},
+					1: {Primary: "node1"},
+					2: {Primary: "node2"},
+					3: {Primary: "node3"},
+				},
+			},
+			incoming: &indexShardPlacement{
+				NumberOfPrimaryShards: 2,
+				RoutingNumShards:      512,
+				ShardToNodes: map[int]*shardNodes{
+					0: {Primary: "node0"},
+					1: {Primary: "node1"},
+				},
+			},
+			wantCount: 2,
+			wantPrims: 2,
+			wantShard: map[int]*shardNodes{
+				0: {Primary: "node0"},
+				1: {Primary: "node1"},
+				2: nil,
+				3: nil,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			c := newIndexSlotCache(indexSlotCacheConfig{})
+			slot := c.getOrCreate("idx")
+			if tt.prior != nil {
+				slot.shardMap.Store(tt.prior)
+			}
+
+			c.updateFromDiscovery(map[string]*indexShardPlacement{"idx": tt.incoming}, 10, time.Now())
+
+			sm := slot.shardMap.Load()
+			require.NotNil(t, sm)
+			require.Equal(t, tt.wantPrims, sm.NumberOfPrimaryShards)
+			require.Len(t, sm.Shards, tt.wantCount)
+
+			for shard, want := range tt.wantShard {
+				if want == nil {
+					require.NotContains(t, sm.Shards, shard, "shard %d should be dropped", shard)
+					continue
+				}
+				got, ok := sm.Shards[shard]
+				require.True(t, ok, "shard %d missing from merged map", shard)
+				require.Equal(t, want.Primary, got.Primary, "shard %d primary", shard)
+				require.Equal(t, want.Replicas, got.Replicas, "shard %d replicas", shard)
+			}
+		})
+	}
+}
+
 func TestIndexSlotShardNodeNameSet(t *testing.T) {
 	t.Parallel()
 
