[BUG] Fix remote shards balancer when filtering throttled nodes

CHANGELOG.md

@@ -241,6 +241,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Fix tracing context propagation for local transport instrumentation ([#11490](https://github.com/opensearch-project/OpenSearch/pull/11490))
 - Fix parsing of single line comments in `lang-painless` ([#11815](https://github.com/opensearch-project/OpenSearch/issues/11815))
 - Fix memory leak issue in ReorganizingLongHash ([#11953](https://github.com/opensearch-project/OpenSearch/issues/11953))
+- [BUG] Fix remote shards balancer when filtering throttled nodes ([#11724](https://github.com/opensearch-project/OpenSearch/pull/11724))
 
 ### Security
 

server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/RemoteShardsBalancer.java

@@ -43,6 +43,8 @@
     private final Logger logger;
     private final RoutingAllocation allocation;
     private final RoutingNodes routingNodes;
+    // indicates if there are any nodes being throttled for allocating any unassigned shards
+    private boolean anyNodesThrottled = false;
 
     public RemoteShardsBalancer(Logger logger, RoutingAllocation allocation) {
         this.logger = logger;
@@ -358,12 +360,16 @@
     }
 
     private void ignoreRemainingShards(Map<String, UnassignedIndexShards> unassignedShardMap) {
+        // If any nodes are throttled during allocation, mark all remaining unassigned shards as THROTTLED
+        final UnassignedInfo.AllocationStatus status = anyNodesThrottled
+            ? UnassignedInfo.AllocationStatus.DECIDERS_THROTTLED
+            : UnassignedInfo.AllocationStatus.DECIDERS_NO;
         for (UnassignedIndexShards indexShards : unassignedShardMap.values()) {
             for (ShardRouting shard : indexShards.getPrimaries()) {
-                routingNodes.unassigned().ignoreShard(shard, UnassignedInfo.AllocationStatus.DECIDERS_NO, allocation.changes());
+                routingNodes.unassigned().ignoreShard(shard, status, allocation.changes());
             }
             for (ShardRouting shard : indexShards.getReplicas()) {
-                routingNodes.unassigned().ignoreShard(shard, UnassignedInfo.AllocationStatus.DECIDERS_NO, allocation.changes());
+                routingNodes.unassigned().ignoreShard(shard, status, allocation.changes());
             }
         }
     }
@@ -424,11 +430,11 @@
     private void tryAllocateUnassignedShard(Queue<RoutingNode> nodeQueue, ShardRouting shard) {
         boolean allocated = false;
         boolean throttled = false;
-        Set<String> nodesCheckedForShard = new HashSet<>();
+        int numNodesToCheck = nodeQueue.size();
         while (nodeQueue.isEmpty() == false) {
             RoutingNode node = nodeQueue.poll();
+            --numNodesToCheck;
             Decision allocateDecision = allocation.deciders().canAllocate(shard, node, allocation);
-            nodesCheckedForShard.add(node.nodeId());
             if (allocateDecision.type() == Decision.Type.YES) {
                 if (logger.isTraceEnabled()) {
                     logger.trace("Assigned shard [{}] to [{}]", shardShortSummary(shard), node.nodeId());
@@ -467,6 +473,10 @@
                     }
                     nodeQueue.offer(node);
                 } else {
+                    if (nodeLevelDecision.type() == Decision.Type.THROTTLE) {
+                        anyNodesThrottled = true;
+                    }
+
                     if (logger.isTraceEnabled()) {
                         logger.trace(
                             "Cannot allocate any shard to node: [{}]. Removing from queue. Node level decisions: [{}],[{}]",
@@ -478,14 +488,14 @@
                 }
 
                 // Break out if all nodes in the queue have been checked for this shard
-                if (nodeQueue.stream().allMatch(rn -> nodesCheckedForShard.contains(rn.nodeId()))) {
+                if (numNodesToCheck == 0) {
                     break;
                 }
             }
         }
 
         if (allocated == false) {
-            UnassignedInfo.AllocationStatus status = throttled
+            UnassignedInfo.AllocationStatus status = (throttled || anyNodesThrottled)
                 ? UnassignedInfo.AllocationStatus.DECIDERS_THROTTLED
                 : UnassignedInfo.AllocationStatus.DECIDERS_NO;
             routingNodes.unassigned().ignoreShard(shard, status, allocation.changes());

server/src/test/java/org/opensearch/cluster/routing/allocation/RemoteShardsBalancerBaseTestCase.java

@@ -229,6 +229,11 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, Routing
                     return Decision.ALWAYS;
                 }
             }
+
+            @Override
+            public Decision canAllocateAnyShardToNode(RoutingNode node, RoutingAllocation allocation) {
+                return throttle ? Decision.THROTTLE : Decision.YES;
+            }
         });
         Collections.shuffle(deciders, random());
         return new AllocationDeciders(deciders);