restore context

Author: dnhatn

Diff for: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java

@@ -98,6 +98,7 @@ void startComputeOnDataNodes(
         Runnable runOnTaskFailure,
         ActionListener<ComputeResponse> outListener
     ) {
+        var storedContext = transportService.getThreadPool().getThreadContext().newStoredContext();
         new DataNodeRequestSender(
             transportService,
             esqlExecutor,
@@ -141,7 +142,7 @@ protected void sendRequest(
                         final Runnable onGroupFailure;
                         final CancellableTask groupTask;
                         if (configuration.allowPartialResults()) {
-                            try {
+                            try (storedContext) {
                                 groupTask = computeService.createGroupTask(
                                     parentTask,
                                     () -> "compute group: data-node [" + node.getName() + "], " + shardIds + " [" + shardIds + "]"