Merge branch 'main' into fix-integ

opensearch-project · Jan 22, 2024 · 81e94b4 · 81e94b4
2 parents fd9c64b + 92d9108
commit 81e94b4
Show file tree

Hide file tree

Showing 38 changed files with 244 additions and 263 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -61,6 +61,7 @@ jobs:
   integTest:
     needs: [spotless, javadoc]
     strategy:
+      fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
         # Don't use 21.0.2 https://github.com/opensearch-project/flow-framework/issues/426

diff --git a/src/main/java/org/opensearch/flowframework/transport/DeprovisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/DeprovisionWorkflowTransportAction.java
@@ -14,6 +14,7 @@
 import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.HandledTransportAction;
+import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.client.Client;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.common.util.concurrent.ThreadContext;
@@ -39,7 +40,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 
 import static org.opensearch.flowframework.common.CommonValue.PROVISIONING_PROGRESS_FIELD;
@@ -160,9 +160,9 @@ private void executeDeprovisionSequence(
                 ProcessNode deprovisionNode = iter.next();
                 ResourceCreated resource = getResourceFromDeprovisionNode(deprovisionNode, resourcesCreated);
                 String resourceNameAndId = getResourceNameAndId(resource);
-                CompletableFuture<WorkflowData> deprovisionFuture = deprovisionNode.execute();
+                PlainActionFuture<WorkflowData> deprovisionFuture = deprovisionNode.execute();
                 try {
-                    deprovisionFuture.join();
+                    deprovisionFuture.get();
                     logger.info("Successful {} for {}", deprovisionNode.id(), resourceNameAndId);
                     // Remove from list so we don't try again
                     iter.remove();

diff --git a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java
@@ -14,6 +14,7 @@
 import org.opensearch.action.get.GetRequest;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.HandledTransportAction;
+import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.client.Client;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.common.util.concurrent.ThreadContext;
@@ -40,7 +41,6 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.CancellationException;
-import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 
 import static org.opensearch.flowframework.common.CommonValue.ERROR_FIELD;
@@ -183,7 +183,7 @@ private void executeWorkflowAsync(String workflowId, List<ProcessNode> workflowS
     private void executeWorkflow(List<ProcessNode> workflowSequence, String workflowId) {
         try {
 
-            List<CompletableFuture<?>> workflowFutureList = new ArrayList<>();
+            List<PlainActionFuture<?>> workflowFutureList = new ArrayList<>();
             for (ProcessNode processNode : workflowSequence) {
                 List<ProcessNode> predecessors = processNode.predecessors();
 
@@ -202,8 +202,8 @@ private void executeWorkflow(List<ProcessNode> workflowSequence, String workflow
                 workflowFutureList.add(processNode.execute());
             }
 
-            // Attempt to join each workflow step future, may throw a CompletionException if any step completes exceptionally
-            workflowFutureList.forEach(CompletableFuture::join);
+            // Attempt to join each workflow step future, may throw a ExecutionException if any step completes exceptionally
+            workflowFutureList.forEach(PlainActionFuture::actionGet);
 
             logger.info("Provisioning completed successfully for workflow {}", workflowId);
             flowFrameworkIndicesHandler.updateFlowFrameworkSystemIndexDoc(

diff --git a/src/main/java/org/opensearch/flowframework/workflow/AbstractRegisterLocalModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/AbstractRegisterLocalModelStep.java
@@ -11,6 +11,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
+import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.flowframework.common.FlowFrameworkSettings;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
@@ -28,7 +29,6 @@
 
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.CompletableFuture;
 import java.util.stream.Stream;
 
 import static org.opensearch.flowframework.common.CommonValue.ALL_CONFIG;
@@ -75,14 +75,14 @@ protected AbstractRegisterLocalModelStep(
     }
 
     @Override
-    public CompletableFuture<WorkflowData> execute(
+    public PlainActionFuture<WorkflowData> execute(
         String currentNodeId,
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs
     ) {
 
-        CompletableFuture<WorkflowData> registerLocalModelFuture = new CompletableFuture<>();
+        PlainActionFuture<WorkflowData> registerLocalModelFuture = PlainActionFuture.newFuture();
 
         try {
             Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
@@ -180,7 +180,7 @@ public CompletableFuture<WorkflowData> execute(
                                         "successfully updated resources created in state index: {}",
                                         deployUpdateResponse.getIndex()
                                     );
-                                    registerLocalModelFuture.complete(
+                                    registerLocalModelFuture.onResponse(
                                         new WorkflowData(
                                             Map.ofEntries(
                                                 Map.entry(resourceName, id),
@@ -192,7 +192,7 @@ public CompletableFuture<WorkflowData> execute(
                                     );
                                 }, deployUpdateException -> {
                                     logger.error("Failed to update simulated deploy step resource", deployUpdateException);
-                                    registerLocalModelFuture.completeExceptionally(
+                                    registerLocalModelFuture.onFailure(
                                         new FlowFrameworkException(
                                             deployUpdateException.getMessage(),
                                             ExceptionsHelper.status(deployUpdateException)
@@ -201,24 +201,22 @@ public CompletableFuture<WorkflowData> execute(
                                 })
                             );
                         } else {
-                            registerLocalModelFuture.complete(
+                            registerLocalModelFuture.onResponse(
                                 new WorkflowData(
                                     Map.ofEntries(Map.entry(resourceName, id), Map.entry(REGISTER_MODEL_STATUS, mlTask.getState().name())),
                                     currentNodeInputs.getWorkflowId(),
                                     currentNodeId
                                 )
                             );
                         }
-                    }, exception -> { registerLocalModelFuture.completeExceptionally(exception); })
+                    }, exception -> { registerLocalModelFuture.onFailure(exception); })
                 );
             }, exception -> {
                 logger.error("Failed to register local model");
-                registerLocalModelFuture.completeExceptionally(
-                    new FlowFrameworkException(exception.getMessage(), ExceptionsHelper.status(exception))
-                );
+                registerLocalModelFuture.onFailure(new FlowFrameworkException(exception.getMessage(), ExceptionsHelper.status(exception)));
             }));
         } catch (FlowFrameworkException e) {
-            registerLocalModelFuture.completeExceptionally(e);
+            registerLocalModelFuture.onFailure(e);
         }
         return registerLocalModelFuture;
     }

diff --git a/src/main/java/org/opensearch/flowframework/workflow/AbstractRetryableWorkflowStep.java b/src/main/java/org/opensearch/flowframework/workflow/AbstractRetryableWorkflowStep.java
@@ -11,6 +11,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
+import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.common.util.concurrent.FutureUtils;
 import org.opensearch.core.action.ActionListener;
@@ -68,7 +69,7 @@ protected AbstractRetryableWorkflowStep(
     protected void retryableGetMlTask(
         String workflowId,
         String nodeId,
-        CompletableFuture<WorkflowData> future,
+        PlainActionFuture<WorkflowData> future,
         String taskId,
         String workflowStep,
         ActionListener<MLTask> mlTaskListener

diff --git a/src/main/java/org/opensearch/flowframework/workflow/CreateConnectorStep.java b/src/main/java/org/opensearch/flowframework/workflow/CreateConnectorStep.java
@@ -11,6 +11,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
+import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
@@ -33,7 +34,6 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
-import java.util.concurrent.CompletableFuture;
 
 import static org.opensearch.flowframework.common.CommonValue.ACTIONS_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.CREDENTIAL_FIELD;
@@ -70,13 +70,13 @@ public CreateConnectorStep(MachineLearningNodeClient mlClient, FlowFrameworkIndi
 
     // TODO: need to add retry conflicts here
     @Override
-    public CompletableFuture<WorkflowData> execute(
+    public PlainActionFuture<WorkflowData> execute(
         String currentNodeId,
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs
     ) {
-        CompletableFuture<WorkflowData> createConnectorFuture = new CompletableFuture<>();
+        PlainActionFuture<WorkflowData> createConnectorFuture = PlainActionFuture.newFuture();
 
         ActionListener<MLCreateConnectorResponse> actionListener = new ActionListener<>() {
 
@@ -93,7 +93,7 @@ public void onResponse(MLCreateConnectorResponse mlCreateConnectorResponse) {
                         mlCreateConnectorResponse.getConnectorId(),
                         ActionListener.wrap(response -> {
                             logger.info("successfully updated resources created in state index: {}", response.getIndex());
-                            createConnectorFuture.complete(
+                            createConnectorFuture.onResponse(
                                 new WorkflowData(
                                     Map.ofEntries(Map.entry(resourceName, mlCreateConnectorResponse.getConnectorId())),
                                     currentNodeInputs.getWorkflowId(),
@@ -102,22 +102,22 @@ public void onResponse(MLCreateConnectorResponse mlCreateConnectorResponse) {
                             );
                         }, exception -> {
                             logger.error("Failed to update new created resource", exception);
-                            createConnectorFuture.completeExceptionally(
+                            createConnectorFuture.onFailure(
                                 new FlowFrameworkException(exception.getMessage(), ExceptionsHelper.status(exception))
                             );
                         })
                     );
 
                 } catch (Exception e) {
                     logger.error("Failed to parse and update new created resource", e);
-                    createConnectorFuture.completeExceptionally(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
+                    createConnectorFuture.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
                 }
             }
 
             @Override
             public void onFailure(Exception e) {
                 logger.error("Failed to create connector");
-                createConnectorFuture.completeExceptionally(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
+                createConnectorFuture.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
             }
         };
 
@@ -171,7 +171,7 @@ public void onFailure(Exception e) {
 
             mlClient.createConnector(mlInput, actionListener);
         } catch (FlowFrameworkException e) {
-            createConnectorFuture.completeExceptionally(e);
+            createConnectorFuture.onFailure(e);
         }
         return createConnectorFuture;
     }

diff --git a/src/main/java/org/opensearch/flowframework/workflow/CreateIndexStep.java b/src/main/java/org/opensearch/flowframework/workflow/CreateIndexStep.java
@@ -13,6 +13,7 @@
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.admin.indices.create.CreateIndexRequest;
 import org.opensearch.action.admin.indices.create.CreateIndexResponse;
+import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.client.Client;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.settings.Settings;
@@ -25,7 +26,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.opensearch.flowframework.common.CommonValue.DEFAULT_MAPPING_OPTION;
@@ -59,13 +59,13 @@ public CreateIndexStep(ClusterService clusterService, Client client, FlowFramewo
     }
 
     @Override
-    public CompletableFuture<WorkflowData> execute(
+    public PlainActionFuture<WorkflowData> execute(
         String currentNodeId,
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs
     ) {
-        CompletableFuture<WorkflowData> createIndexFuture = new CompletableFuture<>();
+        PlainActionFuture<WorkflowData> createIndexFuture = PlainActionFuture.newFuture();
         ActionListener<CreateIndexResponse> actionListener = new ActionListener<>() {
 
             @Override
@@ -80,7 +80,7 @@ public void onResponse(CreateIndexResponse createIndexResponse) {
                         createIndexResponse.index(),
                         ActionListener.wrap(response -> {
                             logger.info("successfully updated resource created in state index: {}", response.getIndex());
-                            createIndexFuture.complete(
+                            createIndexFuture.onResponse(
                                 new WorkflowData(
                                     Map.of(resourceName, createIndexResponse.index()),
                                     currentNodeInputs.getWorkflowId(),
@@ -89,21 +89,21 @@ public void onResponse(CreateIndexResponse createIndexResponse) {
                             );
                         }, exception -> {
                             logger.error("Failed to update new created resource", exception);
-                            createIndexFuture.completeExceptionally(
+                            createIndexFuture.onFailure(
                                 new FlowFrameworkException(exception.getMessage(), ExceptionsHelper.status(exception))
                             );
                         })
                     );
                 } catch (Exception e) {
                     logger.error("Failed to parse and update new created resource", e);
-                    createIndexFuture.completeExceptionally(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
+                    createIndexFuture.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
                 }
             }
 
             @Override
             public void onFailure(Exception e) {
                 logger.error("Failed to create an index", e);
-                createIndexFuture.completeExceptionally(e);
+                createIndexFuture.onFailure(e);
             }
         };
 
@@ -128,7 +128,7 @@ public void onFailure(Exception e) {
             }
         } catch (Exception e) {
             logger.error("Failed to find the correct resource for the workflow step", e);
-            createIndexFuture.completeExceptionally(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
+            createIndexFuture.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
         }
 
         // TODO: