Changing computation of PageToken in ShardStrategy

Signed-off-by: Harsh Garg <[email protected]>
opensearch-project · Oct 3, 2024 · ede0700 · ede0700
1 parent fe6601c
commit ede0700
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 26 deletions.
diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/CatShardsRequest.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/CatShardsRequest.java
@@ -43,6 +43,7 @@ public CatShardsRequest(StreamInput in) throws IOException {
             if (in.readBoolean()) {
                 pageParams = new PageParams(in);
             }
+            requestLimitCheckSupported = in.readBoolean();
         }
     }
 
@@ -60,8 +61,8 @@ public void writeTo(StreamOutput out) throws IOException {
             if (pageParams != null) {
                 pageParams.writeTo(out);
             }
+            out.writeBoolean(requestLimitCheckSupported);
         }
-        this.requestLimitCheckSupported = false;
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/action/pagination/ShardPaginationStrategy.java b/server/src/main/java/org/opensearch/action/pagination/ShardPaginationStrategy.java
@@ -91,36 +91,26 @@ private PageData getPageData(
                 lastAddedIndex = indexMetadata;
                 indices.add(indexName);
             }
-            if (shardCount >= numShardsRequired) {
-                break;
+
+            if (shardCount > numShardsRequired) {
+                return new PageData(
+                    shardRoutings,
+                    indices,
+                    new PageToken(
+                        new ShardStrategyToken(
+                            lastAddedIndex.getIndex().getName(),
+                            shardRoutings.get(shardRoutings.size() - 1).id(),
+                            lastAddedIndex.getCreationDate()
+                        ).generateEncryptedToken(),
+                        DEFAULT_SHARDS_PAGINATED_ENTITY
+                    )
+                );
             }
         }
 
-        if (filteredIndices.isEmpty() || shardRoutings.isEmpty()) {
-            return new PageData(shardRoutings, indices, new PageToken(null, DEFAULT_SHARDS_PAGINATED_ENTITY));
-        }
-        return new PageData(
-            shardRoutings,
-            indices,
-            getResponseToken(
-                lastAddedIndex,
-                filteredIndices.get(filteredIndices.size() - 1).getIndex().getName(),
-                shardRoutings.get(shardRoutings.size() - 1).id()
-            )
-        );
-    }
-
-    private PageToken getResponseToken(IndexMetadata pageEndIndexMetadata, final String lastFilteredIndexName, final int pageEndShardId) {
         // If all the shards of filtered indices list have been included in pageShardRoutings, then no more
         // shards are remaining to be fetched, and the next_token should thus be null.
-        String pageEndIndexName = pageEndIndexMetadata.getIndex().getName();
-        if (pageEndIndexName.equals(lastFilteredIndexName) && pageEndIndexMetadata.getNumberOfShards() == pageEndShardId + 1) {
-            return new PageToken(null, DEFAULT_SHARDS_PAGINATED_ENTITY);
-        }
-        return new PageToken(
-            new ShardStrategyToken(pageEndIndexName, pageEndShardId, pageEndIndexMetadata.getCreationDate()).generateEncryptedToken(),
-            DEFAULT_SHARDS_PAGINATED_ENTITY
-        );
+        return new PageData(shardRoutings, indices, new PageToken(null, DEFAULT_SHARDS_PAGINATED_ENTITY));
     }
 
     private ShardStrategyToken getShardStrategyToken(String requestedToken) {

diff --git a/server/src/test/java/org/opensearch/action/admin/cluster/shards/CatShardsRequestTests.java b/server/src/test/java/org/opensearch/action/admin/cluster/shards/CatShardsRequestTests.java
@@ -30,6 +30,7 @@ public void testSerializationWithDefaultParameters() throws Exception {
                 assertNull(deserialized.getPageParams());
                 assertNull(deserialized.getCancelAfterTimeInterval());
                 assertEquals(0, deserialized.getIndices().length);
+                assertFalse(deserialized.isRequestLimitCheckSupported());
             }
         }
     }
@@ -44,6 +45,7 @@ public void testSerializationWithStringPageParamsNull() throws Exception {
         }
         catShardsRequest.setIndices(indices);
         catShardsRequest.setCancelAfterTimeInterval(TimeValue.timeValueMillis(randomIntBetween(1, 5)));
+        catShardsRequest.setRequestLimitCheckSupported(true);
 
         Version version = Version.CURRENT;
         try (BytesStreamOutput out = new BytesStreamOutput()) {
@@ -58,6 +60,7 @@ public void testSerializationWithStringPageParamsNull() throws Exception {
                 assertArrayEquals(catShardsRequest.getIndices(), deserialized.getIndices());
                 // assert timeout
                 assertEquals(catShardsRequest.getCancelAfterTimeInterval(), deserialized.getCancelAfterTimeInterval());
+                assertTrue(deserialized.isRequestLimitCheckSupported());
             }
         }
     }
@@ -79,6 +82,7 @@ public void testSerializationWithPageParamsSet() throws Exception {
                 assertEquals(catShardsRequest.getPageParams(), deserialized.getPageParams());
                 assertEquals(0, deserialized.getIndices().length);
                 assertNull(deserialized.getCancelAfterTimeInterval());
+                assertFalse(deserialized.isRequestLimitCheckSupported());
             }
         }
     }
@@ -101,6 +105,7 @@ public void testSerializationWithOlderVersionsParametersNotSerialized() throws E
                 assertNull(deserialized.getPageParams());
                 assertNull(deserialized.getIndices());
                 assertNull(deserialized.getCancelAfterTimeInterval());
+                assertFalse(deserialized.isRequestLimitCheckSupported());
             }
         }
     }