From dedaefc1f41a5daff79dc7d30dd2eed135592628 Mon Sep 17 00:00:00 2001 From: "opensearch-trigger-bot[bot]" <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Date: Tue, 23 Jan 2024 10:36:41 -0800 Subject: [PATCH] stash thread context before running forward action (#1904) (#1907) Signed-off-by: Yaliang Wu (cherry picked from commit a14521dd97397c2ec9f6899e5d7d2f450986a5f9) Co-authored-by: Yaliang Wu --- .../TransportDeployModelOnNodeAction.java | 33 +++++++++++-------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/plugin/src/main/java/org/opensearch/ml/action/deploy/TransportDeployModelOnNodeAction.java b/plugin/src/main/java/org/opensearch/ml/action/deploy/TransportDeployModelOnNodeAction.java index 67f58d69a0..bf8c81756b 100644 --- a/plugin/src/main/java/org/opensearch/ml/action/deploy/TransportDeployModelOnNodeAction.java +++ b/plugin/src/main/java/org/opensearch/ml/action/deploy/TransportDeployModelOnNodeAction.java @@ -22,6 +22,7 @@ import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.xcontent.NamedXContentRegistry; @@ -161,13 +162,15 @@ private MLDeployModelNodeResponse createDeployModelNodeResponse(MLDeployModelNod .build(); MLForwardRequest deployModelDoneMessage = new MLForwardRequest(mlForwardInput); - transportService - .sendRequest( - getNodeById(coordinatingNodeId), - MLForwardAction.NAME, - deployModelDoneMessage, - new ActionListenerResponseHandler<>(taskDoneListener, MLForwardResponse::new) - ); + try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { + transportService + .sendRequest( + getNodeById(coordinatingNodeId), + MLForwardAction.NAME, + deployModelDoneMessage, + new ActionListenerResponseHandler<>(taskDoneListener, MLForwardResponse::new) + ); + } }, e -> { MLForwardInput mlForwardInput = MLForwardInput .builder() @@ -179,13 +182,15 @@ private MLDeployModelNodeResponse createDeployModelNodeResponse(MLDeployModelNod .build(); MLForwardRequest deployModelDoneMessage = new MLForwardRequest(mlForwardInput); - transportService - .sendRequest( - getNodeById(coordinatingNodeId), - MLForwardAction.NAME, - deployModelDoneMessage, - new ActionListenerResponseHandler<>(taskDoneListener, MLForwardResponse::new) - ); + try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { + transportService + .sendRequest( + getNodeById(coordinatingNodeId), + MLForwardAction.NAME, + deployModelDoneMessage, + new ActionListenerResponseHandler<>(taskDoneListener, MLForwardResponse::new) + ); + } }) );