Fail fast during catShardsLimit check in cat shards action run

Signed-off-by: Sumit Bansal <[email protected]>
opensearch-project · Sep 19, 2024 · 8a2140f · 8a2140f
1 parent 454e2dd
commit 8a2140f
Showing 1 changed file with 12 additions and 14 deletions.
diff --git a/server/src/main/java/org/opensearch/rest/RequestLimitSettings.java b/server/src/main/java/org/opensearch/rest/RequestLimitSettings.java
@@ -99,8 +99,18 @@ public boolean isCircuitLimitBreached(final ClusterState clusterState, final Blo
                 break;
             case CAT_SHARDS:
                 if (catShardsLimit <= 0) return false;
-                int totalShards = getTotalShards(clusterState);
-                if (totalShards > catShardsLimit) return true;
+                final RoutingTable routingTable = clusterState.getRoutingTable();
+                final Map<String, IndexRoutingTable> indexRoutingTableMap = routingTable.getIndicesRouting();
+                int totalShards = 0;
+                for (final Map.Entry<String, IndexRoutingTable> entry : indexRoutingTableMap.entrySet()) {
+                    for (final Map.Entry<Integer, IndexShardRoutingTable> indexShardRoutingTableEntry : entry.getValue()
+                        .getShards()
+                        .entrySet()) {
+                        totalShards += indexShardRoutingTableEntry.getValue().getShards().size();
+                        // Fail fast if catShardsLimit value is breached and avoid unnecessary computation.
+                        if (totalShards > catShardsLimit) return true;
+                    }
+                }
                 break;
             case CAT_SEGMENTS:
                 if (catSegmentsLimit <= 0) return false;
@@ -126,18 +136,6 @@ private static int getTotalIndices(final ClusterState clusterState) {
         return chainWalk(() -> clusterState.getMetadata().getIndices().size(), 0);
     }
 
-    private static int getTotalShards(final ClusterState clusterState) {
-        final RoutingTable routingTable = clusterState.getRoutingTable();
-        final Map<String, IndexRoutingTable> indexRoutingTableMap = routingTable.getIndicesRouting();
-        int totalShards = 0;
-        for (final Map.Entry<String, IndexRoutingTable> entry : indexRoutingTableMap.entrySet()) {
-            for (final Map.Entry<Integer, IndexShardRoutingTable> indexShardRoutingTableEntry : entry.getValue().getShards().entrySet()) {
-                totalShards += indexShardRoutingTableEntry.getValue().getShards().size();
-            }
-        }
-        return totalShards;
-    }
-
     // TODO: Evaluate if we can move this to common util.
     private static <T> T chainWalk(Supplier<T> supplier, T defaultValue) {
         try {