Implementing pagination for _cat/shards

Signed-off-by: Harsh Garg <[email protected]>

server/src/main/java/org/opensearch/common/Table.java

@@ -43,6 +43,8 @@
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
+import reactor.util.annotation.NonNull;
+
 import static java.util.Collections.emptyMap;
 
 /**
@@ -59,9 +61,19 @@ public class Table {
     private List<Cell> currentCells;
     private boolean inHeaders = false;
     private boolean withTime = false;
+    private PaginationMetadata paginationMetadata = new PaginationMetadata(false, null, null);
+
     public static final String EPOCH = "epoch";
     public static final String TIMESTAMP = "timestamp";
 
+    public Table() {}
+
+    public Table(@Nullable PaginationMetadata paginationMetadata) {
+        if (paginationMetadata != null) {
+            this.paginationMetadata = paginationMetadata;
+        }
+    }
+
     public Table startHeaders() {
         inHeaders = true;
         currentCells = new ArrayList<>();
@@ -230,6 +242,18 @@ public Map<String, String> getAliasMap() {
         return headerAliasMap;
     }
 
+    public boolean isTablePaginated() {
+        return paginationMetadata.isResponsePaginated;
+    }
+
+    public String getNextTokenForTable() {
+        return paginationMetadata.nextToken;
+    }
+
+    public String getPaginatedElementForTable() {
+        return paginationMetadata.paginatedElement;
+    }
+
     /**
      * Cell in a table
      *
@@ -254,4 +278,34 @@ public Cell(Object value, Map<String, String> attr) {
             this.attr = attr;
         }
     }
+
+    /**
+     * Pagination metadata for a table.
+     *
+     * @opensearch.internal
+     */
+    public static class PaginationMetadata {
+
+        /**
+         * boolean denoting whether the table is paginated or not.
+         */
+        public final boolean isResponsePaginated;
+
+        /**
+         * String denoting the element which is being paginated (for e.g. shards, indices..).
+         */
+        public final String paginatedElement;
+
+        /**
+         * String denoting the nextToken of paginated response, which will be used to fetch nextPage (if any).
+         */
+        public final String nextToken;
+
+        public PaginationMetadata(@NonNull boolean isResponsePaginated, @Nullable String paginatedElement, @Nullable String nextToken) {
+            this.isResponsePaginated = isResponsePaginated;
+            assert !isResponsePaginated || paginatedElement != null : "paginatedElement must be specified for a table which is paginated";
+            this.paginatedElement = paginatedElement;
+            this.nextToken = nextToken;
+        }
+    }
 }

server/src/main/java/org/opensearch/rest/action/cat/RestShardsAction.java

@@ -39,6 +39,7 @@
 import org.opensearch.action.admin.indices.stats.IndicesStatsResponse;
 import org.opensearch.action.admin.indices.stats.ShardStats;
 import org.opensearch.client.node.NodeClient;
+import org.opensearch.cluster.routing.IndexShardRoutingTable;
 import org.opensearch.cluster.routing.ShardRouting;
 import org.opensearch.cluster.routing.UnassignedInfo;
 import org.opensearch.common.Table;
@@ -65,11 +66,17 @@
 import org.opensearch.rest.action.RestResponseListener;
 import org.opensearch.search.suggest.completion.CompletionStats;
 
+import java.nio.charset.StandardCharsets;
 import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Base64;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
 import java.util.function.Function;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Arrays.asList;
 import static java.util.Collections.unmodifiableList;
 import static org.opensearch.rest.RestRequest.Method.GET;
