Support cancellation for admin apis

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased 2.x]
 ### Added
+- Support cancellation for cat shards and node stats API.([#13966](https://github.com/opensearch-project/OpenSearch/pull/13966))
 
 ### Dependencies
 - Bump `org.apache.commons:commons-lang3` from 3.14.0 to 3.15.0 ([#14861](https://github.com/opensearch-project/OpenSearch/pull/14861))

server/src/internalClusterTest/java/org/opensearch/action/admin/cluster/shards/TransportCatShardsActionIT.java

@@ -0,0 +1,128 @@
+/*
+ * 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.action.admin.cluster.shards;
+
+import org.opensearch.action.admin.cluster.state.ClusterStateResponse;
+import org.opensearch.action.admin.indices.stats.IndicesStatsResponse;
+import org.opensearch.cluster.metadata.IndexMetadata;
+import org.opensearch.cluster.routing.ShardRouting;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.common.Strings;
+import org.opensearch.core.tasks.TaskCancelledException;
+import org.opensearch.test.InternalTestCluster;
+import org.opensearch.test.OpenSearchIntegTestCase;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+
+import static org.opensearch.cluster.routing.UnassignedInfo.INDEX_DELAYED_NODE_LEFT_TIMEOUT_SETTING;
+import static org.opensearch.common.unit.TimeValue.timeValueMillis;
+import static org.opensearch.search.SearchService.NO_TIMEOUT;
+import static org.junit.Assert.fail;
+
+@OpenSearchIntegTestCase.ClusterScope(numDataNodes = 0, scope = OpenSearchIntegTestCase.Scope.TEST)
+public class TransportCatShardsActionIT extends OpenSearchIntegTestCase {
+
+    public void testCatShardsWithSuccessResponse() throws InterruptedException {
+        internalCluster().startClusterManagerOnlyNodes(1);
+        List<String> nodes = internalCluster().startDataOnlyNodes(3);
+        createIndex(
+            "test",
+            Settings.builder()
+                .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+                .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 2)
+                .put(INDEX_DELAYED_NODE_LEFT_TIMEOUT_SETTING.getKey(), "60m")
+                .build()
+        );
+        ensureGreen("test");
+
+        final CatShardsRequest shardsRequest = new CatShardsRequest();
+        shardsRequest.setCancelAfterTimeInterval(NO_TIMEOUT);
+        shardsRequest.setIndices(Strings.EMPTY_ARRAY);
+        CountDownLatch latch = new CountDownLatch(1);
+        client().execute(CatShardsAction.INSTANCE, shardsRequest, new ActionListener<CatShardsResponse>() {
+            @Override
+            public void onResponse(CatShardsResponse catShardsResponse) {
+                ClusterStateResponse clusterStateResponse = catShardsResponse.getClusterStateResponse();
+                IndicesStatsResponse indicesStatsResponse = catShardsResponse.getIndicesStatsResponse();
+                for (ShardRouting shard : clusterStateResponse.getState().routingTable().allShards()) {
+                    assertEquals("test", shard.getIndexName());
+                    assertNotNull(indicesStatsResponse.asMap().get(shard));
+                }
+                latch.countDown();
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                fail();
+                latch.countDown();
+            }
+        });
+        latch.await();
+    }
+
+    public void testCatShardsWithTimeoutException() throws IOException, AssertionError, InterruptedException {
+        List<String> masterNodes = internalCluster().startClusterManagerOnlyNodes(1);
+        List<String> nodes = internalCluster().startDataOnlyNodes(3);
+        createIndex(
+            "test",
+            Settings.builder()
+                .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+                .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 2)
+                .put(INDEX_DELAYED_NODE_LEFT_TIMEOUT_SETTING.getKey(), "60m")
+                .build()
+        );
+        ensureGreen("test");
+
+        Settings clusterManagerDataPathSettings = internalCluster().dataPathSettings(masterNodes.get(0));
+        // Dropping master node to delay in cluster state call.
+        internalCluster().stopRandomNode(InternalTestCluster.nameFilter(masterNodes.get(0)));
+
+        new Thread(() -> {
+            try {
+                // Ensures the cancellation timeout expires.
+                Thread.sleep(1000);
+                // Starting master node to proceed in cluster state call.
+                internalCluster().startClusterManagerOnlyNode(
+                    Settings.builder().put("node.name", masterNodes.get(0)).put(clusterManagerDataPathSettings).build()
+                );
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+        }).start();
+
+        final CatShardsRequest shardsRequest = new CatShardsRequest();
+        shardsRequest.setCancelAfterTimeInterval(timeValueMillis(1000));
+        shardsRequest.setIndices(Strings.EMPTY_ARRAY);
+        CountDownLatch latch = new CountDownLatch(1);
+        client().execute(CatShardsAction.INSTANCE, shardsRequest, new ActionListener<CatShardsResponse>() {
+            @Override
+            public void onResponse(CatShardsResponse catShardsResponse) {
+                // onResponse should not be called.
+                fail();
+                latch.countDown();
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                boolean timeoutException = (e.getClass() == TaskCancelledException.class);
+                if (e.getCause() != null) {
+                    timeoutException = timeoutException
+                        || (e.getCause().getMessage().contains("The parent task was cancelled, shouldn't start any child tasks"));
+                }
+                assertTrue(timeoutException);
+                latch.countDown();
+            }
+        });
+        latch.await();
+    }
+
+}

server/src/main/java/org/opensearch/action/ActionModule.java

@@ -85,7 +85,9 @@
 import org.opensearch.action.admin.cluster.reroute.TransportClusterRerouteAction;
 import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsAction;
 import org.opensearch.action.admin.cluster.settings.TransportClusterUpdateSettingsAction;
+import org.opensearch.action.admin.cluster.shards.CatShardsAction;
 import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsAction;
+import org.opensearch.action.admin.cluster.shards.TransportCatShardsAction;
 import org.opensearch.action.admin.cluster.shards.TransportClusterSearchShardsAction;
 import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingAction;
 import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.TransportDeleteWeightedRoutingAction;
@@ -646,6 +648,7 @@ public <Request extends ActionRequest, Response extends ActionResponse> void reg
         actions.register(ClusterGetWeightedRoutingAction.INSTANCE, TransportGetWeightedRoutingAction.class);
         actions.register(ClusterDeleteWeightedRoutingAction.INSTANCE, TransportDeleteWeightedRoutingAction.class);
         actions.register(IndicesStatsAction.INSTANCE, TransportIndicesStatsAction.class);
+        actions.register(CatShardsAction.INSTANCE, TransportCatShardsAction.class);
         actions.register(IndicesSegmentsAction.INSTANCE, TransportIndicesSegmentsAction.class);
         actions.register(IndicesShardStoresAction.INSTANCE, TransportIndicesShardStoresAction.class);
         actions.register(CreateIndexAction.INSTANCE, TransportCreateIndexAction.class);

server/src/main/java/org/opensearch/action/admin/cluster/shards/CatShardsAction.java

@@ -0,0 +1,25 @@
+/*
+ * 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.action.admin.cluster.shards;
+
+import org.opensearch.action.ActionType;
+
+/**
+ * Transport action for cat shards
+ *
+ * @opensearch.internal
+ */
+public class CatShardsAction extends ActionType<CatShardsResponse> {
+    public static final CatShardsAction INSTANCE = new CatShardsAction();
+    public static final String NAME = "cluster:monitor/shards";
+
+    private CatShardsAction() {
+        super(NAME, CatShardsResponse::new);
+    }
+}

server/src/main/java/org/opensearch/action/admin/cluster/shards/CatShardsRequest.java

@@ -0,0 +1,76 @@
+/*
+ * 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.action.admin.cluster.shards;
+
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.action.support.clustermanager.ClusterManagerNodeReadRequest;
+import org.opensearch.common.unit.TimeValue;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.tasks.TaskId;
+import org.opensearch.rest.action.admin.cluster.ClusterAdminTask;
+
+import java.io.IOException;
+import java.util.Map;
+
+import static org.opensearch.search.SearchService.NO_TIMEOUT;
+
+/**
+ * A request of _cat/shards.
+ *
+ * @opensearch.api
+ */
+public class CatShardsRequest extends ClusterManagerNodeReadRequest<CatShardsRequest> {
+
+    private String[] indices;
+    private TimeValue cancelAfterTimeInterval = NO_TIMEOUT;
+
+    public CatShardsRequest() {}
+
+    public CatShardsRequest(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        return null;
+    }
+
+    public void setClusterManagerNodeTimeout(TimeValue timeout) {
+        this.clusterManagerNodeTimeout = timeout;
+    }
+
+    public TimeValue getClusterManagerNodeTimeout() {
+        return this.clusterManagerNodeTimeout;
+    }
+
+    public void setLocal(boolean local) {
+        this.local = local;
+    }
+
+    public boolean getLocal() {
+        return this.local;
+    }
+
+    public void setIndices(String[] indices) {
+        this.indices = indices;
+    }
+
+    public String[] getIndices() {
+        return this.indices;
+    }
+
+    public void setCancelAfterTimeInterval(TimeValue timeout) {
+        this.cancelAfterTimeInterval = timeout;
+    }
+
+    @Override
+    public ClusterAdminTask createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+        return new ClusterAdminTask(id, type, action, parentTaskId, headers, this.cancelAfterTimeInterval);
+    }
+}

server/src/main/java/org/opensearch/action/admin/cluster/shards/CatShardsResponse.java

@@ -0,0 +1,56 @@
+/*
+ * 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.action.admin.cluster.shards;
+
+import org.opensearch.action.admin.cluster.state.ClusterStateResponse;
+import org.opensearch.action.admin.indices.stats.IndicesStatsResponse;
+import org.opensearch.core.action.ActionResponse;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+
+/**
+ * A response of a cat shards request.
+ *
+ * @opensearch.api
+ */
+public class CatShardsResponse extends ActionResponse {
+
+    private ClusterStateResponse clusterStateResponse = null;
+
+    private IndicesStatsResponse indicesStatsResponse = null;
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+
+    }
+
+    public CatShardsResponse() {}
+
+    public CatShardsResponse(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    public void setClusterStateResponse(ClusterStateResponse clusterStateResponse) {
+        this.clusterStateResponse = clusterStateResponse;
+    }
+
+    public ClusterStateResponse getClusterStateResponse() {
+        return this.clusterStateResponse;
+    }
+
+    public void setIndicesStatsResponse(IndicesStatsResponse indicesStatsResponse) {
+        this.indicesStatsResponse = indicesStatsResponse;
+    }
+
+    public IndicesStatsResponse getIndicesStatsResponse() {
+        return this.indicesStatsResponse;
+    }
+}