Ensure shard movement is always recorded for hotspot mitigation (elastic#138142)

We should record the movement even if the shard has no write load. This
ensures hot-spot migitation happens only once every ClusterInfo polling
cycle as we agreed.

Resolves: elastic#138137

Author: ywangd

muted-tests.yml

@@ -477,9 +477,6 @@ tests:
 - class: org.elasticsearch.xpack.downsample.DataStreamLifecycleDownsampleIT
   method: testUpdateDownsampleSamplingMode
   issue: https://github.com/elastic/elasticsearch/issues/138135
-- class: org.elasticsearch.indices.mapping.UpdateMappingIntegrationIT
-  method: testUpdateMappingWithConflicts
-  issue: https://github.com/elastic/elasticsearch/issues/138137
 
 # Examples:
 #

server/src/main/java/org/elasticsearch/cluster/routing/ShardMovementWriteLoadSimulator.java

@@ -51,14 +51,17 @@ public void simulateShardStarted(ShardRouting shardRouting) {
                 // This is a shard being relocated
                 simulatedNodeWriteLoadDeltas.addTo(shardRouting.relocatingNodeId(), -1 * writeLoadForShard);
                 simulatedNodeWriteLoadDeltas.addTo(shardRouting.currentNodeId(), writeLoadForShard);
-                nodesWithMovedAwayShard.add(shardRouting.relocatingNodeId());
             } else {
                 // This is a new shard starting, it's unlikely we'll have a write-load value for a new
                 // shard, but we may be able to estimate if the new shard is created as part of a datastream
                 // rollover. See https://elasticco.atlassian.net/browse/ES-12469
                 simulatedNodeWriteLoadDeltas.addTo(shardRouting.currentNodeId(), writeLoadForShard);
             }
         }
+        // Always record the moving shard regardless whether its write load is adjusted
+        if (shardRouting.relocatingNodeId() != null) {
+            nodesWithMovedAwayShard.add(shardRouting.relocatingNodeId());
+        }
     }
 
     /**
@@ -69,15 +72,15 @@ public Map<String, NodeUsageStatsForThreadPools> simulatedNodeUsageStatsForThrea
             originalNodeUsageStatsForThreadPools.size()
         );
         for (Map.Entry<String, NodeUsageStatsForThreadPools> entry : originalNodeUsageStatsForThreadPools.entrySet()) {
-            if (simulatedNodeWriteLoadDeltas.containsKey(entry.getKey())) {
+            if (simulatedNodeWriteLoadDeltas.containsKey(entry.getKey()) || nodesWithMovedAwayShard.contains(entry.getKey())) {
                 var adjustedValue = new NodeUsageStatsForThreadPools(
                     entry.getKey(),
                     Maps.copyMapWithAddedOrReplacedEntry(
                         entry.getValue().threadPoolUsageStatsMap(),
                         ThreadPool.Names.WRITE,
                         replaceWritePoolStats(
                             entry.getValue(),
-                            simulatedNodeWriteLoadDeltas.get(entry.getKey()),
+                            simulatedNodeWriteLoadDeltas.getOrDefault(entry.getKey(), 0.0),
                             nodesWithMovedAwayShard.contains(entry.getKey())
                         )
                     )

server/src/test/java/org/elasticsearch/cluster/routing/ShardMovementWriteLoadSimulatorTests.java

@@ -172,6 +172,39 @@ public void testMovementBetweenNodesWithNoThreadPoolAndWriteLoadStats() {
         assertThat(simulated.containsKey("node_1"), equalTo(originalNode1ThreadPoolStats != null));
     }
 
+    public void testShardWithNoWriteLoadStillResetsQueueLatency() {
+        final ClusterState clusterState = createClusterState();
+        final var allocation = createRoutingAllocationWithShardWriteLoads(
+            clusterState,
+            Set.of(INDICES),
+            Map.of(),
+            randomThreadPoolUsageStats(),
+            randomThreadPoolUsageStats()
+        );
+        final var shardMovementWriteLoadSimulator = new ShardMovementWriteLoadSimulator(allocation);
+
+        // Relocate a random shard from node_0 to node_1
+        final var randomShard = randomFrom(StreamSupport.stream(allocation.routingNodes().node("node_0").spliterator(), false).toList());
+        final var moveShardTuple = allocation.routingNodes().relocateShard(randomShard, "node_1", 0, "testing", NOOP);
+        shardMovementWriteLoadSimulator.simulateShardStarted(moveShardTuple.v2());
+
+        final var simulated = shardMovementWriteLoadSimulator.simulatedNodeUsageStatsForThreadPools();
+        final var threadPoolUsageStats = simulated.get("node_0").threadPoolUsageStatsMap().get("write");
+        assertThat(threadPoolUsageStats.maxThreadPoolQueueLatencyMillis(), equalTo(0L)); // queue latency is reset
+        // No change to write load since shard has no write load
+        assertThat(
+            threadPoolUsageStats.averageThreadPoolUtilization(),
+            equalTo(
+                allocation.clusterInfo()
+                    .getNodeUsageStatsForThreadPools()
+                    .get("node_0")
+                    .threadPoolUsageStatsMap()
+                    .get("write")
+                    .averageThreadPoolUtilization()
+            )
+        );
+    }
+
     public void testUpdateThreadPoolQueueLatencyWithShardMovements() {
         final long originalLatency = randomNonNegativeLong();
 
@@ -206,6 +239,30 @@ private NodeUsageStatsForThreadPools.ThreadPoolUsageStats randomThreadPoolUsageS
     private RoutingAllocation createRoutingAllocationWithRandomisedWriteLoads(
         Set<String> indicesWithNoWriteLoad,
         NodeUsageStatsForThreadPools.ThreadPoolUsageStats... arrayOfNodeThreadPoolStats
+    ) {
+        final ClusterState clusterState = createClusterState();
+        final Map<ShardId, Double> shardWriteLoads = clusterState.metadata()
+            .getProject(ProjectId.DEFAULT)
+            .stream()
+            .filter(index -> indicesWithNoWriteLoad.contains(index.getIndex().getName()) == false)
+            .flatMap(index -> IntStream.range(0, 3).mapToObj(shardNum -> new ShardId(index.getIndex(), shardNum)))
+            .collect(
+                Collectors.toUnmodifiableMap(shardId -> shardId, shardId -> randomBoolean() ? 0.0f : randomDoubleBetween(0.1, 5.0, true))
+            );
+
+        return createRoutingAllocationWithShardWriteLoads(
+            clusterState,
+            indicesWithNoWriteLoad,
+            shardWriteLoads,
+            arrayOfNodeThreadPoolStats
+        );
+    }
+
+    private RoutingAllocation createRoutingAllocationWithShardWriteLoads(
+        ClusterState clusterState,
+        Set<String> indicesWithNoWriteLoad,
+        Map<ShardId, Double> shardWriteLoads,
+        NodeUsageStatsForThreadPools.ThreadPoolUsageStats... arrayOfNodeThreadPoolStats
     ) {
         final Map<String, NodeUsageStatsForThreadPools> nodeUsageStats = new HashMap<>();
         for (int i = 0; i < arrayOfNodeThreadPoolStats.length; i++) {
@@ -216,22 +273,9 @@ private RoutingAllocation createRoutingAllocationWithRandomisedWriteLoads(
             }
         }
 
-        final ClusterState clusterState = createClusterState();
         final ClusterInfo clusterInfo = ClusterInfo.builder()
             .nodeUsageStatsForThreadPools(nodeUsageStats)
-            .shardWriteLoads(
-                clusterState.metadata()
-                    .getProject(ProjectId.DEFAULT)
-                    .stream()
-                    .filter(index -> indicesWithNoWriteLoad.contains(index.getIndex().getName()) == false)
-                    .flatMap(index -> IntStream.range(0, 3).mapToObj(shardNum -> new ShardId(index.getIndex(), shardNum)))
-                    .collect(
-                        Collectors.toUnmodifiableMap(
-                            shardId -> shardId,
-                            shardId -> randomBoolean() ? 0.0f : randomDoubleBetween(0.1, 5.0, true)
-                        )
-                    )
-            )
+            .shardWriteLoads(shardWriteLoads)
             .build();
 
         return new RoutingAllocation(

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceComputerTests.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.DiskUsage;
 import org.elasticsearch.cluster.ESAllocationTestCase;
+import org.elasticsearch.cluster.NodeUsageStatsForThreadPools;
 import org.elasticsearch.cluster.RestoreInProgress;
 import org.elasticsearch.cluster.TestShardRoutingRoleStrategies;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
@@ -99,6 +100,7 @@
 import static org.elasticsearch.cluster.routing.ShardRoutingState.UNASSIGNED;
 import static org.elasticsearch.cluster.routing.TestShardRouting.newShardRouting;
 import static org.elasticsearch.cluster.routing.TestShardRouting.shardRoutingBuilder;
+import static org.elasticsearch.cluster.routing.allocation.WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING;
 import static org.elasticsearch.common.settings.ClusterSettings.createBuiltInClusterSettings;
 import static org.elasticsearch.test.MockLog.assertThatLogger;
 import static org.hamcrest.Matchers.aMapWithSize;
@@ -173,6 +175,58 @@ public void testStopsComputingWhenStale() {
         );
     }
 
+    public void testNoInfiniteLoopBetweenHotspotMitigationAndBalancing() {
+        // This test demonstrates that the computer does not get stuck in an infinite loop when moveShards and balancer moving against
+        // each other. This is done by configuring two shards each on its own node.
+        // - Shard 0 with no write load on node-0 with node load 0.92 and queue latency 15s
+        // - Shard 1 with some write load on node-1 with no node load nor queue latency
+        // 1. MoveShard will want to move shard 0 off node-0 to node-1 for hot-spot mitigation.
+        // 2. Balance will want to move shard 0 back to node-0 to spread the index. Balancer always picks shard 0 because it has
+        // no write load thus write load decider says YES
+        // The computation should stop after one round of moveShards + balance because hot-spot is considered mitigated after
+        // a single shard movement and breaks the loop.
+        final var initialState = createInitialClusterState(2, 2, 0);
+        final var index = initialState.metadata().getProject(ProjectId.DEFAULT).index(TEST_INDEX).getIndex();
+        final RoutingNodes routingNodes = initialState.getRoutingNodes().mutableCopy();
+        final var changes = mock(RoutingChangesObserver.class);
+        for (var iterator = routingNodes.unassigned().iterator(); iterator.hasNext();) {
+            var shardRouting = iterator.next();
+            routingNodes.startShard(
+                iterator.initialize(shardRouting.shardId().id() == 0 ? "node-0" : "node-1", null, 0L, changes),
+                changes,
+                0L
+            );
+        }
+        final var clusterState = rebuildRoutingTable(initialState, routingNodes);
+
+        final var clusterInfo = ClusterInfo.builder()
+            .shardWriteLoads(Map.of(new ShardId(index, 1), 0.004379))
+            .nodeUsageStatsForThreadPools(
+                Map.of(
+                    "node-0",
+                    new NodeUsageStatsForThreadPools(
+                        "node-0",
+                        Map.of("write", new NodeUsageStatsForThreadPools.ThreadPoolUsageStats(4, 0.92f, 15732))
+
+                    ),
+                    "node-1",
+                    new NodeUsageStatsForThreadPools(
+                        "node-1",
+                        Map.of("write", new NodeUsageStatsForThreadPools.ThreadPoolUsageStats(4, 0.0f, 0))
+                    )
+                )
+            )
+            .build();
+
+        final var settings = Settings.builder().put(WRITE_LOAD_DECIDER_ENABLED_SETTING.getKey(), "enabled").build();
+        final var routingAllocation = routingAllocationWithDecidersOf(clusterState, clusterInfo, settings);
+        final var input = new DesiredBalanceInput(42, routingAllocation, List.of());
+        final var computer = createDesiredBalanceComputer(new BalancedShardsAllocator(settings));
+        var balance = computer.compute(DesiredBalance.BECOME_MASTER_INITIAL, input, queue(), ignored -> true);
+        assertThat(balance.getAssignment(new ShardId(index, 0)).nodeIds(), equalTo(Set.of("node-0")));
+        assertThat(balance.getAssignment(new ShardId(index, 1)).nodeIds(), equalTo(Set.of("node-1")));
+    }
+
     public void testIgnoresOutOfScopePrimaries() {
         var desiredBalanceComputer = createDesiredBalanceComputer();
         var clusterState = mutateAllocationStatuses(createInitialClusterState(3));