@@ -106,24 +113,53 @@ protected void documentation(StringBuilder sb) {
 
     @Override
     public RestChannelConsumer doCatRequest(final RestRequest request, final NodeClient client) {
-        final String[] indices = Strings.splitStringByCommaToArray(request.param("index"));
+
+        String[] indices = new String[0];
         final ClusterStateRequest clusterStateRequest = new ClusterStateRequest();
         clusterStateRequest.local(request.paramAsBoolean("local", clusterStateRequest.local()));
         clusterStateRequest.clusterManagerNodeTimeout(
             request.paramAsTime("cluster_manager_timeout", clusterStateRequest.clusterManagerNodeTimeout())
         );
         parseDeprecatedMasterTimeoutParameter(clusterStateRequest, request, deprecationLogger, getName());
-        clusterStateRequest.clear().nodes(true).routingTable(true).indices(indices);
+        if (request.hasParam("nextToken")) {
+            // ToDo: Add validation on the nextToken passed in the request
+            // Need to get the metadata as well
+            request.param("nextToken");
+            clusterStateRequest.clear().nodes(true).routingTable(true).metadata(true);
+        } else {
+            // Only parse the "index" param if the request is not-paginated.
+            indices = Strings.splitStringByCommaToArray(request.param("index"));
+            clusterStateRequest.clear().nodes(true).routingTable(true).indices(indices);
+        }
+
+        String[] finalIndices = indices;
         return channel -> client.admin().cluster().state(clusterStateRequest, new RestActionListener<ClusterStateResponse>(channel) {
             @Override
             public void processResponse(final ClusterStateResponse clusterStateResponse) {
                 IndicesStatsRequest indicesStatsRequest = new IndicesStatsRequest();
                 indicesStatsRequest.all();
-                indicesStatsRequest.indices(indices);
+                final List<ShardRouting> shardRoutingResponseList = new ArrayList<>();
+                final Table.PaginationMetadata paginationMetadata;
+                if (request.hasParam("nextToken")) {
+                    List<String> indicesToBeQueried = new ArrayList<>();
+                    paginationMetadata = new Table.PaginationMetadata(
+                        true,
+                        "shards",
+                        getNextTokenForPaginatedResponse(request, clusterStateResponse, indicesToBeQueried, shardRoutingResponseList)
+                    );
+                    indicesStatsRequest.indices(indicesToBeQueried.toArray(new String[0]));
+                } else {
+                    shardRoutingResponseList.addAll(clusterStateResponse.getState().routingTable().allShards());
+                    indicesStatsRequest.indices(finalIndices);
+                    paginationMetadata = null;
+                }
                 client.admin().indices().stats(indicesStatsRequest, new RestResponseListener<IndicesStatsResponse>(channel) {
                     @Override
                     public RestResponse buildResponse(IndicesStatsResponse indicesStatsResponse) throws Exception {
-                        return RestTable.buildResponse(buildTable(request, clusterStateResponse, indicesStatsResponse), channel);
+                        return RestTable.buildResponse(
+                            buildTable(request, clusterStateResponse, indicesStatsResponse, shardRoutingResponseList, paginationMetadata),
+                            channel
+                        );
                     }
                 });
             }
@@ -132,7 +168,11 @@ public RestResponse buildResponse(IndicesStatsResponse indicesStatsResponse) thr
 
     @Override
     protected Table getTableWithHeader(final RestRequest request) {
-        Table table = new Table();
+        return getTableWithHeader(request, null);
+    }
+
+    protected Table getTableWithHeader(final RestRequest request, Table.PaginationMetadata paginationMetadata) {
+        Table table = new Table(paginationMetadata);
         table.startHeaders()
             .addCell("index", "default:true;alias:i,idx;desc:index name")
             .addCell("shard", "default:true;alias:s,sh;desc:shard name")
@@ -301,10 +341,15 @@ private static <S, T> Object getOrNull(S stats, Function<S, T> accessor, Functio
     }
 
     // package private for testing
-    Table buildTable(RestRequest request, ClusterStateResponse state, IndicesStatsResponse stats) {
-        Table table = getTableWithHeader(request);
-
-        for (ShardRouting shard : state.getState().routingTable().allShards()) {
+    Table buildTable(
+        RestRequest request,
+        ClusterStateResponse state,
+        IndicesStatsResponse stats,
+        List<ShardRouting> shardRoutingList,
+        Table.PaginationMetadata paginationMetadata
+    ) {
+        Table table = getTableWithHeader(request, paginationMetadata);
+        for (ShardRouting shard : shardRoutingList) {
             ShardStats shardStats = stats.asMap().get(shard);
             CommonStats commonStats = null;
             CommitStats commitStats = null;
@@ -453,7 +498,106 @@ Table buildTable(RestRequest request, ClusterStateResponse state, IndicesStatsRe
 
             table.endRow();
         }
-
         return table;
     }
+
+    private List<String> getListOfIndicesSortedByCreateTime(final ClusterStateResponse clusterStateResponse) {
+        List<String> indicesList = new ArrayList<String>(clusterStateResponse.getState().getRoutingTable().getIndicesRouting().keySet());
+        indicesList.sort((index1, index2) -> {
+            Long index1CreationTimeStamp = clusterStateResponse.getState().metadata().indices().get(index1).getCreationDate();
+            Long index2CreationTimeStamp = clusterStateResponse.getState().metadata().indices().get(index2).getCreationDate();
+            if (index1CreationTimeStamp.equals(index2CreationTimeStamp)) {
+                return index1.compareTo(index2);
+            }
+            return (index1CreationTimeStamp - index2CreationTimeStamp) > 0 ? 1 : -1;
+        });
+        return indicesList;
+    }
+
+    private String getNextTokenForPaginatedResponse(
+        final RestRequest request,
+        ClusterStateResponse clusterStateResponse,
+        List<String> indicesToBeQueried,
+        List<ShardRouting> shardRoutingResponseList
+    ) {
+        final long defaultPageSize = (long) clusterStateResponse.getState().nodes().getDataNodes().size() + 1;
+
+        // Get the nextToken provided in the request
+        final String nextTokenInRequest = Objects.equals(request.param("nextToken"), "null")
+            ? null
+            : new String(Base64.getDecoder().decode(request.param("nextToken")), UTF_8);
+
+        List<String> sortedIndicesList = getListOfIndicesSortedByCreateTime(clusterStateResponse);
+
+        // Since all the shards for last ID would have already been sent in the last response,
+        // start iterating from the next shard for current page
+        int newPageStartShardID = nextTokenInRequest == null ? 0 : Integer.parseInt(nextTokenInRequest.split("\\$")[0]) + 1;
+
+        // Since all the shards corresponding to the last processed index might not have been included in the last page,
+        // start iterating from the last index number itself
+        int newPageStartIndexNumber = nextTokenInRequest == null ? 0 : Integer.parseInt(nextTokenInRequest.split("\\$")[1]);
+
+        // Get the number of shards upto the maxPageSize
+        long shardCountSoFar = 0L;
+        int lastProcessedShardNumber = -1;
+        int lastProcessedIndexNumber = -1;
+
+        // ToDo: Handle case when index gets deleted. Select the first index with creationTime just greater than the last index's
+        // creationTime
+        int indexNumberInSortedList = newPageStartIndexNumber;
+        for (; indexNumberInSortedList < sortedIndicesList.size(); indexNumberInSortedList++) {
+            String index = sortedIndicesList.get(indexNumberInSortedList);
+            Map<Integer, IndexShardRoutingTable> indexShards = clusterStateResponse.getState()
+                .getRoutingTable()
+                .getIndicesRouting()
+                .get(index)
+                .getShards();
+            // If all the shards corresponding to the last index were already processed, move to the next Index
+            if (indexNumberInSortedList == newPageStartIndexNumber && (newPageStartShardID > indexShards.size() - 1)) {
+                // ToDo: Add validation that the newPageStartShardID should not be greater than the newPageStartIndexShards.size()
+                newPageStartShardID = 0;
+                continue;
+            }
+            int lastProcessedShardNumberForCurrentIndex = -1;
+            int shardID = (indexNumberInSortedList == newPageStartIndexNumber) ? newPageStartShardID : 0;
+            for (; shardID < indexShards.size(); shardID++) {
+                shardCountSoFar += indexShards.get(shardID).shards().size();
+                if (shardCountSoFar > defaultPageSize) {
+                    break;
+                }
+                shardRoutingResponseList.addAll(indexShards.get(shardID).shards());
+                lastProcessedShardNumberForCurrentIndex = shardID;
+            }
+
+            if (shardCountSoFar > defaultPageSize) {
+                if (lastProcessedShardNumberForCurrentIndex != -1) {
+                    indicesToBeQueried.add(index);
+                    lastProcessedIndexNumber = indexNumberInSortedList;
+                    lastProcessedShardNumber = lastProcessedShardNumberForCurrentIndex;
+                }
+                break;
+            }
+            indicesToBeQueried.add(index);
+            lastProcessedShardNumber = lastProcessedShardNumberForCurrentIndex;
+            lastProcessedIndexNumber = indexNumberInSortedList;
+        }
+
+        // nextToken = "lastProcessedShardNumber$LastProcessedIndexNumber$CreateTimeOfLastProcessedIndex$NameOfLastProcessedIndex"
+        return indexNumberInSortedList >= sortedIndicesList.size()
+            ? null
+            : Base64.getEncoder()
+                .encodeToString(
+                    (lastProcessedShardNumber
+                        + "$"
+                        + (lastProcessedIndexNumber)
+                        + "$"
+                        + clusterStateResponse.getState()
+                            .metadata()
+                            .indices()
+                            .get(sortedIndicesList.get(lastProcessedIndexNumber))
+                            .getCreationDate()
+                        + "$"
+                        + sortedIndicesList.get(lastProcessedIndexNumber)).getBytes(StandardCharsets.UTF_8)
+                );
+    }
 }

server/src/main/java/org/opensearch/rest/action/cat/RestTable.java

@@ -88,7 +88,15 @@ public static RestResponse buildXContentBuilder(Table table, RestChannel channel
         XContentBuilder builder = channel.newBuilder();
         List<DisplayHeader> displayHeaders = buildDisplayHeaders(table, request);
 
-        builder.startArray();
+        if (table.isTablePaginated()) {
+            assert table.getPaginatedElementForTable() != null : "Paginated element is required in-case nextToken is not null";
+            builder.startObject();
+            builder.field("nextToken", table.getNextTokenForTable());
+            builder.startArray(table.getPaginatedElementForTable());
+        } else {
+            builder.startArray();
+        }
+
         List<Integer> rowOrder = getRowOrder(table, request);
         for (Integer row : rowOrder) {
             builder.startObject();
@@ -98,6 +106,11 @@ public static RestResponse buildXContentBuilder(Table table, RestChannel channel
             builder.endObject();
         }
         builder.endArray();
+
+        if (table.isTablePaginated()) {
+            builder.endObject();
+        }
+
         return new BytesRestResponse(RestStatus.OK, builder);
     }
 
@@ -136,6 +149,13 @@ public static RestResponse buildTextPlainResponse(Table table, RestChannel chann
             }
             out.append("\n");
         }
+
+        // Adding a nextToken row, post an empty line, in the response if the table is paginated.
+        if (table.isTablePaginated()) {
+            out.append("\n");
+            out.append("nextToken" + " " + table.getNextTokenForTable());
+            out.append("\n");
+        }
         out.close();
         return new BytesRestResponse(RestStatus.OK, BytesRestResponse.TEXT_CONTENT_TYPE, bytesOut.bytes());
     }

server/src/test/java/org/opensearch/rest/action/cat/RestShardsActionTests.java

@@ -112,7 +112,7 @@ public void testBuildTable() {
         when(state.getState()).thenReturn(clusterState);
 
         final RestShardsAction action = new RestShardsAction();
-        final Table table = action.buildTable(new FakeRestRequest(), state, stats);
+        final Table table = action.buildTable(new FakeRestRequest(), state, stats, state.getState().routingTable().allShards(), null);
 
         // now, verify the table is correct
         List<Table.Cell> headers = table.getHeaders();