check that new shard assignment won't exceed threshold

DiannaHohensee · DiannaHohensee · commit 68be3dfac624 · 2025-08-11T15:31:37.000-07:00
improvements and *Tests are complete
diff --git a/server/src/main/java/org/elasticsearch/cluster/routing/ShardMovementWriteLoadSimulator.java b/server/src/main/java/org/elasticsearch/cluster/routing/ShardMovementWriteLoadSimulator.java
@@ -96,16 +96,34 @@ private static NodeUsageStatsForThreadPools.ThreadPoolUsageStats replaceWritePoo
     ) {
         final NodeUsageStatsForThreadPools.ThreadPoolUsageStats writeThreadPoolStats = value.threadPoolUsageStatsMap()
             .get(ThreadPool.Names.WRITE);
-        var newAverageThreadPoolUtilization = (writeThreadPoolStats.averageThreadPoolUtilization() + (writeLoadDelta / writeThreadPoolStats
-            .totalThreadPoolThreads()));
-        assert newAverageThreadPoolUtilization < 1.00;
         return new NodeUsageStatsForThreadPools.ThreadPoolUsageStats(
             writeThreadPoolStats.totalThreadPoolThreads(),
-            (float) Math.max(
-                (writeThreadPoolStats.averageThreadPoolUtilization() + (writeLoadDelta / writeThreadPoolStats.totalThreadPoolThreads())),
-                0.0
+            updateNodeUtilizationWithShardMovements(
+                writeThreadPoolStats.averageThreadPoolUtilization(),
+                (float) writeLoadDelta,
+                writeThreadPoolStats.totalThreadPoolThreads()
             ),
             writeThreadPoolStats.maxThreadPoolQueueLatencyMillis()
         );
     }
+
+    /**
+     * The {@code nodeUtilization} is the average utilization per thread for some duration of time. The {@code shardWriteLoadDelta} is the
+     * sum of shards' total execution time. Dividing the shards total execution time by the number of threads provides the average
+     * utilization of each write thread for those shards. The change in shard load can then be added to the node utilization.
+     *
+     * @param nodeUtilization The current node-level write load percent utilization.
+     * @param shardWriteLoadDelta The change in shard(s) execution time across all threads. This can be positive or negative depending on
+     *                            whether shards were moved onto the node or off of the node.
+     * @param numberOfWriteThreads The number of threads available in the node's write thread pool.
+     * @return The new node-level write load percent utilization after adding the shard write load delta.
+     */
+    public static float updateNodeUtilizationWithShardMovements(
+        float nodeUtilization,
+        float shardWriteLoadDelta,
+        int numberOfWriteThreads
+    ) {
+        float newNodeUtilization = nodeUtilization + (shardWriteLoadDelta / numberOfWriteThreads);
+        return (float) Math.max(newNodeUtilization, 0.0);
+    }
 }
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
@@ -14,10 +14,12 @@
 import org.elasticsearch.cluster.NodeUsageStatsForThreadPools.ThreadPoolUsageStats;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.routing.RoutingNode;
+import org.elasticsearch.cluster.routing.ShardMovementWriteLoadSimulator;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.cluster.routing.allocation.WriteLoadConstraintSettings;
 import org.elasticsearch.common.settings.ClusterSettings;
