From 8a2140fb60761d8e14dc8adc852a050f4c2ed436 Mon Sep 17 00:00:00 2001
From: Sumit Bansal <sumitsb@amazon.com>
Date: Thu, 19 Sep 2024 18:23:03 +0530
Subject: [PATCH] Fail fast during catShardsLimit check in cat shards action
 run

Signed-off-by: Sumit Bansal <sumitsb@amazon.com>
---
 .../opensearch/rest/RequestLimitSettings.java | 26 +++++++++----------
 1 file changed, 12 insertions(+), 14 deletions(-)

diff --git a/server/src/main/java/org/opensearch/rest/RequestLimitSettings.java b/server/src/main/java/org/opensearch/rest/RequestLimitSettings.java
index 727113def43a6..6522a640eff59 100644
--- 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 {