From a391bd9f1e041ce9807ac09a0aa8c413e605e73f Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Mon, 13 Jul 2026 21:03:05 +0000 Subject: [PATCH 01/11] feat: slow query thread pool Signed-off-by: Simeon Widdis --- .../sql/common/setting/Settings.java | 5 +- .../executor/DirectExecutionDispatcher.java | 26 ++++ .../sql/executor/ExecutionDispatcher.java | 33 +++++ .../opensearch/sql/executor/QueryService.java | 40 +++++- .../executor/OpenSearchQueryManager.java | 1 + .../opensearch/executor/ScriptDetector.java | 37 +++++ .../ThreadPoolExecutionDispatcher.java | 68 +++++++++ .../setting/OpenSearchSettings.java | 14 ++ .../executor/ScriptDetectorTest.java | 90 ++++++++++++ .../ThreadPoolExecutionDispatcherTest.java | 132 ++++++++++++++++++ .../org/opensearch/sql/plugin/SQLPlugin.java | 17 ++- .../plugin/config/OpenSearchPluginModule.java | 12 +- .../plugin/rest/RestUnifiedQueryAction.java | 42 ++++-- 13 files changed, 495 insertions(+), 22 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/executor/DirectExecutionDispatcher.java create mode 100644 core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java diff --git a/common/src/main/java/org/opensearch/sql/common/setting/Settings.java b/common/src/main/java/org/opensearch/sql/common/setting/Settings.java index bf3e65d8741..43087ffcd77 100644 --- a/common/src/main/java/org/opensearch/sql/common/setting/Settings.java +++ b/common/src/main/java/org/opensearch/sql/common/setting/Settings.java @@ -79,7 +79,10 @@ public enum Key { ASYNC_QUERY_EXTERNAL_SCHEDULER_INTERVAL( "plugins.query.executionengine.async_query.external_scheduler.interval"), STREAMING_JOB_HOUSEKEEPER_INTERVAL( - "plugins.query.executionengine.spark.streamingjobs.housekeeper.interval"); + "plugins.query.executionengine.spark.streamingjobs.housekeeper.interval"), + + /** Thread Pool Settings. */ + SQL_SLOW_WORKER_POOL_ENABLED("plugins.sql.slow_worker_pool.enabled"); @Getter private final String keyValue; diff --git a/core/src/main/java/org/opensearch/sql/executor/DirectExecutionDispatcher.java b/core/src/main/java/org/opensearch/sql/executor/DirectExecutionDispatcher.java new file mode 100644 index 00000000000..0b168f53a56 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/executor/DirectExecutionDispatcher.java @@ -0,0 +1,26 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor; + +import org.apache.calcite.rel.RelNode; +import org.opensearch.sql.calcite.CalcitePlanContext; +import org.opensearch.sql.common.response.ResponseListener; + +/** + * Default no-op dispatcher that executes inline on the current thread. Used when slow-pool routing + * is disabled or as a fallback. + */ +public class DirectExecutionDispatcher implements ExecutionDispatcher { + + @Override + public void dispatch( + RelNode plan, + CalcitePlanContext context, + ResponseListener listener, + ExecutionEngine engine) { + engine.execute(plan, context, listener); + } +} diff --git a/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java b/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java new file mode 100644 index 00000000000..cc8f5ee78b9 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor; + +import org.apache.calcite.rel.RelNode; +import org.opensearch.sql.calcite.CalcitePlanContext; +import org.opensearch.sql.common.response.ResponseListener; + +/** + * Dispatches query execution to an appropriate thread pool based on plan characteristics. After + * query analysis and optimization, the dispatcher inspects the plan and routes execution to either + * the fast worker pool (for queries fully pushed to OpenSearch) or the slow worker pool (for + * queries requiring scripts/table scans). + */ +public interface ExecutionDispatcher { + + /** + * Dispatch execution of the given plan. + * + * @param plan the optimized Calcite plan + * @param context the plan context + * @param listener response listener for query results + * @param engine the execution engine to invoke + */ + void dispatch( + RelNode plan, + CalcitePlanContext context, + ResponseListener listener, + ExecutionEngine engine); +} 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 ddb4338bc8f..4fd8dd47336 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -8,7 +8,6 @@ import java.util.List; import java.util.Optional; import javax.annotation.Nullable; -import lombok.AllArgsConstructor; import lombok.Getter; import lombok.RequiredArgsConstructor; import lombok.extern.log4j.Log4j2; @@ -58,7 +57,6 @@ /** The low level interface of core engine. */ @RequiredArgsConstructor -@AllArgsConstructor @Log4j2 public class QueryService { private final Analyzer analyzer; @@ -66,6 +64,37 @@ public class QueryService { private final Planner planner; private DataSourceService dataSourceService; private Settings settings; + private ExecutionDispatcher executionDispatcher = new DirectExecutionDispatcher(); + + public QueryService( + Analyzer analyzer, + ExecutionEngine executionEngine, + Planner planner, + DataSourceService dataSourceService, + Settings settings) { + this( + analyzer, + executionEngine, + planner, + dataSourceService, + settings, + new DirectExecutionDispatcher()); + } + + public QueryService( + Analyzer analyzer, + ExecutionEngine executionEngine, + Planner planner, + DataSourceService dataSourceService, + Settings settings, + ExecutionDispatcher executionDispatcher) { + this.analyzer = analyzer; + this.executionEngine = executionEngine; + this.planner = planner; + this.dataSourceService = dataSourceService; + this.settings = settings; + this.executionDispatcher = executionDispatcher; + } @Getter(lazy = true) private final CalciteRelNodeVisitor relNodeVisitor = new CalciteRelNodeVisitor(dataSourceService); @@ -179,10 +208,13 @@ public void executeWithCalcite( analyzeMetric.set(System.nanoTime() - analyzeStart); - // Wrap execution with EXECUTING stage tracking + // Wrap execution with EXECUTING stage tracking — dispatch via + // ExecutionDispatcher which may route to a slow worker pool StageErrorHandler.executeStageVoid( QueryProcessingStage.EXECUTING, - () -> executionEngine.execute(calcitePlan, context, listener), + () -> + executionDispatcher.dispatch( + calcitePlan, context, listener, executionEngine), "while running the query"); }, QueryService.class); 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 7aaaaa6655e..15245deb01a 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 @@ -32,6 +32,7 @@ public class OpenSearchQueryManager implements QueryManager { private final Settings settings; public static final String SQL_WORKER_THREAD_POOL_NAME = "sql-worker"; + public static final String SQL_SLOW_WORKER_THREAD_POOL_NAME = "sql-slow-worker"; public static final String SQL_BACKGROUND_THREAD_POOL_NAME = "sql_background_io"; private static final ThreadLocal cancellableTask = new ThreadLocal<>(); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java new file mode 100644 index 00000000000..a8a11cfcc22 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java @@ -0,0 +1,37 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.executor; + +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelVisitor; +import org.opensearch.sql.opensearch.storage.scan.AbstractCalciteIndexScan; + +/** + * Inspects a Calcite plan tree to determine whether any scan node requires script evaluation. A + * plan with scripts cannot be fully pushed down to OpenSearch and will perform in-memory + * evaluation, making it a candidate for the slow worker pool. + */ +public final class ScriptDetector { + + private ScriptDetector() {} + + /** Returns true if any scan node in the plan has scripts pushed. */ + public static boolean hasScripts(RelNode plan) { + boolean[] found = {false}; + new RelVisitor() { + @Override + public void visit(RelNode node, int ordinal, RelNode parent) { + if (!found[0] && node instanceof AbstractCalciteIndexScan scan) { + found[0] = scan.isScriptPushed(); + } + if (!found[0]) { + super.visit(node, ordinal, parent); + } + } + }.go(plan); + return found[0]; + } +} 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 new file mode 100644 index 00000000000..95685f2e793 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java @@ -0,0 +1,68 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.executor; + +import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_SLOW_WORKER_THREAD_POOL_NAME; + +import java.util.Map; +import lombok.RequiredArgsConstructor; +import org.apache.calcite.rel.RelNode; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.logging.log4j.ThreadContext; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.sql.calcite.CalcitePlanContext; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.executor.ExecutionDispatcher; +import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.threadpool.ThreadPool; + +/** + * Dispatches query execution to either the fast or slow worker thread pool based on whether the + * plan contains scripts. Plans with scripts require in-memory evaluation and are routed to the slow + * pool so they don't block fast pushdown-only queries. + */ +@RequiredArgsConstructor +public class ThreadPoolExecutionDispatcher implements ExecutionDispatcher { + + private static final Logger LOG = LogManager.getLogger(ThreadPoolExecutionDispatcher.class); + + private final ThreadPool threadPool; + private final Settings settings; + + @Override + public void dispatch( + RelNode plan, + CalcitePlanContext context, + ResponseListener listener, + ExecutionEngine engine) { + if (isSlowPoolEnabled() && ScriptDetector.hasScripts(plan)) { + LOG.debug("Query plan contains scripts, dispatching to slow worker pool"); + Map ctx = ThreadContext.getImmutableContext(); + CancellableTask task = OpenSearchQueryManager.getCancellableTask(); + threadPool.schedule( + () -> { + ThreadContext.putAll(ctx); + OpenSearchQueryManager.setCancellableTask(task); + try { + engine.execute(plan, context, listener); + } finally { + OpenSearchQueryManager.clearCancellableTask(); + } + }, + new TimeValue(0), + SQL_SLOW_WORKER_THREAD_POOL_NAME); + } else { + engine.execute(plan, context, listener); + } + } + + private boolean isSlowPoolEnabled() { + return settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED); + } +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java index b596c7bc47a..d29b551dc3b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java @@ -352,6 +352,13 @@ public class OpenSearchSettings extends Settings { Setting.Property.NodeScope, Setting.Property.Dynamic); + public static final Setting SQL_SLOW_WORKER_POOL_ENABLED_SETTING = + Setting.boolSetting( + Key.SQL_SLOW_WORKER_POOL_ENABLED.getKeyValue(), + true, + Setting.Property.NodeScope, + Setting.Property.Dynamic); + /** Construct OpenSearchSetting. The OpenSearchSetting must be singleton. */ @SuppressWarnings("unchecked") public OpenSearchSettings(ClusterSettings clusterSettings) { @@ -610,6 +617,12 @@ public OpenSearchSettings(ClusterSettings clusterSettings) { Key.FIELD_TYPE_TOLERANCE, FIELD_TYPE_TOLERANCE_SETTING, new Updater(Key.FIELD_TYPE_TOLERANCE)); + register( + settingBuilder, + clusterSettings, + Key.SQL_SLOW_WORKER_POOL_ENABLED, + SQL_SLOW_WORKER_POOL_ENABLED_SETTING, + new Updater(Key.SQL_SLOW_WORKER_POOL_ENABLED)); defaultSettings = settingBuilder.build(); } @@ -703,6 +716,7 @@ public static List> pluginSettings() { .add(SESSION_INACTIVITY_TIMEOUT_MILLIS_SETTING) .add(STREAMING_JOB_HOUSEKEEPER_INTERVAL_SETTING) .add(FIELD_TYPE_TOLERANCE_SETTING) + .add(SQL_SLOW_WORKER_POOL_ENABLED_SETTING) .build(); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java new file mode 100644 index 00000000000..e61064fe190 --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java @@ -0,0 +1,90 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.executor; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.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.List; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelVisitor; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; +import org.opensearch.sql.opensearch.storage.scan.AbstractCalciteIndexScan; + +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) +class ScriptDetectorTest { + + @Test + void returnsFalseForNonScanNode() { + RelNode mockNode = createMockNode(); + assertFalse(ScriptDetector.hasScripts(mockNode)); + } + + @Test + void returnsTrueWhenScanHasScripts() { + AbstractCalciteIndexScan scan = createMockScan(true); + assertTrue(ScriptDetector.hasScripts(scan)); + } + + @Test + void returnsFalseWhenScanHasNoScripts() { + AbstractCalciteIndexScan scan = createMockScan(false); + assertFalse(ScriptDetector.hasScripts(scan)); + } + + @Test + void detectsScriptsInNestedPlan() { + AbstractCalciteIndexScan scan = createMockScan(true); + RelNode parent = createMockNode(scan); + assertTrue(ScriptDetector.hasScripts(parent)); + } + + @Test + void returnsFalseForNestedPlanWithoutScripts() { + AbstractCalciteIndexScan scan = createMockScan(false); + RelNode parent = createMockNode(scan); + assertFalse(ScriptDetector.hasScripts(parent)); + } + + private static RelNode createMockNode(RelNode... children) { + RelNode node = mock(RelNode.class); + List childList = List.of(children); + when(node.getInputs()).thenReturn(childList); + doAnswer( + invocation -> { + RelVisitor visitor = invocation.getArgument(0); + for (int i = 0; i < childList.size(); i++) { + visitor.visit(childList.get(i), i, node); + } + return null; + }) + .when(node) + .childrenAccept(any(RelVisitor.class)); + return node; + } + + private static AbstractCalciteIndexScan createMockScan(boolean scriptPushed) { + AbstractCalciteIndexScan scan = mock(AbstractCalciteIndexScan.class); + when(scan.isScriptPushed()).thenReturn(scriptPushed); + when(scan.getInputs()).thenReturn(List.of()); + doAnswer( + invocation -> { + return null; + }) + .when(scan) + .childrenAccept(any(RelVisitor.class)); + return scan; + } +} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java new file mode 100644 index 00000000000..cd984fd5ba7 --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java @@ -0,0 +1,132 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.executor; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_SLOW_WORKER_THREAD_POOL_NAME; + +import java.util.List; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelVisitor; +import org.junit.jupiter.api.BeforeEach; +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.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.sql.calcite.CalcitePlanContext; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.opensearch.storage.scan.AbstractCalciteIndexScan; +import org.opensearch.threadpool.ThreadPool; + +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) +class ThreadPoolExecutionDispatcherTest { + + @Mock private ThreadPool threadPool; + @Mock private Settings settings; + @Mock private CalcitePlanContext context; + @Mock private ResponseListener listener; + @Mock private ExecutionEngine engine; + + private ThreadPoolExecutionDispatcher dispatcher; + + @BeforeEach + void setUp() { + dispatcher = new ThreadPoolExecutionDispatcher(threadPool, settings); + } + + @Test + void executesInlineWhenNoScripts() { + when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + .thenReturn(true); + RelNode plan = createMockNode(); + + dispatcher.dispatch(plan, context, listener, engine); + + verify(engine).execute(plan, context, listener); + verify(threadPool, never()).schedule(any(), any(TimeValue.class), any()); + } + + @Test + void dispatchesToSlowPoolWhenScriptsDetected() { + when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + .thenReturn(true); + AbstractCalciteIndexScan scan = createMockScan(true); + + dispatcher.dispatch(scan, context, listener, engine); + + verify(threadPool) + .schedule(any(Runnable.class), eq(new TimeValue(0)), eq(SQL_SLOW_WORKER_THREAD_POOL_NAME)); + verify(engine, never()).execute(any(RelNode.class), any(), any(ResponseListener.class)); + } + + @Test + void executesInlineWhenSlowPoolDisabled() { + when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + .thenReturn(false); + AbstractCalciteIndexScan scan = createMockScan(true); + + dispatcher.dispatch(scan, context, listener, engine); + + verify(engine).execute(scan, context, listener); + verify(threadPool, never()).schedule(any(), any(TimeValue.class), any()); + } + + @Test + void scheduledRunnableCallsEngine() { + when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + .thenReturn(true); + doAnswer( + invocation -> { + Runnable task = invocation.getArgument(0); + task.run(); + return null; + }) + .when(threadPool) + .schedule(any(Runnable.class), any(TimeValue.class), any()); + + AbstractCalciteIndexScan scan = createMockScan(true); + dispatcher.dispatch(scan, context, listener, engine); + + verify(engine).execute(scan, context, listener); + } + + private static RelNode createMockNode(RelNode... children) { + RelNode node = mock(RelNode.class); + List childList = List.of(children); + when(node.getInputs()).thenReturn(childList); + doAnswer( + invocation -> { + RelVisitor visitor = invocation.getArgument(0); + for (int i = 0; i < childList.size(); i++) { + visitor.visit(childList.get(i), i, node); + } + return null; + }) + .when(node) + .childrenAccept(any(RelVisitor.class)); + return node; + } + + private static AbstractCalciteIndexScan createMockScan(boolean scriptPushed) { + AbstractCalciteIndexScan scan = mock(AbstractCalciteIndexScan.class); + when(scan.isScriptPushed()).thenReturn(scriptPushed); + when(scan.getInputs()).thenReturn(List.of()); + doAnswer(invocation -> null).when(scan).childrenAccept(any(RelVisitor.class)); + return scan; + } +} 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 e1278aa75e8..320c778ac3f 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -7,6 +7,7 @@ import static org.opensearch.sql.datasource.model.DataSourceMetadata.defaultOpenSearchDataSourceMetadata; import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_BACKGROUND_THREAD_POOL_NAME; +import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_SLOW_WORKER_THREAD_POOL_NAME; import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME; import static org.opensearch.sql.spark.data.constants.SparkConstants.SPARK_REQUEST_BUFFER_INDEX_NAME; @@ -446,10 +447,12 @@ public ScheduledJobParser getJobParser() { @Override public List> getExecutorBuilders(Settings settings) { - // The worker pool is the primary pool where most of the work is done. The background thread - // pool is a separate queue for asynchronous requests to other nodes. We keep them separate to - // prevent deadlocks during async fetches on small node counts. Tasks in the background pool - // should do no work except I/O to other services. + // The worker pool is the primary pool where most of the work is done. The slow-worker pool + // handles queries that require scripts (table scans that can't be pushed to Lucene) so they + // don't starve fast queries. The background thread pool is a separate queue for asynchronous + // requests to other nodes. We keep them separate to prevent deadlocks during async fetches on + // small node counts. Tasks in the background pool should do no work except I/O to other + // services. return List.of( new FixedExecutorBuilder( settings, @@ -457,6 +460,12 @@ public List> getExecutorBuilders(Settings settings) { OpenSearchExecutors.allocatedProcessors(settings), 1000, "thread_pool." + SQL_WORKER_THREAD_POOL_NAME), + new FixedExecutorBuilder( + settings, + SQL_SLOW_WORKER_THREAD_POOL_NAME, + OpenSearchExecutors.allocatedProcessors(settings), + 1000, + "thread_pool." + SQL_SLOW_WORKER_THREAD_POOL_NAME), new FixedExecutorBuilder( settings, SQL_BACKGROUND_THREAD_POOL_NAME, diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java index 057c88c9a02..816f1071310 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java @@ -15,6 +15,7 @@ import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.executor.DelegatingExecutionEngine; +import org.opensearch.sql.executor.ExecutionDispatcher; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.QueryManager; import org.opensearch.sql.executor.QueryService; @@ -26,6 +27,7 @@ import org.opensearch.sql.opensearch.client.OpenSearchNodeClient; import org.opensearch.sql.opensearch.executor.OpenSearchExecutionEngine; import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager; +import org.opensearch.sql.opensearch.executor.ThreadPoolExecutionDispatcher; import org.opensearch.sql.opensearch.executor.protector.ExecutionProtector; import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector; import org.opensearch.sql.opensearch.monitor.OpenSearchMemoryHealthy; @@ -114,13 +116,19 @@ public SQLService sqlService( /** {@link QueryPlanFactory}. */ @Provides public QueryPlanFactory queryPlanFactory( - DataSourceService dataSourceService, ExecutionEngine executionEngine, Settings settings) { + DataSourceService dataSourceService, + ExecutionEngine executionEngine, + Settings settings, + NodeClient nodeClient) { Analyzer analyzer = new Analyzer( new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository); Planner planner = new Planner(LogicalPlanOptimizer.create()); + ExecutionDispatcher executionDispatcher = + new ThreadPoolExecutionDispatcher(nodeClient.threadPool(), settings); QueryService queryService = - new QueryService(analyzer, executionEngine, planner, dataSourceService, settings); + new QueryService( + analyzer, executionEngine, planner, dataSourceService, settings, executionDispatcher); return new QueryPlanFactory(queryService); } } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java index f8214096e41..aca90c8db5e 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java @@ -7,6 +7,7 @@ import static org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; import static org.opensearch.sql.lang.PPLLangSpec.PPL_SPEC; +import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_SLOW_WORKER_THREAD_POOL_NAME; import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME; import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY; @@ -40,12 +41,14 @@ import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit; import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.common.setting.Settings.Key; import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; import org.opensearch.sql.executor.QueryType; import org.opensearch.sql.executor.analytics.AnalyticsExecutionEngine; import org.opensearch.sql.lang.LangSpec; import org.opensearch.sql.monitor.profile.ProfileContext; import org.opensearch.sql.monitor.profile.QueryProfiling; +import org.opensearch.sql.opensearch.executor.ScriptDetector; import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse; import org.opensearch.sql.protocol.response.QueryResult; import org.opensearch.sql.protocol.response.format.ResponseFormatter; @@ -229,18 +232,31 @@ private void doExecute( // string, so apply the equivalent top-level limit here before the system cap. plan = addFetchSizeLimit(plan, planContext, fetchSize); plan = addQuerySizeLimit(plan, planContext); - if (profiling) { - analyticsEngine.executeWithProfile( - plan, - planContext, - queryCtx, - createQueryListener(queryType, profileCtx, closingListener)); + RelNode finalPlan = plan; + Runnable executeTask = + profiling + ? () -> + analyticsEngine.executeWithProfile( + finalPlan, + planContext, + queryCtx, + createQueryListener(queryType, profileCtx, closingListener)) + : () -> + analyticsEngine.execute( + finalPlan, + planContext, + queryCtx, + createQueryListener(queryType, profileCtx, closingListener)); + if (isSlowPoolEnabled() && ScriptDetector.hasScripts(finalPlan)) { + LOG.debug("Query contains scripts, routing to slow worker pool"); + client + .threadPool() + .schedule( + withCurrentContext(executeTask), + new TimeValue(0), + SQL_SLOW_WORKER_THREAD_POOL_NAME); } else { - analyticsEngine.execute( - plan, - planContext, - queryCtx, - createQueryListener(queryType, profileCtx, closingListener)); + executeTask.run(); } } catch (Exception e) { closingListener.onFailure(e); @@ -460,6 +476,10 @@ private static String appendError(String json, Throwable error) { return json.substring(0, json.length() - 1) + ",\"error\":" + err + "}"; } + private boolean isSlowPoolEnabled() { + return pluginSettings.getSettingValue(Key.SQL_SLOW_WORKER_POOL_ENABLED); + } + private static Runnable withCurrentContext(final Runnable task) { final Map currentContext = ThreadContext.getImmutableContext(); return () -> { From 311f7513cd2f584dfd11b4d66667135a8d06f0a3 Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Mon, 13 Jul 2026 22:06:49 +0000 Subject: [PATCH 02/11] handle slow query detection when optimization runs in execution step Signed-off-by: Simeon Widdis --- .../opensearch/executor/ScriptDetector.java | 85 ++++++++++- .../executor/ScriptDetectorTest.java | 132 +++++++++++++++--- .../ThreadPoolExecutionDispatcherTest.java | 15 +- 3 files changed, 204 insertions(+), 28 deletions(-) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java index a8a11cfcc22..a26868910ed 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java @@ -7,25 +7,57 @@ import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelVisitor; +import org.apache.calcite.rel.core.Join; +import org.apache.calcite.rel.core.Project; +import org.apache.calcite.rel.core.Window; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexOver; +import org.apache.calcite.rex.RexVisitorImpl; +import org.apache.calcite.sql.validate.SqlUserDefinedFunction; import org.opensearch.sql.opensearch.storage.scan.AbstractCalciteIndexScan; +import org.opensearch.sql.opensearch.storage.scan.context.AggSpec; +import org.opensearch.sql.opensearch.storage.scan.context.PushDownContext; /** - * Inspects a Calcite plan tree to determine whether any scan node requires script evaluation. A - * plan with scripts cannot be fully pushed down to OpenSearch and will perform in-memory - * evaluation, making it a candidate for the slow worker pool. + * Inspects a Calcite plan tree to determine whether execution will be expensive. Detects: (1) + * user-defined functions (REX_EXTRACT, PARSE, etc.) that become per-document scripts, (2) join + * nodes requiring in-memory merge, and (3) window functions requiring in-memory evaluation. + * + *

Works on both logical plans (before optimization, where PushDownContext is empty) and physical + * plans (after optimization, where PushDownContext tracks scripts). */ public final class ScriptDetector { private ScriptDetector() {} - /** Returns true if any scan node in the plan has scripts pushed. */ + /** + * Returns true if the plan contains patterns indicating expensive execution: UDF calls that + * produce scripts, join nodes, or window functions. + */ public static boolean hasScripts(RelNode plan) { boolean[] found = {false}; new RelVisitor() { @Override public void visit(RelNode node, int ordinal, RelNode parent) { - if (!found[0] && node instanceof AbstractCalciteIndexScan scan) { - found[0] = scan.isScriptPushed(); + if (found[0]) { + return; + } + // Physical plan: check PushDownContext for scripts already detected by optimizer + if (node instanceof AbstractCalciteIndexScan scan) { + found[0] = scanHasScripts(scan); + } + // Logical plan: detect join nodes (always require in-memory processing) + if (!found[0] && node instanceof Join) { + found[0] = true; + } + // Logical plan: detect window rel nodes (eventstats, dedup patterns) + if (!found[0] && node instanceof Window) { + found[0] = true; + } + // Logical plan: check projections for UDFs or window expressions + if (!found[0] && node instanceof Project project) { + found[0] = projectHasExpensiveExpressions(project); } if (!found[0]) { super.visit(node, ordinal, parent); @@ -34,4 +66,45 @@ public void visit(RelNode node, int ordinal, RelNode parent) { }.go(plan); return found[0]; } + + private static boolean scanHasScripts(AbstractCalciteIndexScan scan) { + PushDownContext ctx = scan.getPushDownContext(); + if (ctx.isScriptPushed()) { + return true; + } + if (ctx.isSortExprPushed()) { + return true; + } + AggSpec aggSpec = ctx.getAggSpec(); + return aggSpec != null && aggSpec.getScriptCount() > 0; + } + + private static boolean projectHasExpensiveExpressions(Project project) { + for (RexNode expr : project.getProjects()) { + if (hasExpensiveRex(expr)) { + return true; + } + } + return false; + } + + private static boolean hasExpensiveRex(RexNode expr) { + boolean[] found = {false}; + expr.accept( + new RexVisitorImpl(true) { + @Override + public Void visitCall(RexCall call) { + if (call.getOperator() instanceof SqlUserDefinedFunction) { + found[0] = true; + return null; + } + if (call instanceof RexOver) { + found[0] = true; + return null; + } + return super.visitCall(call); + } + }); + return found[0]; + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java index e61064fe190..e7e186f244f 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java @@ -15,12 +15,25 @@ import java.util.List; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelVisitor; +import org.apache.calcite.rel.core.JoinRelType; +import org.apache.calcite.rel.logical.LogicalJoin; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rel.logical.LogicalWindow; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexOver; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.validate.SqlUserDefinedFunction; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.junit.jupiter.MockitoExtension; import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; import org.opensearch.sql.opensearch.storage.scan.AbstractCalciteIndexScan; +import org.opensearch.sql.opensearch.storage.scan.context.AggSpec; +import org.opensearch.sql.opensearch.storage.scan.context.PushDownContext; @ExtendWith(MockitoExtension.class) @MockitoSettings(strictness = Strictness.LENIENT) @@ -33,29 +46,105 @@ void returnsFalseForNonScanNode() { } @Test - void returnsTrueWhenScanHasScripts() { - AbstractCalciteIndexScan scan = createMockScan(true); + void returnsTrueWhenFilterScriptPushed() { + AbstractCalciteIndexScan scan = createMockScan(true, false, 0); assertTrue(ScriptDetector.hasScripts(scan)); } @Test - void returnsFalseWhenScanHasNoScripts() { - AbstractCalciteIndexScan scan = createMockScan(false); + void returnsTrueWhenAggScriptPresent() { + AbstractCalciteIndexScan scan = createMockScan(false, false, 1); + assertTrue(ScriptDetector.hasScripts(scan)); + } + + @Test + void returnsTrueWhenSortExprPushed() { + AbstractCalciteIndexScan scan = createMockScan(false, true, 0); + assertTrue(ScriptDetector.hasScripts(scan)); + } + + @Test + void returnsFalseWhenNoScripts() { + AbstractCalciteIndexScan scan = createMockScan(false, false, 0); assertFalse(ScriptDetector.hasScripts(scan)); } @Test void detectsScriptsInNestedPlan() { - AbstractCalciteIndexScan scan = createMockScan(true); + AbstractCalciteIndexScan scan = createMockScan(false, false, 3); RelNode parent = createMockNode(scan); assertTrue(ScriptDetector.hasScripts(parent)); } @Test - void returnsFalseForNestedPlanWithoutScripts() { - AbstractCalciteIndexScan scan = createMockScan(false); - RelNode parent = createMockNode(scan); - assertFalse(ScriptDetector.hasScripts(parent)); + void detectsJoinNode() { + LogicalJoin join = mock(LogicalJoin.class); + when(join.getJoinType()).thenReturn(JoinRelType.LEFT); + when(join.getInputs()).thenReturn(List.of()); + doAnswer(invocation -> null).when(join).childrenAccept(any(RelVisitor.class)); + assertTrue(ScriptDetector.hasScripts(join)); + } + + @Test + void detectsWindowNode() { + LogicalWindow window = mock(LogicalWindow.class); + when(window.getInputs()).thenReturn(List.of()); + doAnswer(invocation -> null).when(window).childrenAccept(any(RelVisitor.class)); + assertTrue(ScriptDetector.hasScripts(window)); + } + + @Test + void detectsUdfInProject() { + SqlUserDefinedFunction udf = mock(SqlUserDefinedFunction.class); + RexCall udfCall = mock(RexCall.class); + when(udfCall.getOperator()).thenReturn(udf); + when(udfCall.getOperands()).thenReturn(List.of()); + RelDataType type = mock(RelDataType.class); + when(udfCall.getType()).thenReturn(type); + doAnswer(inv -> inv.>getArgument(0).visitCall(udfCall)) + .when(udfCall) + .accept(any()); + + LogicalProject project = mock(LogicalProject.class); + when(project.getProjects()).thenReturn(List.of(udfCall)); + when(project.getInputs()).thenReturn(List.of()); + doAnswer(invocation -> null).when(project).childrenAccept(any(RelVisitor.class)); + + assertTrue(ScriptDetector.hasScripts(project)); + } + + @Test + void detectsRexOverInProject() { + RexOver rexOver = mock(RexOver.class); + SqlOperator op = mock(SqlOperator.class); + when(rexOver.getOperator()).thenReturn(op); + when(rexOver.getOperands()).thenReturn(List.of()); + RelDataType type = mock(RelDataType.class); + when(rexOver.getType()).thenReturn(type); + doAnswer(inv -> inv.>getArgument(0).visitCall(rexOver)) + .when(rexOver) + .accept(any()); + + LogicalProject project = mock(LogicalProject.class); + when(project.getProjects()).thenReturn(List.of(rexOver)); + when(project.getInputs()).thenReturn(List.of()); + doAnswer(invocation -> null).when(project).childrenAccept(any(RelVisitor.class)); + + assertTrue(ScriptDetector.hasScripts(project)); + } + + @Test + void returnsFalseForSimpleFieldProject() { + RexInputRef fieldRef = mock(RexInputRef.class); + RelDataType type = mock(RelDataType.class); + when(fieldRef.getType()).thenReturn(type); + + LogicalProject project = mock(LogicalProject.class); + when(project.getProjects()).thenReturn(List.of((RexNode) fieldRef)); + when(project.getInputs()).thenReturn(List.of()); + doAnswer(invocation -> null).when(project).childrenAccept(any(RelVisitor.class)); + + assertFalse(ScriptDetector.hasScripts(project)); } private static RelNode createMockNode(RelNode... children) { @@ -75,16 +164,25 @@ private static RelNode createMockNode(RelNode... children) { return node; } - private static AbstractCalciteIndexScan createMockScan(boolean scriptPushed) { + private static AbstractCalciteIndexScan createMockScan( + boolean scriptPushed, boolean sortExprPushed, long aggScriptCount) { AbstractCalciteIndexScan scan = mock(AbstractCalciteIndexScan.class); - when(scan.isScriptPushed()).thenReturn(scriptPushed); + + PushDownContext ctx = mock(PushDownContext.class); + when(ctx.isScriptPushed()).thenReturn(scriptPushed); + when(ctx.isSortExprPushed()).thenReturn(sortExprPushed); + + if (aggScriptCount > 0) { + AggSpec aggSpec = mock(AggSpec.class); + when(aggSpec.getScriptCount()).thenReturn(aggScriptCount); + when(ctx.getAggSpec()).thenReturn(aggSpec); + } else { + when(ctx.getAggSpec()).thenReturn(null); + } + + when(scan.getPushDownContext()).thenReturn(ctx); when(scan.getInputs()).thenReturn(List.of()); - doAnswer( - invocation -> { - return null; - }) - .when(scan) - .childrenAccept(any(RelVisitor.class)); + doAnswer(invocation -> null).when(scan).childrenAccept(any(RelVisitor.class)); return scan; } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java index cd984fd5ba7..5e2d52b7f3a 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java @@ -30,6 +30,7 @@ import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.opensearch.storage.scan.AbstractCalciteIndexScan; +import org.opensearch.sql.opensearch.storage.scan.context.PushDownContext; import org.opensearch.threadpool.ThreadPool; @ExtendWith(MockitoExtension.class) @@ -65,7 +66,7 @@ void executesInlineWhenNoScripts() { void dispatchesToSlowPoolWhenScriptsDetected() { when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) .thenReturn(true); - AbstractCalciteIndexScan scan = createMockScan(true); + AbstractCalciteIndexScan scan = createMockScanWithScripts(); dispatcher.dispatch(scan, context, listener, engine); @@ -78,7 +79,7 @@ void dispatchesToSlowPoolWhenScriptsDetected() { void executesInlineWhenSlowPoolDisabled() { when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) .thenReturn(false); - AbstractCalciteIndexScan scan = createMockScan(true); + AbstractCalciteIndexScan scan = createMockScanWithScripts(); dispatcher.dispatch(scan, context, listener, engine); @@ -99,7 +100,7 @@ void scheduledRunnableCallsEngine() { .when(threadPool) .schedule(any(Runnable.class), any(TimeValue.class), any()); - AbstractCalciteIndexScan scan = createMockScan(true); + AbstractCalciteIndexScan scan = createMockScanWithScripts(); dispatcher.dispatch(scan, context, listener, engine); verify(engine).execute(scan, context, listener); @@ -122,9 +123,13 @@ private static RelNode createMockNode(RelNode... children) { return node; } - private static AbstractCalciteIndexScan createMockScan(boolean scriptPushed) { + private static AbstractCalciteIndexScan createMockScanWithScripts() { AbstractCalciteIndexScan scan = mock(AbstractCalciteIndexScan.class); - when(scan.isScriptPushed()).thenReturn(scriptPushed); + PushDownContext ctx = mock(PushDownContext.class); + when(ctx.isScriptPushed()).thenReturn(true); + when(ctx.isSortExprPushed()).thenReturn(false); + when(ctx.getAggSpec()).thenReturn(null); + when(scan.getPushDownContext()).thenReturn(ctx); when(scan.getInputs()).thenReturn(List.of()); doAnswer(invocation -> null).when(scan).childrenAccept(any(RelVisitor.class)); return scan; From 49e41df2a301f8b087cdfc80ac5b203933e97572 Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Mon, 13 Jul 2026 23:06:32 +0000 Subject: [PATCH 03/11] fixes: assorted context propagation issues Signed-off-by: Simeon Widdis --- .../opensearch/executor/ScriptDetector.java | 18 ++++++---- .../ThreadPoolExecutionDispatcher.java | 21 +++++++++++ .../executor/ScriptDetectorTest.java | 36 +++++++++++++++---- .../plugin/rest/RestUnifiedQueryAction.java | 29 ++++++++++++++- 4 files changed, 91 insertions(+), 13 deletions(-) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java index a26868910ed..aafb2f307b3 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java @@ -5,6 +5,7 @@ package org.opensearch.sql.opensearch.executor; +import java.util.Set; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelVisitor; import org.apache.calcite.rel.core.Join; @@ -14,7 +15,6 @@ import org.apache.calcite.rex.RexNode; import org.apache.calcite.rex.RexOver; import org.apache.calcite.rex.RexVisitorImpl; -import org.apache.calcite.sql.validate.SqlUserDefinedFunction; import org.opensearch.sql.opensearch.storage.scan.AbstractCalciteIndexScan; import org.opensearch.sql.opensearch.storage.scan.context.AggSpec; import org.opensearch.sql.opensearch.storage.scan.context.PushDownContext; @@ -29,6 +29,9 @@ */ public final class ScriptDetector { + private static final Set EXPENSIVE_UDFS = + Set.of("REX_EXTRACT", "REX_EXTRACT_MULTI", "PARSE", "PATTERN_PARSER"); + private ScriptDetector() {} /** @@ -92,13 +95,16 @@ private static boolean hasExpensiveRex(RexNode expr) { boolean[] found = {false}; expr.accept( new RexVisitorImpl(true) { + @Override + public Void visitOver(RexOver over) { + found[0] = true; + return null; + } + @Override public Void visitCall(RexCall call) { - if (call.getOperator() instanceof SqlUserDefinedFunction) { - found[0] = true; - return null; - } - if (call instanceof RexOver) { + String name = call.getOperator().getName(); + if (name != null && EXPENSIVE_UDFS.contains(name)) { found[0] = true; return null; } 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 95685f2e793..8944b51e5f0 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 @@ -8,11 +8,16 @@ import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_SLOW_WORKER_THREAD_POOL_NAME; import java.util.Map; +import java.util.function.Consumer; import lombok.RequiredArgsConstructor; import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.metadata.JaninoRelMetadataProvider; +import org.apache.calcite.rel.metadata.RelMetadataQueryBase; +import org.apache.calcite.runtime.Hook; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.logging.log4j.ThreadContext; +import org.checkerframework.checker.nullness.qual.Nullable; import org.opensearch.common.unit.TimeValue; import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.common.response.ResponseListener; @@ -45,14 +50,30 @@ public void dispatch( LOG.debug("Query plan contains scripts, dispatching to slow worker pool"); Map ctx = ThreadContext.getImmutableContext(); CancellableTask task = OpenSearchQueryManager.getCancellableTask(); + @Nullable JaninoRelMetadataProvider metadataProvider = + RelMetadataQueryBase.THREAD_PROVIDERS.get(); + long currentTime = Hook.CURRENT_TIME.get(-1L); threadPool.schedule( () -> { ThreadContext.putAll(ctx); OpenSearchQueryManager.setCancellableTask(task); + if (metadataProvider != null) { + RelMetadataQueryBase.THREAD_PROVIDERS.set(metadataProvider); + } + Hook.Closeable hookHandle = null; + if (currentTime >= 0) { + hookHandle = + Hook.CURRENT_TIME.addThread( + (Consumer>) h -> h.set(currentTime)); + } try { engine.execute(plan, context, listener); } finally { + if (hookHandle != null) { + hookHandle.close(); + } OpenSearchQueryManager.clearCancellableTask(); + RelMetadataQueryBase.THREAD_PROVIDERS.remove(); } }, new TimeValue(0), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java index e7e186f244f..e44809e40a9 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java @@ -25,7 +25,6 @@ import org.apache.calcite.rex.RexNode; import org.apache.calcite.rex.RexOver; import org.apache.calcite.sql.SqlOperator; -import org.apache.calcite.sql.validate.SqlUserDefinedFunction; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.junit.jupiter.MockitoExtension; @@ -94,10 +93,11 @@ void detectsWindowNode() { } @Test - void detectsUdfInProject() { - SqlUserDefinedFunction udf = mock(SqlUserDefinedFunction.class); + void detectsExpensiveUdfInProject() { + SqlOperator rexExtractOp = mock(SqlOperator.class); + when(rexExtractOp.getName()).thenReturn("REX_EXTRACT"); RexCall udfCall = mock(RexCall.class); - when(udfCall.getOperator()).thenReturn(udf); + when(udfCall.getOperator()).thenReturn(rexExtractOp); when(udfCall.getOperands()).thenReturn(List.of()); RelDataType type = mock(RelDataType.class); when(udfCall.getType()).thenReturn(type); @@ -113,6 +113,30 @@ void detectsUdfInProject() { assertTrue(ScriptDetector.hasScripts(project)); } + @Test + void ignoresCheapUdfInProject() throws Exception { + SqlOperator cheapOp = mock(SqlOperator.class); + when(cheapOp.getName()).thenReturn("NOW"); + RexCall cheapCall = mock(RexCall.class); + when(cheapCall.getOperator()).thenReturn(cheapOp); + RelDataType type = mock(RelDataType.class); + when(cheapCall.getType()).thenReturn(type); + // RexVisitorImpl.visitCall accesses call.operands field directly, set it via reflection + java.lang.reflect.Field operandsField = RexCall.class.getDeclaredField("operands"); + operandsField.setAccessible(true); + operandsField.set(cheapCall, com.google.common.collect.ImmutableList.of()); + doAnswer(inv -> inv.>getArgument(0).visitCall(cheapCall)) + .when(cheapCall) + .accept(any()); + + LogicalProject project = mock(LogicalProject.class); + when(project.getProjects()).thenReturn(List.of((RexNode) cheapCall)); + when(project.getInputs()).thenReturn(List.of()); + doAnswer(invocation -> null).when(project).childrenAccept(any(RelVisitor.class)); + + assertFalse(ScriptDetector.hasScripts(project)); + } + @Test void detectsRexOverInProject() { RexOver rexOver = mock(RexOver.class); @@ -121,9 +145,9 @@ void detectsRexOverInProject() { when(rexOver.getOperands()).thenReturn(List.of()); RelDataType type = mock(RelDataType.class); when(rexOver.getType()).thenReturn(type); - doAnswer(inv -> inv.>getArgument(0).visitCall(rexOver)) + doAnswer(inv -> inv.>getArgument(0).visitOver(rexOver)) .when(rexOver) - .accept(any()); + .accept(any(org.apache.calcite.rex.RexVisitor.class)); LogicalProject project = mock(LogicalProject.class); when(project.getProjects()).thenReturn(List.of(rexOver)); diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java index aca90c8db5e..76f81a50b7b 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java @@ -16,7 +16,12 @@ import com.google.gson.JsonParser; import java.util.Map; import java.util.Optional; +import java.util.function.Consumer; import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.metadata.JaninoRelMetadataProvider; +import org.apache.calcite.rel.metadata.RelMetadataQueryBase; +import org.apache.calcite.runtime.Hook; +import org.apache.calcite.util.Holder; import org.apache.commons.lang3.Strings; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -249,10 +254,32 @@ private void doExecute( createQueryListener(queryType, profileCtx, closingListener)); if (isSlowPoolEnabled() && ScriptDetector.hasScripts(finalPlan)) { LOG.debug("Query contains scripts, routing to slow worker pool"); + JaninoRelMetadataProvider metadataProvider = + RelMetadataQueryBase.THREAD_PROVIDERS.get(); + long currentTime = Hook.CURRENT_TIME.get(-1L); client .threadPool() .schedule( - withCurrentContext(executeTask), + withCurrentContext( + () -> { + if (metadataProvider != null) { + RelMetadataQueryBase.THREAD_PROVIDERS.set(metadataProvider); + } + Hook.Closeable hookHandle = null; + if (currentTime >= 0) { + hookHandle = + Hook.CURRENT_TIME.addThread( + (Consumer>) h -> h.set(currentTime)); + } + try { + executeTask.run(); + } finally { + if (hookHandle != null) { + hookHandle.close(); + } + RelMetadataQueryBase.THREAD_PROVIDERS.remove(); + } + }), new TimeValue(0), SQL_SLOW_WORKER_THREAD_POOL_NAME); } else { From 657657c19ef3dfdc89a2d08000068163e97a1579 Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Tue, 14 Jul 2026 18:46:02 +0000 Subject: [PATCH 04/11] fix remaining integ tests Signed-off-by: Simeon Widdis --- .../executor/ThreadPoolExecutionDispatcher.java | 9 +++++++++ .../sql/plugin/rest/RestUnifiedQueryAction.java | 9 +++++++++ 2 files changed, 18 insertions(+) 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 8944b51e5f0..769ef6bc186 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 @@ -53,6 +53,9 @@ public void dispatch( @Nullable JaninoRelMetadataProvider metadataProvider = RelMetadataQueryBase.THREAD_PROVIDERS.get(); long currentTime = Hook.CURRENT_TIME.get(-1L); + boolean stripNullCols = CalcitePlanContext.stripNullColumns.get(); + String twUnitName = CalcitePlanContext.timewrapUnitName.get(); + String twSeries = CalcitePlanContext.timewrapSeries.get(); threadPool.schedule( () -> { ThreadContext.putAll(ctx); @@ -66,14 +69,20 @@ public void dispatch( Hook.CURRENT_TIME.addThread( (Consumer>) h -> h.set(currentTime)); } + CalcitePlanContext.stripNullColumns.set(stripNullCols); + CalcitePlanContext.timewrapUnitName.set(twUnitName); + CalcitePlanContext.timewrapSeries.set(twSeries); try { engine.execute(plan, context, listener); + } catch (Exception e) { + listener.onFailure(e); } finally { if (hookHandle != null) { hookHandle.close(); } OpenSearchQueryManager.clearCancellableTask(); RelMetadataQueryBase.THREAD_PROVIDERS.remove(); + CalcitePlanContext.clearTimewrapSignals(); } }, new TimeValue(0), diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java index 76f81a50b7b..429b705a1af 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java @@ -257,6 +257,9 @@ private void doExecute( JaninoRelMetadataProvider metadataProvider = RelMetadataQueryBase.THREAD_PROVIDERS.get(); long currentTime = Hook.CURRENT_TIME.get(-1L); + boolean stripNullCols = CalcitePlanContext.stripNullColumns.get(); + String twUnitName = CalcitePlanContext.timewrapUnitName.get(); + String twSeries = CalcitePlanContext.timewrapSeries.get(); client .threadPool() .schedule( @@ -271,13 +274,19 @@ private void doExecute( Hook.CURRENT_TIME.addThread( (Consumer>) h -> h.set(currentTime)); } + CalcitePlanContext.stripNullColumns.set(stripNullCols); + CalcitePlanContext.timewrapUnitName.set(twUnitName); + CalcitePlanContext.timewrapSeries.set(twSeries); try { executeTask.run(); + } catch (Exception e) { + closingListener.onFailure(e); } finally { if (hookHandle != null) { hookHandle.close(); } RelMetadataQueryBase.THREAD_PROVIDERS.remove(); + CalcitePlanContext.clearTimewrapSignals(); } }), new TimeValue(0), From 5b0d7b287a1505b5a57d69982cc0adf1add67edd Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Wed, 15 Jul 2026 16:37:06 +0000 Subject: [PATCH 05/11] add some more thread & security tests Signed-off-by: Simeon Widdis --- .../sql/security/FGACIndexScanningIT.java | 84 +++++++ .../ThreadPoolExecutionDispatcherTest.java | 237 ++++++++++++++++++ 2 files changed, 321 insertions(+) diff --git a/integ-test/src/test/java/org/opensearch/sql/security/FGACIndexScanningIT.java b/integ-test/src/test/java/org/opensearch/sql/security/FGACIndexScanningIT.java index 22c93591241..8bb9718dd61 100644 --- a/integ-test/src/test/java/org/opensearch/sql/security/FGACIndexScanningIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/security/FGACIndexScanningIT.java @@ -656,4 +656,88 @@ public void testRowLevelSecurity(boolean useCalcite) throws IOException { expectedPublicDocs, totalDocs); } + + /** + * Verifies that document-level security is enforced when queries are dispatched to the slow + * worker pool. Queries containing window functions (eventstats) are routed to sql-slow-worker; + * this test ensures the OpenSearch ThreadContext (which carries DLS filters) is correctly + * propagated across that thread pool boundary. + */ + @ParameterizedTest + @ValueSource(booleans = {false, true}) + public void testRowLevelSecurityEnforcedOnSlowPool(boolean useCalcite) throws IOException { + configureEngine(useCalcite); + String engineLabel = useCalcite ? "V3" : "V2"; + + // eventstats creates a Window node, which ScriptDetector flags as expensive, + // routing the query to the sql-slow-worker pool. + String query = + String.format( + "search source=%s | eventstats count() as total_count by security_level" + + " | stats count() by security_level", + SECURE_LOGS); + JSONObject result = executeQueryAsUser(query, LIMITED_USER); + + var datarows = result.getJSONArray("datarows"); + var schema = result.getJSONArray("schema"); + int levelIdx = -1; + for (int i = 0; i < schema.length(); i++) { + String name = schema.getJSONObject(i).getString("name"); + if ("security_level".equals(name)) { + levelIdx = i; + } + } + assertTrue("Expected security_level in schema", levelIdx >= 0); + + for (int i = 0; i < datarows.length(); i++) { + var row = datarows.getJSONArray(i); + String securityLevel = row.getString(levelIdx); + assertFalse( + String.format( + "[%s] SECURITY VIOLATION on slow pool: limited_user saw '%s' documents. " + + "DLS ThreadContext may not be propagated to sql-slow-worker pool.", + engineLabel, securityLevel), + "confidential".equals(securityLevel) || "internal".equals(securityLevel)); + } + } + + /** + * Verifies that field-level security is enforced when queries are dispatched to the slow worker + * pool. The eventstats command creates a Window node that triggers slow pool dispatch; this test + * ensures the restricted field (ssn) remains invisible. + */ + @ParameterizedTest + @ValueSource(booleans = {false, true}) + public void testFieldLevelSecurityEnforcedOnSlowPool(boolean useCalcite) throws IOException { + configureEngine(useCalcite); + + // eventstats creates a Window node, which ScriptDetector flags as expensive, + // routing the query to the sql-slow-worker pool. + // manager_user should still NOT see ssn. + String query = + String.format( + "search source=%s | eventstats avg(salary) as avg_salary by department" + + " | fields name, department, salary, avg_salary | head 10", + EMPLOYEE_RECORDS); + JSONObject result = executeQueryAsUser(query, MANAGER_USER); + + var resultSchema = result.getJSONArray("schema"); + boolean hasSSN = false; + boolean hasName = false; + boolean hasAvgSalary = false; + + for (int i = 0; i < resultSchema.length(); i++) { + String fieldName = resultSchema.getJSONObject(i).getString("name"); + if ("ssn".equals(fieldName)) hasSSN = true; + if ("name".equals(fieldName)) hasName = true; + if ("avg_salary".equals(fieldName)) hasAvgSalary = true; + } + + assertTrue("manager_user should see 'name' field on slow pool", hasName); + assertTrue("manager_user should see computed 'avg_salary' field on slow pool", hasAvgSalary); + assertFalse( + "SECURITY VIOLATION on slow pool: manager_user saw 'ssn' field. " + + "FLS ThreadContext may not be propagated to sql-slow-worker pool.", + hasSSN); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java index 5e2d52b7f3a..f61185e57dc 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java @@ -5,9 +5,13 @@ package org.opensearch.sql.opensearch.executor; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -15,8 +19,13 @@ import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_SLOW_WORKER_THREAD_POOL_NAME; import java.util.List; +import java.util.concurrent.atomic.AtomicReference; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelVisitor; +import org.apache.calcite.rel.metadata.JaninoRelMetadataProvider; +import org.apache.calcite.rel.metadata.RelMetadataQueryBase; +import org.apache.logging.log4j.ThreadContext; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -31,6 +40,7 @@ import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.opensearch.storage.scan.AbstractCalciteIndexScan; import org.opensearch.sql.opensearch.storage.scan.context.PushDownContext; +import org.opensearch.tasks.CancellableTask; import org.opensearch.threadpool.ThreadPool; @ExtendWith(MockitoExtension.class) @@ -50,6 +60,14 @@ void setUp() { dispatcher = new ThreadPoolExecutionDispatcher(threadPool, settings); } + @AfterEach + void tearDown() { + ThreadContext.clearAll(); + OpenSearchQueryManager.clearCancellableTask(); + RelMetadataQueryBase.THREAD_PROVIDERS.remove(); + CalcitePlanContext.clearTimewrapSignals(); + } + @Test void executesInlineWhenNoScripts() { when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) @@ -106,6 +124,225 @@ void scheduledRunnableCallsEngine() { verify(engine).execute(scan, context, listener); } + @Test + void propagatesCancellableTaskToSlowPool() { + when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + .thenReturn(true); + CancellableTask mockTask = mock(CancellableTask.class); + OpenSearchQueryManager.setCancellableTask(mockTask); + + AtomicReference taskOnSlowPool = new AtomicReference<>(); + doAnswer( + invocation -> { + Runnable task = invocation.getArgument(0); + // Simulate running on a different thread — clear the ThreadLocal first + OpenSearchQueryManager.clearCancellableTask(); + task.run(); + taskOnSlowPool.set(OpenSearchQueryManager.getCancellableTask()); + return null; + }) + .when(threadPool) + .schedule(any(Runnable.class), any(TimeValue.class), any()); + + AbstractCalciteIndexScan scan = createMockScanWithScripts(); + dispatcher.dispatch(scan, context, listener, engine); + + // During execution, the task should have been available + // (we check via a side-channel since the finally block clears it) + verify(engine).execute(scan, context, listener); + // After execution, it should be cleaned up + assertNull( + OpenSearchQueryManager.getCancellableTask(), + "CancellableTask should be cleared after execution"); + } + + @Test + void propagatesLog4jThreadContextToSlowPool() { + when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + .thenReturn(true); + ThreadContext.put("request.id", "test-123"); + ThreadContext.put("user", "admin"); + + AtomicReference requestIdOnSlowPool = new AtomicReference<>(); + AtomicReference userOnSlowPool = new AtomicReference<>(); + doAnswer( + invocation -> { + Runnable task = invocation.getArgument(0); + // Simulate a different thread — clear MDC + ThreadContext.clearAll(); + task.run(); + requestIdOnSlowPool.set(ThreadContext.get("request.id")); + userOnSlowPool.set(ThreadContext.get("user")); + return null; + }) + .when(threadPool) + .schedule(any(Runnable.class), any(TimeValue.class), any()); + + AbstractCalciteIndexScan scan = createMockScanWithScripts(); + dispatcher.dispatch(scan, context, listener, engine); + + assertEquals("test-123", requestIdOnSlowPool.get()); + assertEquals("admin", userOnSlowPool.get()); + } + + @Test + void propagatesMetadataProviderToSlowPool() { + when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + .thenReturn(true); + JaninoRelMetadataProvider provider = mock(JaninoRelMetadataProvider.class); + RelMetadataQueryBase.THREAD_PROVIDERS.set(provider); + + AtomicReference providerOnSlowPool = new AtomicReference<>(); + doAnswer( + invocation -> { + Runnable task = invocation.getArgument(0); + RelMetadataQueryBase.THREAD_PROVIDERS.remove(); + task.run(); + providerOnSlowPool.set(RelMetadataQueryBase.THREAD_PROVIDERS.get()); + return null; + }) + .when(threadPool) + .schedule(any(Runnable.class), any(TimeValue.class), any()); + + AbstractCalciteIndexScan scan = createMockScanWithScripts(); + dispatcher.dispatch(scan, context, listener, engine); + + // After finally block, metadata provider should be cleaned up + assertNull( + RelMetadataQueryBase.THREAD_PROVIDERS.get(), + "Metadata provider should be cleaned up after execution"); + } + + @Test + void propagatesTimewrapSignalsToSlowPool() { + when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + .thenReturn(true); + CalcitePlanContext.stripNullColumns.set(true); + CalcitePlanContext.timewrapUnitName.set("HOUR"); + CalcitePlanContext.timewrapSeries.set("timestamp"); + + AtomicReference stripOnSlowPool = new AtomicReference<>(); + AtomicReference unitOnSlowPool = new AtomicReference<>(); + AtomicReference seriesOnSlowPool = new AtomicReference<>(); + doAnswer( + invocation -> { + Runnable task = invocation.getArgument(0); + // Clear thread-locals to simulate new thread + CalcitePlanContext.clearTimewrapSignals(); + CalcitePlanContext.stripNullColumns.set(false); + task.run(); + return null; + }) + .when(threadPool) + .schedule(any(Runnable.class), any(TimeValue.class), any()); + + // Capture the values during engine.execute + doAnswer( + invocation -> { + stripOnSlowPool.set(CalcitePlanContext.stripNullColumns.get()); + unitOnSlowPool.set(CalcitePlanContext.timewrapUnitName.get()); + seriesOnSlowPool.set(CalcitePlanContext.timewrapSeries.get()); + return null; + }) + .when(engine) + .execute(any(RelNode.class), any(), any(ResponseListener.class)); + + AbstractCalciteIndexScan scan = createMockScanWithScripts(); + dispatcher.dispatch(scan, context, listener, engine); + + assertEquals(true, stripOnSlowPool.get()); + assertEquals("HOUR", unitOnSlowPool.get()); + assertEquals("timestamp", seriesOnSlowPool.get()); + } + + @Test + void forwardsExceptionToListenerOnSlowPool() { + when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + .thenReturn(true); + RuntimeException error = new RuntimeException("execution failed"); + doThrow(error).when(engine).execute(any(RelNode.class), any(), any(ResponseListener.class)); + + doAnswer( + invocation -> { + Runnable task = invocation.getArgument(0); + task.run(); + return null; + }) + .when(threadPool) + .schedule(any(Runnable.class), any(TimeValue.class), any()); + + AbstractCalciteIndexScan scan = createMockScanWithScripts(); + dispatcher.dispatch(scan, context, listener, engine); + + verify(listener).onFailure(error); + } + + @Test + void cleansUpThreadLocalsAfterException() { + when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + .thenReturn(true); + CancellableTask mockTask = mock(CancellableTask.class); + OpenSearchQueryManager.setCancellableTask(mockTask); + CalcitePlanContext.timewrapUnitName.set("DAY"); + RelMetadataQueryBase.THREAD_PROVIDERS.set(mock(JaninoRelMetadataProvider.class)); + + doThrow(new RuntimeException("boom")) + .when(engine) + .execute(any(RelNode.class), any(), any(ResponseListener.class)); + + doAnswer( + invocation -> { + Runnable task = invocation.getArgument(0); + task.run(); + return null; + }) + .when(threadPool) + .schedule(any(Runnable.class), any(TimeValue.class), any()); + + AbstractCalciteIndexScan scan = createMockScanWithScripts(); + dispatcher.dispatch(scan, context, listener, engine); + + assertNull(OpenSearchQueryManager.getCancellableTask()); + assertNull(RelMetadataQueryBase.THREAD_PROVIDERS.get()); + // timewrapSignals cleared via clearTimewrapSignals() + assertNull(CalcitePlanContext.timewrapUnitName.get()); + } + + @Test + void cancellableTaskAvailableDuringExecution() { + when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + .thenReturn(true); + CancellableTask mockTask = mock(CancellableTask.class); + OpenSearchQueryManager.setCancellableTask(mockTask); + + AtomicReference taskDuringExecution = new AtomicReference<>(); + doAnswer( + invocation -> { + Runnable task = invocation.getArgument(0); + // Simulate running on a different thread + OpenSearchQueryManager.clearCancellableTask(); + task.run(); + return null; + }) + .when(threadPool) + .schedule(any(Runnable.class), any(TimeValue.class), any()); + + doAnswer( + invocation -> { + taskDuringExecution.set(OpenSearchQueryManager.getCancellableTask()); + return null; + }) + .when(engine) + .execute(any(RelNode.class), any(), any(ResponseListener.class)); + + AbstractCalciteIndexScan scan = createMockScanWithScripts(); + dispatcher.dispatch(scan, context, listener, engine); + + assertNotNull( + taskDuringExecution.get(), "CancellableTask should be available during execution"); + assertEquals(mockTask, taskDuringExecution.get()); + } + private static RelNode createMockNode(RelNode... children) { RelNode node = mock(RelNode.class); List childList = List.of(children); From 46c3cbe9f3322221d5236b6e37dcd1f74fc75099 Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Thu, 16 Jul 2026 16:43:13 +0000 Subject: [PATCH 06/11] code self-review, round 1 Signed-off-by: Simeon Widdis --- .../sql/calcite/utils/CalciteToolsHelper.java | 2 - .../sql/executor/ExecutionDispatcher.java | 14 ++++- .../opensearch/sql/executor/QueryService.java | 7 ++- .../CalcitePPLRelNodeIntegTestCase.java | 4 +- .../executor/OpenSearchExecutionEngine.java | 8 ++- .../ThreadPoolExecutionDispatcher.java | 27 ++++++-- .../org/opensearch/sql/plugin/SQLPlugin.java | 8 ++- .../plugin/rest/RestUnifiedQueryAction.java | 63 +++---------------- .../transport/TransportPPLQueryAction.java | 8 ++- .../rest/RestUnifiedQueryActionTest.java | 3 +- 10 files changed, 72 insertions(+), 72 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java b/core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java index 54b9d4ffbaf..682a3ea17c1 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java +++ b/core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java @@ -516,8 +516,6 @@ private static void enrichErrorsForSpecialCases(ErrorReport.Builder report, SQLE public static PreparedStatement run(CalcitePlanContext context, RelNode rel) { ProfileMetric optimizeTime = QueryProfiling.current().getOrCreateMetric(OPTIMIZE); long startTime = System.nanoTime(); - // Optimize the plan by Calcite's HepPlanner before using VolcanoPlanner in prepareStatement. - rel = CalciteToolsHelper.optimize(rel, context); final RelShuttle shuttle = new RelHomogeneousShuttle() { @Override diff --git a/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java b/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java index cc8f5ee78b9..e497ca9c6d9 100644 --- a/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java +++ b/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java @@ -18,7 +18,7 @@ public interface ExecutionDispatcher { /** - * Dispatch execution of the given plan. + * Dispatch execution of the given plan via the standard ExecutionEngine. * * @param plan the optimized Calcite plan * @param context the plan context @@ -30,4 +30,16 @@ void dispatch( CalcitePlanContext context, ResponseListener listener, ExecutionEngine engine); + + /** + * Dispatch a task to the appropriate thread pool based on plan characteristics. Use this when the + * execution path differs from the standard ExecutionEngine interface (e.g., analytics engine). + * + * @param plan the optimized Calcite plan used for routing decisions + * @param context the plan context + * @param task the execution task to run + */ + default void dispatchTask(RelNode plan, CalcitePlanContext context, Runnable task) { + task.run(); + } } 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 4fd8dd47336..6c283f69d8d 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -35,6 +35,7 @@ import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit; import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit.SystemLimitType; import org.opensearch.sql.calcite.utils.CalciteClassLoaderHelper; +import org.opensearch.sql.calcite.utils.CalciteToolsHelper; import org.opensearch.sql.common.error.ErrorReport; import org.opensearch.sql.common.error.QueryProcessingStage; import org.opensearch.sql.common.error.StageErrorHandler; @@ -208,13 +209,17 @@ public void executeWithCalcite( analyzeMetric.set(System.nanoTime() - analyzeStart); + // Optimize before dispatch so the dispatcher's ScriptDetector + // sees the post-optimization plan for accurate routing. + RelNode optimizedPlan = CalciteToolsHelper.optimize(calcitePlan, context); + // Wrap execution with EXECUTING stage tracking — dispatch via // ExecutionDispatcher which may route to a slow worker pool StageErrorHandler.executeStageVoid( QueryProcessingStage.EXECUTING, () -> executionDispatcher.dispatch( - calcitePlan, context, listener, executionEngine), + optimizedPlan, context, listener, executionEngine), "while running the query"); }, QueryService.class); diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLRelNodeIntegTestCase.java b/integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLRelNodeIntegTestCase.java index 9c7d8c90ac3..744b9ec124b 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLRelNodeIntegTestCase.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLRelNodeIntegTestCase.java @@ -25,6 +25,7 @@ import org.apache.calcite.tools.RelBuilder; import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.calcite.SysLimit; +import org.opensearch.sql.calcite.utils.CalciteToolsHelper; import org.opensearch.sql.calcite.utils.CalciteToolsHelper.OpenSearchRelRunners; import org.opensearch.sql.executor.QueryType; @@ -87,7 +88,8 @@ protected RexNode createStringArray(RexBuilder rexBuilder, String... values) { protected void executeRelNodeAndVerify( CalcitePlanContext planContext, RelNode relNode, ResultVerifier verifier) throws SQLException { - try (PreparedStatement statement = OpenSearchRelRunners.run(planContext, relNode)) { + try (PreparedStatement statement = + OpenSearchRelRunners.run(planContext, CalciteToolsHelper.optimize(relNode, planContext))) { ResultSet resultSet = statement.executeQuery(); verifier.verify(resultSet); } 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 e8c7cfc7c68..17b858ae10a 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 @@ -42,6 +42,7 @@ import org.locationtech.jts.geom.Point; import org.opensearch.sql.ast.statement.ExplainMode; import org.opensearch.sql.calcite.CalcitePlanContext; +import org.opensearch.sql.calcite.utils.CalciteToolsHelper; import org.opensearch.sql.calcite.utils.CalciteToolsHelper.OpenSearchRelRunners; import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory; import org.opensearch.sql.calcite.utils.TimewrapPivot; @@ -259,7 +260,7 @@ public void explain( } })) { // triggers the hook - OpenSearchRelRunners.run(context, rel); + OpenSearchRelRunners.run(context, CalciteToolsHelper.optimize(rel, context)); } if (physicalError.get() != null) { @@ -306,7 +307,7 @@ public void explain( CalcitePlanContext.skipEncoding.set(true); } // triggers the hook - OpenSearchRelRunners.run(context, rel); + OpenSearchRelRunners.run(context, CalciteToolsHelper.optimize(rel, context)); } listener.onResponse( new ExplainResponse( @@ -326,7 +327,8 @@ public void execute( RelNode rel, CalcitePlanContext context, ResponseListener listener) { client.schedule( () -> { - try (PreparedStatement statement = OpenSearchRelRunners.run(context, rel)) { + try (PreparedStatement statement = + OpenSearchRelRunners.run(context, CalciteToolsHelper.optimize(rel, context))) { ProfileMetric metric = QueryProfiling.current().getOrCreateMetric(MetricName.EXECUTE); long execTime = System.nanoTime(); ResultSet result = statement.executeQuery(); 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 769ef6bc186..c8c5365bbd7 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 @@ -46,10 +46,23 @@ public void dispatch( CalcitePlanContext context, ResponseListener listener, ExecutionEngine engine) { - if (isSlowPoolEnabled() && ScriptDetector.hasScripts(plan)) { + dispatchInternal(plan, context, () -> engine.execute(plan, context, listener), listener); + } + + @Override + public void dispatchTask(RelNode plan, CalcitePlanContext context, Runnable task) { + dispatchInternal(plan, context, task, null); + } + + private void dispatchInternal( + RelNode optimizedPlan, + CalcitePlanContext context, + Runnable task, + @Nullable ResponseListener failureListener) { + if (isSlowPoolEnabled() && ScriptDetector.hasScripts(optimizedPlan)) { LOG.debug("Query plan contains scripts, dispatching to slow worker pool"); Map ctx = ThreadContext.getImmutableContext(); - CancellableTask task = OpenSearchQueryManager.getCancellableTask(); + CancellableTask cancellableTask = OpenSearchQueryManager.getCancellableTask(); @Nullable JaninoRelMetadataProvider metadataProvider = RelMetadataQueryBase.THREAD_PROVIDERS.get(); long currentTime = Hook.CURRENT_TIME.get(-1L); @@ -59,7 +72,7 @@ public void dispatch( threadPool.schedule( () -> { ThreadContext.putAll(ctx); - OpenSearchQueryManager.setCancellableTask(task); + OpenSearchQueryManager.setCancellableTask(cancellableTask); if (metadataProvider != null) { RelMetadataQueryBase.THREAD_PROVIDERS.set(metadataProvider); } @@ -73,9 +86,11 @@ public void dispatch( CalcitePlanContext.timewrapUnitName.set(twUnitName); CalcitePlanContext.timewrapSeries.set(twSeries); try { - engine.execute(plan, context, listener); + task.run(); } catch (Exception e) { - listener.onFailure(e); + if (failureListener != null) { + failureListener.onFailure(e); + } } finally { if (hookHandle != null) { hookHandle.close(); @@ -88,7 +103,7 @@ public void dispatch( new TimeValue(0), SQL_SLOW_WORKER_THREAD_POOL_NAME); } else { - engine.execute(plan, context, listener); + task.run(); } } 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 320c778ac3f..92b512eaca1 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -234,7 +234,13 @@ private BiFunction createSqlAnalyticsRout } cached[0] = new RestUnifiedQueryAction( - client, clusterService, executor, contextProvider, pluginSettings); + client, + clusterService, + executor, + contextProvider, + pluginSettings, + new org.opensearch.sql.opensearch.executor.ThreadPoolExecutionDispatcher( + client.threadPool(), pluginSettings)); } return cached[0]; }; diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java index 429b705a1af..26c71060c2e 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java @@ -7,7 +7,6 @@ import static org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; import static org.opensearch.sql.lang.PPLLangSpec.PPL_SPEC; -import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_SLOW_WORKER_THREAD_POOL_NAME; import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME; import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY; @@ -16,12 +15,7 @@ import com.google.gson.JsonParser; import java.util.Map; import java.util.Optional; -import java.util.function.Consumer; import org.apache.calcite.rel.RelNode; -import org.apache.calcite.rel.metadata.JaninoRelMetadataProvider; -import org.apache.calcite.rel.metadata.RelMetadataQueryBase; -import org.apache.calcite.runtime.Hook; -import org.apache.calcite.util.Holder; import org.apache.commons.lang3.Strings; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -46,14 +40,12 @@ import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit; import org.opensearch.sql.common.response.ResponseListener; -import org.opensearch.sql.common.setting.Settings.Key; import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; import org.opensearch.sql.executor.QueryType; import org.opensearch.sql.executor.analytics.AnalyticsExecutionEngine; import org.opensearch.sql.lang.LangSpec; import org.opensearch.sql.monitor.profile.ProfileContext; import org.opensearch.sql.monitor.profile.QueryProfiling; -import org.opensearch.sql.opensearch.executor.ScriptDetector; import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse; import org.opensearch.sql.protocol.response.QueryResult; import org.opensearch.sql.protocol.response.format.ResponseFormatter; @@ -77,18 +69,21 @@ public class RestUnifiedQueryAction { private final ClusterService clusterService; private final org.opensearch.analytics.EngineContextProvider contextProvider; private final org.opensearch.sql.common.setting.Settings pluginSettings; + private final org.opensearch.sql.executor.ExecutionDispatcher executionDispatcher; public RestUnifiedQueryAction( NodeClient client, ClusterService clusterService, QueryPlanExecutor> planExecutor, org.opensearch.analytics.EngineContextProvider contextProvider, - org.opensearch.sql.common.setting.Settings pluginSettings) { + org.opensearch.sql.common.setting.Settings pluginSettings, + org.opensearch.sql.executor.ExecutionDispatcher executionDispatcher) { this.client = client; this.clusterService = clusterService; this.analyticsEngine = new AnalyticsExecutionEngine(planExecutor); this.contextProvider = contextProvider; this.pluginSettings = pluginSettings; + this.executionDispatcher = executionDispatcher; } /** @@ -237,6 +232,9 @@ private void doExecute( // string, so apply the equivalent top-level limit here before the system cap. plan = addFetchSizeLimit(plan, planContext, fetchSize); plan = addQuerySizeLimit(plan, planContext); + plan = + org.opensearch.sql.calcite.utils.CalciteToolsHelper.optimize( + plan, planContext); RelNode finalPlan = plan; Runnable executeTask = profiling @@ -252,48 +250,7 @@ private void doExecute( planContext, queryCtx, createQueryListener(queryType, profileCtx, closingListener)); - if (isSlowPoolEnabled() && ScriptDetector.hasScripts(finalPlan)) { - LOG.debug("Query contains scripts, routing to slow worker pool"); - JaninoRelMetadataProvider metadataProvider = - RelMetadataQueryBase.THREAD_PROVIDERS.get(); - long currentTime = Hook.CURRENT_TIME.get(-1L); - boolean stripNullCols = CalcitePlanContext.stripNullColumns.get(); - String twUnitName = CalcitePlanContext.timewrapUnitName.get(); - String twSeries = CalcitePlanContext.timewrapSeries.get(); - client - .threadPool() - .schedule( - withCurrentContext( - () -> { - if (metadataProvider != null) { - RelMetadataQueryBase.THREAD_PROVIDERS.set(metadataProvider); - } - Hook.Closeable hookHandle = null; - if (currentTime >= 0) { - hookHandle = - Hook.CURRENT_TIME.addThread( - (Consumer>) h -> h.set(currentTime)); - } - CalcitePlanContext.stripNullColumns.set(stripNullCols); - CalcitePlanContext.timewrapUnitName.set(twUnitName); - CalcitePlanContext.timewrapSeries.set(twSeries); - try { - executeTask.run(); - } catch (Exception e) { - closingListener.onFailure(e); - } finally { - if (hookHandle != null) { - hookHandle.close(); - } - RelMetadataQueryBase.THREAD_PROVIDERS.remove(); - CalcitePlanContext.clearTimewrapSignals(); - } - }), - new TimeValue(0), - SQL_SLOW_WORKER_THREAD_POOL_NAME); - } else { - executeTask.run(); - } + executionDispatcher.dispatchTask(finalPlan, planContext, executeTask); } catch (Exception e) { closingListener.onFailure(e); } finally { @@ -512,10 +469,6 @@ private static String appendError(String json, Throwable error) { return json.substring(0, json.length() - 1) + ",\"error\":" + err + "}"; } - private boolean isSlowPoolEnabled() { - return pluginSettings.getSettingValue(Key.SQL_SLOW_WORKER_POOL_ENABLED); - } - private static Runnable withCurrentContext(final Runnable task) { final Map currentContext = ThreadContext.getImmutableContext(); return () -> { 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 678ed58f37f..501a826c0d9 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 @@ -135,7 +135,13 @@ private void buildUnifiedQueryHandlerIfReady() { if (executor != null && contextProvider != null) { this.unifiedQueryHandler = new RestUnifiedQueryAction( - clientRef, clusterServiceRef, executor, contextProvider, pluginSettingsRef); + clientRef, + clusterServiceRef, + executor, + contextProvider, + pluginSettingsRef, + new org.opensearch.sql.opensearch.executor.ThreadPoolExecutionDispatcher( + clientRef.threadPool(), pluginSettingsRef)); } } diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryActionTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryActionTest.java index 111597bb587..de1ce111558 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryActionTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryActionTest.java @@ -54,7 +54,8 @@ public void setUp() { clusterService, executor, mock(EngineContextProvider.class), - mock(org.opensearch.sql.common.setting.Settings.class)); + mock(org.opensearch.sql.common.setting.Settings.class), + new org.opensearch.sql.executor.DirectExecutionDispatcher()); } @Test From 434021f903141c516a3eabf752b48f69a4428e35 Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Fri, 17 Jul 2026 00:02:02 +0000 Subject: [PATCH 07/11] use 2x background threads to account for 2x pools Signed-off-by: Simeon Widdis --- plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 e8493cd57db..12071047ff7 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -480,7 +480,7 @@ public List> getExecutorBuilders(Settings settings) { settings, SQL_BACKGROUND_THREAD_POOL_NAME, settings.getAsInt( - "thread_pool.search.size", OpenSearchExecutors.allocatedProcessors(settings)), + "thread_pool.search.size", 2 * OpenSearchExecutors.allocatedProcessors(settings)), 1000, "thread_pool." + SQL_BACKGROUND_THREAD_POOL_NAME)); } From 169375fe9a2c82164acec65ac0b585784e4a5fc3 Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Fri, 17 Jul 2026 21:31:45 +0000 Subject: [PATCH 08/11] rename slow -> complex, add pool indication header Signed-off-by: Simeon Widdis --- .../sql/common/setting/Settings.java | 2 +- .../sql/calcite/CalcitePlanContext.java | 6 +++++ .../executor/DirectExecutionDispatcher.java | 4 +-- .../sql/executor/ExecutionDispatcher.java | 2 +- .../opensearch/sql/executor/QueryService.java | 2 +- .../sql/security/FGACIndexScanningIT.java | 26 +++++++++--------- .../executor/OpenSearchQueryManager.java | 2 +- .../ThreadPoolExecutionDispatcher.java | 21 ++++++++------- .../setting/OpenSearchSettings.java | 12 ++++----- .../ThreadPoolExecutionDispatcherTest.java | 27 ++++++++++--------- .../org/opensearch/sql/plugin/SQLPlugin.java | 8 +++--- .../sql/plugin/rest/RestPPLQueryAction.java | 8 +++++- 12 files changed, 68 insertions(+), 52 deletions(-) diff --git a/common/src/main/java/org/opensearch/sql/common/setting/Settings.java b/common/src/main/java/org/opensearch/sql/common/setting/Settings.java index 952a8045507..7e37df7512a 100644 --- a/common/src/main/java/org/opensearch/sql/common/setting/Settings.java +++ b/common/src/main/java/org/opensearch/sql/common/setting/Settings.java @@ -84,7 +84,7 @@ public enum Key { "plugins.query.executionengine.spark.streamingjobs.housekeeper.interval"), /** Thread Pool Settings. */ - SQL_SLOW_WORKER_POOL_ENABLED("plugins.sql.slow_worker_pool.enabled"); + SQL_COMPLEX_WORKER_POOL_ENABLED("plugins.sql.complex_worker_pool.enabled"); @Getter private final String keyValue; diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java index 162a4895805..1dd28636de5 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -58,6 +58,11 @@ public class CalcitePlanContext { /** Timewrap series mode: "relative", "short", or "exact". */ public static final ThreadLocal timewrapSeries = new ThreadLocal<>(); + /** + * Thread-local tracking which pool executed this query ("sql-worker" or "sql-complex-worker"). + */ + public static final ThreadLocal executionPool = new ThreadLocal<>(); + /** Thread-local switch that tells whether the current query prefers legacy behavior. */ private static final ThreadLocal legacyPreferredFlag = ThreadLocal.withInitial(() -> true); @@ -229,6 +234,7 @@ public static void clearTimewrapSignals() { stripNullColumns.set(false); timewrapUnitName.set(null); timewrapSeries.set(null); + executionPool.set(null); } public void pushForeachBindings( diff --git a/core/src/main/java/org/opensearch/sql/executor/DirectExecutionDispatcher.java b/core/src/main/java/org/opensearch/sql/executor/DirectExecutionDispatcher.java index 0b168f53a56..5af110fa7e5 100644 --- a/core/src/main/java/org/opensearch/sql/executor/DirectExecutionDispatcher.java +++ b/core/src/main/java/org/opensearch/sql/executor/DirectExecutionDispatcher.java @@ -10,8 +10,8 @@ import org.opensearch.sql.common.response.ResponseListener; /** - * Default no-op dispatcher that executes inline on the current thread. Used when slow-pool routing - * is disabled or as a fallback. + * Default no-op dispatcher that executes inline on the current thread. Used when complex-pool + * routing is disabled or as a fallback. */ public class DirectExecutionDispatcher implements ExecutionDispatcher { diff --git a/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java b/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java index e497ca9c6d9..ec7fa193852 100644 --- a/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java +++ b/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java @@ -12,7 +12,7 @@ /** * Dispatches query execution to an appropriate thread pool based on plan characteristics. After * query analysis and optimization, the dispatcher inspects the plan and routes execution to either - * the fast worker pool (for queries fully pushed to OpenSearch) or the slow worker pool (for + * the fast worker pool (for queries fully pushed to OpenSearch) or the complex worker pool (for * queries requiring scripts/table scans). */ public interface ExecutionDispatcher { 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 06c30821d62..947f999f9eb 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -242,7 +242,7 @@ private void executeCalcitePlan( RelNode optimizedPlan = CalciteToolsHelper.optimize(calcitePlan, context); // Wrap execution with EXECUTING stage tracking — dispatch via - // ExecutionDispatcher which may route to a slow worker pool + // ExecutionDispatcher which may route to a complex worker pool StageErrorHandler.executeStageVoid( QueryProcessingStage.EXECUTING, () -> executionDispatcher.dispatch(optimizedPlan, context, listener, executionEngine), diff --git a/integ-test/src/test/java/org/opensearch/sql/security/FGACIndexScanningIT.java b/integ-test/src/test/java/org/opensearch/sql/security/FGACIndexScanningIT.java index 8bb9718dd61..97559b6032c 100644 --- a/integ-test/src/test/java/org/opensearch/sql/security/FGACIndexScanningIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/security/FGACIndexScanningIT.java @@ -658,19 +658,19 @@ public void testRowLevelSecurity(boolean useCalcite) throws IOException { } /** - * Verifies that document-level security is enforced when queries are dispatched to the slow - * worker pool. Queries containing window functions (eventstats) are routed to sql-slow-worker; + * Verifies that document-level security is enforced when queries are dispatched to the complex + * worker pool. Queries containing window functions (eventstats) are routed to sql-complex-worker; * this test ensures the OpenSearch ThreadContext (which carries DLS filters) is correctly * propagated across that thread pool boundary. */ @ParameterizedTest @ValueSource(booleans = {false, true}) - public void testRowLevelSecurityEnforcedOnSlowPool(boolean useCalcite) throws IOException { + public void testRowLevelSecurityEnforcedOnComplexPool(boolean useCalcite) throws IOException { configureEngine(useCalcite); String engineLabel = useCalcite ? "V3" : "V2"; // eventstats creates a Window node, which ScriptDetector flags as expensive, - // routing the query to the sql-slow-worker pool. + // routing the query to the sql-complex-worker pool. String query = String.format( "search source=%s | eventstats count() as total_count by security_level" @@ -694,8 +694,8 @@ public void testRowLevelSecurityEnforcedOnSlowPool(boolean useCalcite) throws IO String securityLevel = row.getString(levelIdx); assertFalse( String.format( - "[%s] SECURITY VIOLATION on slow pool: limited_user saw '%s' documents. " - + "DLS ThreadContext may not be propagated to sql-slow-worker pool.", + "[%s] SECURITY VIOLATION on complex pool: limited_user saw '%s' documents. " + + "DLS ThreadContext may not be propagated to sql-complex-worker pool.", engineLabel, securityLevel), "confidential".equals(securityLevel) || "internal".equals(securityLevel)); } @@ -703,8 +703,8 @@ public void testRowLevelSecurityEnforcedOnSlowPool(boolean useCalcite) throws IO /** * Verifies that field-level security is enforced when queries are dispatched to the slow worker - * pool. The eventstats command creates a Window node that triggers slow pool dispatch; this test - * ensures the restricted field (ssn) remains invisible. + * pool. The eventstats command creates a Window node that triggers complex pool dispatch; this + * test ensures the restricted field (ssn) remains invisible. */ @ParameterizedTest @ValueSource(booleans = {false, true}) @@ -712,7 +712,7 @@ public void testFieldLevelSecurityEnforcedOnSlowPool(boolean useCalcite) throws configureEngine(useCalcite); // eventstats creates a Window node, which ScriptDetector flags as expensive, - // routing the query to the sql-slow-worker pool. + // routing the query to the sql-complex-worker pool. // manager_user should still NOT see ssn. String query = String.format( @@ -733,11 +733,11 @@ public void testFieldLevelSecurityEnforcedOnSlowPool(boolean useCalcite) throws if ("avg_salary".equals(fieldName)) hasAvgSalary = true; } - assertTrue("manager_user should see 'name' field on slow pool", hasName); - assertTrue("manager_user should see computed 'avg_salary' field on slow pool", hasAvgSalary); + assertTrue("manager_user should see 'name' field on complex pool", hasName); + assertTrue("manager_user should see computed 'avg_salary' field on complex pool", hasAvgSalary); assertFalse( - "SECURITY VIOLATION on slow pool: manager_user saw 'ssn' field. " - + "FLS ThreadContext may not be propagated to sql-slow-worker pool.", + "SECURITY VIOLATION on complex pool: manager_user saw 'ssn' field. " + + "FLS ThreadContext may not be propagated to sql-complex-worker pool.", hasSSN); } } 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 15245deb01a..c391153fca6 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 @@ -32,7 +32,7 @@ public class OpenSearchQueryManager implements QueryManager { private final Settings settings; public static final String SQL_WORKER_THREAD_POOL_NAME = "sql-worker"; - public static final String SQL_SLOW_WORKER_THREAD_POOL_NAME = "sql-slow-worker"; + public static final String SQL_COMPLEX_WORKER_THREAD_POOL_NAME = "sql-complex-worker"; public static final String SQL_BACKGROUND_THREAD_POOL_NAME = "sql_background_io"; private static final ThreadLocal cancellableTask = new ThreadLocal<>(); 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 c8c5365bbd7..793a24b2c4b 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 @@ -5,7 +5,8 @@ package org.opensearch.sql.opensearch.executor; -import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_SLOW_WORKER_THREAD_POOL_NAME; +import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_COMPLEX_WORKER_THREAD_POOL_NAME; +import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME; import java.util.Map; import java.util.function.Consumer; @@ -28,9 +29,9 @@ import org.opensearch.threadpool.ThreadPool; /** - * Dispatches query execution to either the fast or slow worker thread pool based on whether the - * plan contains scripts. Plans with scripts require in-memory evaluation and are routed to the slow - * pool so they don't block fast pushdown-only queries. + * Dispatches query execution to either the fast or complex worker thread pool based on whether the + * plan contains scripts. Plans with scripts require in-memory evaluation and are routed to the + * complex pool so they don't block fast pushdown-only queries. */ @RequiredArgsConstructor public class ThreadPoolExecutionDispatcher implements ExecutionDispatcher { @@ -59,8 +60,8 @@ private void dispatchInternal( CalcitePlanContext context, Runnable task, @Nullable ResponseListener failureListener) { - if (isSlowPoolEnabled() && ScriptDetector.hasScripts(optimizedPlan)) { - LOG.debug("Query plan contains scripts, dispatching to slow worker pool"); + if (isComplexPoolEnabled() && ScriptDetector.hasScripts(optimizedPlan)) { + LOG.debug("Query plan contains scripts, dispatching to complex worker pool"); Map ctx = ThreadContext.getImmutableContext(); CancellableTask cancellableTask = OpenSearchQueryManager.getCancellableTask(); @Nullable JaninoRelMetadataProvider metadataProvider = @@ -73,6 +74,7 @@ private void dispatchInternal( () -> { ThreadContext.putAll(ctx); OpenSearchQueryManager.setCancellableTask(cancellableTask); + CalcitePlanContext.executionPool.set(SQL_COMPLEX_WORKER_THREAD_POOL_NAME); if (metadataProvider != null) { RelMetadataQueryBase.THREAD_PROVIDERS.set(metadataProvider); } @@ -101,13 +103,14 @@ private void dispatchInternal( } }, new TimeValue(0), - SQL_SLOW_WORKER_THREAD_POOL_NAME); + SQL_COMPLEX_WORKER_THREAD_POOL_NAME); } else { + CalcitePlanContext.executionPool.set(SQL_WORKER_THREAD_POOL_NAME); task.run(); } } - private boolean isSlowPoolEnabled() { - return settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED); + private boolean isComplexPoolEnabled() { + return settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java index 2e3d71fac33..65dacbfdf89 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java @@ -363,9 +363,9 @@ public class OpenSearchSettings extends Settings { Setting.Property.NodeScope, Setting.Property.Dynamic); - public static final Setting SQL_SLOW_WORKER_POOL_ENABLED_SETTING = + public static final Setting SQL_COMPLEX_WORKER_POOL_ENABLED_SETTING = Setting.boolSetting( - Key.SQL_SLOW_WORKER_POOL_ENABLED.getKeyValue(), + Key.SQL_COMPLEX_WORKER_POOL_ENABLED.getKeyValue(), true, Setting.Property.NodeScope, Setting.Property.Dynamic); @@ -641,9 +641,9 @@ public OpenSearchSettings(ClusterSettings clusterSettings) { register( settingBuilder, clusterSettings, - Key.SQL_SLOW_WORKER_POOL_ENABLED, - SQL_SLOW_WORKER_POOL_ENABLED_SETTING, - new Updater(Key.SQL_SLOW_WORKER_POOL_ENABLED)); + Key.SQL_COMPLEX_WORKER_POOL_ENABLED, + SQL_COMPLEX_WORKER_POOL_ENABLED_SETTING, + new Updater(Key.SQL_COMPLEX_WORKER_POOL_ENABLED)); defaultSettings = settingBuilder.build(); } @@ -739,7 +739,7 @@ public static List> pluginSettings() { .add(SESSION_INACTIVITY_TIMEOUT_MILLIS_SETTING) .add(STREAMING_JOB_HOUSEKEEPER_INTERVAL_SETTING) .add(FIELD_TYPE_TOLERANCE_SETTING) - .add(SQL_SLOW_WORKER_POOL_ENABLED_SETTING) + .add(SQL_COMPLEX_WORKER_POOL_ENABLED_SETTING) .build(); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java index f61185e57dc..abf806ec810 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java @@ -16,7 +16,7 @@ import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_SLOW_WORKER_THREAD_POOL_NAME; +import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_COMPLEX_WORKER_THREAD_POOL_NAME; import java.util.List; import java.util.concurrent.atomic.AtomicReference; @@ -70,7 +70,7 @@ void tearDown() { @Test void executesInlineWhenNoScripts() { - when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); RelNode plan = createMockNode(); @@ -82,20 +82,21 @@ void executesInlineWhenNoScripts() { @Test void dispatchesToSlowPoolWhenScriptsDetected() { - when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); AbstractCalciteIndexScan scan = createMockScanWithScripts(); dispatcher.dispatch(scan, context, listener, engine); verify(threadPool) - .schedule(any(Runnable.class), eq(new TimeValue(0)), eq(SQL_SLOW_WORKER_THREAD_POOL_NAME)); + .schedule( + any(Runnable.class), eq(new TimeValue(0)), eq(SQL_COMPLEX_WORKER_THREAD_POOL_NAME)); verify(engine, never()).execute(any(RelNode.class), any(), any(ResponseListener.class)); } @Test void executesInlineWhenSlowPoolDisabled() { - when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(false); AbstractCalciteIndexScan scan = createMockScanWithScripts(); @@ -107,7 +108,7 @@ void executesInlineWhenSlowPoolDisabled() { @Test void scheduledRunnableCallsEngine() { - when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); doAnswer( invocation -> { @@ -126,7 +127,7 @@ void scheduledRunnableCallsEngine() { @Test void propagatesCancellableTaskToSlowPool() { - when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); CancellableTask mockTask = mock(CancellableTask.class); OpenSearchQueryManager.setCancellableTask(mockTask); @@ -158,7 +159,7 @@ void propagatesCancellableTaskToSlowPool() { @Test void propagatesLog4jThreadContextToSlowPool() { - when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); ThreadContext.put("request.id", "test-123"); ThreadContext.put("user", "admin"); @@ -187,7 +188,7 @@ void propagatesLog4jThreadContextToSlowPool() { @Test void propagatesMetadataProviderToSlowPool() { - when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); JaninoRelMetadataProvider provider = mock(JaninoRelMetadataProvider.class); RelMetadataQueryBase.THREAD_PROVIDERS.set(provider); @@ -215,7 +216,7 @@ void propagatesMetadataProviderToSlowPool() { @Test void propagatesTimewrapSignalsToSlowPool() { - when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); CalcitePlanContext.stripNullColumns.set(true); CalcitePlanContext.timewrapUnitName.set("HOUR"); @@ -257,7 +258,7 @@ void propagatesTimewrapSignalsToSlowPool() { @Test void forwardsExceptionToListenerOnSlowPool() { - when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); RuntimeException error = new RuntimeException("execution failed"); doThrow(error).when(engine).execute(any(RelNode.class), any(), any(ResponseListener.class)); @@ -279,7 +280,7 @@ void forwardsExceptionToListenerOnSlowPool() { @Test void cleansUpThreadLocalsAfterException() { - when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); CancellableTask mockTask = mock(CancellableTask.class); OpenSearchQueryManager.setCancellableTask(mockTask); @@ -310,7 +311,7 @@ void cleansUpThreadLocalsAfterException() { @Test void cancellableTaskAvailableDuringExecution() { - when(settings.getSettingValue(Settings.Key.SQL_SLOW_WORKER_POOL_ENABLED)) + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); CancellableTask mockTask = mock(CancellableTask.class); OpenSearchQueryManager.setCancellableTask(mockTask); 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 12071047ff7..2f40d8100b3 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -7,7 +7,7 @@ import static org.opensearch.sql.datasource.model.DataSourceMetadata.defaultOpenSearchDataSourceMetadata; import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_BACKGROUND_THREAD_POOL_NAME; -import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_SLOW_WORKER_THREAD_POOL_NAME; +import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_COMPLEX_WORKER_THREAD_POOL_NAME; import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME; import static org.opensearch.sql.spark.data.constants.SparkConstants.SPARK_REQUEST_BUFFER_INDEX_NAME; @@ -457,7 +457,7 @@ public ScheduledJobParser getJobParser() { @Override public List> getExecutorBuilders(Settings settings) { - // The worker pool is the primary pool where most of the work is done. The slow-worker pool + // The worker pool is the primary pool where most of the work is done. The complex-worker pool // handles queries that require scripts (table scans that can't be pushed to Lucene) so they // don't starve fast queries. The background thread pool is a separate queue for asynchronous // requests to other nodes. We keep them separate to prevent deadlocks during async fetches on @@ -472,10 +472,10 @@ public List> getExecutorBuilders(Settings settings) { "thread_pool." + SQL_WORKER_THREAD_POOL_NAME), new FixedExecutorBuilder( settings, - SQL_SLOW_WORKER_THREAD_POOL_NAME, + SQL_COMPLEX_WORKER_THREAD_POOL_NAME, OpenSearchExecutors.allocatedProcessors(settings), 1000, - "thread_pool." + SQL_SLOW_WORKER_THREAD_POOL_NAME), + "thread_pool." + SQL_COMPLEX_WORKER_THREAD_POOL_NAME), new FixedExecutorBuilder( settings, SQL_BACKGROUND_THREAD_POOL_NAME, diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java index b6347bdf8e1..fa0f3fe9b76 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java @@ -23,6 +23,7 @@ import org.opensearch.rest.RestChannel; import org.opensearch.rest.RestRequest; import org.opensearch.rest.action.RestCancellableNodeClient; +import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.common.antlr.SyntaxCheckException; import org.opensearch.sql.common.error.ErrorReport; import org.opensearch.sql.datasources.exceptions.DataSourceClientException; @@ -144,7 +145,12 @@ public void onFailure(Exception e) { private void sendResponse( RestChannel channel, RestStatus status, String contentType, String content) { - channel.sendResponse(new BytesRestResponse(status, contentType, content)); + BytesRestResponse response = new BytesRestResponse(status, contentType, content); + String poolName = CalcitePlanContext.executionPool.get(); + if (poolName != null) { + response.addHeader("X-Sql-Thread-Pool", poolName); + } + channel.sendResponse(response); } private void reportError(final RestChannel channel, final Exception e, final RestStatus status) { From 9c57cd74edff26d886bc8feadf100a946da720a6 Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Fri, 17 Jul 2026 21:36:37 +0000 Subject: [PATCH 09/11] add a failure log for slow pool requests Signed-off-by: Simeon Widdis --- .../sql/opensearch/executor/ThreadPoolExecutionDispatcher.java | 1 + 1 file changed, 1 insertion(+) 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 793a24b2c4b..62a74b90af1 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 @@ -90,6 +90,7 @@ private void dispatchInternal( try { task.run(); } catch (Exception e) { + LOG.error("Exception during task execution on complex pool", e); if (failureListener != null) { failureListener.onFailure(e); } From 7da880759cb07d2cac43b28977f3ac35a64c0e10 Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Fri, 17 Jul 2026 23:03:00 +0000 Subject: [PATCH 10/11] fix profile, add thread pool as part of profile object Signed-off-by: Simeon Widdis --- .../profile/DefaultProfileContext.java | 4 ++- .../sql/monitor/profile/QueryProfile.java | 19 ++++++++++- .../sql/monitor/profile/QueryProfiling.java | 10 ++++++ .../ThreadPoolExecutionDispatcher.java | 33 +++++++++++-------- plugin/build.gradle | 19 +++++++++++ .../sql/plugin/rest/RestPPLQueryAction.java | 8 +---- 6 files changed, 70 insertions(+), 23 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/monitor/profile/DefaultProfileContext.java b/core/src/main/java/org/opensearch/sql/monitor/profile/DefaultProfileContext.java index 63327c2d6dd..f3df41356f2 100644 --- a/core/src/main/java/org/opensearch/sql/monitor/profile/DefaultProfileContext.java +++ b/core/src/main/java/org/opensearch/sql/monitor/profile/DefaultProfileContext.java @@ -9,6 +9,7 @@ import java.util.Map; import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; +import org.opensearch.sql.calcite.CalcitePlanContext; /** Default implementation that records profiling metrics. */ public class DefaultProfileContext implements ProfileContext { @@ -63,7 +64,8 @@ public synchronized QueryProfile finish() { double totalMillis = ProfileUtils.roundToMillis(endNanos - startNanos); Object planSnapshot = enginePlan != null ? enginePlan : (planRoot == null ? null : planRoot.snapshot()); - profile = new QueryProfile(totalMillis, snapshot, planSnapshot); + String threadPool = CalcitePlanContext.executionPool.get(); + profile = new QueryProfile(totalMillis, snapshot, planSnapshot, threadPool); return profile; } } diff --git a/core/src/main/java/org/opensearch/sql/monitor/profile/QueryProfile.java b/core/src/main/java/org/opensearch/sql/monitor/profile/QueryProfile.java index d9d2a785868..71454951557 100644 --- a/core/src/main/java/org/opensearch/sql/monitor/profile/QueryProfile.java +++ b/core/src/main/java/org/opensearch/sql/monitor/profile/QueryProfile.java @@ -24,6 +24,9 @@ public final class QueryProfile { /** Execution-engine-specific plan profile: a {@link PlanNode} tree, or a pre-rendered object. */ private final Object plan; + @SerializedName("thread_pool") + private final String threadPool; + /** * Create a new query profile snapshot. * @@ -31,7 +34,7 @@ public final class QueryProfile { * @param phases metric values keyed by {@link MetricName} */ public QueryProfile(double totalTimeMillis, Map phases) { - this(totalTimeMillis, phases, null); + this(totalTimeMillis, phases, null, null); } /** @@ -42,9 +45,23 @@ public QueryProfile(double totalTimeMillis, Map phases) { * @param plan plan tree profiling output */ public QueryProfile(double totalTimeMillis, Map phases, Object plan) { + this(totalTimeMillis, phases, plan, null); + } + + /** + * Create a new query profile snapshot. + * + * @param totalTimeMillis total elapsed milliseconds for the query (rounded to two decimals) + * @param phases metric values keyed by {@link MetricName} + * @param plan plan tree profiling output + * @param threadPool thread pool name that executed the query + */ + public QueryProfile( + double totalTimeMillis, Map phases, Object plan, String threadPool) { this.summary = new Summary(totalTimeMillis); this.phases = buildPhases(phases); this.plan = plan; + this.threadPool = threadPool; } private Map buildPhases(Map phases) { diff --git a/core/src/main/java/org/opensearch/sql/monitor/profile/QueryProfiling.java b/core/src/main/java/org/opensearch/sql/monitor/profile/QueryProfiling.java index 3ef32dac748..900be175050 100644 --- a/core/src/main/java/org/opensearch/sql/monitor/profile/QueryProfiling.java +++ b/core/src/main/java/org/opensearch/sql/monitor/profile/QueryProfiling.java @@ -57,6 +57,16 @@ public static void clear() { CURRENT.remove(); } + /** + * Set the profiling context for the current thread. Used when propagating context across thread + * boundaries. + * + * @param ctx profiling context to bind + */ + public static void set(ProfileContext ctx) { + CURRENT.set(Objects.requireNonNull(ctx, "ctx")); + } + /** * Run a supplier with the provided profiling context bound to the current thread. * 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 62a74b90af1..7f7fbd14dd1 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 @@ -25,6 +25,8 @@ import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.executor.ExecutionDispatcher; import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.monitor.profile.ProfileContext; +import org.opensearch.sql.monitor.profile.QueryProfiling; import org.opensearch.tasks.CancellableTask; import org.opensearch.threadpool.ThreadPool; @@ -64,6 +66,7 @@ private void dispatchInternal( LOG.debug("Query plan contains scripts, dispatching to complex worker pool"); Map ctx = ThreadContext.getImmutableContext(); CancellableTask cancellableTask = OpenSearchQueryManager.getCancellableTask(); + ProfileContext profileContext = QueryProfiling.current(); @Nullable JaninoRelMetadataProvider metadataProvider = RelMetadataQueryBase.THREAD_PROVIDERS.get(); long currentTime = Hook.CURRENT_TIME.get(-1L); @@ -72,22 +75,23 @@ private void dispatchInternal( String twSeries = CalcitePlanContext.timewrapSeries.get(); threadPool.schedule( () -> { - ThreadContext.putAll(ctx); - OpenSearchQueryManager.setCancellableTask(cancellableTask); - CalcitePlanContext.executionPool.set(SQL_COMPLEX_WORKER_THREAD_POOL_NAME); - if (metadataProvider != null) { - RelMetadataQueryBase.THREAD_PROVIDERS.set(metadataProvider); - } Hook.Closeable hookHandle = null; - if (currentTime >= 0) { - hookHandle = - Hook.CURRENT_TIME.addThread( - (Consumer>) h -> h.set(currentTime)); - } - CalcitePlanContext.stripNullColumns.set(stripNullCols); - CalcitePlanContext.timewrapUnitName.set(twUnitName); - CalcitePlanContext.timewrapSeries.set(twSeries); try { + ThreadContext.putAll(ctx); + OpenSearchQueryManager.setCancellableTask(cancellableTask); + QueryProfiling.set(profileContext); + CalcitePlanContext.executionPool.set(SQL_COMPLEX_WORKER_THREAD_POOL_NAME); + if (metadataProvider != null) { + RelMetadataQueryBase.THREAD_PROVIDERS.set(metadataProvider); + } + if (currentTime >= 0) { + hookHandle = + Hook.CURRENT_TIME.addThread( + (Consumer>) h -> h.set(currentTime)); + } + CalcitePlanContext.stripNullColumns.set(stripNullCols); + CalcitePlanContext.timewrapUnitName.set(twUnitName); + CalcitePlanContext.timewrapSeries.set(twSeries); task.run(); } catch (Exception e) { LOG.error("Exception during task execution on complex pool", e); @@ -101,6 +105,7 @@ private void dispatchInternal( OpenSearchQueryManager.clearCancellableTask(); RelMetadataQueryBase.THREAD_PROVIDERS.remove(); CalcitePlanContext.clearTimewrapSignals(); + QueryProfiling.clear(); } }, new TimeValue(0), diff --git a/plugin/build.gradle b/plugin/build.gradle index fbd2cd1a331..05278e26b35 100644 --- a/plugin/build.gradle +++ b/plugin/build.gradle @@ -305,8 +305,27 @@ def getJobSchedulerPlugin() { }) } +def getWorkloadManagementPlugin() { + provider(new Callable() { + @Override + RegularFile call() throws Exception { + return new RegularFile() { + @Override + File getAsFile() { + // ponytail: assumes ~/workplace/opensearch is built + def wlmPath = System.getProperty('user.home') + '/workplace/opensearch/plugins/workload-management/build/distributions' + return fileTree(wlmPath).matching { + include 'workload-management*.zip' + }.singleFile + } + } + } + }) +} + testClusters.integTest { plugin(getJobSchedulerPlugin()) + plugin(getWorkloadManagementPlugin()) plugin(project.tasks.bundlePlugin.archiveFile) testDistribution = "ARCHIVE" diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java index fa0f3fe9b76..b6347bdf8e1 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java @@ -23,7 +23,6 @@ import org.opensearch.rest.RestChannel; import org.opensearch.rest.RestRequest; import org.opensearch.rest.action.RestCancellableNodeClient; -import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.common.antlr.SyntaxCheckException; import org.opensearch.sql.common.error.ErrorReport; import org.opensearch.sql.datasources.exceptions.DataSourceClientException; @@ -145,12 +144,7 @@ public void onFailure(Exception e) { private void sendResponse( RestChannel channel, RestStatus status, String contentType, String content) { - BytesRestResponse response = new BytesRestResponse(status, contentType, content); - String poolName = CalcitePlanContext.executionPool.get(); - if (poolName != null) { - response.addHeader("X-Sql-Thread-Pool", poolName); - } - channel.sendResponse(response); + channel.sendResponse(new BytesRestResponse(status, contentType, content)); } private void reportError(final RestChannel channel, final Exception e, final RestStatus status) { From 559a2d0c1da2415ab8a9b98acb4d2c14832cf383 Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Fri, 17 Jul 2026 23:13:16 +0000 Subject: [PATCH 11/11] remove leftover build.gradle changes from another branch Signed-off-by: Simeon Widdis --- plugin/build.gradle | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/plugin/build.gradle b/plugin/build.gradle index 05278e26b35..fbd2cd1a331 100644 --- a/plugin/build.gradle +++ b/plugin/build.gradle @@ -305,27 +305,8 @@ def getJobSchedulerPlugin() { }) } -def getWorkloadManagementPlugin() { - provider(new Callable() { - @Override - RegularFile call() throws Exception { - return new RegularFile() { - @Override - File getAsFile() { - // ponytail: assumes ~/workplace/opensearch is built - def wlmPath = System.getProperty('user.home') + '/workplace/opensearch/plugins/workload-management/build/distributions' - return fileTree(wlmPath).matching { - include 'workload-management*.zip' - }.singleFile - } - } - } - }) -} - testClusters.integTest { plugin(getJobSchedulerPlugin()) - plugin(getWorkloadManagementPlugin()) plugin(project.tasks.bundlePlugin.archiveFile) testDistribution = "ARCHIVE"