Refactor code, rebase with pagination PR merge and add relevant changes

Signed-off-by: Sumit Bansal <sumitsb@amazon.com>
sumitasr · Sep 30, 2024 · 5d2e2f9 · 5d2e2f9
1 parent 23b9c64
commit 5d2e2f9
Show file tree

Hide file tree

Showing 15 changed files with 180 additions and 150 deletions.
diff --git a/server/src/main/java/org/opensearch/OpenSearchServerException.java b/server/src/main/java/org/opensearch/OpenSearchServerException.java
@@ -1213,8 +1213,8 @@ public static void registerExceptions() {
         );
         registerExceptionHandle(
             new OpenSearchExceptionHandle(
-                org.opensearch.rest.ResponseLimitBreachedException.class,
-                org.opensearch.rest.ResponseLimitBreachedException::new,
+                org.opensearch.common.breaker.ResponseLimitBreachedException.class,
+                org.opensearch.common.breaker.ResponseLimitBreachedException::new,
                 175,
                 V_2_18_0
             )

diff --git a/server/src/main/java/org/opensearch/action/ActionModule.java b/server/src/main/java/org/opensearch/action/ActionModule.java
@@ -300,6 +300,7 @@
 import org.opensearch.cluster.node.DiscoveryNodes;
 import org.opensearch.common.NamedRegistry;
 import org.opensearch.common.annotation.PublicApi;
+import org.opensearch.common.breaker.ResponseLimitSettings;
 import org.opensearch.common.inject.AbstractModule;
 import org.opensearch.common.inject.TypeLiteral;
 import org.opensearch.common.inject.multibindings.MapBinder;
@@ -325,7 +326,6 @@
 import org.opensearch.plugins.ActionPlugin;
 import org.opensearch.plugins.ActionPlugin.ActionHandler;
 import org.opensearch.rest.NamedRoute;
-import org.opensearch.rest.ResponseLimitSettings;
 import org.opensearch.rest.RestController;
 import org.opensearch.rest.RestHandler;
 import org.opensearch.rest.RestHeaderDefinition;
@@ -992,7 +992,7 @@ public void initRestHandlers(Supplier<DiscoveryNodes> nodesInCluster) {
         registerHandler.accept(new RestTemplatesAction());
 
         // LIST API
-        registerHandler.accept(new RestIndicesListAction());
+        registerHandler.accept(new RestIndicesListAction(responseLimitSettings));
 
         // Point in time API
         registerHandler.accept(new RestCreatePitAction());

diff --git a/...er/src/main/java/org/opensearch/action/admin/cluster/shards/TransportCatShardsAction.java b/...er/src/main/java/org/opensearch/action/admin/cluster/shards/TransportCatShardsAction.java
@@ -16,18 +16,18 @@
 import org.opensearch.action.support.HandledTransportAction;
 import org.opensearch.action.support.TimeoutTaskCancellationUtility;
 import org.opensearch.client.node.NodeClient;
+import org.opensearch.common.breaker.ResponseLimitBreachedException;
+import org.opensearch.common.breaker.ResponseLimitSettings;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.action.NotifyOnceListener;
-import org.opensearch.rest.ResponseLimitBreachedException;
-import org.opensearch.rest.ResponseLimitSettings;
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
 
 import java.util.Objects;
 
-import static org.opensearch.rest.ResponseLimitSettings.LimitEntity.SHARDS;
+import static org.opensearch.common.breaker.ResponseLimitSettings.LimitEntity.SHARDS;
 
 /**
  * Perform cat shards action
@@ -133,7 +133,11 @@ private void validateRequestLimit(
             int limit = responseLimitSettings.getCatShardsResponseLimit();
             if (ResponseLimitSettings.isResponseLimitBreached(clusterStateResponse.getState().getRoutingTable(), SHARDS, limit)) {
                 listener.onFailure(
-                    new ResponseLimitBreachedException("Too many shards requested. Can not request shards beyond {" + limit + "}")
+                    new ResponseLimitBreachedException(
+                        "Too many shards requested. Can not request shards beyond {" + limit + "}",
+                        limit,
+                        SHARDS
+                    )
                 );
             }
         }

diff --git a/server/src/main/java/org/opensearch/common/breaker/ResponseLimitBreachedException.java b/server/src/main/java/org/opensearch/common/breaker/ResponseLimitBreachedException.java
@@ -0,0 +1,66 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.common.breaker;
+
+import org.opensearch.OpenSearchException;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.rest.RestStatus;
+import org.opensearch.core.xcontent.XContentBuilder;
+
+import java.io.IOException;
+
+/**
+ * Thrown when api response breaches threshold limit.
+ *
+ * @opensearch.internal
+ */
+public class ResponseLimitBreachedException extends OpenSearchException {
+
+    private final int responseLimit;
+    private final ResponseLimitSettings.LimitEntity limitEntity;
+
+    public ResponseLimitBreachedException(StreamInput in) throws IOException {
+        super(in);
+        responseLimit = in.readVInt();
+        limitEntity = in.readEnum(ResponseLimitSettings.LimitEntity.class);
+    }
+
+    public ResponseLimitBreachedException(String msg, int responseLimit, ResponseLimitSettings.LimitEntity limitEntity) {
+        super(msg);
+        this.responseLimit = responseLimit;
+        this.limitEntity = limitEntity;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeVInt(responseLimit);
+        out.writeEnum(limitEntity);
+    }
+
+    public int getResponseLimit() {
+        return responseLimit;
+    }
+
+    public ResponseLimitSettings.LimitEntity getLimitEntity() {
+        return limitEntity;
+    }
+
+    @Override
+    public RestStatus status() {
+        return RestStatus.TOO_MANY_REQUESTS;
+    }
+
+    @Override
+    protected void metadataToXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.field("response_limit", responseLimit);
+        builder.field("limit_entity", limitEntity);
+    }
+}
diff --git a/...pensearch/rest/ResponseLimitSettings.java → ...common/breaker/ResponseLimitSettings.java b/...pensearch/rest/ResponseLimitSettings.java → ...common/breaker/ResponseLimitSettings.java
@@ -6,10 +6,11 @@
  * compatible open source license.
  */
 
-package org.opensearch.rest;
+package org.opensearch.common.breaker;
 
 import org.opensearch.cluster.metadata.Metadata;
 import org.opensearch.cluster.routing.IndexRoutingTable;
+import org.opensearch.cluster.routing.IndexShardRoutingTable;
 import org.opensearch.cluster.routing.RoutingTable;
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Setting;
@@ -94,8 +95,9 @@ public static boolean isResponseLimitBreached(final Metadata metadata, final Lim
         if (Objects.requireNonNull(limitEntity) == LimitEntity.INDICES) {
             int indicesCount = getTotalIndicesFromMetadata.apply(metadata);
             return indicesCount > limit;
+        } else {
+            throw new IllegalArgumentException("Unsupported limit entity [" + limitEntity + "]");
         }
-        return false;
     }
 
     /**
@@ -116,17 +118,23 @@ public static boolean isResponseLimitBreached(final RoutingTable routingTable, f
                 if (indicesCount > limit) return true;
                 break;
             case SHARDS:
-                final Map<String, IndexRoutingTable> indexRoutingTableMap = routingTable.getIndicesRouting();
-                int totalShards = 0;
-                for (final Map.Entry<String, IndexRoutingTable> entry : indexRoutingTableMap.entrySet()) {
-                    // In case routing table is corrupted. We will not block actions.
-                    if (Objects.isNull(entry.getValue()) || Objects.isNull(entry.getValue().getShards())) {
-                        return false;
-                    }
-                    totalShards += entry.getValue().getShards().size();
-                    if (totalShards > limit) return true;
-                }
+                if (isShardsLimitBreached(routingTable, limit)) return true;
                 break;
+            default:
+                throw new IllegalArgumentException("Unsupported limit entity [" + limitEntity + "]");
+        }
+        return false;
+    }
+
+    private static boolean isShardsLimitBreached(final RoutingTable routingTable, final int limit) {
+        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 limit value is breached and avoid unnecessary computation.
+                if (totalShards > limit) return true;
+            }
         }
         return false;
     }

diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java
@@ -84,6 +84,7 @@
 import org.opensearch.cluster.service.ClusterManagerTaskThrottler;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.annotation.PublicApi;
+import org.opensearch.common.breaker.ResponseLimitSettings;
 import org.opensearch.common.cache.CacheType;
 import org.opensearch.common.cache.settings.CacheSettings;
 import org.opensearch.common.cache.store.settings.OpenSearchOnHeapCacheSettings;
@@ -155,7 +156,6 @@
 import org.opensearch.repositories.blobstore.BlobStoreRepository;
 import org.opensearch.repositories.fs.FsRepository;
 import org.opensearch.rest.BaseRestHandler;
-import org.opensearch.rest.ResponseLimitSettings;
 import org.opensearch.script.ScriptService;
 import org.opensearch.search.SearchService;
 import org.opensearch.search.aggregations.MultiBucketConsumerService;

diff --git a/server/src/main/java/org/opensearch/rest/ResponseLimitBreachedException.java b/server/src/main/java/org/opensearch/rest/ResponseLimitBreachedException.java
diff --git a/server/src/main/java/org/opensearch/rest/action/cat/AbstractCatAction.java b/server/src/main/java/org/opensearch/rest/action/cat/AbstractCatAction.java
@@ -33,13 +33,13 @@
 
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.common.Table;
+import org.opensearch.common.breaker.ResponseLimitSettings;
 import org.opensearch.common.io.Streams;
 import org.opensearch.common.io.UTF8StreamWriter;
 import org.opensearch.core.common.io.stream.BytesStream;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.BytesRestResponse;
-import org.opensearch.rest.ResponseLimitSettings;
 import org.opensearch.rest.RestRequest;
 
 import java.io.IOException;

diff --git a/server/src/main/java/org/opensearch/rest/action/cat/RestIndicesAction.java b/server/src/main/java/org/opensearch/rest/action/cat/RestIndicesAction.java
@@ -50,6 +50,8 @@
 import org.opensearch.cluster.health.ClusterIndexHealth;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.common.Table;
+import org.opensearch.common.breaker.ResponseLimitBreachedException;
+import org.opensearch.common.breaker.ResponseLimitSettings;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.logging.DeprecationLogger;
 import org.opensearch.common.settings.Settings;
@@ -59,8 +61,6 @@
 import org.opensearch.core.action.ActionResponse;
 import org.opensearch.core.common.Strings;
 import org.opensearch.index.IndexSettings;
-import org.opensearch.rest.ResponseLimitBreachedException;
-import org.opensearch.rest.ResponseLimitSettings;
 import org.opensearch.rest.RestRequest;
 import org.opensearch.rest.RestResponse;
 import org.opensearch.rest.action.RestResponseListener;
@@ -88,7 +88,7 @@
 import static java.util.Arrays.asList;
 import static java.util.Collections.unmodifiableList;
 import static org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest.DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT;
-import static org.opensearch.rest.ResponseLimitSettings.LimitEntity.INDICES;
+import static org.opensearch.common.breaker.ResponseLimitSettings.LimitEntity.INDICES;
 import static org.opensearch.rest.RestRequest.Method.GET;
 
 /**
@@ -253,7 +253,11 @@ private void validateRequestLimit(final ClusterStateResponse clusterStateRespons
             int limit = responseLimitSettings.getCatIndicesResponseLimit();
             if (ResponseLimitSettings.isResponseLimitBreached(clusterStateResponse.getState().getMetadata(), INDICES, limit)) {
                 listener.onFailure(
-                    new ResponseLimitBreachedException("Too many indices requested. Can not request indices beyond {" + limit + "}")
+                    new ResponseLimitBreachedException(
+                        "Too many indices requested. Can not request indices beyond {" + limit + "}",
+                        limit,
+                        INDICES
+                    )
                 );
             }
         }

diff --git a/server/src/main/java/org/opensearch/rest/action/cat/RestSegmentsAction.java b/server/src/main/java/org/opensearch/rest/action/cat/RestSegmentsAction.java
@@ -42,11 +42,11 @@
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.cluster.node.DiscoveryNodes;
 import org.opensearch.common.Table;
+import org.opensearch.common.breaker.ResponseLimitBreachedException;
+import org.opensearch.common.breaker.ResponseLimitSettings;
 import org.opensearch.common.logging.DeprecationLogger;
 import org.opensearch.core.common.Strings;
 import org.opensearch.index.engine.Segment;
-import org.opensearch.rest.ResponseLimitBreachedException;
-import org.opensearch.rest.ResponseLimitSettings;
 import org.opensearch.rest.RestRequest;
 import org.opensearch.rest.RestResponse;
 import org.opensearch.rest.action.RestActionListener;
@@ -58,7 +58,7 @@
 
 import static java.util.Arrays.asList;
 import static java.util.Collections.unmodifiableList;
-import static org.opensearch.rest.ResponseLimitSettings.LimitEntity.INDICES;
+import static org.opensearch.common.breaker.ResponseLimitSettings.LimitEntity.INDICES;
 import static org.opensearch.rest.RestRequest.Method.GET;
 
 /**
@@ -131,7 +131,9 @@ private void validateRequestLimit(final ClusterStateResponse clusterStateRespons
             int limit = responseLimitSettings.getCatSegmentsResponseLimit();
             if (ResponseLimitSettings.isResponseLimitBreached(clusterStateResponse.getState().getRoutingTable(), INDICES, limit)) {
                 throw new ResponseLimitBreachedException(
-                    "Segments from too many indices requested. Can not request indices beyond {" + limit + "}"
+                    "Segments from too many indices requested. Can not request indices beyond {" + limit + "}",
+                    limit,
+                    INDICES
                 );
             }
         }

diff --git a/server/src/main/java/org/opensearch/rest/action/list/RestIndicesListAction.java b/server/src/main/java/org/opensearch/rest/action/list/RestIndicesListAction.java
@@ -9,6 +9,7 @@
 package org.opensearch.rest.action.list;
 
 import org.opensearch.action.admin.cluster.state.ClusterStateResponse;
+import org.opensearch.common.breaker.ResponseLimitSettings;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.rest.RestRequest;
@@ -35,6 +36,10 @@ public class RestIndicesListAction extends RestIndicesAction {
     private static final int MAX_SUPPORTED_LIST_INDICES_PAGE_SIZE = 5000;
     private static final int DEFAULT_LIST_INDICES_PAGE_SIZE = 500;
 
+    public RestIndicesListAction(final ResponseLimitSettings responseLimitSettings) {
+        super(responseLimitSettings);
+    }
+
     @Override
     public List<Route> routes() {
         return unmodifiableList(asList(new Route(GET, "/_list/indices"), new Route(GET, "/_list/indices/{index}")));
@@ -70,6 +75,11 @@ protected PageParams validateAndGetPageParams(RestRequest restRequest) {
         return pageParams;
     }
 
+    @Override
+    public boolean isRequestLimitCheckSupported() {
+        return false;
+    }
+
     protected int defaultPageSize() {
         return DEFAULT_LIST_INDICES_PAGE_SIZE;
     }

diff --git a/server/src/test/java/org/opensearch/ExceptionSerializationTests.java b/server/src/test/java/org/opensearch/ExceptionSerializationTests.java
@@ -65,6 +65,7 @@
 import org.opensearch.cluster.routing.UnsupportedWeightedRoutingStateException;
 import org.opensearch.cluster.service.ClusterManagerThrottlingException;
 import org.opensearch.common.UUIDs;
+import org.opensearch.common.breaker.ResponseLimitBreachedException;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.io.PathUtils;
 import org.opensearch.common.io.stream.BytesStreamOutput;
@@ -106,7 +107,6 @@
 import org.opensearch.indices.replication.common.ReplicationFailedException;
 import org.opensearch.ingest.IngestProcessorException;
 import org.opensearch.repositories.RepositoryException;
-import org.opensearch.rest.ResponseLimitBreachedException;
 import org.opensearch.rest.action.admin.indices.AliasesNotFoundException;
 import org.opensearch.search.SearchContextMissingException;
 import org.opensearch.search.SearchException;