From 4f5510a8902c457ae4a31aeced1d8d7ab59c291b Mon Sep 17 00:00:00 2001
From: Daniel Widdis <widdis@gmail.com>
Date: Thu, 15 Aug 2024 15:18:15 -0700
Subject: [PATCH 1/2] Delay context restoring for initializing master key and
 deleting agent

Signed-off-by: Daniel Widdis <widdis@gmail.com>
---
 .../ml/engine/encryptor/EncryptorImpl.java       |  2 +-
 .../agents/DeleteAgentTransportAction.java       | 16 +++++++---------
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java
index 329459973e..5615236f95 100644
--- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java
+++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java
@@ -187,7 +187,6 @@ private void handleGetDataObjectResponse(
         AtomicReference<Exception> exceptionRef,
         CountDownLatch latch
     ) {
-        context.restore();
         log.debug("Completed Get MASTER_KEY Request, for tenant id:{}", tenantId);
 
         if (throwable != null) {
@@ -195,6 +194,7 @@ private void handleGetDataObjectResponse(
         } else {
             handleGetDataObjectSuccess(response, tenantId, exceptionRef, latch, context);
         }
+        context.restore();
     }
 
     private void handleGetDataObjectFailure(Throwable throwable, AtomicReference<Exception> exceptionRef, CountDownLatch latch) {
diff --git a/plugin/src/main/java/org/opensearch/ml/action/agents/DeleteAgentTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/agents/DeleteAgentTransportAction.java
index 8e702d7186..cbe6cd0e48 100644
--- a/plugin/src/main/java/org/opensearch/ml/action/agents/DeleteAgentTransportAction.java
+++ b/plugin/src/main/java/org/opensearch/ml/action/agents/DeleteAgentTransportAction.java
@@ -99,9 +99,9 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
             sdkClient
                 .getDataObjectAsync(getDataObjectRequest, client.threadPool().executor(GENERAL_THREAD_POOL))
                 .whenComplete((r, throwable) -> {
-                    context.restore();
                     log.debug("Completed Get Agent Request, Agent id:{}", agentId);
                     if (throwable != null) {
+                        context.restore();
                         Exception cause = SdkClientUtils.unwrapAndConvertToException(throwable);
                         if (cause instanceof IndexNotFoundException) {
                             log.info("Failed to get Agent index", cause);
@@ -124,6 +124,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
                                     if (TenantAwareHelper
                                         .validateTenantResource(mlFeatureEnabledSetting, tenantId, mlAgent.getTenantId(), actionListener)) {
                                         if (mlAgent.getIsHidden() && !isSuperAdmin) {
+                                            context.restore();
                                             actionListener
                                                 .onFailure(
                                                     new OpenSearchStatusException(
@@ -144,15 +145,12 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
                                                             .build(),
                                                         client.threadPool().executor(GENERAL_THREAD_POOL)
                                                     )
-                                                    .whenComplete(
-                                                        (response, delThrowable) -> handleDeleteResponse(
-                                                            response,
-                                                            delThrowable,
-                                                            tenantId,
-                                                            actionListener
-                                                        )
-                                                    );
+                                                    .whenComplete((response, delThrowable) -> {
+                                                        context.restore();
+                                                        handleDeleteResponse(response, delThrowable, tenantId, actionListener);
+                                                    });
                                             } catch (Exception e) {
+                                                context.restore();
                                                 log.error("Failed to delete ML agent: {}", agentId, e);
                                                 actionListener.onFailure(e);
                                             }

From 83510046add7a3186b733a3acef880b6fa25fb8a Mon Sep 17 00:00:00 2001
From: Daniel Widdis <widdis@gmail.com>
Date: Tue, 17 Sep 2024 06:19:47 -0700
Subject: [PATCH 2/2] Simplify restoring context with runBefore

Signed-off-by: Daniel Widdis <widdis@gmail.com>
---
 .../agents/DeleteAgentTransportAction.java    | 29 ++++++++++---------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/plugin/src/main/java/org/opensearch/ml/action/agents/DeleteAgentTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/agents/DeleteAgentTransportAction.java
index cbe6cd0e48..f9d3d8a469 100644
--- a/plugin/src/main/java/org/opensearch/ml/action/agents/DeleteAgentTransportAction.java
+++ b/plugin/src/main/java/org/opensearch/ml/action/agents/DeleteAgentTransportAction.java
@@ -95,20 +95,19 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
             .build();
 
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
-
+            ActionListener<DeleteResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
             sdkClient
                 .getDataObjectAsync(getDataObjectRequest, client.threadPool().executor(GENERAL_THREAD_POOL))
                 .whenComplete((r, throwable) -> {
                     log.debug("Completed Get Agent Request, Agent id:{}", agentId);
                     if (throwable != null) {
-                        context.restore();
                         Exception cause = SdkClientUtils.unwrapAndConvertToException(throwable);
                         if (cause instanceof IndexNotFoundException) {
                             log.info("Failed to get Agent index", cause);
-                            actionListener.onFailure(new OpenSearchStatusException("Failed to get agent index", RestStatus.NOT_FOUND));
+                            wrappedListener.onFailure(new OpenSearchStatusException("Failed to get agent index", RestStatus.NOT_FOUND));
                         } else {
                             log.error("Failed to get ML Agent {}", agentId, cause);
-                            actionListener.onFailure(cause);
+                            wrappedListener.onFailure(cause);
                         }
                     } else {
                         try {
@@ -122,10 +121,14 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
                                     ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
                                     MLAgent mlAgent = MLAgent.parse(parser);
                                     if (TenantAwareHelper
-                                        .validateTenantResource(mlFeatureEnabledSetting, tenantId, mlAgent.getTenantId(), actionListener)) {
+                                        .validateTenantResource(
+                                            mlFeatureEnabledSetting,
+                                            tenantId,
+                                            mlAgent.getTenantId(),
+                                            wrappedListener
+                                        )) {
                                         if (mlAgent.getIsHidden() && !isSuperAdmin) {
-                                            context.restore();
-                                            actionListener
+                                            wrappedListener
                                                 .onFailure(
                                                     new OpenSearchStatusException(
                                                         "User doesn't have privilege to perform this operation on this agent",
@@ -146,26 +149,24 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
                                                         client.threadPool().executor(GENERAL_THREAD_POOL)
                                                     )
                                                     .whenComplete((response, delThrowable) -> {
-                                                        context.restore();
-                                                        handleDeleteResponse(response, delThrowable, tenantId, actionListener);
+                                                        handleDeleteResponse(response, delThrowable, tenantId, wrappedListener);
                                                     });
                                             } catch (Exception e) {
-                                                context.restore();
                                                 log.error("Failed to delete ML agent: {}", agentId, e);
-                                                actionListener.onFailure(e);
+                                                wrappedListener.onFailure(e);
                                             }
                                         }
                                     }
                                 } catch (Exception e) {
                                     log.error("Failed to parse ml agent {}", agentId, e);
-                                    actionListener.onFailure(e);
+                                    wrappedListener.onFailure(e);
                                 }
                             } else {
-                                actionListener.onFailure(new OpenSearchStatusException("Fail to find ml agent", RestStatus.NOT_FOUND));
+                                wrappedListener.onFailure(new OpenSearchStatusException("Fail to find ml agent", RestStatus.NOT_FOUND));
                             }
                         } catch (Exception e) {
                             log.error("Failed to delete ML agent: {}", agentId, e);
-                            actionListener.onFailure(e);
+                            wrappedListener.onFailure(e);
                         }
                     }
                 });