From 393b51710df01f85bd31ca2537d0aef22c725acc Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Thu, 10 Sep 2026 15:11:55 -0700 Subject: [PATCH] feat(ppl): Report completed PPL queries to Query Insights Top N Emit a Query Insights record for every completed PPL query so it appears in the Top N overview alongside DSL queries. Each record carries the query source (PPL), the resolved indices, and query-total latency, CPU, and memory from the task resource-tracking framework. Child DSL searches spawned by a PPL query (e.g. a join's per-table scans) are tagged with a parent marker (X-Query-Insights-Parent) so Query Insights can associate them with their parent as sub-queries. The marker is carried on a ThreadLocal across the engine's worker and complex-worker thread hops and re-applied as a task header on the background scan pool. This is reporting only: no phase breakdown and no profiling changes. The profiling subsystem is left as-is for a future change. Known limitation: for multi-scan queries (joins), child capture can be partial under background-pool thread contention; parent reporting is unaffected. Signed-off-by: Kishore Kumaar Natarajan --- .../executor/OpenSearchQueryManager.java | 107 +++++++++- .../ThreadPoolExecutionDispatcher.java | 21 ++ .../storage/scan/BackgroundSearchScanner.java | 75 ++++++- .../executor/OpenSearchQueryManagerTest.java | 89 ++++++++ .../org/opensearch/sql/plugin/SQLPlugin.java | 11 + .../sql/plugin/transport/PPLQueryTask.java | 11 + .../plugin/transport/QueryInsightsMarker.java | 35 ++++ .../transport/QueryInsightsReporter.java | 105 ++++++++++ .../transport/TransportPPLQueryAction.java | 191 +++++++++++++++++- .../transport/ExtractPplIndicesTest.java | 108 ++++++++++ .../plugin/transport/PPLQueryTaskTest.java | 31 +++ 11 files changed, 778 insertions(+), 6 deletions(-) create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsMarker.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/ExtractPplIndicesTest.java diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java index c391153fca6..26f00031b41 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java @@ -5,6 +5,8 @@ package org.opensearch.sql.opensearch.executor; +import com.sun.management.ThreadMXBean; +import java.lang.management.ManagementFactory; import java.util.Map; import lombok.RequiredArgsConstructor; import org.apache.logging.log4j.LogManager; @@ -12,6 +14,9 @@ import org.apache.logging.log4j.ThreadContext; import org.opensearch.OpenSearchTimeoutException; import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.tasks.resourcetracker.ResourceStats; +import org.opensearch.core.tasks.resourcetracker.ResourceStatsType; +import org.opensearch.core.tasks.resourcetracker.ResourceUsageMetric; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.executor.QueryId; import org.opensearch.sql.executor.QueryManager; @@ -27,6 +32,21 @@ public class OpenSearchQueryManager implements QueryManager { private static final Logger LOG = LogManager.getLogger(OpenSearchQueryManager.class); + /** Samples per-thread CPU/memory for resource tracking; null when the JVM bean is unavailable. */ + private static final ThreadMXBean THREAD_MX_BEAN = resolveThreadMXBean(); + + private static ThreadMXBean resolveThreadMXBean() { + try { + java.lang.management.ThreadMXBean bean = ManagementFactory.getThreadMXBean(); + if (bean instanceof ThreadMXBean) { + return (ThreadMXBean) bean; + } + } catch (Exception e) { + LOG.warn("Per-thread resource metrics unavailable; PPL task resource tracking disabled", e); + } + return null; + } + private final NodeClient nodeClient; private final Settings settings; @@ -49,24 +69,53 @@ public static void clearCancellableTask() { cancellableTask.remove(); } + /** + * Query Insights parent marker ({@code PPL::}), propagated across the engine's + * thread hops the same way as {@link #cancellableTask}. Reading it off the OpenSearch + * ThreadContext at each hop is racy (join sides / prefetch batches can land on threads that never + * received the header), so a dedicated ThreadLocal makes child tagging deterministic. + */ + private static final ThreadLocal queryInsightsParentMarker = new ThreadLocal<>(); + + public static void setQueryInsightsParentMarker(String marker) { + queryInsightsParentMarker.set(marker); + } + + public static String getQueryInsightsParentMarker() { + return queryInsightsParentMarker.get(); + } + + public static void clearQueryInsightsParentMarker() { + queryInsightsParentMarker.remove(); + } + @Override public QueryId submit(AbstractPlan queryPlan) { TimeValue timeout = settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT); CancellableTask cancelTask = cancellableTask.get(); cancellableTask.remove(); - schedule(nodeClient, queryPlan::execute, timeout, cancelTask); + // Carry the parent marker to the worker thread alongside the cancellable task. + String parentMarker = queryInsightsParentMarker.get(); + queryInsightsParentMarker.remove(); + schedule(nodeClient, queryPlan::execute, timeout, cancelTask, parentMarker); return queryPlan.getQueryId(); } private void schedule( - NodeClient client, Runnable task, TimeValue timeout, CancellableTask cancelTask) { + NodeClient client, + Runnable task, + TimeValue timeout, + CancellableTask cancelTask, + String parentMarker) { ThreadPool threadPool = client.threadPool(); Runnable wrappedTask = withCurrentContext( () -> { final Thread executionThread = Thread.currentThread(); + // Re-establish on this worker thread so downstream scans can read it. + setQueryInsightsParentMarker(parentMarker); Scheduler.ScheduledCancellable timeoutTask = threadPool.schedule( @@ -81,6 +130,16 @@ private void schedule( setCancellableTask(cancelTask); + // Bracket resource tracking for the inline (non-script) path, which runs its work on + // this thread; the worker hop skips core's automatic TaskAwareRunnable association. + // Script plans hand off to the complex-worker pool, which brackets itself (see + // ThreadPoolExecutionDispatcher). Stop only if start succeeded. + final boolean trackResources = + cancelTask != null && cancelTask.supportsResourceTracking(); + final long trackedThreadId = Thread.currentThread().getId(); + final boolean trackingStarted = + trackResources && startThreadResourceTracking(cancelTask, trackedThreadId); + try { task.run(); timeoutTask.cancel(); @@ -98,7 +157,11 @@ private void schedule( throw e; } finally { + if (trackingStarted) { + stopThreadResourceTracking(cancelTask, trackedThreadId); + } clearCancellableTask(); + clearQueryInsightsParentMarker(); } }); @@ -112,4 +175,44 @@ private Runnable withCurrentContext(final Runnable task) { task.run(); }; } + + /** + * Records the starting CPU/memory snapshot for {@code threadId}. + * + * @return true if tracking started; only then should {@link #stopThreadResourceTracking} be + * called. + */ + static boolean startThreadResourceTracking(CancellableTask task, long threadId) { + try { + task.startThreadResourceTracking( + threadId, ResourceStatsType.WORKER_STATS, currentThreadResourceMetrics(threadId)); + return true; + } catch (Exception e) { + LOG.warn("Failed to start resource tracking for task [{}]", task.getId(), e); + return false; + } + } + + /** Records the final CPU/memory snapshot for {@code threadId}. */ + static void stopThreadResourceTracking(CancellableTask task, long threadId) { + try { + task.stopThreadResourceTracking( + threadId, ResourceStatsType.WORKER_STATS, currentThreadResourceMetrics(threadId)); + } catch (Exception e) { + LOG.warn("Failed to stop resource tracking for task [{}]", task.getId(), e); + } + } + + /** Per-thread memory and CPU usage; empty array when the JVM bean is unavailable. */ + private static ResourceUsageMetric[] currentThreadResourceMetrics(long threadId) { + if (THREAD_MX_BEAN == null) { + return new ResourceUsageMetric[0]; + } + ResourceUsageMetric memory = + new ResourceUsageMetric( + ResourceStats.MEMORY, THREAD_MX_BEAN.getThreadAllocatedBytes(threadId)); + ResourceUsageMetric cpu = + new ResourceUsageMetric(ResourceStats.CPU, THREAD_MX_BEAN.getThreadCpuTime(threadId)); + return new ResourceUsageMetric[] {memory, cpu}; + } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java index 3de8bb85c2b..6bb5c896591 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java @@ -69,6 +69,9 @@ private void dispatchInternal( // Capture thread-local state to propagate across thread boundary Map ctx = ThreadContext.getImmutableContext(); CancellableTask cancellableTask = OpenSearchQueryManager.getCancellableTask(); + // Carry the parent marker across to the complex-worker thread; without it, join scans there + // lose the parent link and go untagged. + String queryInsightsParentMarker = OpenSearchQueryManager.getQueryInsightsParentMarker(); ProfileContext profileContext = QueryProfiling.current(); CalcitePlanContext.ThreadLocalSnapshot snapshot = CalcitePlanContext.snapshotThreadLocals(); @Nullable JaninoRelMetadataProvider metadataProvider = @@ -91,10 +94,14 @@ private void dispatchInternal( ThreadPool.Names.GENERIC); Cancellable cancelPoller = scheduleCancellationPoller(cancellableTask, executionThread); Hook.Closeable hookHandle = null; + boolean trackResources = false; + long trackedThreadId = -1L; + boolean trackingStarted = false; try { // Restore state from caller thread ThreadContext.putAll(ctx); OpenSearchQueryManager.setCancellableTask(cancellableTask); + OpenSearchQueryManager.setQueryInsightsParentMarker(queryInsightsParentMarker); QueryProfiling.set(profileContext); CalcitePlanContext.restoreThreadLocals(snapshot); // Override execution pool to indicate complex pool @@ -106,6 +113,16 @@ private void dispatchInternal( hookHandle = Hook.CURRENT_TIME.addThread((Consumer>) h -> h.set(currentTime)); } + // Script plans run their real work on this complex-worker thread (the sql-worker + // returned once it scheduled us), so bracket resource tracking here. Stop only if + // start succeeded. + trackResources = + cancellableTask != null && cancellableTask.supportsResourceTracking(); + trackedThreadId = Thread.currentThread().getId(); + trackingStarted = + trackResources + && OpenSearchQueryManager.startThreadResourceTracking( + cancellableTask, trackedThreadId); task.run(); } catch (Exception e) { LOG.error("Exception during task execution on complex pool", e); @@ -119,7 +136,11 @@ private void dispatchInternal( if (hookHandle != null) { hookHandle.close(); } + if (trackingStarted) { + OpenSearchQueryManager.stopThreadResourceTracking(cancellableTask, trackedThreadId); + } OpenSearchQueryManager.clearCancellableTask(); + OpenSearchQueryManager.clearQueryInsightsParentMarker(); RelMetadataQueryBase.THREAD_PROVIDERS.remove(); CalcitePlanContext.clearTimewrapSignals(); QueryProfiling.clear(); 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..e83ea45b3df 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 @@ -15,14 +15,17 @@ import javax.annotation.Nullable; import org.opensearch.OpenSearchException; import org.opensearch.OpenSearchSecurityException; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.tasks.TaskCancelledException; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.exception.NonFallbackCalciteException; 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 @@ -64,6 +67,16 @@ * cleanup. */ public class BackgroundSearchScanner { + /** + * Task header linking a child DSL search to its originating SQL/PPL query. Must match {@code + * QueryInsightsMarker.PARENT_HEADER} in the {@code plugin} module; duplicated as a literal + * because {@code opensearch} must not depend on {@code plugin}. OpenSearch only copies this + * header onto a child SearchTask from the thread that issues the search, and background scans run + * on a pool thread without the coordinator's context, so it is re-applied per search (see {@link + * #searchWithParentHeader}). + */ + private static final String QUERY_INSIGHTS_PARENT_HEADER = "X-Query-Insights-Parent"; + private final OpenSearchClient client; @Nullable private final Executor backgroundExecutor; private CompletableFuture nextBatchFuture = null; @@ -106,13 +119,68 @@ public boolean isScanDone() { public void startScanning(OpenSearchRequest request) { if (isAsync()) { ProfileContext ctx = QueryProfiling.current(); + // Capture task + parent marker on the calling thread; the background pool thread does not + // inherit them and re-establishes both around the search. + CancellableTask task = OpenSearchQueryManager.getCancellableTask(); + String parentMarker = currentQueryInsightsParentHeader(); nextBatchFuture = CompletableFuture.supplyAsync( - () -> QueryProfiling.withCurrentContext(ctx, () -> client.search(request)), + () -> + QueryProfiling.withCurrentContext( + ctx, () -> searchWithTask(request, task, parentMarker)), backgroundExecutor); } } + /** + * Parent marker on the calling thread, or null for non-SQL/PPL queries or when recording is off. + */ + @Nullable + private String currentQueryInsightsParentHeader() { + return OpenSearchQueryManager.getQueryInsightsParentMarker(); + } + + /** Binds {@code task} to this pool thread for the search, restoring the previous task after. */ + private OpenSearchResponse searchWithTask( + OpenSearchRequest request, @Nullable CancellableTask task, @Nullable String parentMarker) { + if (task == null) { + return searchWithParentHeader(request, parentMarker); + } + CancellableTask previous = OpenSearchQueryManager.getCancellableTask(); + OpenSearchQueryManager.setCancellableTask(task); + try { + return searchWithParentHeader(request, parentMarker); + } finally { + if (previous != null) { + OpenSearchQueryManager.setCancellableTask(previous); + } else { + OpenSearchQueryManager.clearCancellableTask(); + } + } + } + + /** + * Re-applies the parent-marker header to this pool thread's context so OpenSearch stamps it onto + * the child SearchTask. Uses {@link ThreadContext#newStoredContext} (snapshot + restore) rather + * than {@code stashContext()}: stashing would drop the security user transient and run the search + * unauthenticated. Falls back to a plain search when there is no marker or node client. + */ + private OpenSearchResponse searchWithParentHeader( + OpenSearchRequest request, @Nullable String parentMarker) { + if (parentMarker == null || parentMarker.isEmpty() || client.getNodeClient().isEmpty()) { + return client.search(request); + } + ThreadContext threadContext = client.getNodeClient().get().threadPool().getThreadContext(); + // Already set on this pool thread: leave it be, putHeader would throw on a duplicate. + if (threadContext.getHeader(QUERY_INSIGHTS_PARENT_HEADER) != null) { + return client.search(request); + } + try (ThreadContext.StoredContext ignored = threadContext.newStoredContext(true)) { + threadContext.putHeader(QUERY_INSIGHTS_PARENT_HEADER, parentMarker); + return client.search(request); + } + } + private OpenSearchResponse getCurrentResponse(OpenSearchRequest request) { if (isAsync()) { try { @@ -176,8 +244,11 @@ public SearchBatchResult fetchNextBatch(OpenSearchRequest request) { // Pre-fetch next batch if needed if (!stopIteration && isAsync()) { + CancellableTask task = OpenSearchQueryManager.getCancellableTask(); + String parentMarker = currentQueryInsightsParentHeader(); nextBatchFuture = - CompletableFuture.supplyAsync(() -> client.search(request), backgroundExecutor); + CompletableFuture.supplyAsync( + () -> searchWithTask(request, task, parentMarker), backgroundExecutor); } } else { iterator = Collections.emptyIterator(); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java index 1463cf48fff..830bebb58d4 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java @@ -5,18 +5,26 @@ package org.opensearch.sql.opensearch.executor; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import java.util.Collections; +import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.core.tasks.resourcetracker.ThreadResourceInfo; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; @@ -26,6 +34,7 @@ import org.opensearch.sql.executor.QueryType; import org.opensearch.sql.executor.execution.AbstractPlan; import org.opensearch.sql.executor.execution.QueryPlan; +import org.opensearch.tasks.CancellableTask; import org.opensearch.threadpool.Scheduler; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.client.node.NodeClient; @@ -76,4 +85,84 @@ public void execute() { assertTrue(isRun.get()); } + + @AfterEach + public void clearTask() { + OpenSearchQueryManager.clearCancellableTask(); + } + + @Test + public void tracksResourceUsageWhenTaskSupportsIt() { + TrackingTask task = new TrackingTask(true); + OpenSearchQueryManager.setCancellableTask(task); + + runSubmit(); + + // The worker thread should have been bracketed with start/stop tracking, leaving a completed + // (inactive) resource entry on the coordinator task. + Map> stats = task.getResourceStats(); + assertEquals(1, stats.size()); + List infos = stats.values().iterator().next(); + assertEquals(1, infos.size()); + assertFalse("resource entry should be closed after execution", infos.get(0).isActive()); + } + + @Test + public void skipsResourceTrackingWhenTaskDoesNotSupportIt() { + TrackingTask task = new TrackingTask(false); + OpenSearchQueryManager.setCancellableTask(task); + + runSubmit(); + + assertTrue("no resource entries expected", task.getResourceStats().isEmpty()); + } + + /** Runs a trivial query through the manager with the schedule hook executing tasks inline. */ + private void runSubmit() { + NodeClient nodeClient = mock(NodeClient.class); + ThreadPool threadPool = mock(ThreadPool.class); + Settings settings = mock(Settings.class); + Scheduler.ScheduledCancellable mockScheduledTask = mock(Scheduler.ScheduledCancellable.class); + + when(nodeClient.threadPool()).thenReturn(threadPool); + when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) + .thenReturn(TimeValue.timeValueSeconds(60)); + + AbstractPlan queryPlan = + new QueryPlan(queryId, queryType, plan, queryService, listener) { + @Override + public void execute() {} + }; + + doAnswer( + invocation -> { + Runnable task = invocation.getArgument(0); + task.run(); + return mockScheduledTask; + }) + .when(threadPool) + .schedule(any(), any(), any()); + + new OpenSearchQueryManager(nodeClient, settings).submit(queryPlan); + } + + /** Minimal CancellableTask whose resource-tracking support is configurable. */ + private static class TrackingTask extends CancellableTask { + private final boolean supportsTracking; + + TrackingTask(boolean supportsTracking) { + super(1L, "ppl", "action", "desc", TaskId.EMPTY_TASK_ID, Collections.emptyMap()); + this.supportsTracking = supportsTracking; + } + + @Override + public boolean supportsResourceTracking() { + return supportsTracking; + } + + @Override + public boolean shouldCancelChildrenOnCancellation() { + return true; + } + } } 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 cdbc0d5b3be..36ae2885bed 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -119,6 +119,7 @@ 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.QueryInsightsMarker; import org.opensearch.sql.plugin.transport.TransportPPLQueryAction; import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse; import org.opensearch.sql.prometheus.storage.PrometheusStorageFactory; @@ -324,6 +325,16 @@ public void onFailure(Exception e) { }; } + /** + * Register the Query Insights parent-marker header so OpenSearch copies it from a SQL/PPL query's + * thread context into the child DSL search tasks it spawns (including on remote data nodes). This + * lets Query Insights classify each child's source and associate it with the originating query. + */ + @Override + public Collection getTaskHeaders() { + return List.of(QueryInsightsMarker.PARENT_HEADER); + } + /** Register action and handler so that transportClient can find proxy for action. */ @Override public List> getActions() { diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java index 2df96bdbd12..3f0d935ca67 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java @@ -25,4 +25,15 @@ public PPLQueryTask( public boolean shouldCancelChildrenOnCancellation() { return true; } + + /** + * Enable per-thread CPU/memory accounting for the PPL coordinator task so its resource usage (and + * the PPL engine work attributed to it on the {@code sql-worker} pool) is exposed through the + * tasks API and consumed by Query Insights. Defaults to {@code false} on {@link CancellableTask}, + * which is why it must be overridden here. + */ + @Override + public boolean supportsResourceTracking() { + return true; + } } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsMarker.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsMarker.java new file mode 100644 index 00000000000..78248c65c29 --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsMarker.java @@ -0,0 +1,35 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +/** + * Constants and helpers for the task header that links a child DSL search back to the SQL/PPL query + * that spawned it, for Query Insights. + * + *

The value is {@code ::} (e.g. {@code PPL:node-1:42}). Registered via + * {@code SQLPlugin.getTaskHeaders()} so OpenSearch copies it onto child search tasks, where Query + * Insights reads it to classify each child and link it to its parent. Source-neutral so the same + * mechanism serves PPL now and SQL later. + */ +public final class QueryInsightsMarker { + + /** Task header name carrying {@code ::} of the originating query. */ + public static final String PARENT_HEADER = "X-Query-Insights-Parent"; + + private QueryInsightsMarker() {} + + /** + * Build the header value for a coordinator query. + * + * @param source query source label (e.g. {@code "PPL"}) + * @param nodeId coordinator node id + * @param taskId coordinator task id + * @return the {@code ::} header value + */ + public static String value(String source, String nodeId, long taskId) { + return source + ":" + nodeId + ":" + taskId; + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java new file mode 100644 index 00000000000..0c8adae97cb --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java @@ -0,0 +1,105 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.Version; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.transport.BytesTransportRequest; +import org.opensearch.transport.EmptyTransportResponseHandler; +import org.opensearch.transport.TransportService; + +/** + * Serializes a completed PPL/SQL query record and sends it to Query Insights over the transport + * layer, feeding its {@code addRecord} pipeline (Top N, historical index, roll-up). + * + *

Sent as the core {@link BytesTransportRequest} via {@link TransportService#sendRequest} (not + * {@code client.execute}): the core request type is loaded by both plugin classloaders, avoiding + * the {@link ClassCastException} a plugin-defined type would hit on same-node delivery. + * + *

The action name and wire format are owned by Query Insights and duplicated here as literals + * (SQL must not depend on the Query Insights artifact), so they MUST stay in lock-step with {@code + * ReportQueryBytesAction}. + */ +public final class QueryInsightsReporter { + + private static final Logger LOG = LogManager.getLogger(QueryInsightsReporter.class); + + /** Transport action name. MUST match {@code ReportQueryBytesAction.NAME} in Query Insights. */ + public static final String ACTION_NAME = + "cluster:admin/opensearch/query_insights/report_query_bytes"; + + /** + * Wire format version. MUST match {@code ReportQueryBytesAction.FORMAT_VERSION} in Query + * Insights. The two plugins are unreleased and always built together, so there is a single + * format: any change to the layout below is a coordinated change on both sides, not a + * compatibility boundary. + */ + public static final int FORMAT_VERSION = 1; + + private QueryInsightsReporter() {} + + /** + * Serialize the record and send it to the local node's Query Insights handler. Errors (e.g. Query + * Insights not installed, so the action is unregistered) are logged at debug and ignored. + * + * @param transportService the transport service used to send the request + * @param localNode the local (coordinator) node — the handler is registered on every node, so we + * send to ourselves + * @param querySource the query source label (e.g. {@code "PPL"}) + * @param parentMarker the originating query's marker {@code ::}; used as + * both the record id and the parent marker so child DSL records (tagged with the same value + * via {@code DERIVED_FROM}) roll up into this record + * @param nodeId the coordinator node id + * @param queryText the (prefix-stripped) query text + * @param timestampMillis the record timestamp + * @param latencyMillis end-to-end coordinator latency + * @param cpuNanos coordinator CPU nanos + * @param memoryBytes coordinator memory bytes + * @param indices the resolved index name(s) the query reads from; may be empty + */ + public static void report( + TransportService transportService, + DiscoveryNode localNode, + String querySource, + String parentMarker, + String nodeId, + String queryText, + long timestampMillis, + long latencyMillis, + long cpuNanos, + long memoryBytes, + java.util.List indices) { + try { + final BytesStreamOutput out = new BytesStreamOutput(); + out.writeVInt(FORMAT_VERSION); + out.writeString(querySource == null ? "" : querySource); + out.writeString(parentMarker == null ? "" : parentMarker); + out.writeString(nodeId == null ? "" : nodeId); + out.writeString(queryText == null ? "" : queryText); + out.writeVLong(timestampMillis); + out.writeVLong(Math.max(0L, latencyMillis)); + out.writeVLong(Math.max(0L, cpuNanos)); + out.writeVLong(Math.max(0L, memoryBytes)); + + final java.util.List safeIndices = indices == null ? java.util.List.of() : indices; + out.writeVInt(safeIndices.size()); + for (final String index : safeIndices) { + out.writeString(index == null ? "" : index); + } + + final BytesTransportRequest request = new BytesTransportRequest(out.bytes(), Version.CURRENT); + transportService.sendRequest( + localNode, ACTION_NAME, request, EmptyTransportResponseHandler.INSTANCE_SAME); + } catch (Exception e) { + // Query Insights may not be installed (action unregistered) or the send may fail; never + // affect query execution. + LOG.debug("Failed to report PPL query to Query Insights", e); + } + } +} 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 f052a4efc7f..f61a612a88f 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 @@ -28,7 +28,11 @@ import org.opensearch.common.inject.Inject; import org.opensearch.common.inject.Injector; import org.opensearch.common.inject.ModulesBuilder; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; +import org.opensearch.core.action.NotifyOnceListener; +import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.common.utils.QueryContext; @@ -88,6 +92,7 @@ public class TransportPPLQueryAction private final NodeClient clientRef; private final ClusterService clusterServiceRef; private final org.opensearch.sql.common.setting.Settings pluginSettingsRef; + private final TransportService transportServiceRef; @Inject public TransportPPLQueryAction( @@ -102,6 +107,7 @@ public TransportPPLQueryAction( super(PPLQueryAction.NAME, transportService, actionFilters, TransportPPLQueryRequest::new); this.clientRef = client; this.clusterServiceRef = clusterService; + this.transportServiceRef = transportService; ModulesBuilder modules = new ModulesBuilder(); modules.add(new OpenSearchPluginModule(extensionsHolder.engines(), tracer)); @@ -160,6 +166,177 @@ private void buildUnifiedQueryHandlerIfReady() { } } + /** + * Stamp the parent marker into the thread context (and the engine ThreadLocal) so child DSL + * searches carry it back to this PPL query. Won't overwrite a marker already set by an enclosing + * PPL call. + */ + private void stampQueryInsightsParentHeader(PPLQueryTask pplQueryTask) { + try { + ThreadContext threadContext = clientRef.threadPool().getThreadContext(); + String value = + QueryInsightsMarker.value( + "PPL", clusterServiceRef.localNode().getId(), pplQueryTask.getId()); + if (threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER) == null) { + threadContext.putHeader(QueryInsightsMarker.PARENT_HEADER, value); + } + // The ThreadContext header alone doesn't survive the engine's thread hops, so also carry the + // marker on the ThreadLocal channel for deterministic child tagging. + OpenSearchQueryManager.setQueryInsightsParentMarker(value); + } catch (Exception e) { + LOG.warn("Failed to stamp Query Insights parent header for query association", e); + } + } + + /** + * Toggle for recording PPL queries into Query Insights Top N. Registered by the Query Insights + * plugin (the two plugins share only this key string, not classes), so it is read through the + * registered Setting to honor opensearch.yml, dynamic updates, and the default. + */ + static final String QUERY_INSIGHTS_PPL_ENABLED_KEY = "search.insights.top_queries.ppl.enabled"; + + /** + * True only when Query Insights is installed (its setting is registered) and the toggle is + * enabled. Defaults to disabled on any failure. + */ + private boolean isQueryInsightsRecordingEnabled() { + try { + // Setting is registered iff Query Insights is installed. + Setting setting = + clusterServiceRef.getClusterSettings().get(QUERY_INSIGHTS_PPL_ENABLED_KEY); + if (setting == null) { + return false; + } + Object value = clusterServiceRef.getClusterSettings().get(setting); + return Boolean.TRUE.equals(value); + } catch (Exception e) { + LOG.debug("Failed to evaluate Query Insights PPL recording gate; defaulting to disabled", e); + return false; + } + } + + /** + * Report the finished PPL query to Query Insights once its resource tracking completes, so {@code + * getTotalResourceStats()} is final. + */ + private void registerQueryInsightsReport(PPLQueryTask reportTask) { + try { + boolean registered = + reportTask.addResourceTrackingCompletionListener( + new NotifyOnceListener<>() { + @Override + protected void innerOnResponse(Task task) { + writeQueryInsightsRecord((PPLQueryTask) task); + } + + @Override + protected void innerOnFailure(Exception e) { + // Tracking didn't complete cleanly; nothing to report. + } + }); + if (registered == false) { + LOG.debug( + "Query Insights report listener not registered; resource tracking already complete"); + } + } catch (Exception e) { + LOG.warn("Failed to register Query Insights report listener", e); + } + } + + /** + * Serialize a completed PPL query (resource stats + query text + indices) and send it to Query + * Insights over the transport layer. Errors are logged at debug and ignored. + */ + private void writeQueryInsightsRecord(PPLQueryTask reportTask) { + try { + String nodeId = clusterServiceRef.localNode().getId(); + String queryText = stripPplPrefix(reportTask.getDescription()); + + TaskResourceUsage usage = reportTask.getTotalResourceStats(); + long cpuNanos = usage == null ? 0L : usage.getCpuTimeInNanos(); + long memoryBytes = usage == null ? 0L : usage.getMemoryInBytes(); + + // Wall-clock latency from task start to now (monotonic clock). Used instead of the profile + // total, which collapses to ~0 when finish() resolves on a different thread. + long latencyMillis = + Math.max(0L, (System.nanoTime() - reportTask.getStartTimeNanos()) / 1_000_000L); + + // Must equal the marker stamped into child DSL tasks (child DERIVED_FROM == this) so Query + // Insights can roll child CPU/memory into this parent. + String parentMarker = QueryInsightsMarker.value("PPL", nodeId, reportTask.getId()); + + // No DSL SearchSourceBuilder exists at this layer, so index names come from the query text. + java.util.List indices = extractPplIndices(queryText); + + QueryInsightsReporter.report( + transportServiceRef, + clusterServiceRef.localNode(), + "PPL", + parentMarker, + nodeId, + queryText, + System.currentTimeMillis(), + latencyMillis, + cpuNanos, + memoryBytes, + indices); + } catch (Exception e) { + LOG.debug("Failed to write PPL query to Query Insights", e); + } + } + + /** + * Heuristic extraction of index name(s) from {@code source=}/{@code index=} clauses in the query + * text (including {@code join} sources). Not an AST parse; empty list when nothing matches. + * + * @param queryText the prefix-stripped PPL query text + * @return index name(s), de-duplicated in encounter order + */ + static java.util.List extractPplIndices(String queryText) { + final java.util.LinkedHashSet indices = new java.util.LinkedHashSet<>(); + if (queryText == null || queryText.isEmpty()) { + return new java.util.ArrayList<>(indices); + } + // Value is a quoted span (may contain spaces) or an unquoted run up to the first whitespace or + // pipe, so it stops at the clause boundary. Comma lists without spaces are captured whole and + // split below. Single-level alternation of linear branches — no catastrophic backtracking. + final java.util.regex.Matcher m = + java.util.regex.Pattern.compile( + "(?i)\\b(?:source|index)\\s*=\\s*(\"[^\"]*\"|'[^']*'|`[^`]*`|[^\\s|]+)") + .matcher(queryText); + while (m.find()) { + final String raw = m.group(1).trim(); + // A source clause may list multiple comma-separated indices. + for (String part : raw.split(",")) { + String name = part.trim(); + // Strip surrounding quotes/backticks if present. + if (name.length() >= 2) { + final char c0 = name.charAt(0); + final char cN = name.charAt(name.length() - 1); + if ((c0 == '"' && cN == '"') || (c0 == '\'' && cN == '\'') || (c0 == '`' && cN == '`')) { + name = name.substring(1, name.length() - 1).trim(); + } + } + if (name.isEmpty() == false) { + indices.add(name); + } + } + } + return new java.util.ArrayList<>(indices); + } + + /** Strip the {@code "PPL: "} / {@code "PPL [queryId=...]: "} description prefix. */ + private static String stripPplPrefix(String description) { + if (description == null) { + return ""; + } + int colon = description.indexOf(": "); + if (description.startsWith("PPL") && colon >= 0) { + return description.substring(colon + 2); + } + return description; + } + /** * {@inheritDoc} Transform the request and call super.doExecute() to support call from other * plugins. @@ -183,8 +360,14 @@ protected void doExecute( return; } - if (task instanceof PPLQueryTask pplQueryTask) { - OpenSearchQueryManager.setCancellableTask(pplQueryTask); + final PPLQueryTask reportTask = task instanceof PPLQueryTask ? (PPLQueryTask) task : null; + if (reportTask != null) { + OpenSearchQueryManager.setCancellableTask(reportTask); + // Opt-in: when disabled, skip header stamping and report registration entirely. + if (isQueryInsightsRecordingEnabled()) { + stampQueryInsightsParentHeader(reportTask); + registerQueryInsightsReport(reportTask); + } } Metrics.getInstance().getNumericalMetric(MetricName.PPL_REQ_TOTAL).increment(); Metrics.getInstance().getNumericalMetric(MetricName.PPL_REQ_COUNT_TOTAL).increment(); @@ -274,6 +457,10 @@ protected void doExecute( clearingListener.onFailure(e); } finally { spanScope.close(); + // submit() removes these on the normal path; clear again so an early return/throw before + // submit() doesn't leak them onto this pooled transport thread. + OpenSearchQueryManager.clearQueryInsightsParentMarker(); + OpenSearchQueryManager.clearCancellableTask(); } } diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/ExtractPplIndicesTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/ExtractPplIndicesTest.java new file mode 100644 index 00000000000..a8c8d9fe6c4 --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/ExtractPplIndicesTest.java @@ -0,0 +1,108 @@ +/* + * 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.assertTrue; + +import java.time.Duration; +import java.util.List; +import org.junit.Test; + +/** + * Unit tests for {@link TransportPPLQueryAction#extractPplIndices(String)} — the heuristic that + * pulls index names out of a PPL query string for Query Insights reporting. + */ +public class ExtractPplIndicesTest { + + @Test + public void singleIndex() { + assertEquals( + List.of("employees"), TransportPPLQueryAction.extractPplIndices("source=employees")); + } + + @Test + public void indexKeywordAlias() { + assertEquals(List.of("logs"), TransportPPLQueryAction.extractPplIndices("index=logs")); + } + + @Test + public void caseInsensitiveKeyword() { + assertEquals(List.of("t"), TransportPPLQueryAction.extractPplIndices("SOURCE=t")); + } + + @Test + public void spacesAroundEquals() { + assertEquals( + List.of("employees"), TransportPPLQueryAction.extractPplIndices("source = employees")); + } + + @Test + public void stopsAtPipe() { + // The value must not swallow the trailing "| where ..." clause. + assertEquals( + List.of("employees"), + TransportPPLQueryAction.extractPplIndices( + "source=employees | where dept=\"eng\" | fields name")); + } + + @Test + public void stopsAtWhitespace() { + // A space-separated trailing term must not become part of the index name. + assertEquals( + List.of("employees"), TransportPPLQueryAction.extractPplIndices("source=employees a=1")); + } + + @Test + public void commaSeparatedList() { + assertEquals( + List.of("a", "b", "c"), + TransportPPLQueryAction.extractPplIndices("source=a,b,c | stats count()")); + } + + @Test + public void quotedIndexIsUnquoted() { + assertEquals( + List.of("my index"), + TransportPPLQueryAction.extractPplIndices("source=\"my index\" | fields x")); + } + + @Test + public void deduplicatesInEncounterOrder() { + assertEquals( + List.of("a", "b"), + TransportPPLQueryAction.extractPplIndices("source=a | join source=b | join source=a")); + } + + @Test + public void nullAndEmptyYieldEmpty() { + assertTrue(TransportPPLQueryAction.extractPplIndices(null).isEmpty()); + assertTrue(TransportPPLQueryAction.extractPplIndices("").isEmpty()); + } + + @Test + public void noSourceClauseYieldsEmpty() { + assertTrue(TransportPPLQueryAction.extractPplIndices("search x=1 | head 5").isEmpty()); + } + + /** Regression guard: hostile input must match in linear time (no catastrophic backtracking). */ + @Test + public void hostileInputRunsInLinearTime() { + StringBuilder sb = new StringBuilder("source="); + for (int i = 0; i < 50_000; i++) { + sb.append("a,\""); + } + final String hostile = sb.toString(); + long start = System.nanoTime(); + TransportPPLQueryAction.extractPplIndices(hostile); + Duration elapsed = Duration.ofNanos(System.nanoTime() - start); + assertTrue( + "extractPplIndices should complete quickly on hostile input, took " + + elapsed.toMillis() + + "ms", + elapsed.toMillis() < 2_000); + } +} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java index c9502ac3bbf..7f82cf7bd1e 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java @@ -62,4 +62,35 @@ public void testCooperativeModel() { task.cancel("Test"); assertTrue(task.isCancelled()); } + + private PPLQueryTask newTask() { + return new PPLQueryTask( + 1, + "transport", + "cluster:admin/opensearch/ppl", + "test query", + TaskId.EMPTY_TASK_ID, + Map.of()); + } + + @Test + public void testSupportsResourceTracking() { + // Resource tracking must be enabled so the coordinator task's CPU/memory is exposed through + // the tasks API for Query Insights. + assertTrue(newTask().supportsResourceTracking()); + } + + @Test + public void testQueryInsightsParentHeaderName() { + // Query Insights relies on this exact header name to classify child DSL searches and associate + // them with the originating query; it must match what SQLPlugin.getTaskHeaders() registers. + assertEquals("X-Query-Insights-Parent", QueryInsightsMarker.PARENT_HEADER); + } + + @Test + public void testQueryInsightsParentHeaderValueIsSourcePrefixed() { + // Value format is :: so QI reads both source and parent id from one + // header. SQL will reuse the same helper with source "SQL". + assertEquals("PPL:node-1:42", QueryInsightsMarker.value("PPL", "node-1", 42L)); + } }