Throw CircuitBreakingException for too many shards in snashot status …

…API call (#15688) Signed-off-by: Lakshya Taragi <lakshya.taragi@gmail.com>
opensearch-project · Sep 5, 2024 · f33c786 · f33c786
1 parent 73e14ab
commit f33c786
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 16 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/snapshots/SnapshotStatusApisIT.java b/server/src/internalClusterTest/java/org/opensearch/snapshots/SnapshotStatusApisIT.java
@@ -51,6 +51,7 @@
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.common.util.io.IOUtils;
 import org.opensearch.core.common.Strings;
+import org.opensearch.core.common.breaker.CircuitBreakingException;
 import org.opensearch.core.common.unit.ByteSizeUnit;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.index.IndexNotFoundException;
@@ -602,28 +603,28 @@ public void testSnapshotStatusApiFailureForTooManyShardsAcrossSnapshots() throws
 
         // across a single snapshot
         assertBusy(() -> {
-            TooManyShardsInSnapshotsStatusException exception = expectThrows(
-                TooManyShardsInSnapshotsStatusException.class,
+            CircuitBreakingException exception = expectThrows(
+                CircuitBreakingException.class,
                 () -> client().admin().cluster().prepareSnapshotStatus(repositoryName).setSnapshots(snapshot1).execute().actionGet()
             );
-            assertEquals(exception.status(), RestStatus.REQUEST_ENTITY_TOO_LARGE);
+            assertEquals(exception.status(), RestStatus.TOO_MANY_REQUESTS);
             assertTrue(
                 exception.getMessage().endsWith(" is more than the maximum allowed value of shard count [2] for snapshot status request")
             );
         }, 1, TimeUnit.MINUTES);
 
         // across multiple snapshots
         assertBusy(() -> {
-            TooManyShardsInSnapshotsStatusException exception = expectThrows(
-                TooManyShardsInSnapshotsStatusException.class,
+            CircuitBreakingException exception = expectThrows(
+                CircuitBreakingException.class,
                 () -> client().admin()
                     .cluster()
                     .prepareSnapshotStatus(repositoryName)
                     .setSnapshots(snapshot1, snapshot2)
                     .execute()
                     .actionGet()
             );
-            assertEquals(exception.status(), RestStatus.REQUEST_ENTITY_TOO_LARGE);
+            assertEquals(exception.status(), RestStatus.TOO_MANY_REQUESTS);
             assertTrue(
                 exception.getMessage().endsWith(" is more than the maximum allowed value of shard count [2] for snapshot status request")
             );
@@ -741,8 +742,8 @@ public void testSnapshotStatusFailuresWithIndexFilter() throws Exception {
             updateSettingsRequest.persistentSettings(Settings.builder().put(MAX_SHARDS_ALLOWED_IN_STATUS_API.getKey(), 2));
             assertAcked(client().admin().cluster().updateSettings(updateSettingsRequest).actionGet());
 
-            TooManyShardsInSnapshotsStatusException ex = expectThrows(
-                TooManyShardsInSnapshotsStatusException.class,
+            CircuitBreakingException ex = expectThrows(
+                CircuitBreakingException.class,
                 () -> client().admin()
                     .cluster()
                     .prepareSnapshotStatus(repositoryName)
@@ -751,7 +752,7 @@ public void testSnapshotStatusFailuresWithIndexFilter() throws Exception {
                     .execute()
                     .actionGet()
             );
-            assertEquals(ex.status(), RestStatus.REQUEST_ENTITY_TOO_LARGE);
+            assertEquals(ex.status(), RestStatus.TOO_MANY_REQUESTS);
             assertTrue(ex.getMessage().endsWith(" is more than the maximum allowed value of shard count [2] for snapshot status request"));
 
             logger.info("Reset MAX_SHARDS_ALLOWED_IN_STATUS_API to default value");

diff --git a/.../org/opensearch/action/admin/cluster/snapshots/status/TransportSnapshotsStatusAction.java b/.../org/opensearch/action/admin/cluster/snapshots/status/TransportSnapshotsStatusAction.java
@@ -49,6 +49,8 @@
 import org.opensearch.common.util.set.Sets;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.Strings;
+import org.opensearch.core.common.breaker.CircuitBreaker;
+import org.opensearch.core.common.breaker.CircuitBreakingException;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.util.CollectionUtils;
 import org.opensearch.core.index.shard.ShardId;
@@ -66,7 +68,6 @@
 import org.opensearch.snapshots.SnapshotShardsService;
 import org.opensearch.snapshots.SnapshotState;
 import org.opensearch.snapshots.SnapshotsService;
-import org.opensearch.snapshots.TooManyShardsInSnapshotsStatusException;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportService;
 
@@ -458,13 +459,17 @@ private Map<SnapshotId, SnapshotInfo> snapshotsInfo(
             snapshotsInfoMap.put(snapshotId, snapshotInfo);
         }
         if (totalShardsAcrossSnapshots > maximumAllowedShardCount && request.indices().length == 0) {
-            String message = "Total shard count ["
+            String message = "["
+                + repositoryName
+                + ":"
+                + String.join(", ", request.snapshots())
+                + "]"
+                + " Total shard count ["
                 + totalShardsAcrossSnapshots
                 + "] is more than the maximum allowed value of shard count ["
                 + maximumAllowedShardCount
                 + "] for snapshot status request";
-
-            throw new TooManyShardsInSnapshotsStatusException(repositoryName, message, request.snapshots());
+            throw new CircuitBreakingException(message, CircuitBreaker.Durability.PERMANENT);
         }
         return unmodifiableMap(snapshotsInfoMap);
     }
@@ -520,15 +525,19 @@ private Map<ShardId, IndexShardSnapshotStatus> snapshotShards(
         }
 
         if (totalShardsAcrossIndices > maximumAllowedShardCount && requestedIndexNames.isEmpty() == false && isV2Snapshot == false) {
-            String message = "Total shard count ["
+            String message = "["
+                + repositoryName
+                + ":"
+                + String.join(", ", request.snapshots())
+                + "]"
+                + " Total shard count ["
                 + totalShardsAcrossIndices
                 + "] across the requested indices ["
                 + requestedIndexNames.stream().collect(Collectors.joining(", "))
                 + "] is more than the maximum allowed value of shard count ["
                 + maximumAllowedShardCount
                 + "] for snapshot status request";
-
-            throw new TooManyShardsInSnapshotsStatusException(repositoryName, message, snapshotName);
+            throw new CircuitBreakingException(message, CircuitBreaker.Durability.PERMANENT);
         }
 
         final Map<ShardId, IndexShardSnapshotStatus> shardStatus = new HashMap<>();