From 6e4576f01b37e59b4cc4142004479f444a7b68c5 Mon Sep 17 00:00:00 2001
From: Owais Kazi <owaiskazi19@gmail.com>
Date: Tue, 7 Nov 2023 14:53:29 -0800
Subject: [PATCH] Handled failure for throttling

Signed-off-by: Owais Kazi <owaiskazi19@gmail.com>
---
 .../CreateWorkflowTransportAction.java        | 21 ++++++++++---------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
index be100fbce..80f8f8a2d 100644
--- a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
+++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
@@ -111,7 +111,6 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
         }
 
         if (request.getWorkflowId() == null) {
-
             // Throttle incoming requests
             QueryBuilder query = QueryBuilders.matchAllQuery();
             TimeValue requestTimeOut = request.getRequestTimeout();
@@ -119,13 +118,12 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
 
             SearchRequest searchRequest = new SearchRequest(CommonValue.GLOBAL_CONTEXT_INDEX).source(searchSourceBuilder);
 
-            client.search(
-                searchRequest,
-                ActionListener.wrap(
-                    response -> onSearchGlobalContext(response, listener, request.getMaxWorkflows()),
-                    exception -> listener.onFailure(exception)
-                )
-            );
+            client.search(searchRequest, ActionListener.wrap(response -> {
+                if (!onSearchGlobalContext(response, listener, request.getMaxWorkflows())) {
+                    logger.error("Failed to save use case template");
+                    return;
+                }
+            }, exception -> listener.onFailure(exception)));
 
             // Create new global context and state index entries
             flowFrameworkIndicesHandler.putTemplateToGlobalContext(templateWithUser, ActionListener.wrap(globalContextResponse -> {
@@ -195,12 +193,15 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
      * @param listener ActionListener of the SearchRequest
      * @param maxWorkflow max workflows
      */
-    protected void onSearchGlobalContext(SearchResponse response, ActionListener listener, Integer maxWorkflow) {
+    protected Boolean onSearchGlobalContext(SearchResponse response, ActionListener listener, Integer maxWorkflow) {
         if (response.getHits().getTotalHits().value >= maxWorkflow) {
             String errorMessage = "Maximum workflows limit reached " + maxWorkflow;
             logger.error(errorMessage);
-            listener.onFailure(new FlowFrameworkException(errorMessage, RestStatus.BAD_REQUEST));
+            FlowFrameworkException ffe = new FlowFrameworkException(errorMessage, RestStatus.BAD_REQUEST);
+            listener.onFailure(ffe);
+            return false;
         }
+        return true;
     }
 
     private void validateWorkflows(Template template) throws Exception {