Skip to content

Commit 53029b6

Browse files
xiangfu0xiangfu0
committed
Add timeout overflow break handling at leaf stage
1 parent 62b20d9 commit 53029b6

9 files changed

Lines changed: 481 additions & 18 deletions

File tree

‎pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,8 @@ public int getNumRowsResultSet() {
109109
@JsonProperty(access = JsonProperty.Access.READ_ONLY)
110110
@Override
111111
public boolean isPartialResult() {
112-
return getExceptionsSize() > 0 || isNumGroupsLimitReached() || isMaxRowsInJoinReached();
112+
return getExceptionsSize() > 0 || isNumGroupsLimitReached() || isMaxRowsInJoinReached()
113+
|| isTimeoutOverflowReached();
113114
}
114115

115116
@Override
@@ -152,6 +153,14 @@ public void mergeNumGroupsWarningLimitReached(boolean numGroupsWarningLimitReach
152153
_brokerStats.merge(StatKey.NUM_GROUPS_WARNING_LIMIT_REACHED, numGroupsWarningLimitReached);
153154
}
154155

156+
public boolean isTimeoutOverflowReached() {
157+
return _brokerStats.getBoolean(StatKey.TIMEOUT_OVERFLOW_REACHED);
158+
}
159+
160+
public void mergeTimeoutOverflowReached(boolean timeoutOverflowReached) {
161+
_brokerStats.merge(StatKey.TIMEOUT_OVERFLOW_REACHED, timeoutOverflowReached);
162+
}
163+
155164
@Override
156165
public boolean isMaxRowsInJoinReached() {
157166
return _maxRowsInJoinReached;
@@ -453,7 +462,8 @@ public long merge(long value1, long value2) {
453462
NUM_SEGMENTS_PRUNED_BY_VALUE(StatMap.Type.INT),
454463
GROUPS_TRIMMED(StatMap.Type.BOOLEAN),
455464
NUM_GROUPS_LIMIT_REACHED(StatMap.Type.BOOLEAN),
456-
NUM_GROUPS_WARNING_LIMIT_REACHED(StatMap.Type.BOOLEAN);
465+
NUM_GROUPS_WARNING_LIMIT_REACHED(StatMap.Type.BOOLEAN),
466+
TIMEOUT_OVERFLOW_REACHED(StatMap.Type.BOOLEAN);
457467

458468
private final StatMap.Type _type;
459469

‎pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import org.apache.pinot.spi.utils.CommonConstants;
3535
import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
3636
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.JoinOverFlowMode;
37+
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.TimeoutOverflowMode;
3738
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.WindowOverFlowMode;
3839

3940

@@ -432,6 +433,12 @@ public static WindowOverFlowMode getWindowOverflowMode(Map<String, String> query
432433
return windowOverflowModeStr != null ? WindowOverFlowMode.valueOf(windowOverflowModeStr) : null;
433434
}
434435

436+
@Nullable
437+
public static TimeoutOverflowMode getTimeoutOverflowMode(Map<String, String> queryOptions) {
438+
String timeoutOverflowModeStr = queryOptions.get(QueryOptionKey.TIMEOUT_OVERFLOW_MODE);
439+
return timeoutOverflowModeStr != null ? TimeoutOverflowMode.valueOf(timeoutOverflowModeStr) : null;
440+
}
441+
435442
public static boolean isSkipUnavailableServers(Map<String, String> queryOptions) {
436443
return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.SKIP_UNAVAILABLE_SERVERS));
437444
}

‎pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@
8686
import org.apache.pinot.spi.utils.CommonConstants.Helix;
8787
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner;
8888
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.JoinOverFlowMode;
89+
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.TimeoutOverflowMode;
8990
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.WindowOverFlowMode;
9091
import org.apache.pinot.spi.utils.CommonConstants.Query.Request;
9192
import org.apache.pinot.spi.utils.CommonConstants.Query.Response;
@@ -138,6 +139,8 @@ public class QueryRunner {
138139
@Nullable
139140
private WindowOverFlowMode _windowOverflowMode;
140141
@Nullable
142+
private TimeoutOverflowMode _timeoutOverflowMode;
143+
@Nullable
141144
private PhysicalTimeSeriesServerPlanVisitor _timeSeriesPhysicalPlanVisitor;
142145
private BooleanSupplier _sendStats;
143146

@@ -191,6 +194,10 @@ public void init(PinotConfiguration serverConf, InstanceDataManager instanceData
191194
String windowOverflowModeStr = serverConf.getProperty(MultiStageQueryRunner.KEY_OF_WINDOW_OVERFLOW_MODE);
192195
_windowOverflowMode = windowOverflowModeStr != null ? WindowOverFlowMode.valueOf(windowOverflowModeStr) : null;
193196

197+
String timeoutOverflowModeStr = serverConf.getProperty(MultiStageQueryRunner.KEY_OF_TIMEOUT_OVERFLOW_MODE);
198+
_timeoutOverflowMode =
199+
timeoutOverflowModeStr != null ? TimeoutOverflowMode.valueOf(timeoutOverflowModeStr) : null;
200+
194201
ExecutorService baseExecutorService =
195202
ExecutorServiceUtils.create(serverConf, Server.MULTISTAGE_EXECUTOR_CONFIG_PREFIX, "query-runner-on-" + port,
196203
Server.DEFAULT_MULTISTAGE_EXECUTOR_TYPE);
@@ -508,6 +515,14 @@ private Map<String, String> consolidateMetadata(Map<String, String> customProper
508515
opChainMetadata.put(QueryOptionKey.WINDOW_OVERFLOW_MODE, windowOverflowMode.name());
509516
}
510517

518+
TimeoutOverflowMode timeoutOverflowMode = QueryOptionsUtils.getTimeoutOverflowMode(opChainMetadata);
519+
if (timeoutOverflowMode == null) {
520+
timeoutOverflowMode = _timeoutOverflowMode;
521+
}
522+
if (timeoutOverflowMode != null) {
523+
opChainMetadata.put(QueryOptionKey.TIMEOUT_OVERFLOW_MODE, timeoutOverflowMode.name());
524+
}
525+
511526
return opChainMetadata;
512527
}
513528

‎pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java‎

Lines changed: 67 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@
6666
import org.apache.pinot.spi.exception.TerminationException;
6767
import org.apache.pinot.spi.query.QueryThreadContext;
6868
import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionValue;
69+
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.TimeoutOverflowMode;
6970
import org.slf4j.Logger;
7071
import org.slf4j.LoggerFactory;
7172

@@ -81,6 +82,8 @@ public class LeafOperator extends MultiStageOperator {
8182
ErrorMseBlock.fromError(QueryErrorCode.QUERY_CANCELLATION, "Cancelled while waiting for leaf results");
8283
private static final ErrorMseBlock TIMEOUT_BLOCK =
8384
ErrorMseBlock.fromError(QueryErrorCode.EXECUTION_TIMEOUT, "Timed out waiting for leaf results");
85+
public static final String LEAF_ACTIVE_DEADLINE_METADATA_KEY = "leafActiveDeadlineMs";
86+
public static final String LEAF_PASSIVE_DEADLINE_METADATA_KEY = "leafPassiveDeadlineMs";
8487

8588
// Use a special results block to indicate that this is the last results block
8689
@VisibleForTesting
@@ -98,10 +101,14 @@ public class LeafOperator extends MultiStageOperator {
98101
// Use a limit-sized BlockingQueue to store the results blocks and apply back pressure to the single-stage threads
99102
@VisibleForTesting
100103
final BlockingQueue<BaseResultsBlock> _blockingQueue;
104+
private final TimeoutOverflowMode _timeoutOverflowMode;
101105

102106
@Nullable
103107
private volatile Future<?> _executionFuture;
104108
private volatile boolean _terminated;
109+
private volatile boolean _timeoutOverflowTriggered;
110+
private final long _customActiveDeadlineMs;
111+
private final long _customPassiveDeadlineMs;
105112

106113
public LeafOperator(OpChainExecutionContext context, List<ServerQueryRequest> requests, DataSchema dataSchema,
107114
QueryExecutor queryExecutor, ExecutorService executorService) {
@@ -118,6 +125,12 @@ public LeafOperator(OpChainExecutionContext context, List<ServerQueryRequest> re
118125
Integer maxStreamingPendingBlocks = QueryOptionsUtils.getMaxStreamingPendingBlocks(context.getOpChainMetadata());
119126
_blockingQueue = new ArrayBlockingQueue<>(maxStreamingPendingBlocks != null ? maxStreamingPendingBlocks
120127
: QueryOptionValue.DEFAULT_MAX_STREAMING_PENDING_BLOCKS);
128+
TimeoutOverflowMode timeoutOverflowMode =
129+
QueryOptionsUtils.getTimeoutOverflowMode(context.getOpChainMetadata());
130+
_timeoutOverflowMode = timeoutOverflowMode != null ? timeoutOverflowMode : TimeoutOverflowMode.THROW;
131+
Map<String, String> metadata = context.getOpChainMetadata();
132+
_customActiveDeadlineMs = parseDeadline(metadata.get(LEAF_ACTIVE_DEADLINE_METADATA_KEY));
133+
_customPassiveDeadlineMs = parseDeadline(metadata.get(LEAF_PASSIVE_DEADLINE_METADATA_KEY));
121134
}
122135

123136
public List<ServerQueryRequest> getRequests() {
@@ -168,15 +181,15 @@ protected MseBlock getNextBlock() {
168181
BaseResultsBlock resultsBlock;
169182
try {
170183
// Here we use passive deadline because we end up waiting for the SSE operators which can timeout by their own.
171-
resultsBlock =
172-
_blockingQueue.poll(_context.getPassiveDeadlineMs() - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
184+
long deadlineMs = getWaitDeadlineMs();
185+
resultsBlock = _blockingQueue.poll(Math.max(0L, deadlineMs - System.currentTimeMillis()), TimeUnit.MILLISECONDS);
173186
} catch (InterruptedException e) {
174187
terminateAndClearResultsBlocks();
175188
return replaceWithTerminateExceptionIfAvailable(CANCELLED_BLOCK);
176189
}
177190
if (resultsBlock == null) {
178191
terminateAndClearResultsBlocks();
179-
return replaceWithTerminateExceptionIfAvailable(TIMEOUT_BLOCK);
192+
return handleTimeoutOverflow(TIMEOUT_BLOCK, "waiting for leaf results");
180193
}
181194
// Terminate when there is error block
182195
ErrorMseBlock errorBlock = getErrorBlock();
@@ -201,6 +214,42 @@ private MseBlock replaceWithTerminateExceptionIfAvailable(ErrorMseBlock errorBlo
201214
terminateException.getMessage()) : errorBlock;
202215
}
203216

217+
private MseBlock handleTimeoutOverflow(ErrorMseBlock defaultBlock, String reason) {
218+
if (recordTimeoutOverflow(reason)) {
219+
return SuccessMseBlock.INSTANCE;
220+
}
221+
return replaceWithTerminateExceptionIfAvailable(defaultBlock);
222+
}
223+
224+
private boolean recordTimeoutOverflow(String reason) {
225+
if (_timeoutOverflowMode != TimeoutOverflowMode.BREAK) {
226+
return false;
227+
}
228+
if (!_timeoutOverflowTriggered) {
229+
LOGGER.warn("Leaf operator for table '{}' returning partial results after timeout: {}", _tableName, reason);
230+
}
231+
_timeoutOverflowTriggered = true;
232+
_statMap.merge(StatKey.TIMEOUT_OVERFLOW_REACHED, true);
233+
return true;
234+
}
235+
236+
private long getActiveDeadlineMsOverride() {
237+
return _customActiveDeadlineMs != Long.MAX_VALUE ? _customActiveDeadlineMs : _context.getActiveDeadlineMs();
238+
}
239+
240+
private long getPassiveDeadlineMsOverride() {
241+
return _customPassiveDeadlineMs != Long.MAX_VALUE ? _customPassiveDeadlineMs : _context.getPassiveDeadlineMs();
242+
}
243+
244+
private static long parseDeadline(@Nullable String value) {
245+
return value != null ? Long.parseLong(value) : Long.MAX_VALUE;
246+
}
247+
248+
private long getWaitDeadlineMs() {
249+
return _timeoutOverflowMode == TimeoutOverflowMode.BREAK ? getActiveDeadlineMsOverride()
250+
: getPassiveDeadlineMsOverride();
251+
}
252+
204253
public ExplainedNode explain() {
205254
if (_executionFuture == null) {
206255
_executionFuture = startExecution();
@@ -214,8 +263,9 @@ public ExplainedNode explain() {
214263
BaseResultsBlock resultsBlock;
215264
try {
216265
// Here we use passive deadline because we end up waiting for the SSE operators which can timeout by their own.
217-
resultsBlock =
218-
_blockingQueue.poll(_context.getPassiveDeadlineMs() - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
266+
long deadlineMs = getWaitDeadlineMs();
267+
long waitMs = Math.max(0L, deadlineMs - System.currentTimeMillis());
268+
resultsBlock = _blockingQueue.poll(waitMs, TimeUnit.MILLISECONDS);
219269
} catch (InterruptedException e) {
220270
terminateAndClearResultsBlocks();
221271
checkTerminateException();
@@ -470,8 +520,11 @@ void execute() {
470520
});
471521
}
472522
try {
473-
if (!latch.await(_context.getPassiveDeadlineMs() - System.currentTimeMillis(), TimeUnit.MILLISECONDS)) {
474-
setErrorBlock(TIMEOUT_BLOCK);
523+
long deadlineMs = getWaitDeadlineMs();
524+
if (!latch.await(Math.max(0L, deadlineMs - System.currentTimeMillis()), TimeUnit.MILLISECONDS)) {
525+
if (!recordTimeoutOverflow("waiting for hybrid leaf responses")) {
526+
setErrorBlock(TIMEOUT_BLOCK);
527+
}
475528
}
476529
} catch (InterruptedException e) {
477530
setErrorBlock(CANCELLED_BLOCK);
@@ -510,7 +563,9 @@ private void executeOneRequest(ServerQueryRequest request, @Nullable Runnable on
510563
} catch (InterruptedException e) {
511564
setErrorBlock(CANCELLED_BLOCK);
512565
} catch (TimeoutException e) {
513-
setErrorBlock(TIMEOUT_BLOCK);
566+
if (!recordTimeoutOverflow("adding results block")) {
567+
setErrorBlock(TIMEOUT_BLOCK);
568+
}
514569
} catch (Exception e) {
515570
if (!(e instanceof EarlyTerminationException)) {
516571
LOGGER.warn("Failed to add results block", e);
@@ -526,8 +581,9 @@ void addResultsBlock(BaseResultsBlock resultsBlock)
526581
if (_terminated) {
527582
throw new EarlyTerminationException("Query has been terminated");
528583
}
529-
if (!_blockingQueue.offer(resultsBlock, _context.getPassiveDeadlineMs() - System.currentTimeMillis(),
530-
TimeUnit.MILLISECONDS)) {
584+
long deadlineMs = getWaitDeadlineMs();
585+
long waitMs = Math.max(0L, deadlineMs - System.currentTimeMillis());
586+
if (!_blockingQueue.offer(resultsBlock, waitMs, TimeUnit.MILLISECONDS)) {
531587
throw new TimeoutException("Timed out waiting to add results block");
532588
}
533589
}
@@ -704,6 +760,7 @@ public long merge(long value1, long value2) {
704760
GROUPS_TRIMMED(StatMap.Type.BOOLEAN),
705761
NUM_GROUPS_LIMIT_REACHED(StatMap.Type.BOOLEAN),
706762
NUM_GROUPS_WARNING_LIMIT_REACHED(StatMap.Type.BOOLEAN),
763+
TIMEOUT_OVERFLOW_REACHED(StatMap.Type.BOOLEAN, BrokerResponseNativeV2.StatKey.TIMEOUT_OVERFLOW_REACHED),
707764
NUM_RESIZES(StatMap.Type.INT, null),
708765
RESIZE_TIME_MS(StatMap.Type.LONG, null),
709766
THREAD_CPU_TIME_NS(StatMap.Type.LONG, null),

‎pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,11 +38,17 @@
3838
import org.apache.pinot.core.query.request.ServerQueryRequest;
3939
import org.apache.pinot.core.query.request.context.QueryContext;
4040
import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
41+
import org.apache.pinot.query.mailbox.MailboxService;
42+
import org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
43+
import org.apache.pinot.query.routing.StageMetadata;
4144
import org.apache.pinot.query.routing.VirtualServerAddress;
45+
import org.apache.pinot.query.routing.WorkerMetadata;
4246
import org.apache.pinot.query.runtime.blocks.MseBlock;
4347
import org.apache.pinot.query.runtime.blocks.SuccessMseBlock;
48+
import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
4449
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
4550
import org.apache.pinot.spi.exception.QueryErrorCode;
51+
import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
4652
import org.mockito.Mock;
4753
import org.mockito.MockitoAnnotations;
4854
import org.testng.annotations.AfterClass;
@@ -109,6 +115,39 @@ private List<ServerQueryRequest> mockQueryRequests(int numRequests) {
109115
return queryRequests;
110116
}
111117

118+
private LeafOperator createLeafOperatorWithTimeoutMode(String timeoutOverflowMode, long timeoutMs,
119+
QueryExecutor queryExecutor, DataSchema schema) {
120+
Map<String, String> metadata = new HashMap<>();
121+
metadata.put(QueryOptionKey.TIMEOUT_OVERFLOW_MODE, timeoutOverflowMode);
122+
return new LeafOperator(createExecutionContext(metadata, timeoutMs), mockQueryRequests(1), schema, queryExecutor,
123+
_executorService);
124+
}
125+
126+
private OpChainExecutionContext createExecutionContext(Map<String, String> metadata, long timeoutMs) {
127+
long deadline = System.currentTimeMillis() + timeoutMs;
128+
MailboxService mailboxService = mock(MailboxService.class);
129+
when(mailboxService.getHostname()).thenReturn("localhost");
130+
when(mailboxService.getPort()).thenReturn(1234);
131+
WorkerMetadata workerMetadata = new WorkerMetadata(0, Map.of(), Map.of());
132+
StageMetadata stageMetadata =
133+
new StageMetadata(0, List.of(workerMetadata), Map.of(DispatchablePlanFragment.TABLE_NAME_KEY, "testTable"));
134+
return new OpChainExecutionContext(mailboxService, 0L, "cid", deadline, deadline, "brokerId", metadata,
135+
stageMetadata, workerMetadata, null, true);
136+
}
137+
138+
private QueryExecutor slowQueryExecutor(long sleepMs) {
139+
QueryExecutor queryExecutor = mock(QueryExecutor.class);
140+
when(queryExecutor.execute(any(), any(), any())).thenAnswer(invocation -> {
141+
try {
142+
Thread.sleep(sleepMs);
143+
} catch (InterruptedException e) {
144+
Thread.currentThread().interrupt();
145+
}
146+
return new InstanceResponseBlock(new MetadataResultsBlock());
147+
});
148+
return queryExecutor;
149+
}
150+
112151
@Test
113152
public void shouldReturnDataBlockThenMetadataBlock() {
114153
// Given:
@@ -201,6 +240,39 @@ public void shouldReturnMultipleDataBlockThenMetadataBlock() {
201240
operator.close();
202241
}
203242

243+
@Test
244+
public void shouldBreakOnTimeoutOverflowMode()
245+
throws Exception {
246+
DataSchema schema = new DataSchema(new String[]{"strCol"},
247+
new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING});
248+
LeafOperator operator =
249+
createLeafOperatorWithTimeoutMode("BREAK", 20, slowQueryExecutor(50), schema);
250+
_operatorRef.set(operator);
251+
252+
MseBlock block = operator.nextBlock();
253+
assertSame(block, SuccessMseBlock.INSTANCE);
254+
MultiStageQueryStats stats = operator.calculateStats();
255+
assertTrue(OperatorTestUtil.getStatMap(LeafOperator.StatKey.class, stats)
256+
.getBoolean(LeafOperator.StatKey.TIMEOUT_OVERFLOW_REACHED));
257+
258+
operator.close();
259+
}
260+
261+
@Test
262+
public void shouldThrowOnTimeoutOverflowMode()
263+
throws Exception {
264+
DataSchema schema = new DataSchema(new String[]{"strCol"},
265+
new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING});
266+
LeafOperator operator =
267+
createLeafOperatorWithTimeoutMode("THROW", 20, slowQueryExecutor(50), schema);
268+
_operatorRef.set(operator);
269+
270+
MseBlock block = operator.nextBlock();
271+
assertTrue(block.isError());
272+
273+
operator.close();
274+
}
275+
204276
@Test
205277
public void shouldHandleMultipleRequests() {
206278
// Given:

0 commit comments

Comments
 (0)