From e8aa8777f25ba50893cc470e9ee841a90468de62 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Fri, 14 Nov 2025 19:12:07 -0800 Subject: [PATCH] Add timeout overflow break handling at leaf stage --- .../MultiStageBrokerRequestHandler.java | 29 +++ .../broker/BrokerResponseNativeV2.java | 21 +- .../utils/config/QueryOptionsUtils.java | 19 ++ .../broker/BrokerResponseNativeV2Test.java | 20 ++ .../pinot/query/runtime/QueryRunner.java | 65 +++++ .../runtime/operator/BaseJoinOperator.java | 50 +++- .../query/runtime/operator/LeafOperator.java | 84 ++++++- .../runtime/operator/MultiStageOperator.java | 22 +- .../runtime/operator/LeafOperatorTest.java | 97 +++++++ .../TimeoutOverflowQueryRunnerTest.java | 237 ++++++++++++++++++ .../MockInstanceDataManagerFactory.java | 21 +- .../SlowMockInstanceDataManagerFactory.java | 55 ++++ .../pinot/spi/utils/CommonConstants.java | 7 + 13 files changed, 705 insertions(+), 22 deletions(-) create mode 100644 pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/TimeoutOverflowQueryRunnerTest.java create mode 100644 pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/SlowMockInstanceDataManagerFactory.java diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java index 91e2d2c2c9a0..1725b176fca2 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java @@ -991,6 +991,10 @@ private void fillOldBrokerResponseStats(BrokerResponseNativeV2 brokerResponse, (int) numSegmentsPrunedByBroker); brokerResponse.addBrokerStats(brokerPruningStats); } + + if (isLeafStageMissingStats(queryStageMap, queryStats)) { + brokerResponse.mergeTimeoutOverflowReached(true); + } } catch (Exception e) { LOGGER.warn("Error encountered while collecting multi-stage stats", e); brokerResponse.setStageStats(JsonNodeFactory.instance.objectNode() @@ -1012,6 +1016,31 @@ private void fillOldBrokerResponseStats(BrokerResponseNativeV2 brokerResponse, } } + private boolean isLeafStageMissingStats(Map queryStageMap, + List queryStats) { + for (Map.Entry entry : queryStageMap.entrySet()) { + int stageId = entry.getKey(); + DispatchablePlanFragment fragment = entry.getValue(); + if (!isLeafStage(fragment)) { + continue; + } + MultiStageQueryStats.StageStats.Closed stageStats = + stageId < queryStats.size() ? queryStats.get(stageId) : null; + if (stageStats == null) { + return true; + } + } + return false; + } + + private boolean isLeafStage(DispatchablePlanFragment fragment) { + Map customProperties = fragment.getCustomProperties(); + if (customProperties.containsKey(DispatchablePlanFragment.TABLE_NAME_KEY)) { + return true; + } + return !fragment.getWorkerIdToSegmentsMap().isEmpty(); + } + private BrokerResponse constructMultistageExplainPlan(String sql, String plan, Map extraFields) { BrokerResponseNative brokerResponse = BrokerResponseNative.empty(); int totalFieldCount = extraFields.size() + 2; diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java index 3012c67831d0..05448d8fb203 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java @@ -88,6 +88,7 @@ public class BrokerResponseNativeV2 implements BrokerResponse { private String _requestId; private String _clientRequestId; private String _brokerId; + private boolean _timeoutOverflowReached; private int _numServersQueried; private int _numServersResponded; private long _brokerReduceTimeMs; @@ -126,7 +127,7 @@ public int getNumRowsResultSet() { @Override public boolean isPartialResult() { return getExceptionsSize() > 0 || isNumGroupsLimitReached() || !getEarlyTerminationReasons().isEmpty() - || isMaxRowsInJoinReached() || isMseLiteLeafStageLimitReached(); + || isMaxRowsInJoinReached() || isMseLiteLeafStageLimitReached() || isTimeoutOverflowReached(); } @Override @@ -211,6 +212,18 @@ public List getEarlyTerminationReasons() { return List.copyOf(_brokerStats.getStringSet(StatKey.EARLY_TERMINATION_REASONS)); } + @JsonProperty("timeoutOverflowReached") + public boolean isTimeoutOverflowReached() { + return _timeoutOverflowReached; + } + + public void mergeTimeoutOverflowReached(boolean timeoutOverflowReached) { + if (timeoutOverflowReached) { + _timeoutOverflowReached = true; + } + _brokerStats.merge(StatKey.TIMEOUT_OVERFLOW_REACHED, timeoutOverflowReached); + } + @Override public boolean isMaxRowsInJoinReached() { return _maxRowsInJoinReached; @@ -515,6 +528,9 @@ public boolean getRLSFiltersApplied() { public void addBrokerStats(StatMap brokerStats) { _brokerStats.merge(brokerStats); + if (brokerStats.getBoolean(StatKey.TIMEOUT_OVERFLOW_REACHED)) { + _timeoutOverflowReached = true; + } } // NOTE: The following keys should match the keys in the leaf-stage operator. @@ -550,7 +566,8 @@ public long merge(long value1, long value2) { } }, EARLY_TERMINATION_REASONS(StatMap.Type.STRING_SET), - LITE_MODE_LEAF_STAGE_LIMIT_REACHED(StatMap.Type.BOOLEAN); + LITE_MODE_LEAF_STAGE_LIMIT_REACHED(StatMap.Type.BOOLEAN), + TIMEOUT_OVERFLOW_REACHED(StatMap.Type.BOOLEAN); private final StatMap.Type _type; diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java index 4a39ab5d8e70..b76ffcb76f0c 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java @@ -33,6 +33,7 @@ import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.JoinOverFlowMode; +import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.TimeoutOverflowMode; import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.WindowOverFlowMode; @@ -511,6 +512,24 @@ public static WindowOverFlowMode getWindowOverflowMode(Map query return windowOverflowModeStr != null ? WindowOverFlowMode.valueOf(windowOverflowModeStr) : null; } + @Nullable + public static TimeoutOverflowMode getTimeoutOverflowMode(Map queryOptions) { + String timeoutOverflowModeStr = queryOptions.get(QueryOptionKey.TIMEOUT_OVERFLOW_MODE); + return timeoutOverflowModeStr != null ? TimeoutOverflowMode.valueOf(timeoutOverflowModeStr) : null; + } + + @Nullable + public static Long getLeafTimeoutMs(Map queryOptions) { + String leafTimeoutMs = queryOptions.get(QueryOptionKey.LEAF_TIMEOUT_MS); + return checkedParseLongPositive(QueryOptionKey.LEAF_TIMEOUT_MS, leafTimeoutMs); + } + + @Nullable + public static Long getLeafExtraPassiveTimeoutMs(Map queryOptions) { + String leafExtraPassiveTimeoutMs = queryOptions.get(QueryOptionKey.LEAF_EXTRA_PASSIVE_TIMEOUT_MS); + return checkedParseLong(QueryOptionKey.LEAF_EXTRA_PASSIVE_TIMEOUT_MS, leafExtraPassiveTimeoutMs, 0); + } + public static boolean isSkipUnavailableServers(Map queryOptions) { return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.SKIP_UNAVAILABLE_SERVERS)); } diff --git a/pinot-common/src/test/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2Test.java b/pinot-common/src/test/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2Test.java index 320db22557a0..ba3176f766ef 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2Test.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2Test.java @@ -18,10 +18,13 @@ */ package org.apache.pinot.common.response.broker; +import com.fasterxml.jackson.databind.JsonNode; +import java.io.IOException; import java.util.LinkedHashSet; import java.util.List; import java.util.Set; import org.apache.pinot.common.datatable.StatMap; +import org.apache.pinot.spi.utils.JsonUtils; import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; @@ -47,6 +50,23 @@ public void testEarlyTerminationReasonsMarkPartialResponse() { assertTrue(brokerResponse.isPartialResult()); } + @Test + public void testTimeoutOverflowFlagSerialized() + throws IOException { + BrokerResponseNativeV2 response = new BrokerResponseNativeV2(); + assertFalse(response.isTimeoutOverflowReached()); + assertFalse(response.isPartialResult()); + + response.mergeTimeoutOverflowReached(true); + + assertTrue(response.isTimeoutOverflowReached()); + assertTrue(response.isPartialResult()); + + JsonNode json = JsonUtils.objectToJsonNode(response); + assertTrue(json.get("timeoutOverflowReached").asBoolean()); + assertTrue(json.get("partialResult").asBoolean()); + } + private static Set stringSet(String... values) { return new LinkedHashSet<>(List.of(values)); } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java index e5575d7680f6..e5b134b18608 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java @@ -77,6 +77,7 @@ import org.apache.pinot.spi.utils.CommonConstants.Helix; import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner; import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.JoinOverFlowMode; +import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.TimeoutOverflowMode; import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.WindowOverFlowMode; import org.apache.pinot.spi.utils.CommonConstants.Query.Request; import org.apache.pinot.spi.utils.CommonConstants.Query.Response; @@ -98,6 +99,7 @@ */ public class QueryRunner { private static final Logger LOGGER = LoggerFactory.getLogger(QueryRunner.class); + private static final double DEFAULT_LEAF_TIMEOUT_RATIO = 0.8; private ExecutorService _executorService; private OpChainSchedulerService _opChainScheduler; @@ -130,6 +132,8 @@ public class QueryRunner { @Nullable private WindowOverFlowMode _windowOverflowMode; @Nullable + private TimeoutOverflowMode _timeoutOverflowMode; + @Nullable private PhysicalTimeSeriesServerPlanVisitor _timeSeriesPhysicalPlanVisitor; /// Cluster-level decision on whether to send stats over the mailbox path, driven by the {@code SendStatsPredicate} /// at startup time. May be overridden per-request via the {@code KEY_OF_STATS_REPORTING_MODE} metadata key — @@ -191,6 +195,10 @@ public void init(PinotConfiguration serverConf, String instanceId, @Nullable Ins String windowOverflowModeStr = serverConf.getProperty(MultiStageQueryRunner.KEY_OF_WINDOW_OVERFLOW_MODE); _windowOverflowMode = windowOverflowModeStr != null ? WindowOverFlowMode.valueOf(windowOverflowModeStr) : null; + String timeoutOverflowModeStr = serverConf.getProperty(MultiStageQueryRunner.KEY_OF_TIMEOUT_OVERFLOW_MODE); + _timeoutOverflowMode = + timeoutOverflowModeStr != null ? TimeoutOverflowMode.valueOf(timeoutOverflowModeStr) : null; + ExecutorService baseExecutorService = ExecutorServiceUtils.create(serverConf, Server.MULTISTAGE_EXECUTOR_CONFIG_PREFIX, "query-runner-on-" + port, Server.DEFAULT_MULTISTAGE_EXECUTOR_TYPE); @@ -301,6 +309,7 @@ private void processQueryBlocking(WorkerMetadata workerMetadata, StagePlan stage Map requestMetadata) { StageMetadata stageMetadata = stagePlan.getStageMetadata(); Map opChainMetadata = consolidateMetadata(stageMetadata.getCustomProperties(), requestMetadata); + applyLeafDeadlineOverrides(workerMetadata, opChainMetadata); // The cluster-level _sendStats decision can be overridden per-request by the SubmitWithStream RPC handler via // MultiStageQueryRunner.KEY_OF_STATS_REPORTING_MODE; in stream mode stats travel out-of-band @@ -557,6 +566,14 @@ private Map consolidateMetadata(Map customProper opChainMetadata.put(QueryOptionKey.WINDOW_OVERFLOW_MODE, windowOverflowMode.name()); } + TimeoutOverflowMode timeoutOverflowMode = QueryOptionsUtils.getTimeoutOverflowMode(opChainMetadata); + if (timeoutOverflowMode == null) { + timeoutOverflowMode = _timeoutOverflowMode; + } + if (timeoutOverflowMode != null) { + opChainMetadata.put(QueryOptionKey.TIMEOUT_OVERFLOW_MODE, timeoutOverflowMode.name()); + } + return opChainMetadata; } @@ -604,6 +621,53 @@ public void unregisterOpChainCompletionListener(long requestId) { _opChainScheduler.unregisterCompletionListener(requestId); } + private void applyLeafDeadlineOverrides(WorkerMetadata workerMetadata, Map opChainMetadata) { + if (!workerMetadata.isLeafStageWorker()) { + return; + } + Long leafTimeoutMs = QueryOptionsUtils.getLeafTimeoutMs(opChainMetadata); + Long leafExtraPassiveTimeoutMs = QueryOptionsUtils.getLeafExtraPassiveTimeoutMs(opChainMetadata); + TimeoutOverflowMode timeoutOverflowMode = _timeoutOverflowMode; + String timeoutOverflowModeStr = opChainMetadata.get(QueryOptionKey.TIMEOUT_OVERFLOW_MODE); + if (timeoutOverflowModeStr != null) { + try { + timeoutOverflowMode = TimeoutOverflowMode.valueOf(timeoutOverflowModeStr); + } catch (IllegalArgumentException e) { + LOGGER.warn("Invalid timeout overflow mode: {}", timeoutOverflowModeStr, e); + } + } + QueryThreadContext threadContext = QueryThreadContext.getIfAvailable(); + if (threadContext == null) { + return; + } + QueryExecutionContext executionContext = threadContext.getExecutionContext(); + long startTimeMs = executionContext.getStartTimeMs(); + if (leafTimeoutMs == null && timeoutOverflowMode == TimeoutOverflowMode.BREAK) { + long globalActiveWindowMs = Math.max(1L, executionContext.getActiveDeadlineMs() - startTimeMs); + leafTimeoutMs = Math.max(1L, (long) Math.floor(globalActiveWindowMs * DEFAULT_LEAF_TIMEOUT_RATIO)); + } + if (leafExtraPassiveTimeoutMs == null && timeoutOverflowMode == TimeoutOverflowMode.BREAK) { + long globalPassiveWindowMs = + Math.max(0L, executionContext.getPassiveDeadlineMs() - executionContext.getActiveDeadlineMs()); + leafExtraPassiveTimeoutMs = + Math.max(0L, (long) Math.floor(globalPassiveWindowMs * DEFAULT_LEAF_TIMEOUT_RATIO)); + } + if (leafTimeoutMs == null && leafExtraPassiveTimeoutMs == null) { + return; + } + long activeDeadlineMs = executionContext.getActiveDeadlineMs(); + long passiveDeadlineMs = executionContext.getPassiveDeadlineMs(); + if (leafTimeoutMs != null) { + activeDeadlineMs = Math.min(activeDeadlineMs, startTimeMs + leafTimeoutMs); + } + if (leafExtraPassiveTimeoutMs != null) { + passiveDeadlineMs = Math.min(passiveDeadlineMs, activeDeadlineMs + leafExtraPassiveTimeoutMs); + } + passiveDeadlineMs = Math.max(passiveDeadlineMs, activeDeadlineMs); + opChainMetadata.put(LeafOperator.LEAF_ACTIVE_DEADLINE_METADATA_KEY, Long.toString(activeDeadlineMs)); + opChainMetadata.put(LeafOperator.LEAF_PASSIVE_DEADLINE_METADATA_KEY, Long.toString(passiveDeadlineMs)); + } + public StagePlan explainQuery(WorkerMetadata workerMetadata, StagePlan stagePlan, Map requestMetadata) { if (!workerMetadata.isLeafStageWorker()) { @@ -613,6 +677,7 @@ public StagePlan explainQuery(WorkerMetadata workerMetadata, StagePlan stagePlan StageMetadata stageMetadata = stagePlan.getStageMetadata(); Map opChainMetadata = consolidateMetadata(stageMetadata.getCustomProperties(), requestMetadata); + applyLeafDeadlineOverrides(workerMetadata, opChainMetadata); if (PipelineBreakerExecutor.hasPipelineBreakers(stagePlan)) { //TODO: See https://github.com/apache/pinot/pull/13733#discussion_r1752031714 diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java index 8ec374ca4b5a..64fc3929351f 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java @@ -39,9 +39,11 @@ import org.apache.pinot.query.runtime.operator.operands.TransformOperandFactory; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryException; import org.apache.pinot.spi.utils.BooleanUtils; import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.JoinOverFlowMode; +import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.TimeoutOverflowMode; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -83,10 +85,12 @@ public abstract class BaseJoinOperator extends MultiStageOperator { * BREAK: Break right table build process, continue to perform JOIN operation, results might be partial. */ protected final JoinOverFlowMode _joinOverflowMode; + protected final TimeoutOverflowMode _timeoutOverflowMode; protected boolean _isRightTableBuilt; @Nullable protected MseBlock.Eos _eos; + private boolean _timeoutOverflowTriggered; public BaseJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema, MultiStageOperator rightInput, JoinNode node) { @@ -106,6 +110,8 @@ public BaseJoinOperator(OpChainExecutionContext context, MultiStageOperator left PlanNode.NodeHint nodeHint = node.getNodeHint(); _maxRowsInJoin = getMaxRowsInJoin(metadata, nodeHint); _joinOverflowMode = getJoinOverflowMode(metadata, nodeHint); + TimeoutOverflowMode timeoutOverflowMode = QueryOptionsUtils.getTimeoutOverflowMode(metadata); + _timeoutOverflowMode = timeoutOverflowMode != null ? timeoutOverflowMode : TimeoutOverflowMode.THROW; } /// Constructor that takes the schema for NonEquiEvaluator as an argument @@ -127,6 +133,8 @@ public BaseJoinOperator(OpChainExecutionContext context, MultiStageOperator left PlanNode.NodeHint nodeHint = node.getNodeHint(); _maxRowsInJoin = getMaxRowsInJoin(metadata, nodeHint); _joinOverflowMode = getJoinOverflowMode(metadata, nodeHint); + TimeoutOverflowMode timeoutOverflowMode = QueryOptionsUtils.getTimeoutOverflowMode(metadata); + _timeoutOverflowMode = timeoutOverflowMode != null ? timeoutOverflowMode : TimeoutOverflowMode.THROW; } protected static int getMaxRowsInJoin(Map opChainMetadata, @Nullable PlanNode.NodeHint nodeHint) { @@ -158,6 +166,14 @@ protected static JoinOverFlowMode getJoinOverflowMode(Map contex return joinOverflowMode != null ? joinOverflowMode : DEFAULT_JOIN_OVERFLOW_MODE; } + @Override + protected long getDeadlineMs() { + if (_timeoutOverflowTriggered && _timeoutOverflowMode == TimeoutOverflowMode.BREAK) { + return Long.MAX_VALUE; + } + return super.getDeadlineMs(); + } + @Override public void registerExecution(long time, int numRows, long memoryUsedBytes, long gcTimeMs) { _statMap.merge(StatKey.EXECUTION_TIME_MS, time); @@ -286,6 +302,14 @@ protected MseBlock buildJoinedDataBlock() { protected abstract List buildNonMatchRightRows(); + @Override + protected MseBlock handleException(Exception e) { + if (maybeHandleTimeoutOverflow(e)) { + return SuccessMseBlock.INSTANCE; + } + return super.handleException(e); + } + // TODO: Optimize this to avoid unnecessary object copy. protected Object[] joinRow(@Nullable Object[] leftRow, @Nullable Object[] rightRow) { Object[] resultRow = new Object[_resultColumnSize]; @@ -332,6 +356,29 @@ protected void earlyTerminateLeftInput() { _eos = (MseBlock.Eos) leftBlock; } + private boolean maybeHandleTimeoutOverflow(Exception e) { + if (_timeoutOverflowMode != TimeoutOverflowMode.BREAK) { + return false; + } + if (e instanceof QueryException) { + QueryException queryException = (QueryException) e; + if (queryException.getErrorCode() == QueryErrorCode.EXECUTION_TIMEOUT) { + if (!_timeoutOverflowTriggered) { + logger().warn("Join operator {} returning partial results after timeout: {}", _operatorId, + queryException.getMessage()); + } + _timeoutOverflowTriggered = true; + _statMap.merge(StatKey.TIMEOUT_OVERFLOW_REACHED, true); + _leftInput.earlyTerminate(); + _rightInput.earlyTerminate(); + _eos = SuccessMseBlock.INSTANCE; + onEosProduced(); + return true; + } + } + return false; + } + @Override public StatMap copyStatMaps() { return new StatMap<>(_statMap); @@ -412,7 +459,8 @@ public long merge(long value1, long value2) { /** * Time spent on GC while this operator or its children in the same stage were running. */ - GC_TIME_MS(StatMap.Type.LONG); + GC_TIME_MS(StatMap.Type.LONG), + TIMEOUT_OVERFLOW_REACHED(StatMap.Type.BOOLEAN); private final StatMap.Type _type; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java index 5ec9d9fd2e28..6a89f4f5f843 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java @@ -69,6 +69,7 @@ import org.apache.pinot.spi.exception.TerminationException; import org.apache.pinot.spi.query.QueryThreadContext; import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionValue; +import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.TimeoutOverflowMode; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -84,6 +85,8 @@ public class LeafOperator extends MultiStageOperator { ErrorMseBlock.fromError(QueryErrorCode.QUERY_CANCELLATION, "Cancelled while waiting for leaf results"); private static final ErrorMseBlock TIMEOUT_BLOCK = ErrorMseBlock.fromError(QueryErrorCode.EXECUTION_TIMEOUT, "Timed out waiting for leaf results"); + public static final String LEAF_ACTIVE_DEADLINE_METADATA_KEY = "leafActiveDeadlineMs"; + public static final String LEAF_PASSIVE_DEADLINE_METADATA_KEY = "leafPassiveDeadlineMs"; // Use a special results block to indicate that this is the last results block @VisibleForTesting @@ -103,10 +106,14 @@ public class LeafOperator extends MultiStageOperator { // Use a limit-sized BlockingQueue to store the results blocks and apply back pressure to the single-stage threads @VisibleForTesting final BlockingQueue _blockingQueue; + private final TimeoutOverflowMode _timeoutOverflowMode; @Nullable private volatile Future _executionFuture; private volatile boolean _terminated; + private volatile boolean _timeoutOverflowTriggered; + private final long _customActiveDeadlineMs; + private final long _customPassiveDeadlineMs; public LeafOperator(OpChainExecutionContext context, List requests, DataSchema dataSchema, QueryExecutor queryExecutor, ExecutorService executorService) { @@ -134,6 +141,12 @@ public LeafOperator( _blockingQueue = new ArrayBlockingQueue<>(maxStreamingPendingBlocks != null ? maxStreamingPendingBlocks : QueryOptionValue.DEFAULT_MAX_STREAMING_PENDING_BLOCKS); _pipelineBreakerStats = pipelineBreakerStats; + TimeoutOverflowMode timeoutOverflowMode = + QueryOptionsUtils.getTimeoutOverflowMode(context.getOpChainMetadata()); + _timeoutOverflowMode = timeoutOverflowMode != null ? timeoutOverflowMode : TimeoutOverflowMode.THROW; + Map metadata = context.getOpChainMetadata(); + _customActiveDeadlineMs = parseDeadline(metadata.get(LEAF_ACTIVE_DEADLINE_METADATA_KEY)); + _customPassiveDeadlineMs = parseDeadline(metadata.get(LEAF_PASSIVE_DEADLINE_METADATA_KEY)); } public List getRequests() { @@ -197,15 +210,15 @@ protected MseBlock getNextBlock() { BaseResultsBlock resultsBlock; try { // Here we use passive deadline because we end up waiting for the SSE operators which can timeout by their own. - resultsBlock = - _blockingQueue.poll(_context.getPassiveDeadlineMs() - System.currentTimeMillis(), TimeUnit.MILLISECONDS); + long deadlineMs = getWaitDeadlineMs(); + resultsBlock = _blockingQueue.poll(Math.max(0L, deadlineMs - System.currentTimeMillis()), TimeUnit.MILLISECONDS); } catch (InterruptedException e) { terminateAndClearResultsBlocks(); return replaceWithTerminateExceptionIfAvailable(CANCELLED_BLOCK); } if (resultsBlock == null) { terminateAndClearResultsBlocks(); - return replaceWithTerminateExceptionIfAvailable(TIMEOUT_BLOCK); + return handleTimeoutOverflow(TIMEOUT_BLOCK, "waiting for leaf results"); } // Terminate when there is error block ErrorMseBlock errorBlock = getErrorBlock(); @@ -230,6 +243,42 @@ private MseBlock replaceWithTerminateExceptionIfAvailable(ErrorMseBlock errorBlo terminateException.getMessage()) : errorBlock; } + private MseBlock handleTimeoutOverflow(ErrorMseBlock defaultBlock, String reason) { + if (recordTimeoutOverflow(reason)) { + return SuccessMseBlock.INSTANCE; + } + return replaceWithTerminateExceptionIfAvailable(defaultBlock); + } + + private boolean recordTimeoutOverflow(String reason) { + if (_timeoutOverflowMode != TimeoutOverflowMode.BREAK) { + return false; + } + if (!_timeoutOverflowTriggered) { + LOGGER.warn("Leaf operator for table '{}' returning partial results after timeout: {}", _tableName, reason); + } + _timeoutOverflowTriggered = true; + _statMap.merge(StatKey.TIMEOUT_OVERFLOW_REACHED, true); + return true; + } + + private long getActiveDeadlineMsOverride() { + return _customActiveDeadlineMs != Long.MAX_VALUE ? _customActiveDeadlineMs : _context.getActiveDeadlineMs(); + } + + private long getPassiveDeadlineMsOverride() { + return _customPassiveDeadlineMs != Long.MAX_VALUE ? _customPassiveDeadlineMs : _context.getPassiveDeadlineMs(); + } + + private static long parseDeadline(@Nullable String value) { + return value != null ? Long.parseLong(value) : Long.MAX_VALUE; + } + + private long getWaitDeadlineMs() { + return _timeoutOverflowMode == TimeoutOverflowMode.BREAK ? getActiveDeadlineMsOverride() + : getPassiveDeadlineMsOverride(); + } + public ExplainedNode explain() { if (_executionFuture == null) { _executionFuture = startExecution(); @@ -243,8 +292,9 @@ public ExplainedNode explain() { BaseResultsBlock resultsBlock; try { // Here we use passive deadline because we end up waiting for the SSE operators which can timeout by their own. - resultsBlock = - _blockingQueue.poll(_context.getPassiveDeadlineMs() - System.currentTimeMillis(), TimeUnit.MILLISECONDS); + long deadlineMs = getWaitDeadlineMs(); + long waitMs = Math.max(0L, deadlineMs - System.currentTimeMillis()); + resultsBlock = _blockingQueue.poll(waitMs, TimeUnit.MILLISECONDS); } catch (InterruptedException e) { terminateAndClearResultsBlocks(); checkTerminateException(); @@ -519,8 +569,11 @@ void execute() { }); } try { - if (!latch.await(_context.getPassiveDeadlineMs() - System.currentTimeMillis(), TimeUnit.MILLISECONDS)) { - setErrorBlock(TIMEOUT_BLOCK); + long deadlineMs = getWaitDeadlineMs(); + if (!latch.await(Math.max(0L, deadlineMs - System.currentTimeMillis()), TimeUnit.MILLISECONDS)) { + if (!recordTimeoutOverflow("waiting for hybrid leaf responses")) { + setErrorBlock(TIMEOUT_BLOCK); + } } } catch (InterruptedException e) { setErrorBlock(CANCELLED_BLOCK); @@ -545,7 +598,12 @@ private void executeOneRequest(ServerQueryRequest request, @Nullable Runnable on // SERVER_SEGMENT_MISSING_ERROR are counted as query failure. Map exceptions = instanceResponseBlock.getExceptions(); if (!exceptions.isEmpty()) { - setErrorBlock(ErrorMseBlock.fromMap(QueryErrorCode.fromKeyMap(exceptions))); + Map errorMessages = QueryErrorCode.fromKeyMap(exceptions); + if (errorMessages.size() == 1 && errorMessages.containsKey(QueryErrorCode.EXECUTION_TIMEOUT) + && recordTimeoutOverflow("single-stage execution")) { + return; + } + setErrorBlock(ErrorMseBlock.fromMap(errorMessages)); if (onException != null) { onException.run(); } @@ -559,7 +617,9 @@ private void executeOneRequest(ServerQueryRequest request, @Nullable Runnable on } catch (InterruptedException e) { setErrorBlock(CANCELLED_BLOCK); } catch (TimeoutException e) { - setErrorBlock(TIMEOUT_BLOCK); + if (!recordTimeoutOverflow("adding results block")) { + setErrorBlock(TIMEOUT_BLOCK); + } } catch (Exception e) { if (!(e instanceof EarlyTerminationException)) { LOGGER.warn("Failed to add results block", e); @@ -575,8 +635,9 @@ void addResultsBlock(BaseResultsBlock resultsBlock) if (_terminated) { throw new EarlyTerminationException("Query has been terminated"); } - if (!_blockingQueue.offer(resultsBlock, _context.getPassiveDeadlineMs() - System.currentTimeMillis(), - TimeUnit.MILLISECONDS)) { + long deadlineMs = getWaitDeadlineMs(); + long waitMs = Math.max(0L, deadlineMs - System.currentTimeMillis()); + if (!_blockingQueue.offer(resultsBlock, waitMs, TimeUnit.MILLISECONDS)) { throw new TimeoutException("Timed out waiting to add results block"); } } @@ -754,6 +815,7 @@ public long merge(long value1, long value2) { NUM_GROUPS_LIMIT_REACHED(StatMap.Type.BOOLEAN), NUM_GROUPS_WARNING_LIMIT_REACHED(StatMap.Type.BOOLEAN), LITE_MODE_LEAF_STAGE_LIMIT_REACHED(StatMap.Type.BOOLEAN), + TIMEOUT_OVERFLOW_REACHED(StatMap.Type.BOOLEAN, BrokerResponseNativeV2.StatKey.TIMEOUT_OVERFLOW_REACHED), NUM_RESIZES(StatMap.Type.INT, null), RESIZE_TIME_MS(StatMap.Type.LONG, null), THREAD_CPU_TIME_NS(StatMap.Type.LONG, null), diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java index dcde2f2dff8c..b57b9d9c637b 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java @@ -114,8 +114,13 @@ public MseBlock nextBlock() { checkTermination(); nextBlock = getNextBlock(); } catch (Exception e) { - logger().warn("Operator {}: Exception while processing next block", _operatorId, e); - nextBlock = ErrorMseBlock.fromException(e); + MseBlock handled = handleException(e); + if (handled != null) { + nextBlock = handled; + } else { + logger().warn("Operator {}: Exception while processing next block", _operatorId, e); + nextBlock = ErrorMseBlock.fromException(e); + } } int numRows = nextBlock instanceof MseBlock.Data ? ((MseBlock.Data) nextBlock).getNumRows() : 0; long memoryUsedBytes = resourceSnapshot.getAllocatedBytes(); @@ -133,6 +138,18 @@ public MseBlock nextBlock() { protected abstract MseBlock getNextBlock() throws Exception; + /** + * Gives operator implementations a chance to translate an exception into a result block instead of bubbling the + * error up to the caller. + * + * @return a {@link MseBlock} that should be returned to the caller, or {@code null} to keep the default behavior of + * returning an {@link ErrorMseBlock}. + */ + @Nullable + protected MseBlock handleException(Exception e) { + return null; + } + /** * Signals the operator to terminate early. * @@ -276,6 +293,7 @@ public void mergeInto(BrokerResponseNativeV2 response, StatMap map) { response.mergeMaxRowsInOperator(stats.getLong(HashJoinOperator.StatKey.EMITTED_ROWS)); response.mergeMaxRowsInJoinReached(stats.getBoolean(HashJoinOperator.StatKey.MAX_ROWS_IN_JOIN_REACHED)); response.mergeMaxRowsInJoin(stats.getLong(HashJoinOperator.StatKey.MAX_ROWS_IN_JOIN)); + response.mergeTimeoutOverflowReached(stats.getBoolean(HashJoinOperator.StatKey.TIMEOUT_OVERFLOW_REACHED)); } @Override diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java index 21bea3dbbac7..0e2ce13acdd9 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java @@ -50,13 +50,18 @@ import org.apache.pinot.core.query.request.ServerQueryRequest; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.query.mailbox.MailboxService; +import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; +import org.apache.pinot.query.routing.StageMetadata; import org.apache.pinot.query.routing.VirtualServerAddress; +import org.apache.pinot.query.routing.WorkerMetadata; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.blocks.SuccessMseBlock; import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.apache.pinot.query.runtime.plan.pipeline.PipelineBreakerOperator; import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; import org.apache.pinot.spi.utils.JsonUtils; import org.mockito.Mock; import org.mockito.MockitoAnnotations; @@ -124,6 +129,39 @@ private List mockQueryRequests(int numRequests) { return queryRequests; } + private LeafOperator createLeafOperatorWithTimeoutMode(String timeoutOverflowMode, long timeoutMs, + QueryExecutor queryExecutor, DataSchema schema) { + Map metadata = new HashMap<>(); + metadata.put(QueryOptionKey.TIMEOUT_OVERFLOW_MODE, timeoutOverflowMode); + return new LeafOperator(createExecutionContext(metadata, timeoutMs), mockQueryRequests(1), schema, queryExecutor, + _executorService); + } + + private OpChainExecutionContext createExecutionContext(Map metadata, long timeoutMs) { + long deadline = System.currentTimeMillis() + timeoutMs; + MailboxService mailboxService = mock(MailboxService.class); + when(mailboxService.getHostname()).thenReturn("localhost"); + when(mailboxService.getPort()).thenReturn(1234); + WorkerMetadata workerMetadata = new WorkerMetadata(0, Map.of(), Map.of()); + StageMetadata stageMetadata = + new StageMetadata(0, List.of(workerMetadata), Map.of(DispatchablePlanFragment.TABLE_NAME_KEY, "testTable")); + return new OpChainExecutionContext(mailboxService, 0L, "cid", deadline, deadline, "brokerId", metadata, + stageMetadata, workerMetadata, null, true, true); + } + + private QueryExecutor slowQueryExecutor(long sleepMs) { + QueryExecutor queryExecutor = mock(QueryExecutor.class); + when(queryExecutor.execute(any(), any(), any())).thenAnswer(invocation -> { + try { + Thread.sleep(sleepMs); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return new InstanceResponseBlock(new MetadataResultsBlock()); + }); + return queryExecutor; + } + @Test public void shouldReturnDataBlockThenMetadataBlock() { // Given: @@ -262,6 +300,65 @@ public void shouldReturnMultipleDataBlockThenMetadataBlock() { operator.close(); } + @Test + public void shouldBreakOnTimeoutOverflowMode() + throws Exception { + DataSchema schema = new DataSchema(new String[]{"strCol"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING}); + LeafOperator operator = + createLeafOperatorWithTimeoutMode("BREAK", 20, slowQueryExecutor(50), schema); + _operatorRef.set(operator); + + MseBlock block = operator.nextBlock(); + assertSame(block, SuccessMseBlock.INSTANCE); + MultiStageQueryStats stats = operator.calculateStats(); + assertTrue(OperatorTestUtil.getStatMap(LeafOperator.StatKey.class, stats) + .getBoolean(LeafOperator.StatKey.TIMEOUT_OVERFLOW_REACHED)); + + operator.close(); + } + + @Test + public void shouldEmitPartialResultsWhenTimeoutErrorBlockReceived() + throws Exception { + QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT strCol FROM tbl"); + DataSchema schema = new DataSchema(new String[]{"strCol"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING}); + List dataBlocks = List.of( + new SelectionResultsBlock(schema, Arrays.asList(new Object[]{"foo"}, new Object[]{"bar"}), queryContext)); + InstanceResponseBlock metadataBlock = new InstanceResponseBlock(new MetadataResultsBlock()); + metadataBlock.addException(QueryErrorCode.EXECUTION_TIMEOUT, "Timing out on: LEAF"); + QueryExecutor queryExecutor = mockQueryExecutor(dataBlocks, metadataBlock); + LeafOperator operator = createLeafOperatorWithTimeoutMode("BREAK", 1_000, queryExecutor, schema); + _operatorRef.set(operator); + + MseBlock dataBlock = operator.nextBlock(); + List rows = ((MseBlock.Data) dataBlock).asRowHeap().getRows(); + assertEquals(rows.get(0), new Object[]{"foo"}); + assertEquals(rows.get(1), new Object[]{"bar"}); + assertSame(operator.nextBlock(), SuccessMseBlock.INSTANCE); + MultiStageQueryStats stats = operator.calculateStats(); + assertTrue(OperatorTestUtil.getStatMap(LeafOperator.StatKey.class, stats) + .getBoolean(LeafOperator.StatKey.TIMEOUT_OVERFLOW_REACHED)); + + operator.close(); + } + + @Test + public void shouldThrowOnTimeoutOverflowMode() + throws Exception { + DataSchema schema = new DataSchema(new String[]{"strCol"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING}); + LeafOperator operator = + createLeafOperatorWithTimeoutMode("THROW", 20, slowQueryExecutor(50), schema); + _operatorRef.set(operator); + + MseBlock block = operator.nextBlock(); + assertTrue(block.isError()); + + operator.close(); + } + @Test public void shouldHandleMultipleRequests() { // Given: diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/TimeoutOverflowQueryRunnerTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/TimeoutOverflowQueryRunnerTest.java new file mode 100644 index 000000000000..a207c993a0cc --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/TimeoutOverflowQueryRunnerTest.java @@ -0,0 +1,237 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.queries; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.pinot.common.response.broker.QueryProcessingException; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.query.QueryEnvironmentTestBase; +import org.apache.pinot.query.QueryServerEnclosure; +import org.apache.pinot.query.mailbox.MailboxService; +import org.apache.pinot.query.routing.QueryServerInstance; +import org.apache.pinot.query.runtime.operator.LeafOperator; +import org.apache.pinot.query.runtime.operator.MultiStageOperator; +import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; +import org.apache.pinot.query.service.dispatch.QueryDispatcher; +import org.apache.pinot.query.testutils.MockInstanceDataManagerFactory; +import org.apache.pinot.query.testutils.QueryTestUtils; +import org.apache.pinot.query.testutils.SlowMockInstanceDataManagerFactory; +import org.apache.pinot.spi.config.instance.InstanceType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + + +/** + * Integration tests covering timeout overflow behavior end-to-end. + */ +public class TimeoutOverflowQueryRunnerTest extends QueryRunnerTestBase { + private static final String TABLE_NAME = "slowTable"; + private static final String TABLE_NAME_WITH_TYPE = TABLE_NAME + "_OFFLINE"; + private static final String FAST_TABLE_NAME = "fastTable"; + private static final String FAST_TABLE_WITH_TYPE = FAST_TABLE_NAME + "_OFFLINE"; + private static final long LEAF_TIMEOUT_MS = 1; + private static final long LEAF_EXTRA_PASSIVE_TIMEOUT_MS = 0; + private static final long QUERY_TIMEOUT_MS = 2000; + private static final long QUERY_EXTRA_PASSIVE_TIMEOUT_MS = 1000; + + private QueryServerEnclosure _slowServer; + private QueryServerEnclosure _fastServer; + + @BeforeClass + public void setUp() + throws Exception { + SlowMockInstanceDataManagerFactory slowFactory = new SlowMockInstanceDataManagerFactory("slowServer", 500); + Schema schema = new Schema.SchemaBuilder().addSingleValueDimension("col1", FieldSpec.DataType.STRING, "") + .addSingleValueDimension("col2", FieldSpec.DataType.STRING, "") + .addMetric("col3", FieldSpec.DataType.INT, 0) + .addDateTime("ts", FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:HOURS") + .setSchemaName(TABLE_NAME) + .setEnableColumnBasedNullHandling(true) + .build(); + slowFactory.registerTable(schema, TABLE_NAME_WITH_TYPE); + slowFactory.addSegment(TABLE_NAME_WITH_TYPE, QueryRunnerTest.buildRows(TABLE_NAME_WITH_TYPE)); + slowFactory.addSegment(TABLE_NAME_WITH_TYPE, QueryRunnerTest.buildRows(TABLE_NAME_WITH_TYPE)); + + MockInstanceDataManagerFactory fastFactory = new MockInstanceDataManagerFactory("fastServer"); + fastFactory.registerTable( + new Schema.SchemaBuilder().addSingleValueDimension("col1", FieldSpec.DataType.STRING, "") + .addSingleValueDimension("col2", FieldSpec.DataType.STRING, "") + .addMetric("col3", FieldSpec.DataType.INT, 0) + .addDateTime("ts", FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:HOURS") + .setSchemaName(FAST_TABLE_NAME) + .setEnableColumnBasedNullHandling(true) + .build(), + FAST_TABLE_WITH_TYPE); + fastFactory.addSegment(FAST_TABLE_WITH_TYPE, QueryRunnerTest.buildRows(FAST_TABLE_WITH_TYPE)); + fastFactory.addSegment(FAST_TABLE_WITH_TYPE, QueryRunnerTest.buildRows(FAST_TABLE_WITH_TYPE)); + + _reducerHostname = "localhost"; + _reducerPort = QueryTestUtils.getAvailablePort(); + Map reducerConfig = new HashMap<>(); + reducerConfig.put(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME, _reducerHostname); + reducerConfig.put(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT, _reducerPort); + _mailboxService = new MailboxService(_reducerHostname, _reducerPort, InstanceType.BROKER, + new PinotConfiguration(reducerConfig)); + _mailboxService.start(); + + _slowServer = new QueryServerEnclosure(slowFactory); + _slowServer.start(); + _fastServer = new QueryServerEnclosure(fastFactory); + _fastServer.start(); + + QueryServerInstance slowInstance = + new QueryServerInstance("Server_localhost_" + _slowServer.getPort(), "localhost", _slowServer.getPort(), + _slowServer.getPort()); + QueryServerInstance fastInstance = + new QueryServerInstance("Server_localhost_" + _fastServer.getPort(), "localhost", _fastServer.getPort(), + _fastServer.getPort()); + _servers.put(slowInstance, _slowServer); + _servers.put(fastInstance, _fastServer); + + Map schemaMap = new HashMap<>(); + schemaMap.putAll(slowFactory.getRegisteredSchemaMap()); + schemaMap.putAll(fastFactory.getRegisteredSchemaMap()); + _queryEnvironment = QueryEnvironmentTestBase.getQueryEnvironment(_reducerPort, _slowServer.getPort(), + _fastServer.getPort(), schemaMap, slowFactory.buildTableSegmentNameMap(), + fastFactory.buildTableSegmentNameMap(), null); + } + + @AfterClass + public void tearDownClass() { + if (_mailboxService != null) { + _mailboxService.shutdown(); + } + if (_slowServer != null) { + _slowServer.shutDown(); + } + if (_fastServer != null) { + _fastServer.shutDown(); + } + } + + @Test + public void testTimeoutOverflowBreakReturnsPartialResults() { + String sql = queryPrefix("BREAK") + "SELECT col1 FROM " + TABLE_NAME + " LIMIT 5"; + QueryDispatcher.QueryResult result = queryRunner(sql, false); + Assert.assertNull(result.getProcessingException(), "Expected partial success without exception"); + ResultTable resultTable = result.getResultTable(); + Assert.assertNotNull(resultTable, "Result table should be present even for partial results"); + Assert.assertTrue(hasTimeoutOverflowStats(result), "Leaf stats should record timeout overflow"); + } + + @Test + public void testTimeoutOverflowThrowFailsQuery() { + String sql = queryPrefix("THROW") + "SELECT col1 FROM " + TABLE_NAME + " LIMIT 5"; + QueryDispatcher.QueryResult result = queryRunner(sql, false); + QueryProcessingException processingException = result.getProcessingException(); + Assert.assertNotNull(processingException, "Expected exception when overflow mode is THROW"); + Assert.assertEquals(processingException.getErrorCode(), QueryErrorCode.EXECUTION_TIMEOUT.getId()); + } + + @Test + public void testTimeoutOverflowBreakOnJoin() { + String sql = queryPrefix("BREAK") + + "SELECT slow.col1, fast.col1 FROM " + TABLE_NAME + " slow JOIN " + + FAST_TABLE_NAME + " fast ON slow.col2 = fast.col2"; + QueryDispatcher.QueryResult result = queryRunner(sql, false); + Assert.assertNull(result.getProcessingException(), "Join should return partial results without error"); + Assert.assertNotNull(result.getResultTable()); + Assert.assertTrue(hasTimeoutOverflowStats(result)); + } + + @Test + public void testTimeoutOverflowThrowOnJoin() { + String sql = queryPrefix("THROW") + + "SELECT slow.col1, fast.col1 FROM " + TABLE_NAME + " slow JOIN " + + FAST_TABLE_NAME + " fast ON slow.col2 = fast.col2"; + QueryDispatcher.QueryResult result = queryRunner(sql, false); + Assert.assertNotNull(result.getProcessingException(), "Join should fail when overflow mode THROW is used"); + Assert.assertEquals(result.getProcessingException().getErrorCode(), QueryErrorCode.EXECUTION_TIMEOUT.getId()); + } + + @Test + public void testTimeoutOverflowBreakOnAggregation() { + String sql = queryPrefix("BREAK") + + "SELECT col2, COUNT(*) FROM " + TABLE_NAME + " GROUP BY col2"; + QueryDispatcher.QueryResult result = queryRunner(sql, false); + Assert.assertNull(result.getProcessingException(), "Aggregation should emit partial results on BREAK"); + Assert.assertNotNull(result.getResultTable()); + Assert.assertTrue(hasTimeoutOverflowStats(result)); + } + @Test + public void testTimeoutOverflowBreakUsesDefaultLeafTimeout() { + long timeoutMs = 100; + long extraPassiveTimeoutMs = 50; + String sql = queryPrefixWithoutLeafOverrides("BREAK", timeoutMs, extraPassiveTimeoutMs) + + "SELECT col1 FROM " + TABLE_NAME + " LIMIT 5"; + QueryDispatcher.QueryResult result = queryRunner(sql, false); + Assert.assertNull(result.getProcessingException(), "Expected partial results when relying on default leaf timeout"); + ResultTable resultTable = result.getResultTable(); + Assert.assertNotNull(resultTable); + Assert.assertTrue(hasTimeoutOverflowStats(result)); + } + + private String queryPrefix(String overflowMode) { + return buildQueryPrefix(overflowMode, QUERY_TIMEOUT_MS, QUERY_EXTRA_PASSIVE_TIMEOUT_MS, true); + } + + private String queryPrefixWithoutLeafOverrides(String overflowMode, long timeoutMs, long extraPassiveTimeoutMs) { + return buildQueryPrefix(overflowMode, timeoutMs, extraPassiveTimeoutMs, false); + } + + private String buildQueryPrefix(String overflowMode, long timeoutMs, long extraPassiveTimeoutMs, + boolean includeLeafOverrides) { + StringBuilder builder = new StringBuilder(); + builder.append("SET useMultistageEngine='true';") + .append("SET timeoutOverflowMode='").append(overflowMode).append("';") + .append("SET timeoutMs='").append(timeoutMs).append("';") + .append("SET extraPassiveTimeoutMs='").append(extraPassiveTimeoutMs).append("';"); + if (includeLeafOverrides) { + builder.append("SET leafTimeoutMs='").append(LEAF_TIMEOUT_MS).append("';") + .append("SET leafExtraPassiveTimeoutMs='").append(LEAF_EXTRA_PASSIVE_TIMEOUT_MS).append("';"); + } + return builder.toString(); + } + + private boolean hasTimeoutOverflowStats(QueryDispatcher.QueryResult result) { + AtomicBoolean overflow = new AtomicBoolean(false); + for (MultiStageQueryStats.StageStats.Closed stageStats : result.getQueryStats()) { + stageStats.forEach((type, statMap) -> { + if (type == MultiStageOperator.Type.LEAF) { + @SuppressWarnings("unchecked") + org.apache.pinot.common.datatable.StatMap stats = + (org.apache.pinot.common.datatable.StatMap) statMap; + if (stats.getBoolean(LeafOperator.StatKey.TIMEOUT_OVERFLOW_REACHED)) { + overflow.set(true); + } + } + }); + } + return overflow.get(); + } +} diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/MockInstanceDataManagerFactory.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/MockInstanceDataManagerFactory.java index f3fb219e8c6a..6575a3993d6e 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/MockInstanceDataManagerFactory.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/MockInstanceDataManagerFactory.java @@ -46,6 +46,7 @@ import org.apache.pinot.spi.utils.ReadMode; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.apache.pinot.spi.utils.builder.TableNameBuilder; +import org.mockito.stubbing.Answer; import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.anyMap; @@ -125,7 +126,10 @@ public InstanceDataManager buildInstanceDataManager() { InstanceDataManager instanceDataManager = mock(InstanceDataManager.class); Map tableDataManagers = new HashMap<>(); for (Map.Entry> e : _tableSegmentMap.entrySet()) { - TableDataManager tableDataManager = mockTableDataManager(e.getKey(), e.getValue()); + String tableNameWithType = e.getKey(); + TableConfig tableConfig = createTableConfig(tableNameWithType); + Schema schema = _schemaMap.get(TableNameBuilder.extractRawTableName(tableNameWithType)); + TableDataManager tableDataManager = createTableDataManager(tableNameWithType, e.getValue(), tableConfig, schema); tableDataManagers.put(e.getKey(), tableDataManager); } for (Map.Entry e : tableDataManagers.entrySet()) { @@ -146,27 +150,32 @@ public Map> buildTableSegmentNameMap() { return _tableSegmentNameMap; } - private TableDataManager mockTableDataManager(String tableNameWithType, List segmentList) { + protected TableDataManager createTableDataManager(String tableNameWithType, List segmentList, + TableConfig tableConfig, Schema schema) { TableDataManager tableDataManager = mock(TableDataManager.class); when(tableDataManager.getTableName()).thenReturn(tableNameWithType); - TableConfig tableConfig = createTableConfig(tableNameWithType); - Schema schema = _schemaMap.get(TableNameBuilder.extractRawTableName(tableNameWithType)); when(tableDataManager.getCachedTableConfigAndSchema()).thenReturn(Pair.of(tableConfig, schema)); Map segmentDataManagerMap = segmentList.stream().collect(Collectors.toMap(IndexSegment::getSegmentName, ImmutableSegmentDataManager::new)); // TODO: support optional segments for multi-stage engine, but for now, it's always null. - when(tableDataManager.acquireSegments(anyList(), eq(null), anyList())).thenAnswer(invocation -> { + Answer> acquireSegmentsAnswer = invocation -> { List segments = invocation.getArgument(0); return segments.stream().map(segmentDataManagerMap::get).collect(Collectors.toList()); - }); + }; + when(tableDataManager.acquireSegments(anyList(), eq(null), anyList())).thenAnswer(acquireSegmentsAnswer); when(tableDataManager.getSegmentContexts(anyList(), anyMap())).thenAnswer(invocation -> { List segments = invocation.getArgument(0); return segments.stream().map(SegmentContext::new).collect(Collectors.toList()); }); + customizeTableDataManager(tableDataManager, tableNameWithType, acquireSegmentsAnswer); return tableDataManager; } + protected void customizeTableDataManager(TableDataManager tableDataManager, String tableNameWithType, + Answer> acquireSegmentsAnswer) { + } + private ImmutableSegment buildSegment(String tableNameWithType, File indexDir, List rows) { String segmentName = String.format("%s_%s", tableNameWithType, UUID.randomUUID()); TableConfig tableConfig = createTableConfig(tableNameWithType); diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/SlowMockInstanceDataManagerFactory.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/SlowMockInstanceDataManagerFactory.java new file mode 100644 index 000000000000..c8f35fd7c640 --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/SlowMockInstanceDataManagerFactory.java @@ -0,0 +1,55 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.testutils; + +import java.util.List; +import org.apache.pinot.segment.local.data.manager.SegmentDataManager; +import org.apache.pinot.segment.local.data.manager.TableDataManager; +import org.mockito.stubbing.Answer; + +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.when; + + +/** + * A {@link MockInstanceDataManagerFactory} variant that sleeps when segments are acquired. + *

This allows integration tests to deterministically trigger server-side timeouts.

+ */ +public class SlowMockInstanceDataManagerFactory extends MockInstanceDataManagerFactory { + private final long _sleepMs; + + public SlowMockInstanceDataManagerFactory(String serverName, long sleepMs) { + super(serverName); + _sleepMs = sleepMs; + } + + @Override + protected void customizeTableDataManager(TableDataManager tableDataManager, String tableNameWithType, + Answer> acquireSegmentsAnswer) { + when(tableDataManager.acquireSegments(anyList(), eq(null), anyList())).thenAnswer(invocation -> { + try { + Thread.sleep(_sleepMs); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return acquireSegmentsAnswer.answer(invocation); + }); + } +} diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 3ddfaf79c0f7..c8b420fddb3a 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -741,6 +741,8 @@ public static class QueryOptionKey { /// SQL sets this option. public static final String ENABLE_MATERIALIZED_VIEW_REWRITE = "enableMaterializedViewRewrite"; public static final String EXTRA_PASSIVE_TIMEOUT_MS = "extraPassiveTimeoutMs"; + public static final String LEAF_TIMEOUT_MS = "leafTimeoutMs"; + public static final String LEAF_EXTRA_PASSIVE_TIMEOUT_MS = "leafExtraPassiveTimeoutMs"; public static final String SKIP_UPSERT = "skipUpsert"; public static final String SKIP_UPSERT_VIEW = "skipUpsertView"; public static final String UPSERT_VIEW_FRESHNESS_MS = "upsertViewFreshnessMs"; @@ -936,6 +938,7 @@ public static class QueryOptionKey { // Handle WINDOW Overflow public static final String MAX_ROWS_IN_WINDOW = "maxRowsInWindow"; public static final String WINDOW_OVERFLOW_MODE = "windowOverflowMode"; + public static final String TIMEOUT_OVERFLOW_MODE = "timeoutOverflowMode"; // Indicates the maximum length of the serialized response per server for a query. public static final String MAX_SERVER_RESPONSE_SIZE_BYTES = "maxServerResponseSizeBytes"; @@ -2635,6 +2638,7 @@ public static class MultiStageQueryRunner { */ public static final String KEY_OF_MAX_ROWS_IN_JOIN = "pinot.query.join.max.rows"; public static final String KEY_OF_JOIN_OVERFLOW_MODE = "pinot.query.join.overflow.mode"; + public static final String KEY_OF_TIMEOUT_OVERFLOW_MODE = "pinot.query.timeout.overflow.mode"; /// Specifies the send stats mode used in MSE. /// @@ -2676,6 +2680,9 @@ public static class MultiStageQueryRunner { public enum JoinOverFlowMode { THROW, BREAK } + public enum TimeoutOverflowMode { + THROW, BREAK + } /** * Configuration for window overflow.