Update tests for new update async code

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Jan 9, 2025 · 0583004 · 0583004
1 parent 13d1b83
commit 0583004
Show file tree

Hide file tree

Showing 5 changed files with 190 additions and 203 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java
@@ -251,21 +251,21 @@ public List<ExecutorBuilder<?>> getExecutorBuilders(Settings settings) {
             new ScalingExecutorBuilder(
                 WORKFLOW_THREAD_POOL,
                 1,
-                Math.max(2, OpenSearchExecutors.allocatedProcessors(settings) - 1),
+                Math.max(4, OpenSearchExecutors.allocatedProcessors(settings) - 1),
                 TimeValue.timeValueMinutes(1),
                 FLOW_FRAMEWORK_THREAD_POOL_PREFIX + WORKFLOW_THREAD_POOL
             ),
             new ScalingExecutorBuilder(
                 PROVISION_WORKFLOW_THREAD_POOL,
                 1,
-                Math.max(4, OpenSearchExecutors.allocatedProcessors(settings) - 1),
+                Math.max(8, OpenSearchExecutors.allocatedProcessors(settings) - 1),
                 TimeValue.timeValueMinutes(5),
                 FLOW_FRAMEWORK_THREAD_POOL_PREFIX + PROVISION_WORKFLOW_THREAD_POOL
             ),
             new ScalingExecutorBuilder(
                 DEPROVISION_WORKFLOW_THREAD_POOL,
                 1,
-                Math.max(2, OpenSearchExecutors.allocatedProcessors(settings) - 1),
+                Math.max(4, OpenSearchExecutors.allocatedProcessors(settings) - 1),
                 TimeValue.timeValueMinutes(1),
                 FLOW_FRAMEWORK_THREAD_POOL_PREFIX + DEPROVISION_WORKFLOW_THREAD_POOL
             )

diff --git a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java
@@ -12,6 +12,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.message.ParameterizedMessageFactory;
 import org.opensearch.ExceptionsHelper;
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.DocWriteRequest.OpType;
 import org.opensearch.action.admin.indices.create.CreateIndexRequest;
 import org.opensearch.action.admin.indices.create.CreateIndexResponse;
@@ -47,7 +48,6 @@
 import org.opensearch.flowframework.util.ParseUtils;
 import org.opensearch.flowframework.workflow.WorkflowData;
 import org.opensearch.index.IndexNotFoundException;
-import org.opensearch.index.engine.VersionConflictEngineException;
 import org.opensearch.remote.metadata.client.DeleteDataObjectRequest;
 import org.opensearch.remote.metadata.client.GetDataObjectRequest;
 import org.opensearch.remote.metadata.client.PutDataObjectRequest;
@@ -845,7 +845,7 @@ public void addResourceToStateIndex(
         String resourceId,
         ActionListener<WorkflowData> listener
     ) {
-        addResourceToStateIndex(currentNodeInputs, nodeId, workflowStepName, resourceId, null, listener);
+        addResourceToStateIndex(currentNodeInputs, nodeId, workflowStepName, resourceId, "fakeTenantId", listener);
     }
 
     /**
@@ -892,7 +892,7 @@ public void addResourceToStateIndex(
      */
     @Deprecated
     public void deleteResourceFromStateIndex(String workflowId, ResourceCreated resourceToDelete, ActionListener<WorkflowData> listener) {
-        deleteResourceFromStateIndex(workflowId, null, resourceToDelete, listener);
+        deleteResourceFromStateIndex(workflowId, "fakeTenantId", resourceToDelete, listener);
     }
 
     /**
@@ -967,7 +967,7 @@ private void getAndUpdateResourceInStateDocumentWithRetries(
                 try {
                     GetResponse getResponse = GetResponse.fromXContent(r.parser());
                     handleStateGetResponse(workflowId, tenantId, resource, operation, retries, listener, getResponse);
-                } catch (IOException e) {
+                } catch (Exception e) {
                     logger.error("Failed to parse get response", e);
                     listener.onFailure(new FlowFrameworkException("Failed to parse get response", INTERNAL_SERVER_ERROR));
                 }
@@ -991,17 +991,16 @@ private void handleStateGetResponse(
             listener.onFailure(new FlowFrameworkException("Workflow state not found for " + workflowId, RestStatus.NOT_FOUND));
             return;
         }
-        WorkflowState currentState;
         try {
-            currentState = WorkflowState.parse(getResponse.getSourceAsString());
+            WorkflowState currentState = WorkflowState.parse(getResponse.getSourceAsString());
             List<ResourceCreated> resourcesCreated = new ArrayList<>(currentState.resourcesCreated());
             if (operation == OpType.DELETE) {
                 resourcesCreated.removeIf(r -> r.resourceMap().equals(resource.resourceMap()));
             } else {
                 resourcesCreated.add(resource);
             }
             WorkflowState newState = WorkflowState.builder(currentState).resourcesCreated(resourcesCreated).build();
-            UpdateDataObjectRequest updateRequest2 = UpdateDataObjectRequest.builder()
+            UpdateDataObjectRequest updateRequest = UpdateDataObjectRequest.builder()
                 .index(WORKFLOW_STATE_INDEX)
                 .id(workflowId)
                 .tenantId(tenantId)
@@ -1010,7 +1009,7 @@ private void handleStateGetResponse(
                 .ifPrimaryTerm(getResponse.getPrimaryTerm())
                 .build();
             sdkClient.updateDataObjectAsync(
-                updateRequest2,
+                updateRequest,
                 client.threadPool().executor(operation == OpType.DELETE ? DEPROVISION_WORKFLOW_THREAD_POOL : PROVISION_WORKFLOW_THREAD_POOL)
             ).whenComplete((r, throwable) -> {
                 if (throwable == null) {
@@ -1059,7 +1058,7 @@ private void handleStateUpdateException(
         ActionListener<WorkflowData> listener,
         Exception e
     ) {
-        if (e instanceof VersionConflictEngineException && retries > 0) {
+        if (e instanceof OpenSearchStatusException && ((OpenSearchStatusException) e).status() == RestStatus.CONFLICT && retries > 0) {
             // Retry if we haven't exhausted retries
             getAndUpdateResourceInStateDocumentWithRetries(workflowId, tenantId, newResource, operation, retries - 1, listener);
             return;

diff --git a/src/main/java/org/opensearch/flowframework/transport/DeprovisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/DeprovisionWorkflowTransportAction.java
@@ -246,7 +246,7 @@ private void executeDeprovisionSequence(
                     deprovisionFuture.get();
                     logger.info("Successful {} for {}", deprovisionNode.id(), resourceNameAndId);
                     // Remove from state index resource list
-                    flowFrameworkIndicesHandler.deleteResourceFromStateIndex(workflowId, resource, stateUpdateFuture);
+                    flowFrameworkIndicesHandler.deleteResourceFromStateIndex(workflowId, tenantId, resource, stateUpdateFuture);
                     try {
                         // Wait at most 1 second for state index update.
                         stateUpdateFuture.actionGet(1, TimeUnit.SECONDS);