Don't allow allocation of shards with zero/unknown write load to overloaded nodes (#138148)

nicktindall · web-flow · commit 012df2423706 · 2025-11-18T14:17:13.000+11:00
diff --git a/server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/WriteLoadConstraintDecider.java b/server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/WriteLoadConstraintDecider.java
@@ -59,13 +59,6 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, Routing
             return allocation.decision(Decision.YES, NAME, "Shard is unassigned. Decider takes no action.");
         }
 
-        // Check whether the shard being relocated has any write load estimate. If it does not, then this decider has no opinion.
-        var allShardWriteLoads = allocation.clusterInfo().getShardWriteLoads();
-        var shardWriteLoad = allShardWriteLoads.get(shardRouting.shardId());
-        if (shardWriteLoad == null || shardWriteLoad == 0) {
-            return allocation.decision(Decision.YES, NAME, "Shard has no estimated write load. Decider takes no action.");
-        }
-
         var allNodeUsageStats = allocation.clusterInfo().getNodeUsageStatsForThreadPools();
         var nodeUsageStatsForThreadPools = allNodeUsageStats.get(node.nodeId());
         if (nodeUsageStatsForThreadPools == null) {
@@ -98,6 +91,8 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, Routing
             }
         }
 
+        var allShardWriteLoads = allocation.clusterInfo().getShardWriteLoads();
+        var shardWriteLoad = allShardWriteLoads.getOrDefault(shardRouting.shardId(), 0.0);
         var newWriteThreadPoolUtilization = calculateShardMovementChange(nodeWriteThreadPoolStats, shardWriteLoad);
         if (newWriteThreadPoolUtilization >= nodeWriteThreadPoolLoadThreshold) {
             // The node's write thread pool usage would be raised above the high utilization threshold with assignment of the new shard.
@@ -195,7 +190,6 @@ private Double getShardWriteLoad(RoutingAllocation allocation, ShardRouting shar
      * Returns the percent thread pool utilization change.
      */
     private float calculateShardMovementChange(ThreadPoolUsageStats nodeWriteThreadPoolStats, double shardWriteLoad) {
-        assert shardWriteLoad > 0;
         return ShardMovementWriteLoadSimulator.updateNodeUtilizationWithShardMovements(
             nodeWriteThreadPoolStats.averageThreadPoolUtilization(),
             (float) shardWriteLoad,
diff --git a/server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceComputerTests.java b/server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceComputerTests.java
@@ -223,7 +223,8 @@ public void testNoInfiniteLoopBetweenHotspotMitigationAndBalancing() {
         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")));
+        // This is an unusual edge case that will result in both shards being moved to the non-hot-spotting node
+        assertThat(balance.getAssignment(new ShardId(index, 0)).nodeIds(), equalTo(Set.of("node-1")));
         assertThat(balance.getAssignment(new ShardId(index, 1)).nodeIds(), equalTo(Set.of("node-1")));
     }
 
diff --git a/server/src/test/java/org/elasticsearch/cluster/routing/allocation/decider/WriteLoadConstraintDeciderTests.java b/server/src/test/java/org/elasticsearch/cluster/routing/allocation/decider/WriteLoadConstraintDeciderTests.java
@@ -11,12 +11,15 @@
 
 import org.elasticsearch.action.support.replication.ClusterStateCreationUtils;
 import org.elasticsearch.cluster.ClusterInfo;
+import org.elasticsearch.cluster.ClusterInfoSimulator;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.ESAllocationTestCase;
 import org.elasticsearch.cluster.NodeUsageStatsForThreadPools;
 import org.elasticsearch.cluster.NodeUsageStatsForThreadPools.ThreadPoolUsageStats;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.cluster.routing.RoutingChangesObserver;
 import org.elasticsearch.cluster.routing.RoutingNode;
 import org.elasticsearch.cluster.routing.RoutingNodes;
 import org.elasticsearch.cluster.routing.RoutingNodesHelper;
@@ -25,15 +28,22 @@
 import org.elasticsearch.cluster.routing.TestShardRouting;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.cluster.routing.allocation.WriteLoadConstraintSettings;
+import org.elasticsearch.cluster.routing.allocation.allocator.BalancedShardsAllocator;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.shard.ShardId;
+import org.elasticsearch.snapshots.SnapshotShardSizeInfo;
 import org.elasticsearch.threadpool.ThreadPool;
 
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
 
 import static org.elasticsearch.common.settings.ClusterSettings.createBuiltInClusterSettings;
 import static org.mockito.ArgumentMatchers.any;
@@ -155,14 +165,25 @@ public void testWriteLoadDeciderCanAllocate() {
             "Shard [[test-index][0]] in index [[test-index]] can be assigned to node [*]. The node's utilization would become [*]"
         );
         assertDecisionMatches(
-            "Assigning a new shard without a write load estimate should _not_ be blocked by lack of capacity",
+            "Assigning a new shard without a write load estimate to an over-threshold node should be blocked",
             writeLoadDecider.canAllocate(
                 testHarness.shardRoutingNoWriteLoad,
                 testHarness.exceedingThresholdRoutingNode,
                 testHarness.routingAllocation
             ),
+            Decision.Type.NOT_PREFERRED,
+            "Node [*] with write thread pool utilization [0.99] already exceeds the high utilization threshold of "
+                + "[0.900000]. Cannot allocate shard [[test-index][2]] to node without risking increased write latencies."
+        );
+        assertDecisionMatches(
+            "Assigning a new shard without a write load estimate to an under-threshold node should be allowed",
+            writeLoadDecider.canAllocate(
+                testHarness.shardRoutingNoWriteLoad,
+                testHarness.belowThresholdRoutingNode,
+                testHarness.routingAllocation
+            ),
             Decision.Type.YES,
-            "Shard has no estimated write load. Decider takes no action."
+            "Shard [[test-index][2]] in index [[test-index]] can be assigned to node [*]. The node's utilization would become [*]"
         );
         assertDecisionMatches(
             "Assigning a new shard that would cause the node to exceed capacity should fail",
@@ -247,6 +268,109 @@ public void testWriteLoadDeciderCanRemain() {
         );
     }
 
+    public void testWriteLoadDeciderShouldPreventBalancerMovingShardsBack() {
+        final var indexName = randomIdentifier();
+        final int numThreads = randomIntBetween(1, 10);
+        final float highUtilizationThreshold = randomFloatBetween(0.5f, 0.9f, true);
+        final long highLatencyThreshold = randomLongBetween(1000, 10000);
+        final var settings = Settings.builder()
+            .put(
+                WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING.getKey(),
+                WriteLoadConstraintSettings.WriteLoadDeciderStatus.ENABLED
+            )
+            .put(WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_HIGH_UTILIZATION_THRESHOLD_SETTING.getKey(), highUtilizationThreshold)
+            .put(
+                WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_QUEUE_LATENCY_THRESHOLD_SETTING.getKey(),
+                TimeValue.timeValueMillis(highLatencyThreshold)
+            )
+            .build();
+
+        final var state = ClusterStateCreationUtils.state(2, new String[] { indexName }, 4);
+        final var balancedShardsAllocator = new BalancedShardsAllocator(settings);
+        final var overloadedNode = randomFrom(state.nodes().getAllNodes());
+        final var otherNode = state.nodes().stream().filter(node -> node != overloadedNode).findFirst().orElseThrow();
+        final var clusterInfo = ClusterInfo.builder()
+            .nodeUsageStatsForThreadPools(
+                Map.of(
+                    overloadedNode.getId(),
+                    new NodeUsageStatsForThreadPools(
+                        overloadedNode.getId(),
+                        Map.of(
+                            ThreadPool.Names.WRITE,
+                            new ThreadPoolUsageStats(
+                                numThreads,
+                                randomFloatBetween(highUtilizationThreshold, 1.1f, false),
+                                randomLongBetween(highLatencyThreshold, highLatencyThreshold * 2)
+                            )
+                        )
+                    ),
+                    otherNode.getId(),
+                    new NodeUsageStatsForThreadPools(
+                        otherNode.getId(),
+                        Map.of(
+                            ThreadPool.Names.WRITE,
+                            new ThreadPoolUsageStats(
+                                numThreads,
+                                randomFloatBetween(0.0f, highUtilizationThreshold / 2, true),
+                                randomLongBetween(0, highLatencyThreshold / 2)
+                            )
+                        )
+                    )
+                )
+            )
+            // simulate all zero or missing shard write loads
+            .shardWriteLoads(
+                state.routingTable(ProjectId.DEFAULT)
+                    .allShards()
+                    .filter(ignored -> randomBoolean()) // some write-loads are missing altogether
+                    .collect(Collectors.toMap(ShardRouting::shardId, ignored -> 0.0d))  // the rest are zero
+            )
+            .build();
+
+        final var clusterSettings = createBuiltInClusterSettings(settings);
+        final var writeLoadConstraintDecider = new WriteLoadConstraintDecider(clusterSettings);
+        final var routingAllocation = new RoutingAllocation(
+            new AllocationDeciders(List.of(writeLoadConstraintDecider)),
+            state.getRoutingNodes().mutableCopy(),
+            state,
+            clusterInfo,
+            SnapshotShardSizeInfo.EMPTY,
+            randomLong()
+        );
+
+        // This should move a shard in an attempt to resolve the hot-spot
+        balancedShardsAllocator.allocate(routingAllocation);
+
+        assertEquals(1, routingAllocation.routingNodes().node(overloadedNode.getId()).numberOfOwningShards());
+        assertEquals(3, routingAllocation.routingNodes().node(otherNode.getId()).numberOfOwningShards());
+
+        final var clusterInfoSimulator = new ClusterInfoSimulator(routingAllocation);
+        final var movedShards = new HashSet<ShardRouting>();
+        for (RoutingNode routingNode : routingAllocation.routingNodes()) {
+            movedShards.addAll(routingNode.shardsWithState(ShardRoutingState.INITIALIZING).collect(Collectors.toSet()));
+        }
+        movedShards.forEach(shardRouting -> {
+            routingAllocation.routingNodes().startShard(shardRouting, new RoutingChangesObserver() {
+            }, randomNonNegativeLong());
+            clusterInfoSimulator.simulateShardStarted(shardRouting);
+        });
+
+        // This should run through the balancer without moving any shards back
+        ClusterInfo simulatedClusterInfo = clusterInfoSimulator.getClusterInfo();
+        balancedShardsAllocator.allocate(
+            new RoutingAllocation(
+                routingAllocation.deciders(),
+                routingAllocation.routingNodes(),
+                routingAllocation.getClusterState(),
+                simulatedClusterInfo,
+                routingAllocation.snapshotShardSizeInfo(),
+                randomLong()
+            )
+        );
+        assertEquals(1, routingAllocation.routingNodes().node(overloadedNode.getId()).numberOfOwningShards());
+        assertEquals(3, routingAllocation.routingNodes().node(otherNode.getId()).numberOfOwningShards());
+    }
+
     private void assertDecisionMatches(String description, Decision decision, Decision.Type type, String explanationPattern) {
         assertEquals(description, type, decision.type());
         if (explanationPattern == null) {

Original file line number	Diff line number	Diff line change
`@@ -223,7 +223,8 @@ public void testNoInfiniteLoopBetweenHotspotMitigationAndBalancing() {`
`223`	`223`	`final var input = new DesiredBalanceInput(42, routingAllocation, List.of());`
`224`	`224`	`final var computer = createDesiredBalanceComputer(new BalancedShardsAllocator(settings));`
`225`	`225`	`var balance = computer.compute(DesiredBalance.BECOME_MASTER_INITIAL, input, queue(), ignored -> true);`
`226`		`- assertThat(balance.getAssignment(new ShardId(index, 0)).nodeIds(), equalTo(Set.of("node-0")));`
	`226`	`+ // This is an unusual edge case that will result in both shards being moved to the non-hot-spotting node`
	`227`	`+ assertThat(balance.getAssignment(new ShardId(index, 0)).nodeIds(), equalTo(Set.of("node-1")));`
`227`	`228`	`assertThat(balance.getAssignment(new ShardId(index, 1)).nodeIds(), equalTo(Set.of("node-1")));`
`228`	`229`	`}`
`229`	`230`