+import org.elasticsearch.core.Strings;
 import org.elasticsearch.threadpool.ThreadPool;
 
 /**
@@ -38,57 +40,59 @@ public WriteLoadConstraintDecider(ClusterSettings clusterSettings) {
     @Override
     public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, RoutingAllocation allocation) {
         if (writeLoadConstraintSettings.getWriteLoadConstraintEnabled() != WriteLoadConstraintSettings.WriteLoadDeciderStatus.ENABLED) {
-            return Decision.YES;
+            return Decision.single(Decision.Type.YES, NAME, "Decider is disabled");
         }
 
         // 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 Decision.YES;
+            return Decision.single(Decision.Type.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) {
             // No node-level thread pool usage stats were reported for this node. Let's assume this is OK and that the simulator will handle
             // setting a node-level write load for this node after this shard is assigned.
-            return Decision.YES;
+            return Decision.single(Decision.Type.YES, NAME, "The node has no write load estimate. Decider takes no action.");
         }
 
-        // NOMERGE: should create a utility class (eventually, maybe duplicate code for this class, for now, if simulator does have bug)
-        // to calculate the change node's usage change with assignment of the new shard
         assert nodeUsageStatsForThreadPools.threadPoolUsageStatsMap().isEmpty() == false;
         assert nodeUsageStatsForThreadPools.threadPoolUsageStatsMap().get(ThreadPool.Names.WRITE) != null;
         var nodeWriteThreadPoolStats = nodeUsageStatsForThreadPools.threadPoolUsageStatsMap().get(ThreadPool.Names.WRITE);
         var nodeWriteThreadPoolLoadThreshold = writeLoadConstraintSettings.getWriteThreadPoolHighUtilizationThresholdSetting();
         if (nodeWriteThreadPoolStats.averageThreadPoolUtilization() >= nodeWriteThreadPoolLoadThreshold) {
             // The node's write thread pool usage stats already show high utilization above the threshold for accepting new shards.
-            logger.debug(
-                "The high utilization threshold of {} has already been reached on node {}. Cannot allocate shard {} to node {} "
-                    + "without risking increased write latencies.",
-                nodeWriteThreadPoolLoadThreshold,
+            String debugMsg = Strings.format(
+                "Node [%s] with write thread pool utilization [%f] already exceeds the high utilization threshold of [%f]. Cannot allocate "
+                    + "shard [%s] to node without risking increased write latencies.",
                 node.nodeId(),
-                shardRouting.shardId(),
-                node.nodeId()
+                nodeWriteThreadPoolStats.averageThreadPoolUtilization(),
+                nodeWriteThreadPoolLoadThreshold,
+                shardRouting.shardId()
             );
-            return Decision.NO;
+            logger.debug(debugMsg);
+            return Decision.single(Decision.Type.NO, NAME, debugMsg);
         }
 
-        if (nodeWriteThreadPoolStats.averageThreadPoolUtilization() + calculateShardMovementChange(nodeWriteThreadPoolStats, shardWriteLoad) >= nodeWriteThreadPoolLoadThreshold) {
-            // The node's write thread pool usage would be raised above the high utilization threshold. This could lead to a hot spot on
-            // this node and is undesirable.
-            logger.debug(
-                "The high utilization threshold of {} would be exceeded on node {} if shard {} with estimated write load {} were "
-                    + "assigned to it. Cannot allocate shard {} to node {} without risking increased write latencies.",
+        if (calculateShardMovementChange(nodeWriteThreadPoolStats, shardWriteLoad) >= nodeWriteThreadPoolLoadThreshold) {
+            // The node's write thread pool usage would be raised above the high utilization threshold with assignment of the new shard.
+            // This could lead to a hot spot on this node and is undesirable.
+            String debugMsg = Strings.format(
+                "The high utilization threshold of [%f] would be exceeded on node [%s] with utilization [%f] if shard [%s] with estimated "
+                    + "write load [%f] (execution time [%f] / threads [%d]) were assigned to it. Cannot allocate shard to node without "
+                    + "risking increased write latencies.",
                 nodeWriteThreadPoolLoadThreshold,
                 node.nodeId(),
+                nodeWriteThreadPoolStats.averageThreadPoolUtilization(),
                 shardRouting.shardId(),
+                shardWriteLoad / nodeWriteThreadPoolStats.totalThreadPoolThreads(),
                 shardWriteLoad,
-                shardRouting.shardId(),
-                node.nodeId()
+                nodeWriteThreadPoolStats.totalThreadPoolThreads()
             );
-            return Decision.NO;
+            logger.debug(debugMsg);
+            return Decision.single(Decision.Type.NO, NAME, debugMsg);
         }
 
         return Decision.YES;
@@ -109,8 +113,10 @@ public Decision canRemain(IndexMetadata indexMetadata, ShardRouting shardRouting
      */
     private float calculateShardMovementChange(ThreadPoolUsageStats nodeWriteThreadPoolStats, double shardWriteLoad) {
         assert shardWriteLoad > 0;
-        // NOMERGE: move this into an utility class, should be commonly accessible with the simulator.
-        // TODO: implement..
-        return 0;
+        return ShardMovementWriteLoadSimulator.updateNodeUtilizationWithShardMovements(
+            nodeWriteThreadPoolStats.averageThreadPoolUtilization(),
+            (float) shardWriteLoad,
+            nodeWriteThreadPoolStats.totalThreadPoolThreads()
+        );
     }
 }
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
@@ -36,34 +36,29 @@
 
 public class WriteLoadConstraintDeciderTests extends ESAllocationTestCase {
 
-    // NOMERGE: I wonder if we don't want a node utilization percent, but rather a total node execution time used and a total node execution
-    // time possible. The math would be a whole lot easier on us...
-    public void testWriteLoadDeciderIsDisabled() {
+    public void testWriteLoadDecider() {
         String indexName = "test-index";
 
-        // Set up multiple nodes and an index with multiple shards.
-        // DiscoveryNode discoveryNode1 = newNode("node1");
-        // DiscoveryNode discoveryNode2 = newNode("node2");
-        // ShardId shardId1 = new ShardId(indexName, IndexMetadata.INDEX_UUID_NA_VALUE, 0);
-        // ShardId shardId2 = new ShardId(indexName, IndexMetadata.INDEX_UUID_NA_VALUE, 1);
-
         /**
-         * Create the ClusterState
+         * Create the ClusterState for multiple nodes and multiple index shards.
          */
 
-        ClusterState clusterState = ClusterStateCreationUtils.stateWithAssignedPrimariesAndReplicas(new String[] { indexName }, 3, 0);
-        assertEquals(2, clusterState.nodes().size());
+        ClusterState clusterState = ClusterStateCreationUtils.stateWithAssignedPrimariesAndReplicas(new String[] { indexName }, 3, 1);
+        // The number of data nodes the util method above creates is numberOfReplicas+1.
+        assertEquals(3, clusterState.nodes().size());
         assertEquals(1, clusterState.metadata().getTotalNumberOfIndices());
 
         /**
-         * Fetch the nodes and index shards from the generated ClusterState.
+         * Fetch references to the nodes and index shards from the generated ClusterState, so the ClusterInfo can be created from them.
          */
 
         var discoveryNodeIterator = clusterState.nodes().iterator();
         assertTrue(discoveryNodeIterator.hasNext());
         var exceedingThresholdDiscoveryNode = discoveryNodeIterator.next();
         assertTrue(discoveryNodeIterator.hasNext());
         var belowThresholdDiscoveryNode2 = discoveryNodeIterator.next();
+        assertTrue(discoveryNodeIterator.hasNext());
+        var nearThresholdDiscoveryNode3 = discoveryNodeIterator.next();
         assertFalse(discoveryNodeIterator.hasNext());
 
         var indexIterator = clusterState.metadata().indicesAllProjects().iterator();
@@ -74,10 +69,10 @@ public void testWriteLoadDeciderIsDisabled() {
         assertEquals(3, testIndexMetadata.getNumberOfShards());
         ShardId testShardId1 = new ShardId(testIndex, 0);
         ShardId testShardId2 = new ShardId(testIndex, 1);
-        ShardId testShardId3NoWriteLoad = new ShardId(testIndex, 1);
+        ShardId testShardId3NoWriteLoad = new ShardId(testIndex, 2);
 
         /**
-         * Create the ClusterInfo that includes the node and shard level write load estimates.
+         * Create a ClusterInfo that includes the node and shard level write load estimates for a variety of node capacity situations.
          */
 
         var nodeThreadPoolStatsWithWriteExceedingThreshold = createNodeUsageStatsForThreadPools(
@@ -87,15 +82,18 @@ public void testWriteLoadDeciderIsDisabled() {
             0
         );
         var nodeThreadPoolStatsWithWriteBelowThreshold = createNodeUsageStatsForThreadPools(belowThresholdDiscoveryNode2, 8, 0.50f, 0);
+        var nodeThreadPoolStatsWithWriteNearThreshold = createNodeUsageStatsForThreadPools(nearThresholdDiscoveryNode3, 8, 0.89f, 0);
 
         // Create a map of usage per node.
         var nodeIdToNodeUsageStatsForThreadPools = new HashMap<String, NodeUsageStatsForThreadPools>();
         nodeIdToNodeUsageStatsForThreadPools.put(exceedingThresholdDiscoveryNode.getId(), nodeThreadPoolStatsWithWriteExceedingThreshold);
         nodeIdToNodeUsageStatsForThreadPools.put(belowThresholdDiscoveryNode2.getId(), nodeThreadPoolStatsWithWriteBelowThreshold);
+        nodeIdToNodeUsageStatsForThreadPools.put(nearThresholdDiscoveryNode3.getId(), nodeThreadPoolStatsWithWriteNearThreshold);
 
+        // Create a map of usage per shard.
         var shardIdToWriteLoadEstimate = new HashMap<ShardId, Double>();
-        shardIdToWriteLoadEstimate.put(testShardId1, 1.5);
-        shardIdToWriteLoadEstimate.put(testShardId2, 1.5);
+        shardIdToWriteLoadEstimate.put(testShardId1, 0.5);
+        shardIdToWriteLoadEstimate.put(testShardId2, 0.5);
         shardIdToWriteLoadEstimate.put(testShardId3NoWriteLoad, 0d);
 
         ClusterInfo clusterInfo = ClusterInfo.builder()
@@ -104,7 +102,7 @@ public void testWriteLoadDeciderIsDisabled() {
             .build();
 
         /**
-         * Create the RoutingAllocation
+         * Create the RoutingAllocation from the ClusterState and ClusterInfo above, and set up the other input for the WriteLoadDecider.
          */
 
         var routingAllocation = new RoutingAllocation(
@@ -138,32 +136,36 @@ public void testWriteLoadDeciderIsDisabled() {
             ShardRoutingState.STARTED
         );
 
-        assertTrue(discoveryNodeIterator.hasNext());
         RoutingNode exceedingThresholdRoutingNode = RoutingNodesHelper.routingNode(
             exceedingThresholdDiscoveryNode.getId(),
-            discoveryNodeIterator.next(),
+            exceedingThresholdDiscoveryNode,
             shardRouting1
         );
-        assertTrue(discoveryNodeIterator.hasNext());
         RoutingNode belowThresholdRoutingNode = RoutingNodesHelper.routingNode(
             belowThresholdDiscoveryNode2.getId(),
-            discoveryNodeIterator.next(),
+            belowThresholdDiscoveryNode2,
             shardRouting2
         );
-        assertFalse(discoveryNodeIterator.hasNext());
+        RoutingNode nearThresholdRoutingNode = RoutingNodesHelper.routingNode(
+            nearThresholdDiscoveryNode3.getId(),
+            nearThresholdDiscoveryNode3,
+            new ShardRouting[] {}
+        );
 
         /**
          * Test the write load decider
          */
 
         // The write load decider is disabled by default.
-        var writeLoadDecider = createWriteLoadConstraintDecider(Settings.builder().build());
 
-        assertEquals(Decision.YES, writeLoadDecider.canAllocate(shardRouting2, exceedingThresholdRoutingNode, routingAllocation));
-        assertEquals(Decision.YES, writeLoadDecider.canAllocate(shardRouting1, belowThresholdRoutingNode, routingAllocation));
-        assertEquals(Decision.YES, writeLoadDecider.canAllocate(thirdRoutingNoWriteLoad, exceedingThresholdRoutingNode, routingAllocation));
+        var writeLoadDecider = createWriteLoadConstraintDecider(Settings.builder().build());
+        assertEquals(Decision.Type.YES, writeLoadDecider.canAllocate(shardRouting2, exceedingThresholdRoutingNode, routingAllocation).type());
+        assertEquals(Decision.Type.YES, writeLoadDecider.canAllocate(shardRouting1, belowThresholdRoutingNode, routingAllocation).type());
+        assertEquals(Decision.Type.YES, writeLoadDecider.canAllocate(shardRouting1, nearThresholdRoutingNode, routingAllocation).type());
+        assertEquals(Decision.Type.YES, writeLoadDecider.canAllocate(thirdRoutingNoWriteLoad, exceedingThresholdRoutingNode, routingAllocation).type());
 
         // Check that the answers change when enabled.
+
         writeLoadDecider = createWriteLoadConstraintDecider(
             Settings.builder()
                 .put(
@@ -172,27 +174,26 @@ public void testWriteLoadDeciderIsDisabled() {
                 )
                 .build()
         );
-
-        assertEquals(Decision.NO, writeLoadDecider.canAllocate(shardRouting2, exceedingThresholdRoutingNode, routingAllocation));
-        assertEquals(Decision.YES, writeLoadDecider.canAllocate(shardRouting1, belowThresholdRoutingNode, routingAllocation));
-        assertEquals(Decision.YES, writeLoadDecider.canAllocate(thirdRoutingNoWriteLoad, exceedingThresholdRoutingNode, routingAllocation));
-
-        // NOMERGE: test that adding a shard is rejected if it would overflow the utilization threshold?
-        // Need to implement the logic in the decider, I don't check right now.
-    }
-
-    public void testShardWithNoWriteLoadEstimateIsAlwaysYES() {
-        Settings writeLoadConstraintSettings = Settings.builder()
-            .put(
-                WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING.getKey(),
-                WriteLoadConstraintSettings.WriteLoadDeciderStatus.ENABLED
-            )
-            .build();
-        // TODO
-    }
-
-    public void testShardWithWriteLoadEstimate() {
-        // TODO: test successful re-assignment and rejected re-assignment due to threshold
+        assertEquals(
+            "Assigning a new shard to a node that is above the threshold should fail",
+            Decision.Type.NO,
+            writeLoadDecider.canAllocate(shardRouting2, exceedingThresholdRoutingNode, routingAllocation).type()
+        );
+        assertEquals(
+            "Assigning a new shard to a node that has capacity should succeed",
+            Decision.Type.YES,
+            writeLoadDecider.canAllocate(shardRouting1, belowThresholdRoutingNode, routingAllocation).type()
+        );
+        assertEquals(
+            "Assigning a new shard without a write load estimate should _not_ be blocked by lack of capacity",
+            Decision.Type.YES,
+            writeLoadDecider.canAllocate(thirdRoutingNoWriteLoad, exceedingThresholdRoutingNode, routingAllocation).type()
+        );
+        assertEquals(
+            "Assigning a new shard that would cause the node to exceed capacity should fail",
+            Decision.Type.NO,
+            writeLoadDecider.canAllocate(shardRouting1, nearThresholdRoutingNode, routingAllocation).type()
+        );
     }
 
     private WriteLoadConstraintDecider createWriteLoadConstraintDecider(Settings settings) {
