From 2e1c827cae3d9d34f3d6d15d13a132d11b5627eb Mon Sep 17 00:00:00 2001
From: Daniel Widdis <widdis@gmail.com>
Date: Mon, 5 Feb 2024 10:49:59 -0800
Subject: [PATCH] Make sure retryableGetMlTask evaluates task response at least
 once (#494)

Signed-off-by: Daniel Widdis <widdis@gmail.com>
---
 .../flowframework/workflow/AbstractRetryableWorkflowStep.java | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/main/java/org/opensearch/flowframework/workflow/AbstractRetryableWorkflowStep.java b/src/main/java/org/opensearch/flowframework/workflow/AbstractRetryableWorkflowStep.java
index 63d0d7587..90734bf02 100644
--- 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));
     }