Just switch to _count

Signed-off-by: Mikayla Thompson <thomika@amazon.com>
opensearch-project · Oct 23, 2024 · e48005b · e48005b
1 parent 31ae47d
commit e48005b
Showing 1 changed file with 9 additions and 14 deletions.
diff --git a/...n/java/org/opensearch/migrations/bulkload/workcoordination/OpenSearchWorkCoordinator.java b/...n/java/org/opensearch/migrations/bulkload/workcoordination/OpenSearchWorkCoordinator.java
@@ -397,14 +397,11 @@ public void completeWorkItem(
         }
     }
 
-    private int numWorkItemsArePending(
-        int maxItemsToCheckFor,
+    private int numWorkItemsArePendingInternal(
         Supplier<IWorkCoordinationContexts.IPendingWorkItemsContext> contextSupplier
     ) throws IOException, InterruptedException {
         try (var context = contextSupplier.get()) {
             refresh(context::getRefreshContext);
-            // TODO: Switch this to use _count
-            log.warn("Switch this to use _count");
             final var queryBody = "{\n"
                 + "\"query\": {"
                 + "  \"bool\": {"
@@ -422,33 +419,31 @@ private int numWorkItemsArePending(
                 + "}"
                 + "}";
 
-            var path = INDEX_NAME + "/_search" + (maxItemsToCheckFor <= 0 ? "" : "?size=" + maxItemsToCheckFor);
+            var path = INDEX_NAME + "/_count";
             var response = httpClient.makeJsonRequest(AbstractedHttpClient.POST_METHOD, path, null, queryBody);
-
-            final var resultHitsUpper = objectMapper.readTree(response.getPayloadBytes()).path("hits");
             var statusCode = response.getStatusCode();
             if (statusCode != 200) {
                 throw new IllegalStateException(
-                    "Querying for pending (expired or not) work, "
-                        + "returned an unexpected status code "
-                        + statusCode
-                        + " instead of 200"
+                        "Querying for pending (expired or not) work, "
+                                + "returned an unexpected status code "
+                                + statusCode
+                                + " instead of 200"
                 );
             }
-            return resultHitsUpper.path("total").path("value").asInt();
+            return objectMapper.readTree(response.getPayloadBytes()).path("count").asInt();
         }
     }
 
     @Override
     public int numWorkItemsArePending(Supplier<IWorkCoordinationContexts.IPendingWorkItemsContext> contextSupplier)
         throws IOException, InterruptedException {
-        return numWorkItemsArePending(-1, contextSupplier);
+        return numWorkItemsArePendingInternal(contextSupplier);
     }
 
     @Override
     public boolean workItemsArePending(Supplier<IWorkCoordinationContexts.IPendingWorkItemsContext> contextSupplier)
         throws IOException, InterruptedException {
-        return numWorkItemsArePending(1, contextSupplier) >= 1;
+        return numWorkItemsArePendingInternal(contextSupplier) >= 1;
     }
 
     enum UpdateResult {