Make resilient to metadata changes

nicktindall · nicktindall · commit 19b74750dffc · 2025-11-06T15:41:44.000+11:00
diff --git a/server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/UndesiredAllocationsTracker.java b/server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/UndesiredAllocationsTracker.java
@@ -9,9 +9,6 @@
 
 package org.elasticsearch.cluster.routing.allocation.allocator;
 
-import com.carrotsearch.hppc.ObjectLongHashMap;
-import com.carrotsearch.hppc.ObjectLongMap;
-
 import org.elasticsearch.cluster.routing.RoutingNodes;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
@@ -20,12 +17,13 @@
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.time.TimeProvider;
 import org.elasticsearch.core.TimeValue;
-import org.elasticsearch.core.Tuple;
+import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 
+import java.util.HashMap;
+import java.util.Map;
 import java.util.stream.Collectors;
-import java.util.stream.StreamSupport;
 
 /**
  * Keeps track of a limited number of shards that are currently in undesired allocations. If the
@@ -70,7 +68,7 @@ public class UndesiredAllocationsTracker {
     );
 
     private final TimeProvider timeProvider;
-    private final ObjectLongHashMap<ShardRouting> undesiredAllocations = new ObjectLongHashMap<>();
+    private final Map<String, UndesiredAllocation> undesiredAllocations = new HashMap<>();
     private final FrequencyCappedAction undesiredAllocationDurationLogInterval;
     private volatile TimeValue undesiredAllocationDurationLoggingThreshold;
     private volatile int maxUndesiredAllocationsToTrack;
@@ -94,30 +92,36 @@ public class UndesiredAllocationsTracker {
      */
     public void trackUndesiredAllocation(ShardRouting shardRouting) {
         assert shardRouting.unassigned() == false : "Shouldn't record unassigned shards as undesired allocations";
-        if (undesiredAllocations.size() < maxUndesiredAllocationsToTrack && undesiredAllocations.containsKey(shardRouting) == false) {
-            undesiredAllocations.put(shardRouting, timeProvider.relativeTimeInMillis());
+        if (undesiredAllocations.size() < maxUndesiredAllocationsToTrack) {
+            final var allocationId = shardRouting.allocationId().getId();
+            if (undesiredAllocations.containsKey(allocationId) == false) {
+                undesiredAllocations.put(
+                    allocationId,
+                    new UndesiredAllocation(shardRouting.shardId(), timeProvider.relativeTimeInMillis())
+                );
+            }
         }
     }
 
     /**
      * Remove any tracking of the specified allocation (a no-op if the allocation isn't being tracked)
      */
     public void removeTracking(ShardRouting shardRouting) {
-        undesiredAllocations.remove(shardRouting);
+        if (shardRouting.unassigned() == false) {
+            undesiredAllocations.remove(shardRouting.allocationId().getId());
+        } else {
+            assert false : "Shouldn't remove tracking of unassigned shards";
+        }
     }
 
     /**
      * Clear any {@link ShardRouting} that are no longer present in the routing nodes
      */
     public void cleanup(RoutingNodes routingNodes) {
-        undesiredAllocations.removeAll(shardRouting -> {
-            final var allocationId = shardRouting.allocationId();
-            if (allocationId != null) {
-                return routingNodes.getByAllocationId(shardRouting.shardId(), allocationId.getId()) == null;
-            } else {
-                assert false : "Unassigned shards shouldn't be marked as undesired";
-                return true;
-            }
+        undesiredAllocations.entrySet().removeIf(e -> {
+            final var undesiredAllocation = e.getValue();
+            final var allocationId = e.getKey();
+            return routingNodes.getByAllocationId(undesiredAllocation.shardId(), allocationId) == null;
         });
         shrinkIfOversized();
     }
@@ -140,9 +144,9 @@ public void maybeLogUndesiredShardsWarning(
     ) {
         final long currentTimeMillis = timeProvider.relativeTimeInMillis();
         long earliestUndesiredTimestamp = Long.MAX_VALUE;
-        for (var allocation : undesiredAllocations) {
-            if (allocation.value < earliestUndesiredTimestamp) {
-                earliestUndesiredTimestamp = allocation.value;
+        for (var undesiredAllocation : undesiredAllocations.values()) {
+            if (undesiredAllocation.undesiredSince() < earliestUndesiredTimestamp) {
+                earliestUndesiredTimestamp = undesiredAllocation.undesiredSince();
             }
         }
         if (earliestUndesiredTimestamp < currentTimeMillis
@@ -160,15 +164,22 @@ private void logDecisionsForUndesiredShardsOverThreshold(
     ) {
         final long currentTimeMillis = timeProvider.relativeTimeInMillis();
         final long loggingThresholdTimestamp = currentTimeMillis - undesiredAllocationDurationLoggingThreshold.millis();
-        for (var allocation : undesiredAllocations) {
-            if (allocation.value < loggingThresholdTimestamp) {
-                logUndesiredShardDetails(
-                    allocation.key,
-                    TimeValue.timeValueMillis(currentTimeMillis - allocation.value),
-                    routingNodes,
-                    routingAllocation,
-                    desiredBalance
-                );
+        for (var allocation : undesiredAllocations.entrySet()) {
+            final var undesiredAllocation = allocation.getValue();
+            final var allocationId = allocation.getKey();
+            if (undesiredAllocation.undesiredSince() < loggingThresholdTimestamp) {
+                final var shardRouting = routingNodes.getByAllocationId(undesiredAllocation.shardId(), allocationId);
+                if (shardRouting != null) {
+                    logUndesiredShardDetails(
+                        shardRouting,
+                        TimeValue.timeValueMillis(currentTimeMillis - undesiredAllocation.undesiredSince()),
+                        routingNodes,
+                        routingAllocation,
+                        desiredBalance
+                    );
+                } else {
+                    assert false : undesiredAllocation + " for allocationID " + allocationId + " was not cleaned up";
+                }
             }
         }
     }
@@ -200,19 +211,28 @@ private void logUndesiredShardDetails(
      */
     private void shrinkIfOversized() {
         if (undesiredAllocations.size() > maxUndesiredAllocationsToTrack) {
-            final var newestExcessValues = StreamSupport.stream(undesiredAllocations.spliterator(), false)
-                // we need to take a copy from the cursor because the cursors are re-used, so don't work with #sorted
-                .map(cursor -> new Tuple<>(cursor.key, cursor.value))
-                .sorted((a, b) -> Long.compare(b.v2(), a.v2()))
+            final var newestExcessAllocationIds = undesiredAllocations.entrySet()
+                .stream()
+                .sorted((a, b) -> Long.compare(b.getValue().undesiredSince(), a.getValue().undesiredSince()))
                 .limit(undesiredAllocations.size() - maxUndesiredAllocationsToTrack)
-                .map(Tuple::v1)
+                .map(Map.Entry::getKey)
                 .collect(Collectors.toSet());
-            undesiredAllocations.removeAll(newestExcessValues::contains);
+            undesiredAllocations.keySet().removeAll(newestExcessAllocationIds);
         }
     }
 
     // visible for testing
-    ObjectLongMap<ShardRouting> getUndesiredAllocations() {
-        return undesiredAllocations.clone();
+    Map<String, UndesiredAllocation> getUndesiredAllocations() {
+        return Map.copyOf(undesiredAllocations);
     }
+
+    /**
+     * Rather than storing the {@link ShardRouting}, we store a map of allocationId -> {@link UndesiredAllocation}
+     * this is because the allocation ID will persist as long as a shard stays on the same node, but the
+     * {@link ShardRouting} changes for a variety of reasons even when the shard doesn't move.
+     *
+     * @param shardId The shard ID
+     * @param undesiredSince The timestamp when the shard was first observed in an undesired allocation
+     */
+    record UndesiredAllocation(ShardId shardId, long undesiredSince) {}
 }
diff --git a/server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/UndesiredAllocationsTrackerTests.java b/server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/UndesiredAllocationsTrackerTests.java
@@ -31,7 +31,6 @@
 
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
-import java.util.stream.StreamSupport;
 
 public class UndesiredAllocationsTrackerTests extends ESTestCase {
 
@@ -70,9 +69,14 @@ public void testShardsArePrunedWhenRemovedFromRoutingTable() {
             RoutingNodes.immutable(stateWithIndexRemoved.globalRoutingTable(), stateWithIndexRemoved.nodes())
         );
         assertEquals((numberOfIndices - 1) * primaryShards, undesiredAllocationsTracker.getUndesiredAllocations().size());
-        undesiredAllocationsTracker.getUndesiredAllocations()
-            .iterator()
-            .forEachRemaining(olc -> assertNotNull(stateWithIndexRemoved.routingTable(ProjectId.DEFAULT).index(olc.key.index())));
+        assertTrue(
+            undesiredAllocationsTracker.getUndesiredAllocations()
+                .values()
+                .stream()
+                .allMatch(
+                    allocation -> stateWithIndexRemoved.routingTable(ProjectId.DEFAULT).index(allocation.shardId().getIndex()) != null
+                )
+        );
     }
 
     public void testNewestRecordsArePurgedWhenLimitIsDecreased() {
@@ -110,8 +114,10 @@ public void testNewestRecordsArePurgedWhenLimitIsDecreased() {
         // We should purge the most recent entries in #cleanup
         undesiredAllocationsTracker.cleanup(routingNodes);
         assertEquals(reducedMaximum, undesiredAllocationsTracker.getUndesiredAllocations().size());
-        final var remainingShardIds = StreamSupport.stream(undesiredAllocationsTracker.getUndesiredAllocations().spliterator(), false)
-            .map(olc -> olc.key.shardId().id())
+        final var remainingShardIds = undesiredAllocationsTracker.getUndesiredAllocations()
+            .values()
+            .stream()
+            .map(allocation -> allocation.shardId().id())
             .collect(Collectors.toSet());
         assertEquals(IntStream.range(0, reducedMaximum).boxed().collect(Collectors.toSet()), remainingShardIds);
     }
@@ -133,12 +139,43 @@ public void testCannotTrackMoreShardsThanTheLimit() {
 
         // Only the first {maxToTrack} shards should be tracked
         assertEquals(maxToTrack, undesiredAllocationsTracker.getUndesiredAllocations().size());
-        final var trackedShardIds = StreamSupport.stream(undesiredAllocationsTracker.getUndesiredAllocations().spliterator(), false)
-            .map(olc -> olc.key.shardId().id())
+        final var trackedShardIds = undesiredAllocationsTracker.getUndesiredAllocations()
+            .values()
+            .stream()
+            .map(allocation -> allocation.shardId().id())
             .collect(Collectors.toSet());
         assertEquals(IntStream.range(0, maxToTrack).boxed().collect(Collectors.toSet()), trackedShardIds);
     }
 
+    public void testUndesiredAllocationsAreIdentifiableDespiteMetadataChanges() {
+        final var clusterSettings = ClusterSettings.createBuiltInClusterSettings(
+            Settings.builder().put(UndesiredAllocationsTracker.MAX_UNDESIRED_ALLOCATIONS_TO_TRACK.getKey(), randomIntBetween(1, 10)).build()
+        );
+        final var advancingTimeProvider = new AdvancingTimeProvider();
+        final var undesiredAllocationsTracker = new UndesiredAllocationsTracker(clusterSettings, advancingTimeProvider);
+        final var index = new Index(randomIdentifier(), randomIdentifier());
+
+        ShardRouting shardRouting = createAssignedRouting(index, 0);
+
+        undesiredAllocationsTracker.trackUndesiredAllocation(shardRouting);
+        assertEquals(1, undesiredAllocationsTracker.getUndesiredAllocations().size());
+
+        // move to started
+        shardRouting = shardRouting.moveToStarted(randomNonNegativeLong());
+        undesiredAllocationsTracker.trackUndesiredAllocation(shardRouting);
+        assertEquals(1, undesiredAllocationsTracker.getUndesiredAllocations().size());
+
+        // start a relocation
+        shardRouting = shardRouting.relocate(randomIdentifier(), randomNonNegativeLong());
+        undesiredAllocationsTracker.trackUndesiredAllocation(shardRouting);
+        assertEquals(1, undesiredAllocationsTracker.getUndesiredAllocations().size());
+
+        // cancel that relocation
+        shardRouting = shardRouting.cancelRelocation();
+        undesiredAllocationsTracker.removeTracking(shardRouting);
+        assertEquals(0, undesiredAllocationsTracker.getUndesiredAllocations().size());
+    }
+
     private ClusterState removeRandomIndex(ClusterState state) {
         RoutingTable originalRoutingTable = state.routingTable(ProjectId.DEFAULT);
         RoutingTable updatedRoutingTable = RoutingTable.builder(originalRoutingTable)
