Skip to content

feat(ppl): Report completed PPL queries to Query Insights Top N #5760

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
KishoreKicha14 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(ppl): Report completed PPL queries to Query Insights Top N #5760

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 105 additions & 2 deletions opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,18 @@

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;
import org.apache.logging.log4j.Logger;
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;
Expand All @@ -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;
Expand All @@ -49,24 +69,53 @@ public static void clearCancellableTask() {
cancellableTask.remove();
}

/**
* Query Insights parent marker ({@code PPL:<nodeId>:<taskId>}), 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<String> 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(
Expand All @@ -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();
Expand All @@ -98,7 +157,11 @@ private void schedule(

throw e;
} finally {
if (trackingStarted) {
stopThreadResourceTracking(cancelTask, trackedThreadId);
}
clearCancellableTask();
clearQueryInsightsParentMarker();
}
});

Expand All @@ -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};
}
}
21 changes: 21 additions & 0 deletions ...h/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,9 @@ private void dispatchInternal(
// Capture thread-local state to propagate across thread boundary
Map<String, String> 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 =
Expand All @@ -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
Expand All @@ -106,6 +113,16 @@ private void dispatchInternal(
hookHandle =
Hook.CURRENT_TIME.addThread((Consumer<Holder<Long>>) 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);
Expand All @@ -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();
Expand Down
75 changes: 73 additions & 2 deletions ...rch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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<OpenSearchResponse> nextBatchFuture = null;
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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();
Expand Down
Loading
Loading