Handled failure for throttling

Signed-off-by: Owais Kazi <[email protected]>
opensearch-project · Nov 7, 2023 · 6e4576f · 6e4576f
1 parent 722f09d
commit 6e4576f
Showing 1 changed file with 11 additions and 10 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
@@ -111,21 +111,19 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
         }
 
         if (request.getWorkflowId() == null) {
-
             // Throttle incoming requests
             QueryBuilder query = QueryBuilders.matchAllQuery();
             TimeValue requestTimeOut = request.getRequestTimeout();
             SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder().query(query).size(0).timeout(requestTimeOut);
 
             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 {