Fix EXECUTE profiling metric by recording inside AnalyticsExecutionEngine

ahkcs · ahkcs · commit ba3efd771ab3 · 2026-03-30T16:45:06.000-07:00
Move EXECUTE metric recording into AnalyticsExecutionEngine.execute(),
between the actual execution (planExecutor + row conversion) and the
listener.onResponse() call. This ensures the metric is written before
SimpleJsonResponseFormatter calls QueryProfiling.finish() to snapshot.

Previously context.measure() was used in RestUnifiedQueryAction, but
finish() was called inside the listener callback (synchronously) before
measure()'s finally block could write the metric, resulting in 0ms.

Add IT assertion that execute phase time_ms &gt; 0 to catch this bug.

Signed-off-by: Kai Huang &lt;kaihuang@amazon.com&gt;
Signed-off-by: Kai Huang &lt;ahkcs@amazon.com&gt;
diff --git a/core/src/main/java/org/opensearch/sql/executor/analytics/AnalyticsExecutionEngine.java b/core/src/main/java/org/opensearch/sql/executor/analytics/AnalyticsExecutionEngine.java
@@ -24,6 +24,9 @@
 import org.opensearch.sql.executor.ExecutionContext;
 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.QueryProfiling;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 
 /**
@@ -67,12 +70,20 @@ public void explain(PhysicalPlan plan, ResponseListener<ExplainResponse> listene
   public void execute(
       RelNode plan, CalcitePlanContext context, ResponseListener<QueryResponse> listener) {
     try {
+      // Record EXECUTE metric before calling listener, because the listener's onResponse
+      // triggers SimpleJsonResponseFormatter which calls QueryProfiling.finish() to snapshot
+      // all metrics. The metric must be written before that snapshot.
+      ProfileMetric execMetric = QueryProfiling.current().getOrCreateMetric(MetricName.EXECUTE);
+      long execStart = System.nanoTime();
+
       Iterable<Object[]> rows = planExecutor.execute(plan, null);
 
       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);
diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/AnalyticsPPLIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/AnalyticsPPLIT.java
@@ -131,6 +131,13 @@ public void testProfileResponseIncludesProfilingData() throws IOException {
     JSONObject result = new JSONObject(getResponseBody(response, true));
 
     assertTrue("Response should have 'profile' field when profile=true", result.has("profile"));
+    JSONObject profile = result.getJSONObject("profile");
+    assertTrue("Profile should have 'phases' field", profile.has("phases"));
+    JSONObject phases = profile.getJSONObject("phases");
+    assertTrue("Phases should have 'execute' field", phases.has("execute"));
+    double executeTime = phases.getJSONObject("execute").getDouble("time_ms");
+    assertTrue(
+        "Execute phase should have non-zero time, got " + executeTime + "ms", executeTime > 0);
   }
 
   // --- Error handling tests ---
diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java
@@ -31,7 +31,6 @@
 import org.opensearch.sql.executor.analytics.AnalyticsExecutionEngine;
 import org.opensearch.sql.executor.analytics.QueryPlanExecutor;
 import org.opensearch.sql.lang.LangSpec;
-import org.opensearch.sql.monitor.profile.MetricName;
 import org.opensearch.sql.plugin.rest.analytics.stub.StubSchemaProvider;
 import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse;
 import org.opensearch.sql.ppl.domain.PPLQueryRequest;
@@ -123,14 +122,7 @@ private void doExecute(
       CalcitePlanContext planContext = context.getPlanContext();
       plan = addQuerySizeLimit(plan, planContext);
 
-      RelNode finalPlan = plan;
-      context.measure(
-          MetricName.EXECUTE,
-          () -> {
-            analyticsEngine.execute(
-                finalPlan, planContext, createQueryListener(queryType, listener));
-            return null;
-          });
+      analyticsEngine.execute(plan, planContext, createQueryListener(queryType, listener));
     } catch (Exception e) {
       listener.onFailure(e);
     }
@@ -148,13 +140,7 @@ private void doExplain(
       CalcitePlanContext planContext = context.getPlanContext();
       plan = addQuerySizeLimit(plan, planContext);
 
-      RelNode finalPlan = plan;
-      context.measure(
-          MetricName.EXECUTE,
-          () -> {
-            analyticsEngine.explain(finalPlan, pplRequest.mode(), planContext, listener);
-            return null;
-          });
+      analyticsEngine.explain(plan, pplRequest.mode(), planContext, listener);
     } catch (Exception e) {
       listener.onFailure(e);
     }
