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 5473aa8812e..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 @@ -81,7 +81,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_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/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/DirectExecutionDispatcher.java b/core/src/main/java/org/opensearch/sql/executor/DirectExecutionDispatcher.java new file mode 100644 index 00000000000..5af110fa7e5 --- /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 complex-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..ec7fa193852 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/executor/ExecutionDispatcher.java @@ -0,0 +1,45 @@ +/* + * 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 complex worker pool (for + * queries requiring scripts/table scans). + */ +public interface ExecutionDispatcher { + + /** + * Dispatch execution of the given plan via the standard ExecutionEngine. + * + * @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); + + /** + * 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 b97a679cbd3..947f999f9eb 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; @@ -42,6 +41,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; @@ -64,7 +64,6 @@ /** The low level interface of core engine. */ @RequiredArgsConstructor -@AllArgsConstructor @Log4j2 public class QueryService { private final Analyzer analyzer; @@ -72,6 +71,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); @@ -207,9 +237,15 @@ private void executeCalcitePlan( CalcitePlanContext context, ResponseListener listener) { try { + // 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 complex worker pool StageErrorHandler.executeStageVoid( QueryProcessingStage.EXECUTING, - () -> executionEngine.execute(calcitePlan, context, listener), + () -> executionDispatcher.dispatch(optimizedPlan, context, listener, executionEngine), "while running the query"); } catch (RuntimeException e) { ArithmeticException overflow = findArithmeticOverflow(e); 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/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/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..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 @@ -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 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 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-complex-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 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)); + } + } + + /** + * 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 complex 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-complex-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 complex pool", hasName); + assertTrue("manager_user should see computed 'avg_salary' field on complex pool", hasAvgSalary); + assertFalse( + "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/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/OpenSearchQueryManager.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java index 7aaaaa6655e..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,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_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/ScriptDetector.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java new file mode 100644 index 00000000000..aafb2f307b3 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ScriptDetector.java @@ -0,0 +1,116 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +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; +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.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 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 static final Set EXPENSIVE_UDFS = + Set.of("REX_EXTRACT", "REX_EXTRACT_MULTI", "PARSE", "PATTERN_PARSER"); + + private ScriptDetector() {} + + /** + * 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]) { + 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); + } + } + }.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 visitOver(RexOver over) { + found[0] = true; + return null; + } + + @Override + public Void visitCall(RexCall call) { + String name = call.getOperator().getName(); + if (name != null && EXPENSIVE_UDFS.contains(name)) { + found[0] = true; + return null; + } + return super.visitCall(call); + } + }); + 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..7f7fbd14dd1 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java @@ -0,0 +1,122 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.executor; + +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; +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; +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; + +/** + * 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 { + + 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) { + 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 (isComplexPoolEnabled() && ScriptDetector.hasScripts(optimizedPlan)) { + 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); + boolean stripNullCols = CalcitePlanContext.stripNullColumns.get(); + String twUnitName = CalcitePlanContext.timewrapUnitName.get(); + String twSeries = CalcitePlanContext.timewrapSeries.get(); + threadPool.schedule( + () -> { + Hook.Closeable hookHandle = null; + 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); + if (failureListener != null) { + failureListener.onFailure(e); + } + } finally { + if (hookHandle != null) { + hookHandle.close(); + } + OpenSearchQueryManager.clearCancellableTask(); + RelMetadataQueryBase.THREAD_PROVIDERS.remove(); + CalcitePlanContext.clearTimewrapSignals(); + QueryProfiling.clear(); + } + }, + new TimeValue(0), + SQL_COMPLEX_WORKER_THREAD_POOL_NAME); + } else { + CalcitePlanContext.executionPool.set(SQL_WORKER_THREAD_POOL_NAME); + task.run(); + } + } + + 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 ce2bdd4960b..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,6 +363,13 @@ public class OpenSearchSettings extends Settings { Setting.Property.NodeScope, Setting.Property.Dynamic); + public static final Setting SQL_COMPLEX_WORKER_POOL_ENABLED_SETTING = + Setting.boolSetting( + Key.SQL_COMPLEX_WORKER_POOL_ENABLED.getKeyValue(), + true, + Setting.Property.NodeScope, + Setting.Property.Dynamic); + /** Construct OpenSearchSetting. The OpenSearchSetting must be singleton. */ @SuppressWarnings("unchecked") public OpenSearchSettings(ClusterSettings clusterSettings) { @@ -631,6 +638,12 @@ public OpenSearchSettings(ClusterSettings clusterSettings) { Key.FIELD_TYPE_TOLERANCE, FIELD_TYPE_TOLERANCE_SETTING, new Updater(Key.FIELD_TYPE_TOLERANCE)); + register( + settingBuilder, + clusterSettings, + Key.SQL_COMPLEX_WORKER_POOL_ENABLED, + SQL_COMPLEX_WORKER_POOL_ENABLED_SETTING, + new Updater(Key.SQL_COMPLEX_WORKER_POOL_ENABLED)); defaultSettings = settingBuilder.build(); } @@ -726,6 +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_COMPLEX_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..e44809e40a9 --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ScriptDetectorTest.java @@ -0,0 +1,212 @@ +/* + * 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.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.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) +class ScriptDetectorTest { + + @Test + void returnsFalseForNonScanNode() { + RelNode mockNode = createMockNode(); + assertFalse(ScriptDetector.hasScripts(mockNode)); + } + + @Test + void returnsTrueWhenFilterScriptPushed() { + AbstractCalciteIndexScan scan = createMockScan(true, false, 0); + assertTrue(ScriptDetector.hasScripts(scan)); + } + + @Test + 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(false, false, 3); + RelNode parent = createMockNode(scan); + assertTrue(ScriptDetector.hasScripts(parent)); + } + + @Test + 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 detectsExpensiveUdfInProject() { + SqlOperator rexExtractOp = mock(SqlOperator.class); + when(rexExtractOp.getName()).thenReturn("REX_EXTRACT"); + RexCall udfCall = mock(RexCall.class); + when(udfCall.getOperator()).thenReturn(rexExtractOp); + 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 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); + 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).visitOver(rexOver)) + .when(rexOver) + .accept(any(org.apache.calcite.rex.RexVisitor.class)); + + 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) { + 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, boolean sortExprPushed, long aggScriptCount) { + AbstractCalciteIndexScan scan = mock(AbstractCalciteIndexScan.class); + + 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 -> 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..abf806ec810 --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java @@ -0,0 +1,375 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +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; +import static org.mockito.Mockito.when; +import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_COMPLEX_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; +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.sql.opensearch.storage.scan.context.PushDownContext; +import org.opensearch.tasks.CancellableTask; +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); + } + + @AfterEach + void tearDown() { + ThreadContext.clearAll(); + OpenSearchQueryManager.clearCancellableTask(); + RelMetadataQueryBase.THREAD_PROVIDERS.remove(); + CalcitePlanContext.clearTimewrapSignals(); + } + + @Test + void executesInlineWhenNoScripts() { + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_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_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_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_COMPLEX_WORKER_POOL_ENABLED)) + .thenReturn(false); + AbstractCalciteIndexScan scan = createMockScanWithScripts(); + + 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_COMPLEX_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 = createMockScanWithScripts(); + dispatcher.dispatch(scan, context, listener, engine); + + verify(engine).execute(scan, context, listener); + } + + @Test + void propagatesCancellableTaskToSlowPool() { + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_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_COMPLEX_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_COMPLEX_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_COMPLEX_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_COMPLEX_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_COMPLEX_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_COMPLEX_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); + 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 createMockScanWithScripts() { + AbstractCalciteIndexScan scan = mock(AbstractCalciteIndexScan.class); + 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; + } +} 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 d9437fded5b..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,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_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; @@ -237,7 +238,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]; }; @@ -450,10 +457,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 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 + // small node counts. Tasks in the background pool should do no work except I/O to other + // services. return List.of( new FixedExecutorBuilder( settings, @@ -461,11 +470,17 @@ public List> getExecutorBuilders(Settings settings) { OpenSearchExecutors.allocatedProcessors(settings), 1000, "thread_pool." + SQL_WORKER_THREAD_POOL_NAME), + new FixedExecutorBuilder( + settings, + SQL_COMPLEX_WORKER_THREAD_POOL_NAME, + OpenSearchExecutors.allocatedProcessors(settings), + 1000, + "thread_pool." + SQL_COMPLEX_WORKER_THREAD_POOL_NAME), new FixedExecutorBuilder( 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)); } 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 5685d539541..609b4a71099 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 @@ -69,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; } /** @@ -230,19 +233,25 @@ 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)); - } else { - analyticsEngine.execute( - plan, - planContext, - queryCtx, - createQueryListener(queryType, profileCtx, closingListener)); - } + plan = + org.opensearch.sql.calcite.utils.CalciteToolsHelper.optimize( + plan, planContext); + RelNode finalPlan = plan; + Runnable executeTask = + profiling + ? () -> + analyticsEngine.executeWithProfile( + finalPlan, + planContext, + queryCtx, + createQueryListener(queryType, profileCtx, closingListener)) + : () -> + analyticsEngine.execute( + finalPlan, + planContext, + queryCtx, + createQueryListener(queryType, profileCtx, closingListener)); + executionDispatcher.dispatchTask(finalPlan, planContext, executeTask); } catch (Exception e) { closingListener.onFailure(e); } finally { 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 0cf87f0604e..516c31940c3 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