Make sure retryableGetMlTask evaluates task response at least once (#494

) Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Feb 5, 2024 · 2e1c827 · 2e1c827
1 parent 281a0b2
commit 2e1c827
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/AbstractRetryableWorkflowStep.java b/src/main/java/org/opensearch/flowframework/workflow/AbstractRetryableWorkflowStep.java
@@ -75,7 +75,7 @@ protected void retryableGetMlTask(
         ActionListener<MLTask> mlTaskListener
     ) {
         CompletableFuture.runAsync(() -> {
-            while (!future.isDone()) {
+            do {
                 mlClient.getTask(taskId, ActionListener.wrap(response -> {
                     switch (response.getState()) {
                         case COMPLETED:
@@ -128,7 +128,7 @@ protected void retryableGetMlTask(
                     FutureUtils.cancel(future);
                     Thread.currentThread().interrupt();
                 }
-            }
+            } while (!future.isDone());
         }, threadPool.executor(WORKFLOW_THREAD_POOL));
     }