This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 66b7f21344a Preserve query error codes during mailbox cancellation
(#19682)
66b7f21344a is described below
commit 66b7f21344abb25f1b5cc1e9edff002343ccd648
Author: Xiang Fu <[email protected]>
AuthorDate: Mon Sep 28 16:13:12 2026 +0700
Preserve query error codes during mailbox cancellation (#19682)
Reproduce: Run QueryKillingIntegrationTest.testCpuBasedServerQueryKilling
with MSE; the broker can report 503 for a server CPU kill that should be 245.
---
.../pinot/query/mailbox/GrpcSendingMailbox.java | 14 +++++---
.../query/mailbox/InMemorySendingMailbox.java | 13 ++++---
.../query/mailbox/GrpcSendingMailboxTest.java | 42 ++++++++++++++++++++++
.../query/mailbox/InMemorySendingMailboxTest.java | 38 ++++++++++++++++++++
4 files changed, 98 insertions(+), 9 deletions(-)
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java
index 90ae69c8d0e..74eabb3219a 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java
@@ -134,8 +134,8 @@ public class GrpcSendingMailbox implements SendingMailbox {
// and must always reach the receiver, so they bypass the back-pressure
gate. Bypassing also disables the
// cooperative termination poll inside [#awaitReady]; without it, a
terminate signal raised while the sender is
// mid-way through pushing an error EOS would unwind [#sendInternal] with
a TerminationException, leave
- // [#_senderSideClosed] false, and let [#cancel] run and overwrite the
original error code with
- // QUERY_CANCELLATION on the receiver side.
+ // [#_senderSideClosed] false, and let [#cancel] run and replace the
original error code with
+ // QUERY_CANCELLATION if its throwable does not carry a QueryErrorCode.
//
// Race against a concurrent [#cancel]: both paths perform an early
`isTerminated()` check before taking
// [#_readyLock], so two threads can both pass that check and enter their
respective lock sections. The cancel
@@ -294,10 +294,14 @@ public class GrpcSendingMailbox implements SendingMailbox
{
// sendContent()'s acquisition nests safely.
_readyLock.lock();
try {
- String msg = t != null ? t.getMessage() : "Unknown";
+ QueryErrorCode errorCode = QueryErrorCode.fromThrowable(t,
QueryErrorCode.QUERY_CANCELLATION);
+ String msg = t != null ? t.getMessage() : null;
+ if (msg == null) {
+ msg = "Unknown";
+ }
// NOTE: DO NOT use onError() because it will terminate the stream, and
receiver might not get the callback
- MseBlock errorBlock = ErrorMseBlock.fromError(
- QueryErrorCode.QUERY_CANCELLATION, "Cancelled by sender with
exception: " + msg);
+ MseBlock errorBlock = ErrorMseBlock.fromError(errorCode,
+ errorCode == QueryErrorCode.QUERY_CANCELLATION ? "Cancelled by
sender with exception: " + msg : msg);
processAndSend(errorBlock, List.of(), /* bypassReady */ true);
_contentObserver.onCompleted();
} catch (Exception e) {
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java
index c9cef3552d3..72a737ac89f 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java
@@ -24,7 +24,7 @@ import org.apache.pinot.query.runtime.blocks.ErrorMseBlock;
import org.apache.pinot.query.runtime.blocks.MseBlock;
import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
import org.apache.pinot.segment.spi.memory.DataBuffer;
-import org.apache.pinot.spi.exception.QueryCancelledException;
+import org.apache.pinot.spi.exception.QueryErrorCode;
import org.apache.pinot.spi.query.QueryThreadContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -119,9 +119,14 @@ public class InMemorySendingMailbox implements
SendingMailbox {
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
- _receivingMailbox.setErrorBlock(
- ErrorMseBlock.fromException(new QueryCancelledException(
- "Cancelled by sender with exception: " + t.getMessage())),
List.of());
+ QueryErrorCode errorCode = QueryErrorCode.fromThrowable(t,
QueryErrorCode.QUERY_CANCELLATION);
+ String msg = t != null ? t.getMessage() : null;
+ if (msg == null) {
+ msg = "Unknown";
+ }
+ _receivingMailbox.setErrorBlock(ErrorMseBlock.fromError(errorCode,
+ errorCode == QueryErrorCode.QUERY_CANCELLATION ? "Cancelled by sender
with exception: " + msg : msg),
+ List.of());
}
@Override
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/GrpcSendingMailboxTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/GrpcSendingMailboxTest.java
index 9241f45b754..bda3ea2a8e8 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/GrpcSendingMailboxTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/GrpcSendingMailboxTest.java
@@ -26,6 +26,7 @@ import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
+import java.util.Map;
import java.util.Random;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.ExecutorService;
@@ -50,19 +51,60 @@ import
org.apache.pinot.segment.spi.memory.CompoundDataBuffer;
import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.segment.spi.memory.PinotByteBuffer;
import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.exception.QueryException;
import org.apache.pinot.spi.exception.TerminationException;
import org.apache.pinot.spi.query.QueryThreadContext;
+import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import org.testng.Assert;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;
public class GrpcSendingMailboxTest {
+ @Test(dataProvider = "cancellationErrors")
+ @SuppressWarnings("unchecked")
+ public void cancelPreservesQueryErrorCode(Exception exception,
QueryErrorCode expectedCode, String expectedMessage)
+ throws IOException {
+ ClientCallStreamObserver<MailboxContent> observer =
mock(ClientCallStreamObserver.class);
+ GrpcSendingMailbox mailbox = new GrpcSendingMailbox("test-mailbox",
mock(ChannelManager.class), "localhost", 0,
+ Long.MAX_VALUE, new StatMap<>(MailboxSendOperator.StatKey.class), 4 *
1024 * 1024, false) {
+ @Override
+ ClientCallStreamObserver<MailboxContent> getContentObserver() {
+ return observer;
+ }
+ };
+
+ mailbox.cancel(exception);
+
+ ArgumentCaptor<MailboxContent> content =
ArgumentCaptor.forClass(MailboxContent.class);
+ verify(observer).onNext(content.capture());
+ verify(observer).onCompleted();
+ DataBlock errorBlock =
DataBlockUtils.deserialize(List.of(content.getValue().getPayload().asReadOnlyByteBuffer()));
+ assertEquals(errorBlock.getExceptions(), Map.of(expectedCode.getId(),
expectedMessage));
+ }
+
+ @DataProvider(name = "cancellationErrors")
+ public Object[][] cancellationErrors() {
+ return new Object[][]{
+ {new QueryException(QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED,
"CPU limit"),
+ QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED, "CPU limit"},
+ {new TerminationException(QueryErrorCode.EXECUTION_TIMEOUT, "execution
timed out"),
+ QueryErrorCode.EXECUTION_TIMEOUT, "execution timed out"},
+ {new QueryException(QueryErrorCode.EXECUTION_TIMEOUT, (String) null),
+ QueryErrorCode.EXECUTION_TIMEOUT, "Unknown"},
+ {new RuntimeException("ordinary cancellation"),
QueryErrorCode.QUERY_CANCELLATION,
+ "Cancelled by sender with exception: ordinary cancellation"},
+ {null, QueryErrorCode.QUERY_CANCELLATION, "Cancelled by sender with
exception: Unknown"}
+ };
+ }
+
@Test
public void sendDataThrowsWhenQueryTerminated() {
ChannelManager channelManager = Mockito.mock(ChannelManager.class);
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/InMemorySendingMailboxTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/InMemorySendingMailboxTest.java
index c550d37af5b..543c86a146f 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/InMemorySendingMailboxTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/InMemorySendingMailboxTest.java
@@ -19,20 +19,58 @@
package org.apache.pinot.query.mailbox;
import java.util.List;
+import java.util.Map;
import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.query.runtime.blocks.ErrorMseBlock;
import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock;
import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.exception.QueryException;
import org.apache.pinot.spi.exception.TerminationException;
import org.apache.pinot.spi.query.QueryThreadContext;
import org.mockito.Mockito;
import org.testng.Assert;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+
public class InMemorySendingMailboxTest {
+ @Test(dataProvider = "cancellationErrors")
+ public void cancelPreservesQueryErrorCode(Exception exception,
QueryErrorCode expectedCode, String expectedMessage) {
+ ReceivingMailbox receivingMailbox = new ReceivingMailbox("test-mailbox",
1);
+ receivingMailbox.registeredReader(mock(ReceivingMailbox.Reader.class));
+ MailboxService mailboxService = mock(MailboxService.class);
+
when(mailboxService.getReceivingMailbox("test-mailbox")).thenReturn(receivingMailbox);
+ InMemorySendingMailbox mailbox = new
InMemorySendingMailbox("test-mailbox", mailboxService, Long.MAX_VALUE,
+ new StatMap<>(MailboxSendOperator.StatKey.class));
+
+ mailbox.cancel(exception);
+
+ ErrorMseBlock errorBlock = (ErrorMseBlock)
receivingMailbox.poll().getBlock();
+ assertEquals(errorBlock.getErrorMessages(), Map.of(expectedCode,
expectedMessage));
+ }
+
+ @DataProvider(name = "cancellationErrors")
+ public Object[][] cancellationErrors() {
+ return new Object[][]{
+ {new QueryException(QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED,
"CPU limit"),
+ QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED, "CPU limit"},
+ {new TerminationException(QueryErrorCode.EXECUTION_TIMEOUT, "execution
timed out"),
+ QueryErrorCode.EXECUTION_TIMEOUT, "execution timed out"},
+ {new QueryException(QueryErrorCode.EXECUTION_TIMEOUT, (String) null),
+ QueryErrorCode.EXECUTION_TIMEOUT, "Unknown"},
+ {new RuntimeException("ordinary cancellation"),
QueryErrorCode.QUERY_CANCELLATION,
+ "Cancelled by sender with exception: ordinary cancellation"},
+ {null, QueryErrorCode.QUERY_CANCELLATION, "Cancelled by sender with
exception: Unknown"}
+ };
+ }
+
@Test
public void sendDataThrowsWhenQueryTerminated() {
MailboxService mailboxService = Mockito.mock(MailboxService.class);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
