From ab6c19ed9d5cf14c5f00e4d3190ec02c1dd6257b Mon Sep 17 00:00:00 2001
From: Bhavana Ramaram <rbhavna@amazon.com>
Date: Wed, 4 Oct 2023 17:07:05 -0700
Subject: [PATCH] fix reindex issue in TransportUpdateModelGroupAction

Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com>
---
 .../TransportUpdateModelGroupAction.java       | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/plugin/src/main/java/org/opensearch/ml/action/model_group/TransportUpdateModelGroupAction.java b/plugin/src/main/java/org/opensearch/ml/action/model_group/TransportUpdateModelGroupAction.java
index a0d87b3d64..68768eb54b 100644
--- a/plugin/src/main/java/org/opensearch/ml/action/model_group/TransportUpdateModelGroupAction.java
+++ b/plugin/src/main/java/org/opensearch/ml/action/model_group/TransportUpdateModelGroupAction.java
@@ -91,6 +91,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLUpda
         User user = RestActionUtils.getUserContext(client);
         GetRequest getModelGroupRequest = new GetRequest(ML_MODEL_GROUP_INDEX).id(modelGroupId);
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
+            ActionListener<MLUpdateModelGroupResponse> wrappedListener = ActionListener.runBefore(listener, () -> context.restore());
             client.get(getModelGroupRequest, ActionListener.wrap(modelGroup -> {
                 if (modelGroup.isExists()) {
                     try (
@@ -104,20 +105,20 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLUpda
                         } else {
                             validateSecurityDisabledOrModelAccessControlDisabled(updateModelGroupInput);
                         }
-                        updateModelGroup(modelGroupId, modelGroup.getSource(), updateModelGroupInput, listener, user);
+                        updateModelGroup(modelGroupId, modelGroup.getSource(), updateModelGroupInput, wrappedListener, user);
                     } catch (Exception e) {
                         log.error("Failed to parse ml model group" + modelGroup.getId(), e);
-                        listener.onFailure(e);
+                        wrappedListener.onFailure(e);
                     }
                 } else {
-                    listener.onFailure(new OpenSearchStatusException("Failed to find model group", RestStatus.NOT_FOUND));
+                    wrappedListener.onFailure(new OpenSearchStatusException("Failed to find model group", RestStatus.NOT_FOUND));
                 }
             }, e -> {
                 if (e instanceof IndexNotFoundException) {
-                    listener.onFailure(new MLResourceNotFoundException("Fail to find model group"));
+                    wrappedListener.onFailure(new MLResourceNotFoundException("Fail to find model group"));
                 } else {
                     logException("Failed to get model group", e, log);
-                    listener.onFailure(e);
+                    wrappedListener.onFailure(e);
                 }
             }));
         } catch (Exception e) {
@@ -188,15 +189,16 @@ private void updateModelGroup(String modelGroupId, Map<String, Object> source, A
         UpdateRequest updateModelGroupRequest = new UpdateRequest();
         updateModelGroupRequest.index(ML_MODEL_GROUP_INDEX).id(modelGroupId).doc(source);
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
+            ActionListener<MLUpdateModelGroupResponse> wrappedListener = ActionListener.runBefore(listener, () -> context.restore());
             client
                 .update(
                     updateModelGroupRequest,
-                    ActionListener.wrap(r -> { listener.onResponse(new MLUpdateModelGroupResponse("Updated")); }, e -> {
+                    ActionListener.wrap(r -> { wrappedListener.onResponse(new MLUpdateModelGroupResponse("Updated")); }, e -> {
                         if (e instanceof IndexNotFoundException) {
-                            listener.onFailure(new MLResourceNotFoundException("Fail to find model group"));
+                            wrappedListener.onFailure(new MLResourceNotFoundException("Fail to find model group"));
                         } else {
                             log.error("Failed to update model group", e, log);
-                            listener.onFailure(new MLValidationException("Failed to update Model Group"));
+                            wrappedListener.onFailure(new MLValidationException("Failed to update Model Group"));
                         }
                     })
                 );