Tighten on when THROTTLE decision can be returned

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java

@@ -704,12 +704,18 @@ private boolean balanceByWeights(NodeSorter sorter) {
                             highIdx = relevantNodes - 1;
 
                             if (routingNodes.getRelocatingShardCount() > 0) {
-                                // ES-12955: Check routingNodes.getRelocatingShardCount() > 0 in case the first relocation is a THROTTLE.
-                                // This should rarely happen since in most cases, we don't throttle unless there is an existing relocation.
-                                // But it can happen in production for frozen indices when the cache is still being prepared. It can also
-                                // happen in tests because we have decider like RandomAllocationDecider that can randomly return THROTTLE
-                                // when there is no existing relocation.
+                                // Check routingNodes.getRelocatingShardCount() > 0 in case the first relocation is a THROTTLE.
                                 shardBalanced = true;
+                            } else {
+                                // A THROTTLE decision can happen when:
+                                // 1. Not simulating
+                                // 2. Or, a partial searchable snapshot index due to HasFrozenCacheAllocationDecider
+                                assert allocation.isSimulating() == false || indexMetadata.isPartialSearchableSnapshot()
+                                    : "unexpected THROTTLE decision (simulation="
+                                        + allocation.isSimulating()
+                                        + ") when balancing index ["
+                                        + index
+                                        + "]";
                             }
                             if (completeEarlyOnShardAssignmentChange && shardBalanced) {
                                 return true;
@@ -834,6 +840,20 @@ public boolean moveShards() {
                 } else if (moveDecision.isDecisionTaken() && moveDecision.canRemain() == false) {
                     logger.trace("[{}][{}] can't move", shardRouting.index(), shardRouting.id());
                 }
+
+                // A THROTTLE allocation decision can happen when:
+                // 1. Not simulating
+                // 2. Or, there are prior shard movements
+                // 3. Or, a partial searchable snapshot index due to HasFrozenCacheAllocationDecider
+                assert moveDecision.getAllocationDecision() != AllocationDecision.THROTTLED
+                    || (allocation.isSimulating() == false || shardMoved || indexMetadata(index).isPartialSearchableSnapshot())
+                    : "unexpected allocation decision ["
+                        + moveDecision.getAllocationDecision()
+                        + "] (simulation="
+                        + allocation.isSimulating()
+                        + ") with no prior shard movements when moving shard ["
+                        + shardRouting
+                        + "]";
             }
 
             // If we get here, attempt to move one of the best not-preferred shards that we identified earlier
@@ -1267,7 +1287,18 @@ private boolean allocateUnassigned() {
                             assert allocationDecision.getAllocationStatus() == AllocationStatus.DECIDERS_THROTTLED;
                             final long shardSize = getExpectedShardSize(shard, ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE, allocation);
                             minNode.addShard(projectIndex(shard), shard.initialize(minNode.getNodeId(), null, shardSize));
-                            // If we see a throttle decision in simulation, there must be other shards that got assigned before it.
+                            // If we see a THROTTLE decision, it's either:
+                            // 1. Not simulating
+                            // 2. Or, there is shard assigned before this one
+                            // 3. Or, a partial searchable snapshot index due to HasFrozenCacheAllocationDecider
+                            assert allocation.isSimulating() == false
+                                || shardAssignmentChanged
+                                || indexMetadata(index).isPartialSearchableSnapshot()
+                                : "unexpected THROTTLE decision (simulation="
+                                    + allocation.isSimulating()
+                                    + ") with no prior assignment when allocating unassigned shard ["
+                                    + shard
+                                    + "]";
                             assert allocation.isSimulating() == false || shardAssignmentChanged
                                 : "shard " + shard + " was throttled but no other shards were assigned";
                         } else {

server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ConcurrentRebalanceAllocationDecider.java

@@ -70,6 +70,8 @@ public Decision canRebalance(ShardRouting shardRouting, RoutingAllocation alloca
     public Decision canRebalance(RoutingAllocation allocation) {
         int relocatingShards = allocation.routingNodes().getRelocatingShardCount();
         if (allocation.isSimulating() && relocatingShards >= 2) {
+            // This branch should no longer run after https://github.com/elastic/elasticsearch/pull/134786
+            assert false : "allocation simulation should have returned earlier and not hit throttling";
             // BalancedShardAllocator is prone to perform unnecessary moves when cluster_concurrent_rebalance is set to high values (>2).
             // (See https://github.com/elastic/elasticsearch/issues/87279)
             // Above allocator is used in DesiredBalanceComputer. Since we do not move actual shard data during calculation

server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ThrottlingAllocationDecider.java

@@ -152,6 +152,9 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, Routing
             // Allocating a shard to this node will increase the incoming recoveries
             int currentInRecoveries = allocation.routingNodes().getIncomingRecoveries(node.nodeId());
             if (currentInRecoveries >= concurrentIncomingRecoveries) {
+                // See also https://github.com/elastic/elasticsearch/pull/134786
+                assert allocation.isSimulating() == false || concurrentIncomingRecoveries == 0
+                    : "allocation simulation should have returned earlier and not hitting throttling";
                 return allocation.decision(
                     THROTTLE,
                     NAME,
@@ -169,6 +172,9 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, Routing
                 }
                 int primaryNodeOutRecoveries = allocation.routingNodes().getOutgoingRecoveries(primaryShard.currentNodeId());
                 if (primaryNodeOutRecoveries >= concurrentOutgoingRecoveries) {
+                    // See also https://github.com/elastic/elasticsearch/pull/134786
+                    assert allocation.isSimulating() == false || concurrentOutgoingRecoveries == 0
+                        : "allocation simulation should have returned earlier and not hitting throttling";
                     return allocation.decision(
                         THROTTLE,
                         NAME,

server/src/test/java/org/elasticsearch/cluster/routing/allocation/RandomAllocationDeciderTests.java

@@ -206,29 +206,33 @@ public RandomAllocationDecider(Random random) {
 
         @Override
         public Decision canRebalance(ShardRouting shardRouting, RoutingAllocation allocation) {
-            return getRandomDecision();
+            return getRandomDecision(allocation.isSimulating() == false || allocation.routingNodes().getRelocatingShardCount() > 0);
         }
 
-        private Decision getRandomDecision() {
+        private Decision getRandomDecision(boolean canThrottle) {
             if (alwaysSayYes) {
                 return Decision.YES;
             }
             return switch (random.nextInt(10)) {
                 case 9, 8, 7, 6, 5 -> Decision.NO;
-                case 4 -> Decision.THROTTLE;
+                case 4 -> canThrottle ? Decision.THROTTLE : Decision.YES;
                 case 3, 2, 1 -> Decision.YES;
                 default -> Decision.ALWAYS;
             };
         }
 
         @Override
         public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, RoutingAllocation allocation) {
-            return getRandomDecision();
+            return getRandomDecision(
+                allocation.isSimulating() == false
+                    || allocation.routingNodes().getIncomingRecoveries(node.nodeId()) > 0
+                    || allocation.routingNodes().getOutgoingRecoveries(node.nodeId()) > 0
+            );
         }
 
         @Override
         public Decision canRemain(IndexMetadata indexMetadata, ShardRouting shardRouting, RoutingNode node, RoutingAllocation allocation) {
-            return getRandomDecision();
+            return getRandomDecision(false); // throttle does not make sense for canRemain
         }
 
     }