From c50c1283d7ffc6771dca84f4c968f5b3acf8c351 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 4 Aug 2026 17:25:00 -0700 Subject: [PATCH 1/4] feat: Integrate SQL/PPL with query-insights plugin Add query source tracking headers (x-query-source, x-original-query, x-query-execution-id, x-query-phases) to SQL and PPL execution paths so query-insights can identify and track SQL/PPL queries separately from DSL queries. - Add QueryPhaseTracker for tracking parse/analyze/plan phases - Tag thread context with SQL/PPL source headers in transport actions - Add writePhaseHeader() to Calcite execution path for PPL queries - Register all tracking headers as task headers for propagation Signed-off-by: Kishore Kumaar Natarajan --- .../sql/common/utils/QueryPhaseTracker.java | 184 ++++++++++++ .../sql/common/utils/QuerySourceHeaders.java | 23 ++ .../common/utils/QueryPhaseTrackerTest.java | 271 ++++++++++++++++++ .../opensearch/sql/executor/QueryService.java | 27 +- .../sql/legacy/plugin/RestSqlAction.java | 22 ++ .../executor/OpenSearchExecutionEngine.java | 26 ++ .../org/opensearch/sql/plugin/SQLPlugin.java | 9 + .../transport/TransportPPLQueryAction.java | 21 ++ .../org/opensearch/sql/ppl/PPLService.java | 5 + ...2025-07-15-162000-pr-sql-query-insights.md | 99 +++++++ .../org/opensearch/sql/sql/SQLService.java | 5 + 11 files changed, 686 insertions(+), 6 deletions(-) create mode 100644 common/src/main/java/org/opensearch/sql/common/utils/QueryPhaseTracker.java create mode 100644 common/src/main/java/org/opensearch/sql/common/utils/QuerySourceHeaders.java create mode 100644 common/src/test/java/org/opensearch/sql/common/utils/QueryPhaseTrackerTest.java create mode 100644 semantic-review/2025-07-15-162000-pr-sql-query-insights.md diff --git a/common/src/main/java/org/opensearch/sql/common/utils/QueryPhaseTracker.java b/common/src/main/java/org/opensearch/sql/common/utils/QueryPhaseTracker.java new file mode 100644 index 00000000000..7c53d3e8bb5 --- /dev/null +++ b/common/src/main/java/org/opensearch/sql/common/utils/QueryPhaseTracker.java @@ -0,0 +1,184 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.common.utils; + +import java.lang.management.ManagementFactory; +import java.lang.management.ThreadMXBean; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.StringJoiner; + +/** + * Tracks timing and resource usage for SQL/PPL query execution phases. Thread-local instance + * collects phase durations (nanos) as the query flows through parse → analyze → plan → execute. The + * result is serialized into a compact header value for query-insights to consume. + * + *

Phases that complete on the REST thread (parse) are stored in the Log4j ThreadContext via + * {@link #persist()}, which is propagated to the sql-worker thread by OpenSearchQueryManager. On + * the worker thread, {@link #startOrRestore()} restores prior phases and continues tracking. + * + *

CPU time uses {@code ThreadMXBean.getCurrentThreadCpuTime()}. Memory allocation tracking uses + * {@code com.sun.management.ThreadMXBean.getThreadAllocatedBytes()} which is best-effort and may be + * unavailable on non-HotSpot JVMs (e.g., IBM J9, GraalVM native). + */ +public final class QueryPhaseTracker { + + private static final String LOG4J_KEY = "_sql_phase_tracker"; + private static final ThreadLocal CURRENT = new ThreadLocal<>(); + private static final ThreadMXBean THREAD_MX = ManagementFactory.getThreadMXBean(); + private static final boolean CPU_SUPPORTED = THREAD_MX.isCurrentThreadCpuTimeSupported(); + private static final com.sun.management.ThreadMXBean SUN_THREAD_MX = getSunThreadMXBean(); + + private final Map phases = new LinkedHashMap<>(); + private final Map cpuPhases = new LinkedHashMap<>(); + private final Map memPhases = new LinkedHashMap<>(); + private String activePhase; + private long activeStart; + private long activeCpuStart; + private long activeMemStart; + private long overallStart; + + private static com.sun.management.ThreadMXBean getSunThreadMXBean() { + try { + if (THREAD_MX instanceof com.sun.management.ThreadMXBean sun) { + return sun; + } + } catch (NoClassDefFoundError | UnsupportedOperationException e) { + // com.sun.management not available on this JVM + } + return null; + } + + private QueryPhaseTracker() { + this.overallStart = System.nanoTime(); + } + + public static QueryPhaseTracker start() { + QueryPhaseTracker tracker = new QueryPhaseTracker(); + CURRENT.set(tracker); + return tracker; + } + + /** Restore from Log4j ThreadContext (cross-thread transfer) or create fresh. */ + public static QueryPhaseTracker startOrRestore() { + QueryPhaseTracker tracker = new QueryPhaseTracker(); + String stored = org.apache.logging.log4j.ThreadContext.get(LOG4J_KEY); + if (stored != null && !stored.isEmpty()) { + try { + for (String part : stored.split(",")) { + // Format: "phase:wallNanos|cpu:cpuNanos|mem:memBytes" + String[] segments = part.split("\\|"); + String[] kv = segments[0].split(":", 2); + if (kv.length == 2) { + tracker.phases.put(kv[0], Long.parseLong(kv[1])); + for (int i = 1; i < segments.length; i++) { + String[] metric = segments[i].split(":", 2); + if (metric.length == 2 && "cpu".equals(metric[0])) { + tracker.cpuPhases.put(kv[0], Long.parseLong(metric[1])); + } else if (metric.length == 2 && "mem".equals(metric[0])) { + tracker.memPhases.put(kv[0], Long.parseLong(metric[1])); + } + } + } + } + } catch (NumberFormatException e) { + // Malformed data in ThreadContext — start fresh + tracker.phases.clear(); + tracker.cpuPhases.clear(); + tracker.memPhases.clear(); + } + } + CURRENT.set(tracker); + return tracker; + } + + public static QueryPhaseTracker current() { + return CURRENT.get(); + } + + /** Returns true if no phases have been recorded (no prior SQL/PPL context). */ + public boolean isEmpty() { + return phases.isEmpty() && activePhase == null; + } + + public static void clear() { + CURRENT.remove(); + org.apache.logging.log4j.ThreadContext.remove(LOG4J_KEY); + } + + /** Persist current phases to Log4j ThreadContext for cross-thread propagation. */ + public void persist() { + org.apache.logging.log4j.ThreadContext.put(LOG4J_KEY, serialize()); + } + + public void addCompletedPhase(String name, long nanos) { + phases.put(name, nanos); + } + + public void addCompletedPhase(String name, long nanos, long cpuNanos, long memBytes) { + phases.put(name, nanos); + if (cpuNanos > 0) cpuPhases.put(name, cpuNanos); + if (memBytes > 0) memPhases.put(name, memBytes); + } + + public void beginPhase(String name) { + endCurrentPhase(); + activePhase = name; + activeStart = System.nanoTime(); + activeCpuStart = CPU_SUPPORTED ? THREAD_MX.getCurrentThreadCpuTime() : 0; + activeMemStart = + SUN_THREAD_MX != null + ? SUN_THREAD_MX.getThreadAllocatedBytes(Thread.currentThread().getId()) + : 0; + } + + public void endCurrentPhase() { + if (activePhase != null) { + long elapsed = System.nanoTime() - activeStart; + phases.merge(activePhase, elapsed, Long::sum); + if (CPU_SUPPORTED) { + long cpuElapsed = THREAD_MX.getCurrentThreadCpuTime() - activeCpuStart; + cpuPhases.merge(activePhase, cpuElapsed, Long::sum); + } + if (SUN_THREAD_MX != null) { + long memElapsed = + SUN_THREAD_MX.getThreadAllocatedBytes(Thread.currentThread().getId()) - activeMemStart; + memPhases.merge(activePhase, memElapsed, Long::sum); + } + activePhase = null; + } + } + + public void endAll() { + endCurrentPhase(); + long total = System.nanoTime() - overallStart; + phases.put("total", total); + } + + /** + * Serialize to compact format with all metrics per phase: + * "parse:1234|cpu:1000|mem:5000,analyze:5678|cpu:4000|mem:20000,total:18943" Time and CPU in + * nanoseconds, memory in bytes. + */ + public String serialize() { + StringJoiner joiner = new StringJoiner(","); + for (Map.Entry entry : phases.entrySet()) { + String phase = entry.getKey(); + StringBuilder sb = new StringBuilder(); + sb.append(phase).append(':').append(entry.getValue()); + Long cpu = cpuPhases.get(phase); + if (cpu != null) { + sb.append("|cpu:").append(cpu); + } + Long mem = memPhases.get(phase); + if (mem != null) { + sb.append("|mem:").append(mem); + } + joiner.add(sb.toString()); + } + return joiner.toString(); + } +} diff --git a/common/src/main/java/org/opensearch/sql/common/utils/QuerySourceHeaders.java b/common/src/main/java/org/opensearch/sql/common/utils/QuerySourceHeaders.java new file mode 100644 index 00000000000..7075870d19f --- /dev/null +++ b/common/src/main/java/org/opensearch/sql/common/utils/QuerySourceHeaders.java @@ -0,0 +1,23 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.common.utils; + +/** + * Constants for thread context headers used to propagate query source metadata to downstream + * components like query-insights. + */ +public final class QuerySourceHeaders { + + public static final String QUERY_SOURCE_HEADER = "x-query-source"; + public static final String ORIGINAL_QUERY_HEADER = "x-original-query"; + public static final String QUERY_EXECUTION_ID_HEADER = "x-query-execution-id"; + public static final String QUERY_PHASES_HEADER = "x-query-phases"; + + /** Maximum number of characters stored in the {@link #ORIGINAL_QUERY_HEADER}. */ + public static final int MAX_ORIGINAL_QUERY_LENGTH = 4096; + + private QuerySourceHeaders() {} +} diff --git a/common/src/test/java/org/opensearch/sql/common/utils/QueryPhaseTrackerTest.java b/common/src/test/java/org/opensearch/sql/common/utils/QueryPhaseTrackerTest.java new file mode 100644 index 00000000000..9ac49665665 --- /dev/null +++ b/common/src/test/java/org/opensearch/sql/common/utils/QueryPhaseTrackerTest.java @@ -0,0 +1,271 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.common.utils; + +import static org.junit.jupiter.api.Assertions.*; + +import org.apache.logging.log4j.ThreadContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** Unit tests for {@link QueryPhaseTracker}. */ +public class QueryPhaseTrackerTest { + + @AfterEach + void cleanup() { + QueryPhaseTracker.clear(); + } + + @Test + public void testStartCreatesTrackerAndSetsCurrent() { + assertNull(QueryPhaseTracker.current()); + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + assertNotNull(tracker); + assertSame(tracker, QueryPhaseTracker.current()); + } + + @Test + public void testBeginPhaseAndEndCurrentPhase() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + tracker.beginPhase("parse"); + // Simulate some work + consumeTime(); + tracker.endCurrentPhase(); + + String serialized = tracker.serialize(); + assertTrue( + serialized.startsWith("parse:"), + "Expected serialized to start with 'parse:': " + serialized); + // The nanos value should be > 0 + String wallNanos = serialized.split("\\|")[0].split(":")[1]; + assertTrue(Long.parseLong(wallNanos) > 0, "Wall-clock nanos should be positive"); + } + + @Test + public void testMultiplePhasesInSequence() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + + tracker.beginPhase("parse"); + consumeTime(); + tracker.endCurrentPhase(); + + tracker.beginPhase("analyze"); + consumeTime(); + tracker.endCurrentPhase(); + + tracker.beginPhase("plan"); + consumeTime(); + tracker.endCurrentPhase(); + + String serialized = tracker.serialize(); + // All three phases should be present in order + assertTrue(serialized.contains("parse:"), "Missing parse phase: " + serialized); + assertTrue(serialized.contains("analyze:"), "Missing analyze phase: " + serialized); + assertTrue(serialized.contains("plan:"), "Missing plan phase: " + serialized); + + // Verify ordering (LinkedHashMap preserves insertion order) + int parseIdx = serialized.indexOf("parse:"); + int analyzeIdx = serialized.indexOf("analyze:"); + int planIdx = serialized.indexOf("plan:"); + assertTrue(parseIdx < analyzeIdx, "parse should come before analyze"); + assertTrue(analyzeIdx < planIdx, "analyze should come before plan"); + } + + @Test + public void testEndAllAddsTotalPhase() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + + tracker.beginPhase("parse"); + consumeTime(); + tracker.endCurrentPhase(); + + tracker.endAll(); + + String serialized = tracker.serialize(); + assertTrue(serialized.contains("total:"), "Missing total phase: " + serialized); + + // total should be the last entry + String[] parts = serialized.split(","); + assertTrue(parts[parts.length - 1].startsWith("total:"), "total should be last: " + serialized); + } + + @Test + public void testEndAllFinishesActivePhase() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + tracker.beginPhase("execute"); + consumeTime(); + // Don't call endCurrentPhase — endAll should do it + tracker.endAll(); + + String serialized = tracker.serialize(); + assertTrue(serialized.contains("execute:"), "Missing execute phase: " + serialized); + assertTrue(serialized.contains("total:"), "Missing total phase: " + serialized); + } + + @Test + public void testSerializeFormat() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + // Use addCompletedPhase to have deterministic values + tracker.addCompletedPhase("parse", 1000L, 800L, 5000L); + tracker.addCompletedPhase("analyze", 2000L, 1500L, 10000L); + + String serialized = tracker.serialize(); + // Expected: "parse:1000|cpu:800|mem:5000,analyze:2000|cpu:1500|mem:10000" + assertEquals("parse:1000|cpu:800|mem:5000,analyze:2000|cpu:1500|mem:10000", serialized); + } + + @Test + public void testSerializeFormatWithoutCpuAndMem() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + // Overload with only nanos + tracker.addCompletedPhase("parse", 1000L); + tracker.addCompletedPhase("analyze", 2000L); + + String serialized = tracker.serialize(); + // No cpu/mem segments + assertEquals("parse:1000,analyze:2000", serialized); + } + + @Test + public void testSerializeFormatWithZeroCpuAndMem() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + // cpuNanos=0 and memBytes=0 should be omitted + tracker.addCompletedPhase("parse", 1000L, 0L, 0L); + + String serialized = tracker.serialize(); + assertEquals("parse:1000", serialized); + } + + @Test + public void testPersistStoresToThreadContext() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + tracker.addCompletedPhase("parse", 1234L, 1000L, 5000L); + + tracker.persist(); + + String stored = ThreadContext.get("_sql_phase_tracker"); + assertNotNull(stored); + assertEquals("parse:1234|cpu:1000|mem:5000", stored); + } + + @Test + public void testStartOrRestoreRestoresPhasesFromThreadContext() { + // Simulate a prior thread persisting data + ThreadContext.put( + "_sql_phase_tracker", "parse:1234|cpu:1000|mem:5000,analyze:5678|cpu:4000|mem:20000"); + + QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); + assertNotNull(tracker); + assertSame(tracker, QueryPhaseTracker.current()); + + String serialized = tracker.serialize(); + assertTrue(serialized.contains("parse:1234"), "parse phase not restored: " + serialized); + assertTrue(serialized.contains("cpu:1000"), "parse cpu not restored: " + serialized); + assertTrue(serialized.contains("mem:5000"), "parse mem not restored: " + serialized); + assertTrue(serialized.contains("analyze:5678"), "analyze phase not restored: " + serialized); + assertTrue(serialized.contains("cpu:4000"), "analyze cpu not restored: " + serialized); + assertTrue(serialized.contains("mem:20000"), "analyze mem not restored: " + serialized); + } + + @Test + public void testStartOrRestoreWithEmptyThreadContext() { + // No prior data + ThreadContext.remove("_sql_phase_tracker"); + + QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); + assertNotNull(tracker); + // Should produce an empty serialization (no phases yet) + assertEquals("", tracker.serialize()); + } + + @Test + public void testStartOrRestoreCanContinueWithNewPhases() { + ThreadContext.put("_sql_phase_tracker", "parse:1000"); + + QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); + tracker.beginPhase("analyze"); + consumeTime(); + tracker.endCurrentPhase(); + + String serialized = tracker.serialize(); + assertTrue(serialized.contains("parse:1000"), "Restored parse missing: " + serialized); + assertTrue(serialized.contains("analyze:"), "New analyze phase missing: " + serialized); + } + + @Test + public void testClearRemovesFromThreadLocalAndThreadContext() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + tracker.addCompletedPhase("parse", 1000L); + tracker.persist(); + + // Verify presence before clear + assertNotNull(QueryPhaseTracker.current()); + assertNotNull(ThreadContext.get("_sql_phase_tracker")); + + QueryPhaseTracker.clear(); + + assertNull(QueryPhaseTracker.current()); + assertNull(ThreadContext.get("_sql_phase_tracker")); + } + + @Test + public void testAddCompletedPhaseWallOnly() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + tracker.addCompletedPhase("custom", 9999L); + + String serialized = tracker.serialize(); + assertEquals("custom:9999", serialized); + } + + @Test + public void testAddCompletedPhaseWithAllMetrics() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + tracker.addCompletedPhase("execute", 50000L, 30000L, 100000L); + + String serialized = tracker.serialize(); + assertEquals("execute:50000|cpu:30000|mem:100000", serialized); + } + + @Test + public void testBeginPhaseImplicitlyEndsCurrentPhase() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + tracker.beginPhase("parse"); + consumeTime(); + // Calling beginPhase again should end "parse" first + tracker.beginPhase("analyze"); + consumeTime(); + tracker.endCurrentPhase(); + + String serialized = tracker.serialize(); + assertTrue(serialized.contains("parse:"), "parse should have been ended: " + serialized); + assertTrue(serialized.contains("analyze:"), "analyze missing: " + serialized); + } + + @Test + public void testEndCurrentPhaseWithNoActivePhaseIsNoop() { + QueryPhaseTracker tracker = QueryPhaseTracker.start(); + // Should not throw + tracker.endCurrentPhase(); + assertEquals("", tracker.serialize()); + } + + @Test + public void testQuerySourceHeadersConstants() { + assertEquals("x-query-source", QuerySourceHeaders.QUERY_SOURCE_HEADER); + assertEquals("x-original-query", QuerySourceHeaders.ORIGINAL_QUERY_HEADER); + assertEquals("x-query-execution-id", QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER); + assertEquals("x-query-phases", QuerySourceHeaders.QUERY_PHASES_HEADER); + assertEquals(4096, QuerySourceHeaders.MAX_ORIGINAL_QUERY_LENGTH); + } + + /** Burn a small amount of wall-clock time to ensure non-zero nanos. */ + private void consumeTime() { + long start = System.nanoTime(); + while (System.nanoTime() - start < 1_000_000) { + // spin for ~1ms + } + } +} diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java index 858ba0598e6..3b231e9a485 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -61,6 +61,7 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.common.utils.QueryContext; +import org.opensearch.sql.common.utils.QueryPhaseTracker; import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.exception.CalciteUnsupportedException; import org.opensearch.sql.exception.NonFallbackCalciteException; @@ -201,6 +202,8 @@ public void executeWithCalcite( CalcitePlanContext.run( () -> { try { + QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); + tracker.beginPhase("analyze"); ProfileContext profileContext = QueryProfiling.activate(QueryContext.isProfileEnabled()); ProfileMetric analyzeMetric = profileContext.getOrCreateMetric(MetricName.ANALYZE); @@ -220,6 +223,7 @@ public void executeWithCalcite( () -> analyze(plan, context), "while preparing and validating the query plan"); + tracker.beginPhase("plan"); // Wrap plan conversion with PLAN_CONVERSION stage tracking RelNode calcitePlan = StageErrorHandler.executeStage( @@ -229,6 +233,10 @@ public void executeWithCalcite( convertToCalcitePlan(relNode, context), context), "while converting the query to an executable plan"); + analyzeMetric.set(System.nanoTime() - analyzeStart); + tracker.endCurrentPhase(); + tracker.endAll(); + executeCalcitePlan(calcitePlan, context, listener, analyzeMetric, analyzeStart); }, QueryService.class); @@ -795,11 +803,14 @@ public void executeWithLegacy( ResponseListener listener, Optional calciteFailure) { try { - executePlan(analyze(plan, queryType), PlanContext.emptyPlanContext(), listener); + QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); + tracker.beginPhase("analyze"); + LogicalPlan analyzed = analyze(plan, queryType); + tracker.beginPhase("plan"); + executePlan(analyzed, PlanContext.emptyPlanContext(), listener); } catch (Exception e) { if (calciteFailure.isPresent()) { // This happens if Calcite fell back to V2 due to some issue, and then V2 also failed. - // Prefer the Calcite error. // https://github.com/opensearch-project/sql/issues/5060 propagateCalciteError(calciteFailure.get(), listener); } else { @@ -855,16 +866,20 @@ public void executePlan( PlanContext planContext, ResponseListener listener) { try { + PhysicalPlan physicalPlan = plan(plan); + QueryPhaseTracker tracker = QueryPhaseTracker.current(); + if (tracker != null) { + tracker.endCurrentPhase(); + tracker.endAll(); + } planContext .getSplit() .ifPresentOrElse( - split -> executionEngine.execute(plan(plan), new ExecutionContext(split), listener), + split -> executionEngine.execute(physicalPlan, new ExecutionContext(split), listener), () -> executionEngine.execute( - plan(plan), + physicalPlan, ExecutionContext.querySizeLimit( - // For pagination, querySizeLimit shouldn't take effect. - // See {@link PaginationWindowIT::testQuerySizeLimitDoesNotEffectPageSize} plan instanceof LogicalPaginate ? null : SysLimit.fromSettings(settings).querySizeLimit()), diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java index 4064b73d4a4..9173abe5c89 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java @@ -34,6 +34,7 @@ import org.opensearch.sql.common.antlr.SyntaxCheckException; import org.opensearch.sql.common.error.ErrorReport; import org.opensearch.sql.common.utils.QueryContext; +import org.opensearch.sql.common.utils.QuerySourceHeaders; import org.opensearch.sql.exception.ExpressionEvaluationException; import org.opensearch.sql.exception.SemanticCheckException; import org.opensearch.sql.legacy.antlr.OpenSearchLegacySqlAnalyzer; @@ -153,6 +154,27 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli request.params(), sqlRequest.cursor()); + // Tag the thread context so query-insights can identify this as a SQL-derived query. + client + .threadPool() + .getThreadContext() + .putHeader(QuerySourceHeaders.QUERY_SOURCE_HEADER, "sql"); + client + .threadPool() + .getThreadContext() + .putHeader( + QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER, java.util.UUID.randomUUID().toString()); + String queryText = sqlRequest.getSql(); + if (queryText != null) { + if (queryText.length() > QuerySourceHeaders.MAX_ORIGINAL_QUERY_LENGTH) { + queryText = queryText.substring(0, QuerySourceHeaders.MAX_ORIGINAL_QUERY_LENGTH); + } + client + .threadPool() + .getThreadContext() + .putHeader(QuerySourceHeaders.ORIGINAL_QUERY_HEADER, queryText); + } + // Route to analytics engine for non-Lucene (e.g., Parquet-backed) indices. // The router returns true and sends the response directly if it handled the request. final SQLQueryRequest finalRequest = newSqlRequest; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java index 483f2684d61..6c0a1172f19 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java @@ -40,6 +40,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.locationtech.jts.geom.Point; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.sql.ast.statement.ExplainMode; import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.calcite.utils.CalciteToolsHelper; @@ -51,6 +52,8 @@ import org.opensearch.sql.common.error.ErrorReport; import org.opensearch.sql.common.error.ResourceLimitExceededException; import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.common.utils.QueryPhaseTracker; +import org.opensearch.sql.common.utils.QuerySourceHeaders; import org.opensearch.sql.data.model.ExprTupleValue; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.data.model.ExprValueUtils; @@ -107,6 +110,7 @@ public void execute( ExecutionContext context, ResponseListener listener) { PhysicalPlan plan = executionProtector.protect(physicalPlan); + writePhaseHeader(); client.schedule( () -> { try { @@ -160,6 +164,27 @@ public ExplainResponseNode visitTableScan( }); } + private void writePhaseHeader() { + QueryPhaseTracker tracker = QueryPhaseTracker.current(); + if (tracker != null) { + try { + client + .getNodeClient() + .ifPresent( + nc -> { + ThreadContext tc = nc.threadPool().getThreadContext(); + String header = tc.getHeader(QuerySourceHeaders.QUERY_PHASES_HEADER); + if (header == null) { + tc.putHeader(QuerySourceHeaders.QUERY_PHASES_HEADER, tracker.serialize()); + } + }); + } catch (Exception e) { + // Best-effort — don't fail the query if phase header can't be written + } + QueryPhaseTracker.clear(); + } + } + private Hook.Closeable getPhysicalPlanInHook( AtomicReference physical, SqlExplainLevel level) { return Hook.PLAN_BEFORE_IMPLEMENTATION.addThread( @@ -328,6 +353,7 @@ public void explain( @Override public void execute( RelNode rel, CalcitePlanContext context, ResponseListener listener) { + writePhaseHeader(); client.schedule( () -> { try (PreparedStatement statement = OpenSearchRelRunners.run(context, rel)) { diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java index 31f14e3411b..c2d20a8a303 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -507,6 +507,15 @@ public List> getSettings() { .build(); } + @Override + public Collection getTaskHeaders() { + return List.of( + org.opensearch.sql.common.utils.QuerySourceHeaders.QUERY_SOURCE_HEADER, + org.opensearch.sql.common.utils.QuerySourceHeaders.ORIGINAL_QUERY_HEADER, + org.opensearch.sql.common.utils.QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER, + org.opensearch.sql.common.utils.QuerySourceHeaders.QUERY_PHASES_HEADER); + } + @Override public ScriptEngine getScriptEngine(Settings settings, Collection> contexts) { return new CompoundedScriptEngine(); diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java index 772f1ec123f..a3d80c104da 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java @@ -14,6 +14,7 @@ import java.util.Locale; import java.util.Map; import java.util.Optional; +import java.util.UUID; import java.util.function.Supplier; import org.apache.calcite.rel.RelNode; import org.apache.logging.log4j.LogManager; @@ -31,6 +32,7 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.common.utils.QueryContext; +import org.opensearch.sql.common.utils.QuerySourceHeaders; import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.datasources.service.DataSourceServiceImpl; import org.opensearch.sql.executor.AnalyzeResponse; @@ -179,6 +181,25 @@ protected void doExecute( // in order to use PPL service, we need to convert TransportPPLQueryRequest to PPLQueryRequest PPLQueryRequest transformedRequest = transportRequest.toPPLQueryRequest(); + + // Tag the thread context so query-insights can identify this as a PPL-derived query. + org.opensearch.common.util.concurrent.ThreadContext threadContext = + clientRef.threadPool().getThreadContext(); + if (threadContext.getHeader(QuerySourceHeaders.QUERY_SOURCE_HEADER) == null) { + threadContext.putHeader(QuerySourceHeaders.QUERY_SOURCE_HEADER, "ppl"); + } + if (threadContext.getHeader(QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER) == null) { + threadContext.putHeader( + QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER, UUID.randomUUID().toString()); + } + if (threadContext.getHeader(QuerySourceHeaders.ORIGINAL_QUERY_HEADER) == null + && transformedRequest.getRequest() != null) { + String pplQueryText = transformedRequest.getRequest(); + if (pplQueryText.length() > QuerySourceHeaders.MAX_ORIGINAL_QUERY_LENGTH) { + pplQueryText = pplQueryText.substring(0, QuerySourceHeaders.MAX_ORIGINAL_QUERY_LENGTH); + } + threadContext.putHeader(QuerySourceHeaders.ORIGINAL_QUERY_HEADER, pplQueryText); + } QueryContext.setProfile(transformedRequest.profile()); ActionListener clearingListener = wrapWithProfilingClear(listener); diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java index 7f117e7bb0b..29c15744212 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -185,6 +185,9 @@ private AbstractPlan plan( ResponseListener queryListener, ResponseListener explainListener) { // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) + org.opensearch.sql.common.utils.QueryPhaseTracker tracker = + org.opensearch.sql.common.utils.QueryPhaseTracker.start(); + tracker.beginPhase("parse"); ParseTree cst = parser.parse(request.getRequest()); Statement statement = cst.accept( @@ -202,6 +205,8 @@ private AbstractPlan plan( : null) .explainMode(request.getExplainMode()) .build())); + tracker.endCurrentPhase(); + tracker.persist(); log.info( "[{}] Incoming request {}", diff --git a/semantic-review/2025-07-15-162000-pr-sql-query-insights.md b/semantic-review/2025-07-15-162000-pr-sql-query-insights.md new file mode 100644 index 00000000000..e366551097e --- /dev/null +++ b/semantic-review/2025-07-15-162000-pr-sql-query-insights.md @@ -0,0 +1,99 @@ +# Query phase tracking and source-header propagation for query-insights + +A new `QueryPhaseTracker` collects wall-clock, CPU, and memory-allocation timings per query execution phase (parse, analyze, plan, execute) and serializes them into thread-context headers so the query-insights plugin can attribute costs to SQL/PPL queries. The entry points (`RestSqlAction`, `TransportPPLQueryAction`) stamp source-identification headers, and `SQLPlugin.getTaskHeaders` registers them for cross-node transport. The tracker uses Log4j `ThreadContext` as a shuttle to cross from the REST thread to the sql-worker thread. + +Watch for: +- **Unconditional `putHeader` in `RestSqlAction`** (confirmed) — unlike the PPL path, the SQL REST handler calls `putHeader` without a null-guard, which will throw `IllegalArgumentException` if the header is already present. +- **`Thread.currentThread().getId()` deprecation** (confirmed) — deprecated since Java 19, replaced by `threadId()`. The project targets Java 21. +- **ThreadLocal leak on exception paths** (likely) — if `executeWithCalcite` throws before reaching `endAll()` / `writePhaseHeader()`, the tracker remains in the ThreadLocal and the Log4j key is never cleaned. +- **Missing imports / FQN usage in PPLService and SQLService** (confirmed) — both files use fully-qualified `org.opensearch.sql.common.utils.QueryPhaseTracker` inline instead of an import statement, inconsistent with the rest of the codebase. + +## High-level view + +The header-writing entry points differ in safety: `TransportPPLQueryAction` guards each `putHeader` with a null-check on `getHeader`, while `RestSqlAction` writes unconditionally. OpenSearch's `ThreadContext.putHeader` throws if the key already exists, so the SQL path is fragile in any scenario where the handler runs more than once per request context (retries, plugin chaining). + +The lifecycle management has a gap in the Calcite path: if an exception escapes between `beginPhase("analyze")` and `endAll()`, neither `endAll()` nor `clear()` is reached, leaking the ThreadLocal and Log4j entry on the pooled thread. + +`writePhaseHeader` in `OpenSearchExecutionEngine` is best-effort (catch-all around `putHeader`), which is the right call for observability plumbing — a failure to write metrics should never fail the query. + +

+Issues (6) + +1. **Unconditional putHeader in RestSqlAction** — wrap each `putHeader` call with a `getHeader == null` guard, matching the PPL pattern, to prevent `IllegalArgumentException` if the header already exists. +2. **Deprecated `Thread.currentThread().getId()`** — replace with `Thread.currentThread().threadId()` (available since Java 19; project targets 21). +3. **ThreadLocal leak on exception in executeWithCalcite** — add a `finally` block (or catch) that calls `QueryPhaseTracker.clear()` so the tracker and Log4j key are cleaned on failure paths. +4. **FQN usage instead of imports in PPLService/SQLService** — add proper import statements for `QueryPhaseTracker` to `PPLService.java` and `SQLService.java` for consistency and readability. +5. **Unused `isEmpty()` method** — `isEmpty()` is public but never called anywhere in the diff or existing code. Either document its intended consumer or remove dead code. +6. **No `tracker.endAll()` on legacy V2 exception path** — in `executeWithLegacy`, the tracker begins the "plan" phase but `endAll()` only fires inside `executePlan` if `current()` is non-null. If `plan(plan)` throws before that point, phases are left dangling. + +
+ +
+Details + +## Unconditional putHeader in RestSqlAction vs guarded PPL path + +In `RestSqlAction` (lines 161-175), headers are set without checking whether they already exist: + +```java +client.threadPool().getThreadContext() + .putHeader(QuerySourceHeaders.QUERY_SOURCE_HEADER, "sql"); +client.threadPool().getThreadContext() + .putHeader(QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER, UUID.randomUUID().toString()); +``` + +OpenSearch's `ThreadContext.putHeader` throws `IllegalArgumentException` if the key is already present. The PPL transport action correctly guards with `if (threadContext.getHeader(...) == null)`. The SQL path should follow the same pattern. + +## Deprecated Thread.currentThread().getId() + +`QueryPhaseTracker` calls `Thread.currentThread().getId()` on lines 134 and 148 to pass to `getThreadAllocatedBytes(long)`. Since Java 19, `Thread.getId()` is deprecated in favour of `Thread.threadId()`. With the project targeting Java 21, this will produce deprecation warnings and should be: + +```java +SUN_THREAD_MX.getThreadAllocatedBytes(Thread.currentThread().threadId()) +``` + +## ThreadLocal and Log4j key lifecycle on exception paths + +In `QueryService.executeWithCalcite`, the tracker is created at the top of the try block: + +```java +QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); +tracker.beginPhase("analyze"); +``` + +If `StageErrorHandler.executeStage` throws (e.g. `CalciteUnsupportedException`), control jumps to the catch block which calls `executeWithLegacy`. That path creates a *new* tracker via `startOrRestore()`, which overwrites `CURRENT` — so the first tracker's ThreadLocal slot is released. However, the Log4j `ThreadContext` key `_sql_phase_tracker` from the initial `persist()` (set on the REST thread in `SQLService`) is never cleaned if the whole request fails before reaching `writePhaseHeader`. The fix: call `QueryPhaseTracker.clear()` in a `finally` block at the outermost scope of `executeWithCalcite`. + +## FQN usage in PPLService and SQLService + +Both `PPLService.java` and `SQLService.java` use the fully-qualified class name inline: + +```java +org.opensearch.sql.common.utils.QueryPhaseTracker tracker = + org.opensearch.sql.common.utils.QueryPhaseTracker.start(); +``` + +Every other file in this diff uses an import statement. Add `import org.opensearch.sql.common.utils.QueryPhaseTracker;` to each file and use the short name. + +## Serialization format delimiter assumption + +The format `phase:wallNanos|cpu:cpuNanos|mem:memBytes` uses `:` and `|` as delimiters without escaping. Currently all phase names are hardcoded safe strings ("parse", "analyze", "plan", "total"), so this is not a bug today. Adding a defensive check in `beginPhase` (e.g. `assert !name.contains(":") && !name.contains("|")`) guards against future misuse. + +
+ +
+File map + +| File | Change | +|------|--------| +| `common/.../QueryPhaseTracker.java` | New. Thread-local phase tracker with wall/CPU/mem metrics and Log4j shuttle. | +| `common/.../QuerySourceHeaders.java` | New. Constants for x-query-source/original-query/execution-id/phases headers. | +| `common/.../QueryPhaseTrackerTest.java` | New. Unit tests covering lifecycle, serialization, cross-thread restore. | +| `core/.../QueryService.java` | Integrates tracker into `executeWithCalcite` and `executeWithLegacy` paths. | +| `legacy/.../RestSqlAction.java` | Stamps source-identification headers on the REST thread for SQL queries. | +| `opensearch/.../OpenSearchExecutionEngine.java` | `writePhaseHeader` writes serialized phases into thread-context header. | +| `plugin/.../SQLPlugin.java` | Registers query-insights headers via `getTaskHeaders`. | +| `plugin/.../TransportPPLQueryAction.java` | Stamps source-identification headers for PPL queries (guarded). | +| `ppl/.../PPLService.java` | Starts tracker and tracks parse phase for PPL. | +| `sql/.../SQLService.java` | Starts tracker and tracks parse phase for SQL. | + +
diff --git a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java index 9b4cf8c1a37..c15b8cf5037 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java +++ b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java @@ -99,6 +99,9 @@ private AbstractPlan plan( explainListener); } else { // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) + org.opensearch.sql.common.utils.QueryPhaseTracker tracker = + org.opensearch.sql.common.utils.QueryPhaseTracker.start(); + tracker.beginPhase("parse"); ParseTree cst = parser.parse(request.getQuery()); Statement statement = cst.accept( @@ -109,6 +112,8 @@ private AbstractPlan plan( .fetchSize(request.getFetchSize()) .format(request.getFormat()) .build())); + tracker.endCurrentPhase(); + tracker.persist(); return queryExecutionFactory.create(statement, queryListener, explainListener); } From bd9913aae8dd9f045c9fb222cf3b696e7c66a83b Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 18 Aug 2026 01:07:01 -0700 Subject: [PATCH 2/4] feat: link SQL/PPL DSL search tasks to their originating query via parent task Replace the query-source request headers with parent-task linkage so every DSL search task the SQL/PPL engines spawn references its originating query. query-insights can then look up the coordinator task (and its original request) instead of relying on x-query-* headers, avoiding raw query text on the wire, getTaskHeaders() wiring, and thread-context propagation. - SQL: dispatch execution through a local SqlQueryAction/TransportSqlQueryAction that registers a SqlQueryTask coordinator task (via NodeClient.executeLocally), binds it as the cancellable task, and completes on the REST channel response. Wired into RestSqlAction through a SqlCoordinatorTaskDispatcher seam (legacy module cannot depend on the plugin transport action directly). - PPL v3: propagate the coordinator task to the background prefetch threads in BackgroundSearchScanner so applyParentTask stamps the parent on prefetched DSL searches (previously lost on the sql_background_io pool). - Remove QuerySourceHeaders, QueryPhaseTracker, getTaskHeaders(), and all x-query-* header writes from the SQL/PPL execution paths. Note: SQL coordinator-task lifecycle (task-manager registration visibility, no-leak cleanup, and fallback/cursor/explain behavior) still needs live-cluster integration testing. Unit tests pass and all modules compile. Signed-off-by: Kishore Natarajan --- .../sql/common/utils/QueryPhaseTracker.java | 184 ------------ .../sql/common/utils/QuerySourceHeaders.java | 23 -- .../common/utils/QueryPhaseTrackerTest.java | 271 ------------------ .../opensearch/sql/executor/QueryService.java | 27 +- .../sql/legacy/plugin/RestSqlAction.java | 51 ++-- .../plugin/SqlCoordinatorTaskDispatcher.java | 37 +++ .../executor/OpenSearchExecutionEngine.java | 26 -- .../storage/scan/BackgroundSearchScanner.java | 35 ++- .../org/opensearch/sql/plugin/SQLPlugin.java | 44 ++- .../sql/plugin/transport/SqlQueryAction.java | 22 ++ .../sql/plugin/transport/SqlQueryTask.java | 33 +++ .../transport/TransportPPLQueryAction.java | 20 -- .../transport/TransportSqlQueryAction.java | 134 +++++++++ .../transport/TransportSqlQueryRequest.java | 83 ++++++ .../transport/TransportSqlQueryResponse.java | 30 ++ .../org/opensearch/sql/ppl/PPLService.java | 5 - ...2025-07-15-162000-pr-sql-query-insights.md | 99 ------- .../org/opensearch/sql/sql/SQLService.java | 5 - 18 files changed, 436 insertions(+), 693 deletions(-) delete mode 100644 common/src/main/java/org/opensearch/sql/common/utils/QueryPhaseTracker.java delete mode 100644 common/src/main/java/org/opensearch/sql/common/utils/QuerySourceHeaders.java delete mode 100644 common/src/test/java/org/opensearch/sql/common/utils/QueryPhaseTrackerTest.java create mode 100644 legacy/src/main/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcher.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryAction.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryTask.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequest.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryResponse.java delete mode 100644 semantic-review/2025-07-15-162000-pr-sql-query-insights.md diff --git a/common/src/main/java/org/opensearch/sql/common/utils/QueryPhaseTracker.java b/common/src/main/java/org/opensearch/sql/common/utils/QueryPhaseTracker.java deleted file mode 100644 index 7c53d3e8bb5..00000000000 --- a/common/src/main/java/org/opensearch/sql/common/utils/QueryPhaseTracker.java +++ /dev/null @@ -1,184 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.common.utils; - -import java.lang.management.ManagementFactory; -import java.lang.management.ThreadMXBean; -import java.util.LinkedHashMap; -import java.util.Map; -import java.util.StringJoiner; - -/** - * Tracks timing and resource usage for SQL/PPL query execution phases. Thread-local instance - * collects phase durations (nanos) as the query flows through parse → analyze → plan → execute. The - * result is serialized into a compact header value for query-insights to consume. - * - *

Phases that complete on the REST thread (parse) are stored in the Log4j ThreadContext via - * {@link #persist()}, which is propagated to the sql-worker thread by OpenSearchQueryManager. On - * the worker thread, {@link #startOrRestore()} restores prior phases and continues tracking. - * - *

CPU time uses {@code ThreadMXBean.getCurrentThreadCpuTime()}. Memory allocation tracking uses - * {@code com.sun.management.ThreadMXBean.getThreadAllocatedBytes()} which is best-effort and may be - * unavailable on non-HotSpot JVMs (e.g., IBM J9, GraalVM native). - */ -public final class QueryPhaseTracker { - - private static final String LOG4J_KEY = "_sql_phase_tracker"; - private static final ThreadLocal CURRENT = new ThreadLocal<>(); - private static final ThreadMXBean THREAD_MX = ManagementFactory.getThreadMXBean(); - private static final boolean CPU_SUPPORTED = THREAD_MX.isCurrentThreadCpuTimeSupported(); - private static final com.sun.management.ThreadMXBean SUN_THREAD_MX = getSunThreadMXBean(); - - private final Map phases = new LinkedHashMap<>(); - private final Map cpuPhases = new LinkedHashMap<>(); - private final Map memPhases = new LinkedHashMap<>(); - private String activePhase; - private long activeStart; - private long activeCpuStart; - private long activeMemStart; - private long overallStart; - - private static com.sun.management.ThreadMXBean getSunThreadMXBean() { - try { - if (THREAD_MX instanceof com.sun.management.ThreadMXBean sun) { - return sun; - } - } catch (NoClassDefFoundError | UnsupportedOperationException e) { - // com.sun.management not available on this JVM - } - return null; - } - - private QueryPhaseTracker() { - this.overallStart = System.nanoTime(); - } - - public static QueryPhaseTracker start() { - QueryPhaseTracker tracker = new QueryPhaseTracker(); - CURRENT.set(tracker); - return tracker; - } - - /** Restore from Log4j ThreadContext (cross-thread transfer) or create fresh. */ - public static QueryPhaseTracker startOrRestore() { - QueryPhaseTracker tracker = new QueryPhaseTracker(); - String stored = org.apache.logging.log4j.ThreadContext.get(LOG4J_KEY); - if (stored != null && !stored.isEmpty()) { - try { - for (String part : stored.split(",")) { - // Format: "phase:wallNanos|cpu:cpuNanos|mem:memBytes" - String[] segments = part.split("\\|"); - String[] kv = segments[0].split(":", 2); - if (kv.length == 2) { - tracker.phases.put(kv[0], Long.parseLong(kv[1])); - for (int i = 1; i < segments.length; i++) { - String[] metric = segments[i].split(":", 2); - if (metric.length == 2 && "cpu".equals(metric[0])) { - tracker.cpuPhases.put(kv[0], Long.parseLong(metric[1])); - } else if (metric.length == 2 && "mem".equals(metric[0])) { - tracker.memPhases.put(kv[0], Long.parseLong(metric[1])); - } - } - } - } - } catch (NumberFormatException e) { - // Malformed data in ThreadContext — start fresh - tracker.phases.clear(); - tracker.cpuPhases.clear(); - tracker.memPhases.clear(); - } - } - CURRENT.set(tracker); - return tracker; - } - - public static QueryPhaseTracker current() { - return CURRENT.get(); - } - - /** Returns true if no phases have been recorded (no prior SQL/PPL context). */ - public boolean isEmpty() { - return phases.isEmpty() && activePhase == null; - } - - public static void clear() { - CURRENT.remove(); - org.apache.logging.log4j.ThreadContext.remove(LOG4J_KEY); - } - - /** Persist current phases to Log4j ThreadContext for cross-thread propagation. */ - public void persist() { - org.apache.logging.log4j.ThreadContext.put(LOG4J_KEY, serialize()); - } - - public void addCompletedPhase(String name, long nanos) { - phases.put(name, nanos); - } - - public void addCompletedPhase(String name, long nanos, long cpuNanos, long memBytes) { - phases.put(name, nanos); - if (cpuNanos > 0) cpuPhases.put(name, cpuNanos); - if (memBytes > 0) memPhases.put(name, memBytes); - } - - public void beginPhase(String name) { - endCurrentPhase(); - activePhase = name; - activeStart = System.nanoTime(); - activeCpuStart = CPU_SUPPORTED ? THREAD_MX.getCurrentThreadCpuTime() : 0; - activeMemStart = - SUN_THREAD_MX != null - ? SUN_THREAD_MX.getThreadAllocatedBytes(Thread.currentThread().getId()) - : 0; - } - - public void endCurrentPhase() { - if (activePhase != null) { - long elapsed = System.nanoTime() - activeStart; - phases.merge(activePhase, elapsed, Long::sum); - if (CPU_SUPPORTED) { - long cpuElapsed = THREAD_MX.getCurrentThreadCpuTime() - activeCpuStart; - cpuPhases.merge(activePhase, cpuElapsed, Long::sum); - } - if (SUN_THREAD_MX != null) { - long memElapsed = - SUN_THREAD_MX.getThreadAllocatedBytes(Thread.currentThread().getId()) - activeMemStart; - memPhases.merge(activePhase, memElapsed, Long::sum); - } - activePhase = null; - } - } - - public void endAll() { - endCurrentPhase(); - long total = System.nanoTime() - overallStart; - phases.put("total", total); - } - - /** - * Serialize to compact format with all metrics per phase: - * "parse:1234|cpu:1000|mem:5000,analyze:5678|cpu:4000|mem:20000,total:18943" Time and CPU in - * nanoseconds, memory in bytes. - */ - public String serialize() { - StringJoiner joiner = new StringJoiner(","); - for (Map.Entry entry : phases.entrySet()) { - String phase = entry.getKey(); - StringBuilder sb = new StringBuilder(); - sb.append(phase).append(':').append(entry.getValue()); - Long cpu = cpuPhases.get(phase); - if (cpu != null) { - sb.append("|cpu:").append(cpu); - } - Long mem = memPhases.get(phase); - if (mem != null) { - sb.append("|mem:").append(mem); - } - joiner.add(sb.toString()); - } - return joiner.toString(); - } -} diff --git a/common/src/main/java/org/opensearch/sql/common/utils/QuerySourceHeaders.java b/common/src/main/java/org/opensearch/sql/common/utils/QuerySourceHeaders.java deleted file mode 100644 index 7075870d19f..00000000000 --- a/common/src/main/java/org/opensearch/sql/common/utils/QuerySourceHeaders.java +++ /dev/null @@ -1,23 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.common.utils; - -/** - * Constants for thread context headers used to propagate query source metadata to downstream - * components like query-insights. - */ -public final class QuerySourceHeaders { - - public static final String QUERY_SOURCE_HEADER = "x-query-source"; - public static final String ORIGINAL_QUERY_HEADER = "x-original-query"; - public static final String QUERY_EXECUTION_ID_HEADER = "x-query-execution-id"; - public static final String QUERY_PHASES_HEADER = "x-query-phases"; - - /** Maximum number of characters stored in the {@link #ORIGINAL_QUERY_HEADER}. */ - public static final int MAX_ORIGINAL_QUERY_LENGTH = 4096; - - private QuerySourceHeaders() {} -} diff --git a/common/src/test/java/org/opensearch/sql/common/utils/QueryPhaseTrackerTest.java b/common/src/test/java/org/opensearch/sql/common/utils/QueryPhaseTrackerTest.java deleted file mode 100644 index 9ac49665665..00000000000 --- a/common/src/test/java/org/opensearch/sql/common/utils/QueryPhaseTrackerTest.java +++ /dev/null @@ -1,271 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.common.utils; - -import static org.junit.jupiter.api.Assertions.*; - -import org.apache.logging.log4j.ThreadContext; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.Test; - -/** Unit tests for {@link QueryPhaseTracker}. */ -public class QueryPhaseTrackerTest { - - @AfterEach - void cleanup() { - QueryPhaseTracker.clear(); - } - - @Test - public void testStartCreatesTrackerAndSetsCurrent() { - assertNull(QueryPhaseTracker.current()); - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - assertNotNull(tracker); - assertSame(tracker, QueryPhaseTracker.current()); - } - - @Test - public void testBeginPhaseAndEndCurrentPhase() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - tracker.beginPhase("parse"); - // Simulate some work - consumeTime(); - tracker.endCurrentPhase(); - - String serialized = tracker.serialize(); - assertTrue( - serialized.startsWith("parse:"), - "Expected serialized to start with 'parse:': " + serialized); - // The nanos value should be > 0 - String wallNanos = serialized.split("\\|")[0].split(":")[1]; - assertTrue(Long.parseLong(wallNanos) > 0, "Wall-clock nanos should be positive"); - } - - @Test - public void testMultiplePhasesInSequence() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - - tracker.beginPhase("parse"); - consumeTime(); - tracker.endCurrentPhase(); - - tracker.beginPhase("analyze"); - consumeTime(); - tracker.endCurrentPhase(); - - tracker.beginPhase("plan"); - consumeTime(); - tracker.endCurrentPhase(); - - String serialized = tracker.serialize(); - // All three phases should be present in order - assertTrue(serialized.contains("parse:"), "Missing parse phase: " + serialized); - assertTrue(serialized.contains("analyze:"), "Missing analyze phase: " + serialized); - assertTrue(serialized.contains("plan:"), "Missing plan phase: " + serialized); - - // Verify ordering (LinkedHashMap preserves insertion order) - int parseIdx = serialized.indexOf("parse:"); - int analyzeIdx = serialized.indexOf("analyze:"); - int planIdx = serialized.indexOf("plan:"); - assertTrue(parseIdx < analyzeIdx, "parse should come before analyze"); - assertTrue(analyzeIdx < planIdx, "analyze should come before plan"); - } - - @Test - public void testEndAllAddsTotalPhase() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - - tracker.beginPhase("parse"); - consumeTime(); - tracker.endCurrentPhase(); - - tracker.endAll(); - - String serialized = tracker.serialize(); - assertTrue(serialized.contains("total:"), "Missing total phase: " + serialized); - - // total should be the last entry - String[] parts = serialized.split(","); - assertTrue(parts[parts.length - 1].startsWith("total:"), "total should be last: " + serialized); - } - - @Test - public void testEndAllFinishesActivePhase() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - tracker.beginPhase("execute"); - consumeTime(); - // Don't call endCurrentPhase — endAll should do it - tracker.endAll(); - - String serialized = tracker.serialize(); - assertTrue(serialized.contains("execute:"), "Missing execute phase: " + serialized); - assertTrue(serialized.contains("total:"), "Missing total phase: " + serialized); - } - - @Test - public void testSerializeFormat() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - // Use addCompletedPhase to have deterministic values - tracker.addCompletedPhase("parse", 1000L, 800L, 5000L); - tracker.addCompletedPhase("analyze", 2000L, 1500L, 10000L); - - String serialized = tracker.serialize(); - // Expected: "parse:1000|cpu:800|mem:5000,analyze:2000|cpu:1500|mem:10000" - assertEquals("parse:1000|cpu:800|mem:5000,analyze:2000|cpu:1500|mem:10000", serialized); - } - - @Test - public void testSerializeFormatWithoutCpuAndMem() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - // Overload with only nanos - tracker.addCompletedPhase("parse", 1000L); - tracker.addCompletedPhase("analyze", 2000L); - - String serialized = tracker.serialize(); - // No cpu/mem segments - assertEquals("parse:1000,analyze:2000", serialized); - } - - @Test - public void testSerializeFormatWithZeroCpuAndMem() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - // cpuNanos=0 and memBytes=0 should be omitted - tracker.addCompletedPhase("parse", 1000L, 0L, 0L); - - String serialized = tracker.serialize(); - assertEquals("parse:1000", serialized); - } - - @Test - public void testPersistStoresToThreadContext() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - tracker.addCompletedPhase("parse", 1234L, 1000L, 5000L); - - tracker.persist(); - - String stored = ThreadContext.get("_sql_phase_tracker"); - assertNotNull(stored); - assertEquals("parse:1234|cpu:1000|mem:5000", stored); - } - - @Test - public void testStartOrRestoreRestoresPhasesFromThreadContext() { - // Simulate a prior thread persisting data - ThreadContext.put( - "_sql_phase_tracker", "parse:1234|cpu:1000|mem:5000,analyze:5678|cpu:4000|mem:20000"); - - QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); - assertNotNull(tracker); - assertSame(tracker, QueryPhaseTracker.current()); - - String serialized = tracker.serialize(); - assertTrue(serialized.contains("parse:1234"), "parse phase not restored: " + serialized); - assertTrue(serialized.contains("cpu:1000"), "parse cpu not restored: " + serialized); - assertTrue(serialized.contains("mem:5000"), "parse mem not restored: " + serialized); - assertTrue(serialized.contains("analyze:5678"), "analyze phase not restored: " + serialized); - assertTrue(serialized.contains("cpu:4000"), "analyze cpu not restored: " + serialized); - assertTrue(serialized.contains("mem:20000"), "analyze mem not restored: " + serialized); - } - - @Test - public void testStartOrRestoreWithEmptyThreadContext() { - // No prior data - ThreadContext.remove("_sql_phase_tracker"); - - QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); - assertNotNull(tracker); - // Should produce an empty serialization (no phases yet) - assertEquals("", tracker.serialize()); - } - - @Test - public void testStartOrRestoreCanContinueWithNewPhases() { - ThreadContext.put("_sql_phase_tracker", "parse:1000"); - - QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); - tracker.beginPhase("analyze"); - consumeTime(); - tracker.endCurrentPhase(); - - String serialized = tracker.serialize(); - assertTrue(serialized.contains("parse:1000"), "Restored parse missing: " + serialized); - assertTrue(serialized.contains("analyze:"), "New analyze phase missing: " + serialized); - } - - @Test - public void testClearRemovesFromThreadLocalAndThreadContext() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - tracker.addCompletedPhase("parse", 1000L); - tracker.persist(); - - // Verify presence before clear - assertNotNull(QueryPhaseTracker.current()); - assertNotNull(ThreadContext.get("_sql_phase_tracker")); - - QueryPhaseTracker.clear(); - - assertNull(QueryPhaseTracker.current()); - assertNull(ThreadContext.get("_sql_phase_tracker")); - } - - @Test - public void testAddCompletedPhaseWallOnly() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - tracker.addCompletedPhase("custom", 9999L); - - String serialized = tracker.serialize(); - assertEquals("custom:9999", serialized); - } - - @Test - public void testAddCompletedPhaseWithAllMetrics() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - tracker.addCompletedPhase("execute", 50000L, 30000L, 100000L); - - String serialized = tracker.serialize(); - assertEquals("execute:50000|cpu:30000|mem:100000", serialized); - } - - @Test - public void testBeginPhaseImplicitlyEndsCurrentPhase() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - tracker.beginPhase("parse"); - consumeTime(); - // Calling beginPhase again should end "parse" first - tracker.beginPhase("analyze"); - consumeTime(); - tracker.endCurrentPhase(); - - String serialized = tracker.serialize(); - assertTrue(serialized.contains("parse:"), "parse should have been ended: " + serialized); - assertTrue(serialized.contains("analyze:"), "analyze missing: " + serialized); - } - - @Test - public void testEndCurrentPhaseWithNoActivePhaseIsNoop() { - QueryPhaseTracker tracker = QueryPhaseTracker.start(); - // Should not throw - tracker.endCurrentPhase(); - assertEquals("", tracker.serialize()); - } - - @Test - public void testQuerySourceHeadersConstants() { - assertEquals("x-query-source", QuerySourceHeaders.QUERY_SOURCE_HEADER); - assertEquals("x-original-query", QuerySourceHeaders.ORIGINAL_QUERY_HEADER); - assertEquals("x-query-execution-id", QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER); - assertEquals("x-query-phases", QuerySourceHeaders.QUERY_PHASES_HEADER); - assertEquals(4096, QuerySourceHeaders.MAX_ORIGINAL_QUERY_LENGTH); - } - - /** Burn a small amount of wall-clock time to ensure non-zero nanos. */ - private void consumeTime() { - long start = System.nanoTime(); - while (System.nanoTime() - start < 1_000_000) { - // spin for ~1ms - } - } -} diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java index 3b231e9a485..858ba0598e6 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -61,7 +61,6 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.common.utils.QueryContext; -import org.opensearch.sql.common.utils.QueryPhaseTracker; import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.exception.CalciteUnsupportedException; import org.opensearch.sql.exception.NonFallbackCalciteException; @@ -202,8 +201,6 @@ public void executeWithCalcite( CalcitePlanContext.run( () -> { try { - QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); - tracker.beginPhase("analyze"); ProfileContext profileContext = QueryProfiling.activate(QueryContext.isProfileEnabled()); ProfileMetric analyzeMetric = profileContext.getOrCreateMetric(MetricName.ANALYZE); @@ -223,7 +220,6 @@ public void executeWithCalcite( () -> analyze(plan, context), "while preparing and validating the query plan"); - tracker.beginPhase("plan"); // Wrap plan conversion with PLAN_CONVERSION stage tracking RelNode calcitePlan = StageErrorHandler.executeStage( @@ -233,10 +229,6 @@ public void executeWithCalcite( convertToCalcitePlan(relNode, context), context), "while converting the query to an executable plan"); - analyzeMetric.set(System.nanoTime() - analyzeStart); - tracker.endCurrentPhase(); - tracker.endAll(); - executeCalcitePlan(calcitePlan, context, listener, analyzeMetric, analyzeStart); }, QueryService.class); @@ -803,14 +795,11 @@ public void executeWithLegacy( ResponseListener listener, Optional calciteFailure) { try { - QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); - tracker.beginPhase("analyze"); - LogicalPlan analyzed = analyze(plan, queryType); - tracker.beginPhase("plan"); - executePlan(analyzed, PlanContext.emptyPlanContext(), listener); + executePlan(analyze(plan, queryType), PlanContext.emptyPlanContext(), listener); } catch (Exception e) { if (calciteFailure.isPresent()) { // This happens if Calcite fell back to V2 due to some issue, and then V2 also failed. + // Prefer the Calcite error. // https://github.com/opensearch-project/sql/issues/5060 propagateCalciteError(calciteFailure.get(), listener); } else { @@ -866,20 +855,16 @@ public void executePlan( PlanContext planContext, ResponseListener listener) { try { - PhysicalPlan physicalPlan = plan(plan); - QueryPhaseTracker tracker = QueryPhaseTracker.current(); - if (tracker != null) { - tracker.endCurrentPhase(); - tracker.endAll(); - } planContext .getSplit() .ifPresentOrElse( - split -> executionEngine.execute(physicalPlan, new ExecutionContext(split), listener), + split -> executionEngine.execute(plan(plan), new ExecutionContext(split), listener), () -> executionEngine.execute( - physicalPlan, + plan(plan), ExecutionContext.querySizeLimit( + // For pagination, querySizeLimit shouldn't take effect. + // See {@link PaginationWindowIT::testQuerySizeLimitDoesNotEffectPageSize} plan instanceof LogicalPaginate ? null : SysLimit.fromSettings(settings).querySizeLimit()), diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java index 9173abe5c89..ae92a93b071 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java @@ -34,7 +34,6 @@ import org.opensearch.sql.common.antlr.SyntaxCheckException; import org.opensearch.sql.common.error.ErrorReport; import org.opensearch.sql.common.utils.QueryContext; -import org.opensearch.sql.common.utils.QuerySourceHeaders; import org.opensearch.sql.exception.ExpressionEvaluationException; import org.opensearch.sql.exception.SemanticCheckException; import org.opensearch.sql.legacy.antlr.OpenSearchLegacySqlAnalyzer; @@ -91,14 +90,26 @@ public class RestSqlAction extends BaseRestHandler { */ private final BiFunction analyticsRouter; + /** Runs the SQL execution under a coordinator task so DSL searches link back to the SQL query. */ + private final SqlCoordinatorTaskDispatcher coordinatorTaskDispatcher; + public RestSqlAction( Settings settings, Injector injector, BiFunction analyticsRouter) { + this(settings, injector, analyticsRouter, SqlCoordinatorTaskDispatcher.PASSTHROUGH); + } + + public RestSqlAction( + Settings settings, + Injector injector, + BiFunction analyticsRouter, + SqlCoordinatorTaskDispatcher coordinatorTaskDispatcher) { super(); this.allowExplicitIndex = MULTI_ALLOW_EXPLICIT_INDEX.get(settings); this.newSqlQueryHandler = new RestSQLQueryAction(injector); this.analyticsRouter = analyticsRouter; + this.coordinatorTaskDispatcher = coordinatorTaskDispatcher; } @Override @@ -154,35 +165,21 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli request.params(), sqlRequest.cursor()); - // Tag the thread context so query-insights can identify this as a SQL-derived query. - client - .threadPool() - .getThreadContext() - .putHeader(QuerySourceHeaders.QUERY_SOURCE_HEADER, "sql"); - client - .threadPool() - .getThreadContext() - .putHeader( - QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER, java.util.UUID.randomUUID().toString()); - String queryText = sqlRequest.getSql(); - if (queryText != null) { - if (queryText.length() > QuerySourceHeaders.MAX_ORIGINAL_QUERY_LENGTH) { - queryText = queryText.substring(0, QuerySourceHeaders.MAX_ORIGINAL_QUERY_LENGTH); - } - client - .threadPool() - .getThreadContext() - .putHeader(QuerySourceHeaders.ORIGINAL_QUERY_HEADER, queryText); - } - // Route to analytics engine for non-Lucene (e.g., Parquet-backed) indices. // The router returns true and sends the response directly if it handled the request. final SQLQueryRequest finalRequest = newSqlRequest; - return channel -> { - if (!analyticsRouter.apply(finalRequest, channel)) { - delegateToV2Engine(request, client, sqlRequest, finalRequest, format, channel); - } - }; + // Run under a coordinator task so every DSL search the SQL engine issues carries a parent + // reference back to this SQL query (used by query-insights to correlate and recover source). + return channel -> + coordinatorTaskDispatcher.dispatch( + client, + sqlRequest.getSql(), + channel, + ch -> { + if (!analyticsRouter.apply(finalRequest, ch)) { + delegateToV2Engine(request, client, sqlRequest, finalRequest, format, ch); + } + }); } catch (Exception e) { return channel -> handleException(channel, e); } diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcher.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcher.java new file mode 100644 index 00000000000..93913738b38 --- /dev/null +++ b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcher.java @@ -0,0 +1,37 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.legacy.plugin; + +import java.util.function.Consumer; +import org.opensearch.rest.RestChannel; +import org.opensearch.transport.client.node.NodeClient; + +/** + * Runs a SQL execution under a coordinator task so the DSL search tasks it spawns carry a parent + * reference back to the originating SQL query. + * + *

The concrete implementation (wired in the plugin module) dispatches through a local transport + * action that registers the coordinator task; this seam exists because {@code RestSqlAction} lives + * in the {@code legacy} module and cannot depend on the plugin-module transport action directly. + * The default implementation simply runs the work with no coordinator task, preserving behavior for + * callers/tests that do not supply one. + */ +@FunctionalInterface +public interface SqlCoordinatorTaskDispatcher { + + /** Default: run the work directly without establishing a coordinator task. */ + SqlCoordinatorTaskDispatcher PASSTHROUGH = + (client, query, channel, work) -> work.accept(channel); + + /** + * @param client node client used to dispatch the local coordinator-task action + * @param query original SQL query text (recorded on the coordinator task) + * @param channel REST channel the execution writes its response to + * @param work the SQL execution to run under the coordinator task + */ + void dispatch( + NodeClient client, String query, RestChannel channel, Consumer work); +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java index 6c0a1172f19..483f2684d61 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java @@ -40,7 +40,6 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.locationtech.jts.geom.Point; -import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.sql.ast.statement.ExplainMode; import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.calcite.utils.CalciteToolsHelper; @@ -52,8 +51,6 @@ import org.opensearch.sql.common.error.ErrorReport; import org.opensearch.sql.common.error.ResourceLimitExceededException; import org.opensearch.sql.common.response.ResponseListener; -import org.opensearch.sql.common.utils.QueryPhaseTracker; -import org.opensearch.sql.common.utils.QuerySourceHeaders; import org.opensearch.sql.data.model.ExprTupleValue; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.data.model.ExprValueUtils; @@ -110,7 +107,6 @@ public void execute( ExecutionContext context, ResponseListener listener) { PhysicalPlan plan = executionProtector.protect(physicalPlan); - writePhaseHeader(); client.schedule( () -> { try { @@ -164,27 +160,6 @@ public ExplainResponseNode visitTableScan( }); } - private void writePhaseHeader() { - QueryPhaseTracker tracker = QueryPhaseTracker.current(); - if (tracker != null) { - try { - client - .getNodeClient() - .ifPresent( - nc -> { - ThreadContext tc = nc.threadPool().getThreadContext(); - String header = tc.getHeader(QuerySourceHeaders.QUERY_PHASES_HEADER); - if (header == null) { - tc.putHeader(QuerySourceHeaders.QUERY_PHASES_HEADER, tracker.serialize()); - } - }); - } catch (Exception e) { - // Best-effort — don't fail the query if phase header can't be written - } - QueryPhaseTracker.clear(); - } - } - private Hook.Closeable getPhysicalPlanInHook( AtomicReference physical, SqlExplainLevel level) { return Hook.PLAN_BEFORE_IMPLEMENTATION.addThread( @@ -353,7 +328,6 @@ public void explain( @Override public void execute( RelNode rel, CalcitePlanContext context, ResponseListener listener) { - writePhaseHeader(); client.schedule( () -> { try (PreparedStatement statement = OpenSearchRelRunners.run(context, rel)) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java index 3aa347b70fa..dbc38f4590a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java @@ -21,8 +21,10 @@ import org.opensearch.sql.monitor.profile.ProfileContext; import org.opensearch.sql.monitor.profile.QueryProfiling; import org.opensearch.sql.opensearch.client.OpenSearchClient; +import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; +import org.opensearch.tasks.CancellableTask; /** * Utility class for asynchronously scanning an index. This lets us send background requests to the @@ -106,13 +108,40 @@ public boolean isScanDone() { public void startScanning(OpenSearchRequest request) { if (isAsync()) { ProfileContext ctx = QueryProfiling.current(); + // Capture the coordinator task on the calling thread so the background search thread can + // re-establish it. Without this, applyParentTask() in OpenSearchNodeClient sees a null task + // on the background pool and the prefetched DSL search task loses its parent (SQL/PPL) link. + CancellableTask task = OpenSearchQueryManager.getCancellableTask(); nextBatchFuture = CompletableFuture.supplyAsync( - () -> QueryProfiling.withCurrentContext(ctx, () -> client.search(request)), + () -> QueryProfiling.withCurrentContext(ctx, () -> searchWithTask(request, task)), backgroundExecutor); } } + /** + * Runs {@code client.search} with the given coordinator task bound to this (pooled) thread's + * ThreadLocal, so parent-task linkage is applied to the outgoing DSL search request. Restores the + * thread's previous task afterward to keep the shared background pool clean. + */ + private OpenSearchResponse searchWithTask( + OpenSearchRequest request, @Nullable CancellableTask task) { + if (task == null) { + return client.search(request); + } + CancellableTask previous = OpenSearchQueryManager.getCancellableTask(); + OpenSearchQueryManager.setCancellableTask(task); + try { + return client.search(request); + } finally { + if (previous != null) { + OpenSearchQueryManager.setCancellableTask(previous); + } else { + OpenSearchQueryManager.clearCancellableTask(); + } + } + } + private OpenSearchResponse getCurrentResponse(OpenSearchRequest request) { if (isAsync()) { try { @@ -176,8 +205,10 @@ public SearchBatchResult fetchNextBatch(OpenSearchRequest request) { // Pre-fetch next batch if needed if (!stopIteration && isAsync()) { + CancellableTask task = OpenSearchQueryManager.getCancellableTask(); nextBatchFuture = - CompletableFuture.supplyAsync(() -> client.search(request), backgroundExecutor); + CompletableFuture.supplyAsync( + () -> searchWithTask(request, task), backgroundExecutor); } } else { iterator = Collections.emptyIterator(); diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java index c2d20a8a303..1755457a491 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -101,6 +101,7 @@ import org.opensearch.sql.legacy.esdomain.LocalClusterState; import org.opensearch.sql.legacy.metrics.Metrics; import org.opensearch.sql.legacy.plugin.RestSqlAction; +import org.opensearch.sql.legacy.plugin.SqlCoordinatorTaskDispatcher; import org.opensearch.sql.legacy.plugin.RestSqlStatsAction; import org.opensearch.sql.opensearch.client.OpenSearchNodeClient; import org.opensearch.sql.opensearch.setting.OpenSearchSettings; @@ -119,7 +120,11 @@ import org.opensearch.sql.plugin.rest.RestQuerySettingsAction; import org.opensearch.sql.plugin.rest.RestUnifiedQueryAction; import org.opensearch.sql.plugin.transport.PPLQueryAction; +import org.opensearch.sql.plugin.transport.SqlQueryAction; import org.opensearch.sql.plugin.transport.TransportPPLQueryAction; +import org.opensearch.sql.plugin.transport.TransportSqlQueryAction; +import org.opensearch.sql.plugin.transport.TransportSqlQueryRequest; +import org.opensearch.sql.plugin.transport.TransportSqlQueryResponse; import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse; import org.opensearch.sql.prometheus.storage.PrometheusStorageFactory; import org.opensearch.sql.protocol.response.format.JsonResponseFormatter; @@ -217,7 +222,8 @@ public List getRestHandlers( return Arrays.asList( new RestPPLQueryAction(), new RestPPLGrammarAction(), - new RestSqlAction(settings, injector, createSqlAnalyticsRouter()), + new RestSqlAction( + settings, injector, createSqlAnalyticsRouter(), createSqlCoordinatorTaskDispatcher()), new RestSqlStatsAction(settings, restController), new RestPPLStatsAction(settings, restController), new RestQuerySettingsAction(settings, restController), @@ -323,6 +329,30 @@ public void onFailure(Exception e) { }; } + /** + * Dispatcher that runs each SQL execution under a coordinator {@link SqlQueryTask} via a local + * transport action, so the DSL search tasks the SQL engine spawns reference the SQL query as + * their parent. + */ + private SqlCoordinatorTaskDispatcher createSqlCoordinatorTaskDispatcher() { + return (nodeClient, query, channel, work) -> + nodeClient.executeLocally( + SqlQueryAction.INSTANCE, + new TransportSqlQueryRequest(query, work, channel), + new ActionListener() { + @Override + public void onResponse(TransportSqlQueryResponse response) { + // The SQL result was already written to the REST channel by the execution path. + } + + @Override + public void onFailure(Exception e) { + // Only reached when execution failed before sending any response. + RestSqlAction.handleException(channel, e); + } + }); + } + /** Register action and handler so that transportClient can find proxy for action. */ @Override public List> getActions() { @@ -330,6 +360,9 @@ public void onFailure(Exception e) { new ActionHandler<>( new ActionType<>(PPLQueryAction.NAME, TransportPPLQueryResponse::new), TransportPPLQueryAction.class), + new ActionHandler<>( + new ActionType<>(SqlQueryAction.NAME, TransportSqlQueryResponse::new), + TransportSqlQueryAction.class), new ActionHandler<>( new ActionType<>( TransportCreateDataSourceAction.NAME, CreateDataSourceActionResponse::new), @@ -507,15 +540,6 @@ public List> getSettings() { .build(); } - @Override - public Collection getTaskHeaders() { - return List.of( - org.opensearch.sql.common.utils.QuerySourceHeaders.QUERY_SOURCE_HEADER, - org.opensearch.sql.common.utils.QuerySourceHeaders.ORIGINAL_QUERY_HEADER, - org.opensearch.sql.common.utils.QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER, - org.opensearch.sql.common.utils.QuerySourceHeaders.QUERY_PHASES_HEADER); - } - @Override public ScriptEngine getScriptEngine(Settings settings, Collection> contexts) { return new CompoundedScriptEngine(); diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryAction.java new file mode 100644 index 00000000000..9a7dc816f20 --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryAction.java @@ -0,0 +1,22 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import org.opensearch.action.ActionType; + +/** + * Internal action used to create a coordinator task for a SQL query. Not exposed as a public REST + * API; it is dispatched locally from {@code RestSqlAction} via {@code NodeClient.executeLocally} so + * the transport framework registers a {@link SqlQueryTask} for the duration of the query. + */ +public class SqlQueryAction extends ActionType { + public static final String NAME = "cluster:admin/opensearch/sql"; + public static final SqlQueryAction INSTANCE = new SqlQueryAction(); + + private SqlQueryAction() { + super(NAME, TransportSqlQueryResponse::new); + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryTask.java new file mode 100644 index 00000000000..0b01dd13fef --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryTask.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import java.util.Map; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.tasks.CancellableTask; + +/** + * Coordinator task for a SQL query. Registering this task in the {@code TaskManager} gives every DSL + * search task spawned by the SQL engine a parent reference back to the originating SQL query, so + * downstream consumers (e.g. query-insights) can correlate the DSL searches with their SQL source + * and look up the original request from the task description. + */ +public class SqlQueryTask extends CancellableTask { + public SqlQueryTask( + long id, + String type, + String action, + String description, + TaskId parentTaskId, + Map headers) { + super(id, type, action, description, parentTaskId, headers); + } + + @Override + public boolean shouldCancelChildrenOnCancellation() { + return true; + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java index a3d80c104da..3edd9a40908 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java @@ -14,7 +14,6 @@ import java.util.Locale; import java.util.Map; import java.util.Optional; -import java.util.UUID; import java.util.function.Supplier; import org.apache.calcite.rel.RelNode; import org.apache.logging.log4j.LogManager; @@ -32,7 +31,6 @@ import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.common.utils.QueryContext; -import org.opensearch.sql.common.utils.QuerySourceHeaders; import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.datasources.service.DataSourceServiceImpl; import org.opensearch.sql.executor.AnalyzeResponse; @@ -182,24 +180,6 @@ protected void doExecute( // in order to use PPL service, we need to convert TransportPPLQueryRequest to PPLQueryRequest PPLQueryRequest transformedRequest = transportRequest.toPPLQueryRequest(); - // Tag the thread context so query-insights can identify this as a PPL-derived query. - org.opensearch.common.util.concurrent.ThreadContext threadContext = - clientRef.threadPool().getThreadContext(); - if (threadContext.getHeader(QuerySourceHeaders.QUERY_SOURCE_HEADER) == null) { - threadContext.putHeader(QuerySourceHeaders.QUERY_SOURCE_HEADER, "ppl"); - } - if (threadContext.getHeader(QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER) == null) { - threadContext.putHeader( - QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER, UUID.randomUUID().toString()); - } - if (threadContext.getHeader(QuerySourceHeaders.ORIGINAL_QUERY_HEADER) == null - && transformedRequest.getRequest() != null) { - String pplQueryText = transformedRequest.getRequest(); - if (pplQueryText.length() > QuerySourceHeaders.MAX_ORIGINAL_QUERY_LENGTH) { - pplQueryText = pplQueryText.substring(0, QuerySourceHeaders.MAX_ORIGINAL_QUERY_LENGTH); - } - threadContext.putHeader(QuerySourceHeaders.ORIGINAL_QUERY_HEADER, pplQueryText); - } QueryContext.setProfile(transformedRequest.profile()); ActionListener clearingListener = wrapWithProfilingClear(listener); diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java new file mode 100644 index 00000000000..e9802706d11 --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java @@ -0,0 +1,134 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import java.io.IOException; +import java.util.concurrent.atomic.AtomicBoolean; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.common.inject.Inject; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.xcontent.MediaType; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.rest.RestChannel; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.RestResponse; +import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.Task; +import org.opensearch.transport.TransportService; + +/** + * Establishes a coordinator {@link SqlQueryTask} for the duration of a SQL query. The transport + * framework registers the task (via {@link TransportSqlQueryRequest#createTask}) before {@link + * #doExecute} runs and unregisters it when the response listener completes. We bind the task to the + * executing thread's {@code OpenSearchQueryManager} ThreadLocal so the existing {@code + * applyParentTask} logic in {@code OpenSearchNodeClient} stamps it as the parent on every DSL search + * the SQL engine issues. The listener is completed once the wrapped REST channel emits its response, + * keeping the task alive across the (possibly asynchronous) execution. + */ +public class TransportSqlQueryAction + extends HandledTransportAction { + + @Inject + public TransportSqlQueryAction(TransportService transportService, ActionFilters actionFilters) { + super(SqlQueryAction.NAME, transportService, actionFilters, TransportSqlQueryRequest::new); + } + + @Override + protected void doExecute( + Task task, + TransportSqlQueryRequest request, + ActionListener listener) { + if (task instanceof CancellableTask cancellableTask) { + OpenSearchQueryManager.setCancellableTask(cancellableTask); + } + + AtomicBoolean completed = new AtomicBoolean(false); + RestChannel channel = new CompletionSignalingChannel(request.getChannel(), listener, completed); + try { + request.getWork().accept(channel); + } catch (Exception e) { + // Only surface here if the channel hasn't already reported completion (success or error). + if (completed.compareAndSet(false, true)) { + listener.onFailure(e); + } + } + } + + /** + * Delegates all channel behavior to the real channel, and completes the transport listener the + * first time a response is sent — which is what triggers unregistration of the coordinator task. + */ + private static final class CompletionSignalingChannel implements RestChannel { + private final RestChannel delegate; + private final ActionListener listener; + private final AtomicBoolean completed; + + CompletionSignalingChannel( + RestChannel delegate, + ActionListener listener, + AtomicBoolean completed) { + this.delegate = delegate; + this.listener = listener; + this.completed = completed; + } + + @Override + public void sendResponse(RestResponse response) { + try { + delegate.sendResponse(response); + } finally { + if (completed.compareAndSet(false, true)) { + listener.onResponse(new TransportSqlQueryResponse()); + } + } + } + + @Override + public XContentBuilder newBuilder() throws IOException { + return delegate.newBuilder(); + } + + @Override + public XContentBuilder newErrorBuilder() throws IOException { + return delegate.newErrorBuilder(); + } + + @Override + public XContentBuilder newBuilder(MediaType mediaType, boolean useFiltering) throws IOException { + return delegate.newBuilder(mediaType, useFiltering); + } + + @Override + public XContentBuilder newBuilder( + MediaType mediaType, MediaType responseContentType, boolean useFiltering) + throws IOException { + return delegate.newBuilder(mediaType, responseContentType, useFiltering); + } + + @Override + public BytesStreamOutput bytesOutput() { + return delegate.bytesOutput(); + } + + @Override + public RestRequest request() { + return delegate.request(); + } + + @Override + public boolean detailedErrorsEnabled() { + return delegate.detailedErrorsEnabled(); + } + + @Override + public boolean detailedErrorStackTraceEnabled() { + return delegate.detailedErrorStackTraceEnabled(); + } + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequest.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequest.java new file mode 100644 index 00000000000..12d8b969577 --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequest.java @@ -0,0 +1,83 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import java.io.IOException; +import java.util.Map; +import java.util.function.Consumer; +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.rest.RestChannel; +import org.opensearch.tasks.Task; + +/** + * Request for {@link SqlQueryAction}. Dispatched only locally via {@code + * NodeClient.executeLocally}, so the {@link #work} and {@link #channel} references are carried + * in-process (never serialized). The {@link #query} text becomes the coordinator task description so + * consumers can recover the original SQL from the task. + */ +public class TransportSqlQueryRequest extends ActionRequest { + + /** Truncate very long queries so the task description stays bounded. */ + static final int MAX_DESCRIPTION_LENGTH = 4096; + + private final String query; + + /** The SQL execution to run under the coordinator task. Transient — local dispatch only. */ + private final transient Consumer work; + + /** The REST channel the execution writes its response to. Transient — local dispatch only. */ + private final transient RestChannel channel; + + public TransportSqlQueryRequest(String query, Consumer work, RestChannel channel) { + this.query = query == null ? "" : query; + this.work = work; + this.channel = channel; + } + + public TransportSqlQueryRequest(StreamInput in) throws IOException { + super(in); + this.query = in.readString(); + this.work = null; + this.channel = null; + } + + public Consumer getWork() { + return work; + } + + public RestChannel getChannel() { + return channel; + } + + @Override + public Task createTask( + long id, String type, String action, TaskId parentTaskId, Map headers) { + return new SqlQueryTask(id, type, action, getDescription(), parentTaskId, headers); + } + + @Override + public String getDescription() { + if (query.length() > MAX_DESCRIPTION_LENGTH) { + return query.substring(0, MAX_DESCRIPTION_LENGTH); + } + return query; + } + + @Override + public ActionRequestValidationException validate() { + return null; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(query); + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryResponse.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryResponse.java new file mode 100644 index 00000000000..4d1a891122f --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryResponse.java @@ -0,0 +1,30 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import java.io.IOException; +import org.opensearch.core.action.ActionResponse; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; + +/** + * Completion signal for {@link SqlQueryAction}. The actual SQL response is written directly to the + * REST channel by the execution path; this response only marks that the coordinator task can be + * unregistered, so it carries no payload. + */ +public class TransportSqlQueryResponse extends ActionResponse { + + public TransportSqlQueryResponse() {} + + public TransportSqlQueryResponse(StreamInput in) throws IOException { + super(in); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + // No payload: the SQL result is delivered via the REST channel, not this transport response. + } +} diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java index 29c15744212..7f117e7bb0b 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -185,9 +185,6 @@ private AbstractPlan plan( ResponseListener queryListener, ResponseListener explainListener) { // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) - org.opensearch.sql.common.utils.QueryPhaseTracker tracker = - org.opensearch.sql.common.utils.QueryPhaseTracker.start(); - tracker.beginPhase("parse"); ParseTree cst = parser.parse(request.getRequest()); Statement statement = cst.accept( @@ -205,8 +202,6 @@ private AbstractPlan plan( : null) .explainMode(request.getExplainMode()) .build())); - tracker.endCurrentPhase(); - tracker.persist(); log.info( "[{}] Incoming request {}", diff --git a/semantic-review/2025-07-15-162000-pr-sql-query-insights.md b/semantic-review/2025-07-15-162000-pr-sql-query-insights.md deleted file mode 100644 index e366551097e..00000000000 --- a/semantic-review/2025-07-15-162000-pr-sql-query-insights.md +++ /dev/null @@ -1,99 +0,0 @@ -# Query phase tracking and source-header propagation for query-insights - -A new `QueryPhaseTracker` collects wall-clock, CPU, and memory-allocation timings per query execution phase (parse, analyze, plan, execute) and serializes them into thread-context headers so the query-insights plugin can attribute costs to SQL/PPL queries. The entry points (`RestSqlAction`, `TransportPPLQueryAction`) stamp source-identification headers, and `SQLPlugin.getTaskHeaders` registers them for cross-node transport. The tracker uses Log4j `ThreadContext` as a shuttle to cross from the REST thread to the sql-worker thread. - -Watch for: -- **Unconditional `putHeader` in `RestSqlAction`** (confirmed) — unlike the PPL path, the SQL REST handler calls `putHeader` without a null-guard, which will throw `IllegalArgumentException` if the header is already present. -- **`Thread.currentThread().getId()` deprecation** (confirmed) — deprecated since Java 19, replaced by `threadId()`. The project targets Java 21. -- **ThreadLocal leak on exception paths** (likely) — if `executeWithCalcite` throws before reaching `endAll()` / `writePhaseHeader()`, the tracker remains in the ThreadLocal and the Log4j key is never cleaned. -- **Missing imports / FQN usage in PPLService and SQLService** (confirmed) — both files use fully-qualified `org.opensearch.sql.common.utils.QueryPhaseTracker` inline instead of an import statement, inconsistent with the rest of the codebase. - -## High-level view - -The header-writing entry points differ in safety: `TransportPPLQueryAction` guards each `putHeader` with a null-check on `getHeader`, while `RestSqlAction` writes unconditionally. OpenSearch's `ThreadContext.putHeader` throws if the key already exists, so the SQL path is fragile in any scenario where the handler runs more than once per request context (retries, plugin chaining). - -The lifecycle management has a gap in the Calcite path: if an exception escapes between `beginPhase("analyze")` and `endAll()`, neither `endAll()` nor `clear()` is reached, leaking the ThreadLocal and Log4j entry on the pooled thread. - -`writePhaseHeader` in `OpenSearchExecutionEngine` is best-effort (catch-all around `putHeader`), which is the right call for observability plumbing — a failure to write metrics should never fail the query. - -

-Issues (6) - -1. **Unconditional putHeader in RestSqlAction** — wrap each `putHeader` call with a `getHeader == null` guard, matching the PPL pattern, to prevent `IllegalArgumentException` if the header already exists. -2. **Deprecated `Thread.currentThread().getId()`** — replace with `Thread.currentThread().threadId()` (available since Java 19; project targets 21). -3. **ThreadLocal leak on exception in executeWithCalcite** — add a `finally` block (or catch) that calls `QueryPhaseTracker.clear()` so the tracker and Log4j key are cleaned on failure paths. -4. **FQN usage instead of imports in PPLService/SQLService** — add proper import statements for `QueryPhaseTracker` to `PPLService.java` and `SQLService.java` for consistency and readability. -5. **Unused `isEmpty()` method** — `isEmpty()` is public but never called anywhere in the diff or existing code. Either document its intended consumer or remove dead code. -6. **No `tracker.endAll()` on legacy V2 exception path** — in `executeWithLegacy`, the tracker begins the "plan" phase but `endAll()` only fires inside `executePlan` if `current()` is non-null. If `plan(plan)` throws before that point, phases are left dangling. - -
- -
-Details - -## Unconditional putHeader in RestSqlAction vs guarded PPL path - -In `RestSqlAction` (lines 161-175), headers are set without checking whether they already exist: - -```java -client.threadPool().getThreadContext() - .putHeader(QuerySourceHeaders.QUERY_SOURCE_HEADER, "sql"); -client.threadPool().getThreadContext() - .putHeader(QuerySourceHeaders.QUERY_EXECUTION_ID_HEADER, UUID.randomUUID().toString()); -``` - -OpenSearch's `ThreadContext.putHeader` throws `IllegalArgumentException` if the key is already present. The PPL transport action correctly guards with `if (threadContext.getHeader(...) == null)`. The SQL path should follow the same pattern. - -## Deprecated Thread.currentThread().getId() - -`QueryPhaseTracker` calls `Thread.currentThread().getId()` on lines 134 and 148 to pass to `getThreadAllocatedBytes(long)`. Since Java 19, `Thread.getId()` is deprecated in favour of `Thread.threadId()`. With the project targeting Java 21, this will produce deprecation warnings and should be: - -```java -SUN_THREAD_MX.getThreadAllocatedBytes(Thread.currentThread().threadId()) -``` - -## ThreadLocal and Log4j key lifecycle on exception paths - -In `QueryService.executeWithCalcite`, the tracker is created at the top of the try block: - -```java -QueryPhaseTracker tracker = QueryPhaseTracker.startOrRestore(); -tracker.beginPhase("analyze"); -``` - -If `StageErrorHandler.executeStage` throws (e.g. `CalciteUnsupportedException`), control jumps to the catch block which calls `executeWithLegacy`. That path creates a *new* tracker via `startOrRestore()`, which overwrites `CURRENT` — so the first tracker's ThreadLocal slot is released. However, the Log4j `ThreadContext` key `_sql_phase_tracker` from the initial `persist()` (set on the REST thread in `SQLService`) is never cleaned if the whole request fails before reaching `writePhaseHeader`. The fix: call `QueryPhaseTracker.clear()` in a `finally` block at the outermost scope of `executeWithCalcite`. - -## FQN usage in PPLService and SQLService - -Both `PPLService.java` and `SQLService.java` use the fully-qualified class name inline: - -```java -org.opensearch.sql.common.utils.QueryPhaseTracker tracker = - org.opensearch.sql.common.utils.QueryPhaseTracker.start(); -``` - -Every other file in this diff uses an import statement. Add `import org.opensearch.sql.common.utils.QueryPhaseTracker;` to each file and use the short name. - -## Serialization format delimiter assumption - -The format `phase:wallNanos|cpu:cpuNanos|mem:memBytes` uses `:` and `|` as delimiters without escaping. Currently all phase names are hardcoded safe strings ("parse", "analyze", "plan", "total"), so this is not a bug today. Adding a defensive check in `beginPhase` (e.g. `assert !name.contains(":") && !name.contains("|")`) guards against future misuse. - -
- -
-File map - -| File | Change | -|------|--------| -| `common/.../QueryPhaseTracker.java` | New. Thread-local phase tracker with wall/CPU/mem metrics and Log4j shuttle. | -| `common/.../QuerySourceHeaders.java` | New. Constants for x-query-source/original-query/execution-id/phases headers. | -| `common/.../QueryPhaseTrackerTest.java` | New. Unit tests covering lifecycle, serialization, cross-thread restore. | -| `core/.../QueryService.java` | Integrates tracker into `executeWithCalcite` and `executeWithLegacy` paths. | -| `legacy/.../RestSqlAction.java` | Stamps source-identification headers on the REST thread for SQL queries. | -| `opensearch/.../OpenSearchExecutionEngine.java` | `writePhaseHeader` writes serialized phases into thread-context header. | -| `plugin/.../SQLPlugin.java` | Registers query-insights headers via `getTaskHeaders`. | -| `plugin/.../TransportPPLQueryAction.java` | Stamps source-identification headers for PPL queries (guarded). | -| `ppl/.../PPLService.java` | Starts tracker and tracks parse phase for PPL. | -| `sql/.../SQLService.java` | Starts tracker and tracks parse phase for SQL. | - -
diff --git a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java index c15b8cf5037..9b4cf8c1a37 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java +++ b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java @@ -99,9 +99,6 @@ private AbstractPlan plan( explainListener); } else { // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) - org.opensearch.sql.common.utils.QueryPhaseTracker tracker = - org.opensearch.sql.common.utils.QueryPhaseTracker.start(); - tracker.beginPhase("parse"); ParseTree cst = parser.parse(request.getQuery()); Statement statement = cst.accept( @@ -112,8 +109,6 @@ private AbstractPlan plan( .fetchSize(request.getFetchSize()) .format(request.getFormat()) .build())); - tracker.endCurrentPhase(); - tracker.persist(); return queryExecutionFactory.create(statement, queryListener, explainListener); } From 06af1718de57ef989da5c1876bb83da9e15ba4b0 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 18 Aug 2026 02:44:34 -0700 Subject: [PATCH 3/4] test: add unit tests for SQL coordinator task and dispatcher - SqlQueryTaskTest: child-cancellation flag and cancellation behavior - TransportSqlQueryRequestTest: createTask type, description text + truncation + null handling, work/channel carriage, validate - TransportSqlQueryActionTest: listener completion on response, cancellable task binding during execution, failure before response, no double completion when work throws after response, single completion on repeated sendResponse - SqlCoordinatorTaskDispatcherTest: PASSTHROUGH runs work with the channel Also apply spotless formatting to the coordinator-task sources. Signed-off-by: Kishore Natarajan --- .../plugin/SqlCoordinatorTaskDispatcher.java | 6 +- .../SqlCoordinatorTaskDispatcherTest.java | 38 +++++ .../org/opensearch/sql/plugin/SQLPlugin.java | 4 +- .../sql/plugin/transport/SqlQueryTask.java | 8 +- .../transport/TransportSqlQueryAction.java | 11 +- .../transport/TransportSqlQueryRequest.java | 4 +- .../plugin/transport/SqlQueryTaskTest.java | 34 +++++ .../TransportSqlQueryActionTest.java | 142 ++++++++++++++++++ .../TransportSqlQueryRequestTest.java | 66 ++++++++ 9 files changed, 296 insertions(+), 17 deletions(-) create mode 100644 legacy/src/test/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcherTest.java create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/SqlQueryTaskTest.java create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryActionTest.java create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequestTest.java diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcher.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcher.java index 93913738b38..d2db1814de1 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcher.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcher.java @@ -23,8 +23,7 @@ public interface SqlCoordinatorTaskDispatcher { /** Default: run the work directly without establishing a coordinator task. */ - SqlCoordinatorTaskDispatcher PASSTHROUGH = - (client, query, channel, work) -> work.accept(channel); + SqlCoordinatorTaskDispatcher PASSTHROUGH = (client, query, channel, work) -> work.accept(channel); /** * @param client node client used to dispatch the local coordinator-task action @@ -32,6 +31,5 @@ public interface SqlCoordinatorTaskDispatcher { * @param channel REST channel the execution writes its response to * @param work the SQL execution to run under the coordinator task */ - void dispatch( - NodeClient client, String query, RestChannel channel, Consumer work); + void dispatch(NodeClient client, String query, RestChannel channel, Consumer work); } diff --git a/legacy/src/test/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcherTest.java b/legacy/src/test/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcherTest.java new file mode 100644 index 00000000000..a1f4ac1662b --- /dev/null +++ b/legacy/src/test/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcherTest.java @@ -0,0 +1,38 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.legacy.plugin; + +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; + +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import org.junit.Test; +import org.opensearch.rest.RestChannel; +import org.opensearch.transport.client.node.NodeClient; + +public class SqlCoordinatorTaskDispatcherTest { + + @Test + public void passthroughRunsWorkWithGivenChannel() { + NodeClient client = mock(NodeClient.class); + RestChannel channel = mock(RestChannel.class); + AtomicBoolean ran = new AtomicBoolean(false); + AtomicReference received = new AtomicReference<>(); + Consumer work = + ch -> { + ran.set(true); + received.set(ch); + }; + + SqlCoordinatorTaskDispatcher.PASSTHROUGH.dispatch(client, "SELECT 1", channel, work); + + assertTrue(ran.get()); + assertSame(channel, received.get()); + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java index 1755457a491..f23d3e61617 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -101,8 +101,8 @@ import org.opensearch.sql.legacy.esdomain.LocalClusterState; import org.opensearch.sql.legacy.metrics.Metrics; import org.opensearch.sql.legacy.plugin.RestSqlAction; -import org.opensearch.sql.legacy.plugin.SqlCoordinatorTaskDispatcher; import org.opensearch.sql.legacy.plugin.RestSqlStatsAction; +import org.opensearch.sql.legacy.plugin.SqlCoordinatorTaskDispatcher; import org.opensearch.sql.opensearch.client.OpenSearchNodeClient; import org.opensearch.sql.opensearch.setting.OpenSearchSettings; import org.opensearch.sql.opensearch.storage.OpenSearchDataSourceFactory; @@ -122,10 +122,10 @@ import org.opensearch.sql.plugin.transport.PPLQueryAction; import org.opensearch.sql.plugin.transport.SqlQueryAction; import org.opensearch.sql.plugin.transport.TransportPPLQueryAction; +import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse; import org.opensearch.sql.plugin.transport.TransportSqlQueryAction; import org.opensearch.sql.plugin.transport.TransportSqlQueryRequest; import org.opensearch.sql.plugin.transport.TransportSqlQueryResponse; -import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse; import org.opensearch.sql.prometheus.storage.PrometheusStorageFactory; import org.opensearch.sql.protocol.response.format.JsonResponseFormatter; import org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style; diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryTask.java index 0b01dd13fef..608742040f1 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryTask.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryTask.java @@ -10,10 +10,10 @@ import org.opensearch.tasks.CancellableTask; /** - * Coordinator task for a SQL query. Registering this task in the {@code TaskManager} gives every DSL - * search task spawned by the SQL engine a parent reference back to the originating SQL query, so - * downstream consumers (e.g. query-insights) can correlate the DSL searches with their SQL source - * and look up the original request from the task description. + * Coordinator task for a SQL query. Registering this task in the {@code TaskManager} gives every + * DSL search task spawned by the SQL engine a parent reference back to the originating SQL query, + * so downstream consumers (e.g. query-insights) can correlate the DSL searches with their SQL + * source and look up the original request from the task description. */ public class SqlQueryTask extends CancellableTask { public SqlQueryTask( diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java index e9802706d11..ef21d36b9bf 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java @@ -10,10 +10,10 @@ import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.HandledTransportAction; import org.opensearch.common.inject.Inject; +import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.action.ActionListener; import org.opensearch.core.xcontent.MediaType; import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.rest.RestChannel; import org.opensearch.rest.RestRequest; import org.opensearch.rest.RestResponse; @@ -27,9 +27,9 @@ * framework registers the task (via {@link TransportSqlQueryRequest#createTask}) before {@link * #doExecute} runs and unregisters it when the response listener completes. We bind the task to the * executing thread's {@code OpenSearchQueryManager} ThreadLocal so the existing {@code - * applyParentTask} logic in {@code OpenSearchNodeClient} stamps it as the parent on every DSL search - * the SQL engine issues. The listener is completed once the wrapped REST channel emits its response, - * keeping the task alive across the (possibly asynchronous) execution. + * applyParentTask} logic in {@code OpenSearchNodeClient} stamps it as the parent on every DSL + * search the SQL engine issues. The listener is completed once the wrapped REST channel emits its + * response, keeping the task alive across the (possibly asynchronous) execution. */ public class TransportSqlQueryAction extends HandledTransportAction { @@ -100,7 +100,8 @@ public XContentBuilder newErrorBuilder() throws IOException { } @Override - public XContentBuilder newBuilder(MediaType mediaType, boolean useFiltering) throws IOException { + public XContentBuilder newBuilder(MediaType mediaType, boolean useFiltering) + throws IOException { return delegate.newBuilder(mediaType, useFiltering); } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequest.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequest.java index 12d8b969577..1d48b6e0319 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequest.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequest.java @@ -19,8 +19,8 @@ /** * Request for {@link SqlQueryAction}. Dispatched only locally via {@code * NodeClient.executeLocally}, so the {@link #work} and {@link #channel} references are carried - * in-process (never serialized). The {@link #query} text becomes the coordinator task description so - * consumers can recover the original SQL from the task. + * in-process (never serialized). The {@link #query} text becomes the coordinator task description + * so consumers can recover the original SQL from the task. */ public class TransportSqlQueryRequest extends ActionRequest { diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/SqlQueryTaskTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/SqlQueryTaskTest.java new file mode 100644 index 00000000000..a63af969dda --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/SqlQueryTaskTest.java @@ -0,0 +1,34 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import java.util.Map; +import org.junit.Test; +import org.opensearch.core.tasks.TaskId; + +public class SqlQueryTaskTest { + + private SqlQueryTask newTask() { + return new SqlQueryTask( + 1, "transport", SqlQueryAction.NAME, "SELECT 1", TaskId.EMPTY_TASK_ID, Map.of()); + } + + @Test + public void testShouldCancelChildrenReturnsTrue() { + assertTrue(newTask().shouldCancelChildrenOnCancellation()); + } + + @Test + public void testCancellation() { + SqlQueryTask task = newTask(); + assertFalse(task.isCancelled()); + task.cancel("Test"); + assertTrue(task.isCancelled()); + } +} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryActionTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryActionTest.java new file mode 100644 index 00000000000..f51a4e222ee --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryActionTest.java @@ -0,0 +1,142 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import static org.junit.Assert.assertSame; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +import java.util.HashSet; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.rest.RestChannel; +import org.opensearch.rest.RestResponse; +import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager; +import org.opensearch.tasks.Task; +import org.opensearch.transport.TransportService; + +public class TransportSqlQueryActionTest { + + private TransportSqlQueryAction action; + + @Before + public void setUp() { + TransportService transportService = mock(TransportService.class); + action = new TransportSqlQueryAction(transportService, new ActionFilters(new HashSet<>())); + } + + @After + public void tearDown() { + // doExecute binds the task to a ThreadLocal and does not clear it on the caller thread. + OpenSearchQueryManager.clearCancellableTask(); + } + + private SqlQueryTask newTask() { + return new SqlQueryTask( + 1, "transport", SqlQueryAction.NAME, "SELECT 1", TaskId.EMPTY_TASK_ID, Map.of()); + } + + @SuppressWarnings("unchecked") + private ActionListener mockListener() { + return mock(ActionListener.class); + } + + @Test + public void completesListenerOnceWhenWorkSendsResponse() { + RestChannel delegate = mock(RestChannel.class); + RestResponse response = mock(RestResponse.class); + Consumer work = ch -> ch.sendResponse(response); + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + ActionListener listener = mockListener(); + + action.doExecute(newTask(), request, listener); + + verify(delegate, times(1)).sendResponse(response); + verify(listener, times(1)).onResponse(any(TransportSqlQueryResponse.class)); + verify(listener, never()).onFailure(any()); + } + + @Test + public void bindsCancellableTaskDuringExecution() { + SqlQueryTask task = newTask(); + RestChannel delegate = mock(RestChannel.class); + RestResponse response = mock(RestResponse.class); + AtomicReference seen = new AtomicReference<>(); + Consumer work = + ch -> { + seen.set(OpenSearchQueryManager.getCancellableTask()); + ch.sendResponse(response); + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + + action.doExecute(task, request, mockListener()); + + assertSame(task, seen.get()); + } + + @Test + public void reportsFailureWhenWorkThrowsBeforeResponse() { + RestChannel delegate = mock(RestChannel.class); + RuntimeException boom = new RuntimeException("boom"); + Consumer work = + ch -> { + throw boom; + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + ActionListener listener = mockListener(); + + action.doExecute(newTask(), request, listener); + + verify(listener, times(1)).onFailure(boom); + verify(listener, never()).onResponse(any()); + } + + @Test + public void doesNotDoubleCompleteWhenWorkThrowsAfterResponse() { + RestChannel delegate = mock(RestChannel.class); + RestResponse response = mock(RestResponse.class); + Consumer work = + ch -> { + ch.sendResponse(response); + throw new RuntimeException("late failure after response already sent"); + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + ActionListener listener = mockListener(); + + action.doExecute(newTask(), request, listener); + + verify(listener, times(1)).onResponse(any(TransportSqlQueryResponse.class)); + verify(listener, never()).onFailure(any()); + } + + @Test + public void secondSendResponseDoesNotCompleteListenerTwice() { + RestChannel delegate = mock(RestChannel.class); + RestResponse first = mock(RestResponse.class); + RestResponse second = mock(RestResponse.class); + Consumer work = + ch -> { + ch.sendResponse(first); + ch.sendResponse(second); + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + ActionListener listener = mockListener(); + + action.doExecute(newTask(), request, listener); + + verify(listener, times(1)).onResponse(any(TransportSqlQueryResponse.class)); + } +} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequestTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequestTest.java new file mode 100644 index 00000000000..c4e62830fe8 --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequestTest.java @@ -0,0 +1,66 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.mockito.Mockito.mock; + +import java.util.Map; +import java.util.function.Consumer; +import org.junit.Test; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.rest.RestChannel; +import org.opensearch.tasks.Task; + +public class TransportSqlQueryRequestTest { + + @Test + public void testCreateTaskReturnsSqlQueryTask() { + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", ch -> {}, null); + Task task = + request.createTask(1, "transport", SqlQueryAction.NAME, TaskId.EMPTY_TASK_ID, Map.of()); + assertNotNull(task); + assertEquals(SqlQueryTask.class, task.getClass()); + } + + @Test + public void testDescriptionIsQueryText() { + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", ch -> {}, null); + assertEquals("SELECT 1", request.getDescription()); + } + + @Test + public void testDescriptionTruncatedToMaxLength() { + String longQuery = "x".repeat(TransportSqlQueryRequest.MAX_DESCRIPTION_LENGTH + 100); + TransportSqlQueryRequest request = new TransportSqlQueryRequest(longQuery, ch -> {}, null); + assertEquals( + TransportSqlQueryRequest.MAX_DESCRIPTION_LENGTH, request.getDescription().length()); + } + + @Test + public void testNullQueryBecomesEmptyDescription() { + TransportSqlQueryRequest request = new TransportSqlQueryRequest(null, ch -> {}, null); + assertEquals("", request.getDescription()); + } + + @Test + public void testCarriesWorkAndChannel() { + Consumer work = ch -> {}; + RestChannel channel = mock(RestChannel.class); + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, channel); + assertSame(work, request.getWork()); + assertSame(channel, request.getChannel()); + } + + @Test + public void testValidateReturnsNull() { + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", ch -> {}, null); + assertNull(request.validate()); + } +} From d12593c1f7b78d5705ab9c62ad526b562e7af1ac Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 18 Aug 2026 03:58:44 -0700 Subject: [PATCH 4/4] fix: clear coordinator task from thread-local after SQL execution Restore the initiating thread's prior cancellable task (or clear it) in a finally block in TransportSqlQueryAction.doExecute so pooled transport workers don't retain a stale task reference. Safe because async workers capture the task synchronously during work.accept before it returns. Expand TransportSqlQueryActionTest to cover the thread-local clear (on both success and failure) and all CompletionSignalingChannel delegation methods. Signed-off-by: Kishore Natarajan --- .../transport/TransportSqlQueryAction.java | 10 +++ .../TransportSqlQueryActionTest.java | 86 +++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java index ef21d36b9bf..73d470cab7e 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java @@ -44,6 +44,7 @@ protected void doExecute( Task task, TransportSqlQueryRequest request, ActionListener listener) { + CancellableTask previous = OpenSearchQueryManager.getCancellableTask(); if (task instanceof CancellableTask cancellableTask) { OpenSearchQueryManager.setCancellableTask(cancellableTask); } @@ -57,6 +58,15 @@ protected void doExecute( if (completed.compareAndSet(false, true)) { listener.onFailure(e); } + } finally { + // Restore the initiating thread's prior task so this pooled transport worker does not retain + // a stale reference. Any async workers continuing the execution have already captured the + // task synchronously during work.accept above, so clearing here does not race them. + if (previous != null) { + OpenSearchQueryManager.setCancellableTask(previous); + } else { + OpenSearchQueryManager.clearCancellableTask(); + } } } diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryActionTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryActionTest.java index f51a4e222ee..baa69193b6d 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryActionTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryActionTest.java @@ -5,13 +5,19 @@ package org.opensearch.sql.plugin.transport; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import java.io.IOException; import java.util.HashSet; import java.util.Map; import java.util.concurrent.atomic.AtomicReference; @@ -20,9 +26,13 @@ import org.junit.Before; import org.junit.Test; import org.opensearch.action.support.ActionFilters; +import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.action.ActionListener; import org.opensearch.core.tasks.TaskId; +import org.opensearch.core.xcontent.MediaType; +import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.rest.RestChannel; +import org.opensearch.rest.RestRequest; import org.opensearch.rest.RestResponse; import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager; import org.opensearch.tasks.Task; @@ -122,6 +132,82 @@ public void doesNotDoubleCompleteWhenWorkThrowsAfterResponse() { verify(listener, never()).onFailure(any()); } + @Test + public void clearsCancellableTaskAfterExecution() { + RestChannel delegate = mock(RestChannel.class); + RestResponse response = mock(RestResponse.class); + Consumer work = ch -> ch.sendResponse(response); + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + + action.doExecute(newTask(), request, mockListener()); + + // The pooled transport thread must not retain the task after doExecute returns. + assertNull(OpenSearchQueryManager.getCancellableTask()); + } + + @Test + public void clearsCancellableTaskEvenWhenWorkThrows() { + RestChannel delegate = mock(RestChannel.class); + Consumer work = + ch -> { + throw new RuntimeException("boom"); + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + + action.doExecute(newTask(), request, mockListener()); + + assertNull(OpenSearchQueryManager.getCancellableTask()); + } + + @Test + public void wrappedChannelDelegatesAllMethods() throws IOException { + RestChannel delegate = mock(RestChannel.class); + XContentBuilder builder = mock(XContentBuilder.class); + XContentBuilder errorBuilder = mock(XContentBuilder.class); + BytesStreamOutput bytesOutput = new BytesStreamOutput(); + RestRequest restRequest = mock(RestRequest.class); + MediaType mediaType = mock(MediaType.class); + when(delegate.newBuilder()).thenReturn(builder); + when(delegate.newErrorBuilder()).thenReturn(errorBuilder); + when(delegate.newBuilder(any(MediaType.class), anyBoolean())).thenReturn(builder); + when(delegate.newBuilder(any(MediaType.class), any(MediaType.class), anyBoolean())) + .thenReturn(builder); + when(delegate.bytesOutput()).thenReturn(bytesOutput); + when(delegate.request()).thenReturn(restRequest); + when(delegate.detailedErrorsEnabled()).thenReturn(true); + when(delegate.detailedErrorStackTraceEnabled()).thenReturn(false); + + AtomicReference failure = new AtomicReference<>(); + Consumer work = + ch -> { + try { + assertSame(builder, ch.newBuilder()); + assertSame(errorBuilder, ch.newErrorBuilder()); + assertSame(builder, ch.newBuilder(mediaType, true)); + assertSame(builder, ch.newBuilder(mediaType, mediaType, true)); + assertSame(bytesOutput, ch.bytesOutput()); + assertSame(restRequest, ch.request()); + assertTrue(ch.detailedErrorsEnabled()); + assertFalse(ch.detailedErrorStackTraceEnabled()); + } catch (AssertionError e) { + failure.set(e); + } catch (IOException e) { + failure.set(new AssertionError(e)); + } + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + + action.doExecute(newTask(), request, mockListener()); + + if (failure.get() != null) { + throw failure.get(); + } + verify(delegate, times(1)).newBuilder(); + verify(delegate, times(1)).newErrorBuilder(); + verify(delegate, times(1)).bytesOutput(); + verify(delegate, times(1)).request(); + } + @Test public void secondSendResponseDoesNotCompleteListenerTwice() { RestChannel delegate = mock(RestChannel.class);