fix(analytics): preserve profile context across async listener thread

QueryPlanExecutor (analytics-framework 3.7+) fires its listener on a
worker thread separate from the one that activated QueryProfiling.
Without explicit handoff, the listener-side formatter reads an empty
profile context, and the synchronous UnifiedQueryContext.close() via
try-with-resources fires before async work finishes — closing the
Calcite connection while it is still in use and clearing profile state
before the listener can read it.

- AnalyticsExecutionEngine.execute: capture the active ProfileContext
  and restore it via QueryProfiling.withCurrentContext inside the
  response callback so metric reads/writes see the same context.
- RestUnifiedQueryAction.execute: defer UnifiedQueryContext.close to
  listener completion via ActionListener.runAfter, wrapped in a
  named wrapWithContextClose helper.
- Add regression test simulating async listener delivery.

Refs: #5317
Signed-off-by: Chen Dai <daichen@amazon.com>

Author: dai-chen

core/src/main/java/org/opensearch/sql/executor/analytics/AnalyticsExecutionEngine.java

@@ -33,7 +33,7 @@
 import org.opensearch.sql.executor.ExecutionEngine;
 import org.opensearch.sql.executor.pagination.Cursor;
 import org.opensearch.sql.monitor.profile.MetricName;
-import org.opensearch.sql.monitor.profile.ProfileMetric;
+import org.opensearch.sql.monitor.profile.ProfileContext;
 import org.opensearch.sql.monitor.profile.QueryProfiling;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 
@@ -81,7 +81,7 @@ public void execute(
     // to a worker pool and results arrive on the listener. Record the execute metric in the
     // listener callback, before delegating to the user-supplied listener, so the metric snapshot
     // taken by SimpleJsonResponseFormatter sees the correct value.
-    ProfileMetric execMetric = QueryProfiling.current().getOrCreateMetric(MetricName.EXECUTE);
+    ProfileContext profileCtx = QueryProfiling.current();
     long execStart = System.nanoTime();
 
     planExecutor.execute(
@@ -90,15 +90,22 @@ public void execute(
         new ActionListener<>() {
           @Override
           public void onResponse(Iterable<Object[]> rows) {
-            try {
-              List<RelDataTypeField> fields = plan.getRowType().getFieldList();
-              List<ExprValue> results = convertRows(rows, fields);
-              Schema schema = buildSchema(fields);
-              execMetric.set(System.nanoTime() - execStart);
-              listener.onResponse(new QueryResponse(schema, results, Cursor.None));
-            } catch (Exception e) {
-              listener.onFailure(e);
-            }
+            QueryProfiling.withCurrentContext(
+                profileCtx,
+                () -> {
+                  try {
+                    List<RelDataTypeField> fields = plan.getRowType().getFieldList();
+                    List<ExprValue> results = convertRows(rows, fields);
+                    Schema schema = buildSchema(fields);
+                    profileCtx
+                        .getOrCreateMetric(MetricName.EXECUTE)
+                        .set(System.nanoTime() - execStart);
+                    listener.onResponse(new QueryResponse(schema, results, Cursor.None));
+                  } catch (Exception e) {
+                    listener.onFailure(e);
+                  }
+                  return null;
+                });
           }
 
           @Override

core/src/test/java/org/opensearch/sql/executor/analytics/AnalyticsExecutionEngineTest.java

@@ -7,6 +7,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
@@ -37,6 +38,8 @@
 import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse;
 import org.opensearch.sql.executor.ExecutionEngine.QueryResponse;
+import org.opensearch.sql.monitor.profile.ProfileContext;
+import org.opensearch.sql.monitor.profile.QueryProfiling;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 
 class AnalyticsExecutionEngineTest {
@@ -330,6 +333,49 @@ void physicalPlanExplain_callsOnFailure() {
             + errorRef.get().getMessage());
   }
 
+  @Test
+  @SuppressWarnings("unchecked")
+  void executeRelNode_profilePreservedOnAsyncListener() throws Exception {
+    ProfileContext expected = QueryProfiling.activate(true);
+
+    RelNode relNode = mockRelNode("col", SqlTypeName.VARCHAR);
+    Iterable<Object[]> rows = Collections.singletonList(new Object[] {"value"});
+
+    // Fire the executor's listener on a different thread to simulate async dispatch
+    doAnswer(
+            inv -> {
+              ActionListener<Iterable<Object[]>> al = inv.getArgument(2);
+              Thread t = new Thread(() -> al.onResponse(rows));
+              t.start();
+              t.join();
+              return null;
+            })
+        .when(mockExecutor)
+        .execute(eq(relNode), any(), any(ActionListener.class));
+
+    AtomicReference<ProfileContext> seen = new AtomicReference<>();
+    engine.execute(
+        relNode,
+        mockContext,
+        new ResponseListener<>() {
+          @Override
+          public void onResponse(QueryResponse response) {
+            seen.set(QueryProfiling.current());
+          }
+
+          @Override
+          public void onFailure(Exception e) {
+            throw new AssertionError(e);
+          }
+        });
+
+    try {
+      assertSame(expected, seen.get(), "Profile context not restored on async listener thread");
+    } finally {
+      QueryProfiling.clear();
+    }
+  }
+
   // --- helpers ---
 
   private QueryResponse executeAndCapture(RelNode relNode) {

plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java

@@ -118,15 +118,18 @@ public void execute(
         .schedule(
             withCurrentContext(
                 () -> {
-                  try (UnifiedQueryContext context = buildContext(queryType, profiling)) {
+                  UnifiedQueryContext context = buildContext(queryType, profiling);
+                  ActionListener<TransportPPLQueryResponse> closingListener =
+                      wrapWithContextClose(context, listener);
+                  try {
                     UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context);
                     RelNode plan = planner.plan(query);
                     CalcitePlanContext planContext = context.getPlanContext();
                     plan = addQuerySizeLimit(plan, planContext);
                     analyticsEngine.execute(
-                        plan, planContext, createQueryListener(queryType, listener));
+                        plan, planContext, createQueryListener(queryType, closingListener));
                   } catch (Exception e) {
-                    listener.onFailure(e);
+                    closingListener.onFailure(e);
                   }
                 }),
             new TimeValue(0),
@@ -256,4 +259,17 @@ private static Runnable withCurrentContext(final Runnable task) {
       task.run();
     };
   }
+
+  private static ActionListener<TransportPPLQueryResponse> wrapWithContextClose(
+      UnifiedQueryContext context, ActionListener<TransportPPLQueryResponse> delegate) {
+    return ActionListener.runAfter(
+        delegate,
+        () -> {
+          try {
+            context.close();
+          } catch (Exception e) {
+            LOG.warn("Failed to close query context", e);
+          }
+        });
+  }
 }