You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by "gortiz (via GitHub)" <gi...@apache.org> on 2025/09/26 08:16:05 UTC
[PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
gortiz opened a new pull request, #16903:
URL: https://github.com/apache/pinot/pull/16903
(no comment)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2427009023
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java:
##########
@@ -49,6 +50,7 @@ public class MailboxContentObserver implements StreamObserver<MailboxContent> {
private final MailboxService _mailboxService;
private final StreamObserver<MailboxStatus> _responseObserver;
private final List<ByteBuffer> _mailboxBuffers = Collections.synchronizedList(new ArrayList<>());
+ private boolean _closedStream = false;
Review Comment:
> It doesn't have to be volatile given StreamObservers should be thread-compatible. Given we never manually call methods on this class but instead it is GRPC the one that calls it, it is safe to assume thread-compatibility is honored.
I've been confusing about what thread-compatible stands for. Does that mean we don't need to sync on memory accessed by different threads? I.e. we don't need `_mailboxBuffers` to be synchronized, and `_mailbox` to be `volatile`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #16903:
URL: https://github.com/apache/pinot/pull/16903#issuecomment-3352051857
## [Codecov](https://app.codecov.io/gh/apache/pinot/pull/16903?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `0%` with `196 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 33.57%. Comparing base ([`9baf132`](https://app.codecov.io/gh/apache/pinot/commit/9baf132185cdccea83acb6fef453530bd4edb4ef?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`bc8be71`](https://app.codecov.io/gh/apache/pinot/commit/bc8be711fbc12ffb81de5f18762c260669b0874e?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 5 commits behind head on master.
| [Files with missing lines](https://app.codecov.io/gh/apache/pinot/pull/16903?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [...g/apache/pinot/query/mailbox/ReceivingMailbox.java](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&filepath=pinot-query-runtime%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fquery%2Fmailbox%2FReceivingMailbox.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3QtcXVlcnktcnVudGltZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvbWFpbGJveC9SZWNlaXZpbmdNYWlsYm94LmphdmE=) | 0.00% | [144 Missing :warning: ](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [.../query/mailbox/channel/MailboxContentObserver.java](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&filepath=pinot-query-runtime%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fquery%2Fmailbox%2Fchannel%2FMailboxContentObserver.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3QtcXVlcnktcnVudGltZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvbWFpbGJveC9jaGFubmVsL01haWxib3hDb250ZW50T2JzZXJ2ZXIuamF2YQ==) | 0.00% | [31 Missing :warning: ](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...he/pinot/query/mailbox/InMemorySendingMailbox.java](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&filepath=pinot-query-runtime%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fquery%2Fmailbox%2FInMemorySendingMailbox.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3QtcXVlcnktcnVudGltZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvbWFpbGJveC9Jbk1lbW9yeVNlbmRpbmdNYWlsYm94LmphdmE=) | 0.00% | [14 Missing :warning: ](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...org/apache/pinot/query/mailbox/MailboxService.java](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&filepath=pinot-query-runtime%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fquery%2Fmailbox%2FMailboxService.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3QtcXVlcnktcnVudGltZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvbWFpbGJveC9NYWlsYm94U2VydmljZS5qYXZh) | 0.00% | [3 Missing :warning: ](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...ot/query/runtime/operator/MailboxSendOperator.java](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&filepath=pinot-query-runtime%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fquery%2Fruntime%2Foperator%2FMailboxSendOperator.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3QtcXVlcnktcnVudGltZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvcnVudGltZS9vcGVyYXRvci9NYWlsYm94U2VuZE9wZXJhdG9yLmphdmE=) | 0.00% | [3 Missing :warning: ](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...apache/pinot/query/mailbox/GrpcSendingMailbox.java](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&filepath=pinot-query-runtime%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fquery%2Fmailbox%2FGrpcSendingMailbox.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3QtcXVlcnktcnVudGltZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvbWFpbGJveC9HcnBjU2VuZGluZ01haWxib3guamF2YQ==) | 0.00% | [1 Missing :warning: ](https://app.codecov.io/gh/apache/pinot/pull/16903?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
> :exclamation: There is a different number of reports uploaded between BASE (9baf132) and HEAD (bc8be71). Click for more details.
>
> <details><summary>HEAD has 32 uploads less than BASE</summary>
>
>| Flag | BASE (9baf132) | HEAD (bc8be71) |
>|------|------|------|
>|java-21|5|1|
>|unittests1|2|0|
>|unittests|4|1|
>|temurin|9|1|
>|java-11|4|0|
>|unittests2|2|1|
>|integration|5|0|
>|integration1|2|0|
>|custom-integration1|2|0|
>|integration2|1|0|
></details>
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## master #16903 +/- ##
=============================================
- Coverage 63.51% 33.57% -29.94%
+ Complexity 1411 704 -707
=============================================
Files 3073 3073
Lines 180742 180831 +89
Branches 27628 27642 +14
=============================================
- Hits 114802 60718 -54084
- Misses 57084 115118 +58034
+ Partials 8856 4995 -3861
```
| [Flag](https://app.codecov.io/gh/apache/pinot/pull/16903/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [custom-integration1](https://app.codecov.io/gh/apache/pinot/pull/16903/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [integration](https://app.codecov.io/gh/apache/pinot/pull/16903/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [integration1](https://app.codecov.io/gh/apache/pinot/pull/16903/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [integration2](https://app.codecov.io/gh/apache/pinot/pull/16903/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [java-11](https://app.codecov.io/gh/apache/pinot/pull/16903/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [java-21](https://app.codecov.io/gh/apache/pinot/pull/16903/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `33.57% <0.00%> (-29.91%)` | :arrow_down: |
| [temurin](https://app.codecov.io/gh/apache/pinot/pull/16903/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `33.57% <0.00%> (-29.94%)` | :arrow_down: |
| [unittests](https://app.codecov.io/gh/apache/pinot/pull/16903/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `33.57% <0.00%> (-29.94%)` | :arrow_down: |
| [unittests1](https://app.codecov.io/gh/apache/pinot/pull/16903/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [unittests2](https://app.codecov.io/gh/apache/pinot/pull/16903/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `33.57% <0.00%> (-0.05%)` | :arrow_down: |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
</details>
[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/apache/pinot/pull/16903?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2425547898
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java:
##########
@@ -49,6 +50,7 @@ public class MailboxContentObserver implements StreamObserver<MailboxContent> {
private final MailboxService _mailboxService;
private final StreamObserver<MailboxStatus> _responseObserver;
private final List<ByteBuffer> _mailboxBuffers = Collections.synchronizedList(new ArrayList<>());
+ private boolean _closedStream = false;
Review Comment:
> I think this needs to be volatile
It doesn't have to be volatile given StreamObservers should be thread-compatible. Given we never manually call methods on this class but instead it is GRPC the one that calls it, it is safe to assume thread-compatibility is honored.
> Do you find places where we close the stream multiple times?
It could happen if an incorrect sender sends a block after an EOS.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2402340254
##########
pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/ResourceBasedQueriesTest.java:
##########
@@ -60,6 +60,7 @@
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner;
import org.apache.pinot.spi.utils.JsonUtils;
import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+import org.assertj.core.api.Assertions;
Review Comment:
This is subjective
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2414727719
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/SendingMailbox.java:
##########
@@ -42,6 +42,7 @@ public interface SendingMailbox {
* and they should <b>not</b> acquire any resources when they are created. This method should throw if there was an
* error sending the data, since that would allow {@link BlockExchange} to exit early.
*/
+ // TODO: Remove throws as we don't throw these exceptions anymore
Review Comment:
We can do it as a separate PR, but this should be done immediately after this PR
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2409006519
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -71,31 +79,42 @@ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
sendPrivate(block, serializedStats);
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new QueryException(QueryErrorCode.INTERNAL, "Interrupted while sending data to mailbox: " + _id, e);
+ } catch (TimeoutException e) {
+ throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
+ String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
+ }
+ _isEarlyTerminated = status != ReceivingMailbox.ReceivingMailboxStatus.SUCCESS;
Review Comment:
Should this be handled under `WAITING_EOS`? That is the definition of early termination
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -71,31 +79,42 @@ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
sendPrivate(block, serializedStats);
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
Review Comment:
Do we need to interrupt again? Would it cause other exceptions being thrown from upstream?
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -42,7 +42,15 @@ public class InMemorySendingMailbox implements SendingMailbox {
private final long _deadlineMs;
private ReceivingMailbox _receivingMailbox;
+ /**
Review Comment:
(minor) Use markdown documentation?
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java:
##########
@@ -55,11 +55,14 @@ public class MailboxService {
CacheBuilder.newBuilder().expireAfterAccess(DANGLING_RECEIVING_MAILBOX_EXPIRY_SECONDS, TimeUnit.SECONDS)
.removalListener((RemovalListener<String, ReceivingMailbox>) notification -> {
if (notification.wasEvicted()) {
- int numPendingBlocks = notification.getValue().getNumPendingBlocks();
+ ReceivingMailbox receivingMailbox = notification.getValue();
+ int numPendingBlocks = receivingMailbox.getNumPendingBlocks();
if (numPendingBlocks > 0) {
LOGGER.warn("Evicting dangling receiving mailbox: {} with {} pending blocks", notification.getKey(),
numPendingBlocks);
}
+ // In case there is a leak, we should cancel the mailbox to unblock any waiters and release resources.
+ receivingMailbox.cancel();
Review Comment:
This is a core fix
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -268,7 +216,23 @@ public interface Reader {
}
public enum ReceivingMailboxStatus {
- SUCCESS, FIRST_ERROR, ERROR, TIMEOUT, CANCELLED, EARLY_TERMINATED
+ /// The block was successfully added to the mailbox.
+ ///
+ /// More blocks can be sent.
+ SUCCESS,
+ /// The block is rejected because downstream has early terminated and now is only waiting for EOS in order to
+ /// get the stats.
+ WAITING_EOS,
+ /// The received message is the last block the mailbox will ever read.
+ ///
+ /// This happens for example when an EOS block is added to the mailbox.
+ ///
+ /// No more blocks can be sent.
+ LAST_BLOCK,
+ /// The mailbox has been closed for write. Only EOS blocks are expected.
Review Comment:
Does it still expect EOS?
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -315,4 +279,377 @@ public List<DataBuffer> getSerializedStats() {
return _serializedStats;
}
}
+
+ /// The state of the queue.
+ ///
+ /// ```
+ /// +-------------------+ offerEos +-------------------+
+ /// | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
+ /// +-------------------+ +-------------------+
+ /// | |
+ /// | earlyTerminate | poll -- when all pending data is read
+ /// v v
+ /// +-------------------+ offerEos +-------------------+
+ /// | WAITING_EOS | -----------> | FULL_CLOSED |
+ /// +-------------------+ +-------------------+
+ /// ```
+ private enum State {
+ /// The queue is open for both read and write.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue, or null if the queue is empty.
+ /// - [#offer] accepts both data and EOS blocks.
+ ///
+ /// Transitions to [State#UPSTREAM_FINISHED] when an EOS block is offered or to [State#WAITING_EOS] when
+ /// [#earlyTerminate()] is called.
+ FULL_OPEN,
+ /// The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
+ ///
+ /// - [#poll()] returns null.
+ /// - [#offer] rejects all data blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when an EOS block is offered.
+ WAITING_EOS,
+ /// The upstream has indicated that no more data will be sent.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue and then the EOS block.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when the EOS block is read by [#poll()].
+ UPSTREAM_FINISHED,
+ /// The queue is closed for both read and write.
+ ///
+ /// - [#poll()] always returns the EOS block, which is always not null.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// No transitions out of this state.
+ FULL_CLOSED
+ }
+
+ /// This is a special bounded blocking queue implementation similar to ArrayBlockingQueue, but:
+ /// - Only accepts a single reader (aka downstream).
+ /// - Only accepts a multiple concurrent writers (aka upstream)
+ /// - Can be [closed for write][#closeForWrite(MseBlock.Eos, List)].
+ /// - Can be [#earlyTerminate()]d.
+ ///
+ /// Read the [State] enum to understand the different states and their transitions.
+ ///
+ /// All methods of this class are thread-safe and may block, although only [#offer] should block for a long time.
+ @ThreadSafe
+ private static class CancellableBlockingQueue {
+ private final String _id;
+ @Nullable
+ private volatile Reader _reader;
+ /// This is set when the queue is in [State#FULL_CLOSED] or [State#UPSTREAM_FINISHED].
+ @Nullable
+ @GuardedBy("_lock")
+ private MseBlockWithStats _eos;
+ /// The current state of the queue.
+ ///
+ /// All changes to this field must be done by calling [#changeState(State, String)] in order to log the state
+ /// transitions.
+ @GuardedBy("_lock")
+ private State _state = State.FULL_OPEN;
+ /// The items in the queue.
+ ///
+ /// This is a circular array where [#_putIndex] is the index to add the next item and [#_takeIndex] is the index to
+ /// take the next item from. Only data blocks are stored in this array, the EOS block is stored in [#_eos].
+ ///
+ /// Like in normal blocking queues, elements are added when upstream threads call [#offer] and removed when the
+ /// downstream thread calls [#poll]. Unlike normal blocking queues, elements will be [removed][#drainDataBlocks()]
+ /// when transitioning to [State#WAITING_EOS] or [State#FULL_CLOSED].
+ @GuardedBy("_lock")
+ private final MseBlock.Data[] _dataBlocks;
+ @GuardedBy("_lock")
+ private int _takeIndex;
+ @GuardedBy("_lock")
+ private int _putIndex;
+ @GuardedBy("_lock")
+ private int _count;
+ /// Threads waiting to add more data to the queue.
+ ///
+ /// This is used to prevent the following situation:
+ /// 1. The queue is full.
+ /// 2. Thread A tries to add data. Thread A will be blocked waiting for space in the queue.
+ /// 3. Thread B adds an EOS block, which will transition the queue to [State#UPSTREAM_FINISHED].
+ /// 4. Thread C reads data from the queue in a loop, the scheduler doesn't give time to Thread A.
+ /// 5. Thread C consumes all data from the queue and then reads the EOS block.
+ /// 6. Finally Thread A is unblocked and adds data to the queue, even though the queue is already closed for write
+ ///
+ /// As a result the block from A will be lost. Instead, we use this counter to return null in [#poll] when the
+ /// queue is empty but there are still threads trying to add data to the queue.
+ @GuardedBy("_lock")
+ private int _pendingData;
+ private final ReentrantLock _lock = new ReentrantLock();
+ private final Condition _notFull = _lock.newCondition();
+
+ public CancellableBlockingQueue(String id, int capacity) {
+ _id = id;
+ _dataBlocks = new MseBlock.Data[capacity];
+ }
+
+ /// Notifies the downstream that there is data to read.
+ private void notifyReader() {
+ Reader reader = _reader;
+ if (reader != null) {
+ LOGGER.debug("Notifying reader");
+ reader.blockReadyToRead();
+ } else {
+ LOGGER.debug("No reader to notify");
+ }
+ }
+
+ /// Offers a successful or erroneous EOS block into the queue, returning the status of the operation.
+ ///
+ /// This method never blocks for long, as it doesn't need to wait for space in the queue.
+ public ReceivingMailboxStatus offerEos(MseBlock.Eos block, List<DataBuffer> stats) {
+ ReentrantLock lock = _lock;
+ lock.lock();
+ try {
+ switch (_state) {
+ case FULL_CLOSED:
+ case UPSTREAM_FINISHED:
+ // The queue is closed for write. Always reject the block.
+ LOGGER.debug("Mailbox: {} is already closed for write, ignoring the late {} block", _id, block);
+ return ReceivingMailboxStatus.ALREADY_TERMINATED;
+ case WAITING_EOS:
+ // We got the EOS block we expected. Close the queue for both read and write.
+ changeState(State.FULL_CLOSED, "received EOS block");
+ _eos = new MseBlockWithStats(block, stats);
+ notifyReader();
+ return ReceivingMailboxStatus.LAST_BLOCK;
+ case FULL_OPEN:
+ changeState(State.UPSTREAM_FINISHED, "received EOS block");
+ _eos = new MseBlockWithStats(block, stats);
+ notifyReader();
+ if (block.isError()) {
+ drainDataBlocks();
+ _notFull.signal();
+ }
+ return ReceivingMailboxStatus.LAST_BLOCK;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ /// Offers a data block into the queue within the timeout specified, returning the status of the operation.
+ public ReceivingMailboxStatus offerData(MseBlock.Data block, long timeout, TimeUnit timeUnit)
+ throws InterruptedException, TimeoutException {
+ ReentrantLock lock = _lock;
+ lock.lockInterruptibly();
+ try {
+ while (true) {
+ switch (_state) {
+ case FULL_CLOSED:
+ case UPSTREAM_FINISHED:
+ // The queue is closed for write. Always reject the block.
+ LOGGER.debug("Mailbox: {} is already closed for write, ignoring the late data block", _id);
+ return ReceivingMailboxStatus.ALREADY_TERMINATED;
+ case WAITING_EOS:
+ // The downstream is not interested in reading more data.
+ LOGGER.debug("Mailbox: {} is not interesting in late data block", _id);
+ return ReceivingMailboxStatus.WAITING_EOS;
+ case FULL_OPEN:
+ if (offerDataToBuffer(block, timeout, timeUnit)) {
+ notifyReader();
+ return ReceivingMailboxStatus.SUCCESS;
+ }
+ // otherwise transitioned to FULL_CLOSED or WAITING_EOS while waiting for space in the queue
+ // and we need to re-evaluate the state
+ break;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ /// Offers a data block into the queue within the timeout specified, returning true if the block was added
+ /// successfully.
+ ///
+ /// This method can only be called while the queue is in the FULL_OPEN state and the lock is held.
+ ///
+ /// This method can time out, in which case we automatically transition to the [State#FULL_CLOSED] state.
+ /// But instead of returning false, we throw a [TimeoutException]. This is because the caller may want to
+ /// distinguish between a timeout and other reasons for not being able to add the block to the queue in order to
+ /// report different error messages.
+ ///
+ /// @return true if the block was added successfully, false if the state changed while waiting.
+ /// @throws InterruptedException if the thread is interrupted while waiting for space in the queue.
+ /// @throws TimeoutException if the timeout specified elapsed before space was available in the queue.
+ @GuardedBy("_lock")
+ private boolean offerDataToBuffer(MseBlock.Data block, long timeout, TimeUnit timeUnit)
+ throws InterruptedException, TimeoutException {
+
+ assert _state == State.FULL_OPEN;
+
+ long nanos = timeUnit.toNanos(timeout);
+ MseBlock.Data[] items = _dataBlocks;
+ _pendingData++;
+ try {
+ while (_count == items.length && nanos > 0L) {
+ nanos = _notFull.awaitNanos(nanos);
+
+ switch (_state) {
+ case FULL_OPEN: // we are in the same state, continue waiting for space
+ break;
+ case FULL_CLOSED:
+ case WAITING_EOS:
+ // The queue is closed and the reader is not interested in reading more data.
+ return false;
+ case UPSTREAM_FINISHED:
+ // Another thread offered the EOS while we were waiting for space.
+ assert _eos != null;
+ if (_eos._block.isSuccess()) { // If closed with EOS, the reader is still interested in reading our block
+ continue;
+ }
+ // if closed with an error, the reader is not interested in reading our block
+ return false;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ }
+ if (nanos <= 0L) { // timed out
+ String errorMessage = "Timed out while waiting for receive operator to consume data from mailbox: " + _id;
+ ErrorMseBlock timeoutBlock = ErrorMseBlock.fromError(QueryErrorCode.EXECUTION_TIMEOUT, errorMessage);
+ changeState(State.FULL_CLOSED, "timed out while waiting to offer data block");
+ drainDataBlocks();
+ _eos = new MseBlockWithStats(timeoutBlock, List.of());
+ notifyReader();
+ throw new TimeoutException(errorMessage);
+ }
+ items[_putIndex] = block;
+ if (++_putIndex == items.length) {
+ _putIndex = 0;
+ }
+ _count++;
+ return true;
+ } finally {
+ _pendingData--;
+ }
+ }
+
+ /// Returns the first block from the queue, or `null` if there is no block in the queue. The returned block will be
+ /// an error block if the queue has been cancelled or has encountered an error.
+ ///
+ /// This method may block briefly while acquiring the lock, but it doesn't actually require waiting for data in the
+ /// queue.
+ @Nullable
+ public MseBlockWithStats poll() {
+ Preconditions.checkState(_reader != null, "A reader must be registered");
+ ReentrantLock lock = _lock;
+ lock.lock();
+ try {
+ switch (_state) {
+ case FULL_CLOSED:
+ // The queue is closed for both read and write. Always return the error block.
+ assert _eos != null;
+ return _eos;
+ case WAITING_EOS:
+ // The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
+ // Polls returns null and only EOS blocks are accepted by offer.
+ assert _eos == null;
+ return null;
+ case UPSTREAM_FINISHED:
+ // The upstream has indicated that no more data will be sent. Poll returns pending blocks and then the EOS
+ // block.
+ if (_count == 0) {
+ if (_pendingData > 0) {
+ // There are still threads trying to add data to the queue. We should wait for them to finish.
+ LOGGER.debug("Mailbox: {} has pending {} data blocks, waiting for them to finish", _id, _pendingData);
+ return null;
+ } else {
+ changeState(State.FULL_CLOSED, "read all data blocks");
+ return _eos;
+ }
+ }
+ break;
+ case FULL_OPEN:
+ if (_count == 0) {
+ assert _eos == null;
+ return null;
+ }
+ break;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ assert _count > 0 : "if we reach here, there must be data in the queue";
+ MseBlock.Data[] items = _dataBlocks;
+ MseBlock.Data block = items[_takeIndex];
+ assert block != null : "data block in the queue must not be null";
+ items[_takeIndex] = null;
+ if (++_takeIndex == items.length) {
+ _takeIndex = 0;
+ }
+ _count--;
+ _notFull.signal();
+
+ return new MseBlockWithStats(block, List.of());
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ @GuardedBy("_lock")
+ private void changeState(State newState, String desc) {
+ LOGGER.debug("Mailbox: {} {}, transitioning from {} to {}", _id, desc, _state, newState);
+ _state = newState;
+ }
+
+ @GuardedBy("_lock")
+ private void drainDataBlocks() {
+ Arrays.fill(_dataBlocks, null);
+ _count = 0;
Review Comment:
Do you need to signal `_notFull` here?
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -315,4 +279,377 @@ public List<DataBuffer> getSerializedStats() {
return _serializedStats;
}
}
+
+ /// The state of the queue.
+ ///
+ /// ```
+ /// +-------------------+ offerEos +-------------------+
+ /// | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
+ /// +-------------------+ +-------------------+
+ /// | |
+ /// | earlyTerminate | poll -- when all pending data is read
+ /// v v
+ /// +-------------------+ offerEos +-------------------+
+ /// | WAITING_EOS | -----------> | FULL_CLOSED |
+ /// +-------------------+ +-------------------+
+ /// ```
+ private enum State {
+ /// The queue is open for both read and write.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue, or null if the queue is empty.
+ /// - [#offer] accepts both data and EOS blocks.
+ ///
+ /// Transitions to [State#UPSTREAM_FINISHED] when an EOS block is offered or to [State#WAITING_EOS] when
+ /// [#earlyTerminate()] is called.
+ FULL_OPEN,
+ /// The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
+ ///
+ /// - [#poll()] returns null.
+ /// - [#offer] rejects all data blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when an EOS block is offered.
+ WAITING_EOS,
+ /// The upstream has indicated that no more data will be sent.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue and then the EOS block.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when the EOS block is read by [#poll()].
+ UPSTREAM_FINISHED,
+ /// The queue is closed for both read and write.
+ ///
+ /// - [#poll()] always returns the EOS block, which is always not null.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// No transitions out of this state.
+ FULL_CLOSED
+ }
+
+ /// This is a special bounded blocking queue implementation similar to ArrayBlockingQueue, but:
+ /// - Only accepts a single reader (aka downstream).
+ /// - Only accepts a multiple concurrent writers (aka upstream)
+ /// - Can be [closed for write][#closeForWrite(MseBlock.Eos, List)].
+ /// - Can be [#earlyTerminate()]d.
+ ///
+ /// Read the [State] enum to understand the different states and their transitions.
+ ///
+ /// All methods of this class are thread-safe and may block, although only [#offer] should block for a long time.
+ @ThreadSafe
+ private static class CancellableBlockingQueue {
Review Comment:
This is the core data structure backing the new mailbox mechanism. Can you add some concurrent test to ensure it works as expected?
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/SendingMailbox.java:
##########
@@ -42,6 +42,7 @@ public interface SendingMailbox {
* and they should <b>not</b> acquire any resources when they are created. This method should throw if there was an
* error sending the data, since that would allow {@link BlockExchange} to exit early.
*/
+ // TODO: Remove throws as we don't throw these exceptions anymore
Review Comment:
Let's directly make the change, and document the expected behavior
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -71,31 +79,42 @@ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
sendPrivate(block, serializedStats);
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new QueryException(QueryErrorCode.INTERNAL, "Interrupted while sending data to mailbox: " + _id, e);
+ } catch (TimeoutException e) {
+ throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
+ String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
+ }
+ _isEarlyTerminated = status != ReceivingMailbox.ReceivingMailboxStatus.SUCCESS;
switch (status) {
case SUCCESS:
+ case WAITING_EOS:
+ break;
+ case LAST_BLOCK:
+ _isTerminated = true;
break;
- case CANCELLED:
- throw new QueryCancelledException(String.format("Mailbox: %s already cancelled from upstream", _id));
- case ERROR:
- throw new QueryException(QueryErrorCode.INTERNAL, String.format(
- "Mailbox: %s already errored out (received error block before)", _id));
- case TIMEOUT:
- throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
- String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
- case EARLY_TERMINATED:
- _isEarlyTerminated = true;
+ case ALREADY_TERMINATED:
+ if (_isTerminated) {
+ LOGGER.debug("Local mailbox received a late message once the stream was closed. This can happen due to "
+ + "race condition between sending the last block and closing the stream on the sender side");
+ } else {
+ throw new QueryException(QueryErrorCode.INTERNAL, String.format(
Review Comment:
Do we want to throw exception here, or simply set `isTerminated` to true? Do you see a path where we can hit this?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2419188611
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -71,31 +79,42 @@ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
sendPrivate(block, serializedStats);
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new QueryException(QueryErrorCode.INTERNAL, "Interrupted while sending data to mailbox: " + _id, e);
+ } catch (TimeoutException e) {
+ throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
+ String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
+ }
+ _isEarlyTerminated = status != ReceivingMailbox.ReceivingMailboxStatus.SUCCESS;
switch (status) {
case SUCCESS:
+ case WAITING_EOS:
+ break;
+ case LAST_BLOCK:
+ _isTerminated = true;
break;
- case CANCELLED:
- throw new QueryCancelledException(String.format("Mailbox: %s already cancelled from upstream", _id));
- case ERROR:
- throw new QueryException(QueryErrorCode.INTERNAL, String.format(
- "Mailbox: %s already errored out (received error block before)", _id));
- case TIMEOUT:
- throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
- String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
- case EARLY_TERMINATED:
- _isEarlyTerminated = true;
+ case ALREADY_TERMINATED:
+ if (_isTerminated) {
+ LOGGER.debug("Local mailbox received a late message once the stream was closed. This can happen due to "
+ + "race condition between sending the last block and closing the stream on the sender side");
+ } else {
+ throw new QueryException(QueryErrorCode.INTERNAL, String.format(
Review Comment:
I've changed this code to log in the same way we log in GRPC mailboxes. Now we log a message saying this should not happen
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2402194481
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -315,4 +298,377 @@ public List<DataBuffer> getSerializedStats() {
return _serializedStats;
}
}
+
+ /// The state of the queue.
+ ///
+ /// ```
+ /// +-------------------+ offerEos +-------------------+
+ /// | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
+ /// +-------------------+ +-------------------+
+ /// | |
+ /// | earlyTerminate | poll -- when all pending data is read
+ /// v v
+ /// +-------------------+ offerEos +-------------------+
+ /// | WAITING_EOS | -----------> | FULL_CLOSED |
+ /// +-------------------+ +-------------------+
+ /// ```
Review Comment:
The documentation uses /// for Javadoc comments instead of the standard /** */ format. According to the project's coding standards, Javadoc comments should use either /** */ or /// as specified in JEP-467, but consistency within the file should be maintained. The rest of the file uses /** */ format.
```suggestion
/**
* The state of the queue.
*
* <pre>
* +-------------------+ offerEos +-------------------+
* | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
* +-------------------+ +-------------------+
* | |
* | earlyTerminate | poll -- when all pending data is read
* v v
* +-------------------+ offerEos +-------------------+
* | WAITING_EOS | -----------> | FULL_CLOSED |
* +-------------------+ +-------------------+
* </pre>
*/
```
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -315,4 +298,377 @@ public List<DataBuffer> getSerializedStats() {
return _serializedStats;
}
}
+
+ /// The state of the queue.
+ ///
+ /// ```
+ /// +-------------------+ offerEos +-------------------+
+ /// | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
+ /// +-------------------+ +-------------------+
+ /// | |
+ /// | earlyTerminate | poll -- when all pending data is read
+ /// v v
+ /// +-------------------+ offerEos +-------------------+
+ /// | WAITING_EOS | -----------> | FULL_CLOSED |
+ /// +-------------------+ +-------------------+
+ /// ```
+ private enum State {
+ /// The queue is open for both read and write.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue, or null if the queue is empty.
+ /// - [#offer] accepts both data and EOS blocks.
+ ///
+ /// Transitions to [State#UPSTREAM_FINISHED] when an EOS block is offered or to [State#WAITING_EOS] when
+ /// [#earlyTerminate()] is called.
+ FULL_OPEN,
+ /// The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
+ ///
+ /// - [#poll()] returns null.
+ /// - [#offer] rejects all data blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when an EOS block is offered.
+ WAITING_EOS,
+ /// The upstream has indicated that no more data will be sent.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue and then the EOS block.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when the EOS block is read by [#poll()].
+ UPSTREAM_FINISHED,
+ /// The queue is closed for both read and write.
+ ///
+ /// - [#poll()] always returns the EOS block, which is always not null.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// No transitions out of this state.
+ FULL_CLOSED
+ }
+
+ /// This is a special bounded blocking queue implementation similar to ArrayBlockingQueue, but:
+ /// - Only accepts a single reader (aka downstream).
+ /// - Only accepts a multiple concurrent writers (aka upstream)
+ /// - Can be [closed for write][#closeForWrite(MseBlock.Eos, List)].
+ /// - Can be [#earlyTerminate()]d.
+ ///
+ /// Read the [State] enum to understand the different states and their transitions.
+ ///
+ /// All methods of this class are thread-safe and may block, although only [#offer] should block for a long time.
Review Comment:
Similar documentation style inconsistency. The comment uses /// format while the rest of the class uses /** */ format for documentation.
```suggestion
/**
* +-------------------+ offerEos +-------------------+
* | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
* +-------------------+ +-------------------+
* | |
* | earlyTerminate | poll -- when all pending data is read
* v v
* +-------------------+ offerEos +-------------------+
* | WAITING_EOS | -----------> | FULL_CLOSED |
* +-------------------+ +-------------------+
* ```
*/
private enum State {
/**
* The queue is open for both read and write.
*
* - [#poll()] returns the pending blocks in the queue, or null if the queue is empty.
* - [#offer] accepts both data and EOS blocks.
*
* Transitions to [State#UPSTREAM_FINISHED] when an EOS block is offered or to [State#WAITING_EOS] when
* [#earlyTerminate()] is called.
*/
FULL_OPEN,
/**
* The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
*
* - [#poll()] returns null.
* - [#offer] rejects all data blocks.
*
* Transitions to [State#FULL_CLOSED] when an EOS block is offered.
*/
WAITING_EOS,
/**
* The upstream has indicated that no more data will be sent.
*
* - [#poll()] returns the pending blocks in the queue and then the EOS block.
* - [#offer] rejects all blocks.
*
* Transitions to [State#FULL_CLOSED] when the EOS block is read by [#poll()].
*/
UPSTREAM_FINISHED,
/**
* The queue is closed for both read and write.
*
* - [#poll()] always returns the EOS block, which is always not null.
* - [#offer] rejects all blocks.
*
* No transitions out of this state.
*/
FULL_CLOSED
}
/**
* This is a special bounded blocking queue implementation similar to ArrayBlockingQueue, but:
* - Only accepts a single reader (aka downstream).
* - Only accepts a multiple concurrent writers (aka upstream)
* - Can be [closed for write][#closeForWrite(MseBlock.Eos, List)].
* - Can be [#earlyTerminate()]d.
*
* Read the [State] enum to understand the different states and their transitions.
*
* All methods of this class are thread-safe and may block, although only [#offer] should block for a long time.
*/
```
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -315,4 +298,377 @@ public List<DataBuffer> getSerializedStats() {
return _serializedStats;
}
}
+
+ /// The state of the queue.
+ ///
+ /// ```
+ /// +-------------------+ offerEos +-------------------+
+ /// | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
+ /// +-------------------+ +-------------------+
+ /// | |
+ /// | earlyTerminate | poll -- when all pending data is read
+ /// v v
+ /// +-------------------+ offerEos +-------------------+
+ /// | WAITING_EOS | -----------> | FULL_CLOSED |
+ /// +-------------------+ +-------------------+
+ /// ```
+ private enum State {
+ /// The queue is open for both read and write.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue, or null if the queue is empty.
+ /// - [#offer] accepts both data and EOS blocks.
+ ///
+ /// Transitions to [State#UPSTREAM_FINISHED] when an EOS block is offered or to [State#WAITING_EOS] when
+ /// [#earlyTerminate()] is called.
+ FULL_OPEN,
+ /// The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
+ ///
+ /// - [#poll()] returns null.
+ /// - [#offer] rejects all data blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when an EOS block is offered.
+ WAITING_EOS,
+ /// The upstream has indicated that no more data will be sent.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue and then the EOS block.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when the EOS block is read by [#poll()].
+ UPSTREAM_FINISHED,
+ /// The queue is closed for both read and write.
+ ///
+ /// - [#poll()] always returns the EOS block, which is always not null.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// No transitions out of this state.
+ FULL_CLOSED
+ }
+
+ /// This is a special bounded blocking queue implementation similar to ArrayBlockingQueue, but:
+ /// - Only accepts a single reader (aka downstream).
+ /// - Only accepts a multiple concurrent writers (aka upstream)
+ /// - Can be [closed for write][#closeForWrite(MseBlock.Eos, List)].
+ /// - Can be [#earlyTerminate()]d.
+ ///
+ /// Read the [State] enum to understand the different states and their transitions.
+ ///
+ /// All methods of this class are thread-safe and may block, although only [#offer] should block for a long time.
+ @ThreadSafe
+ private static class CancellableBlockingQueue {
+ private final String _id;
+ @Nullable
+ private volatile Reader _reader;
+ /// This is set when the queue is in [State#FULL_CLOSED] or [State#UPSTREAM_FINISHED].
+ @Nullable
+ @GuardedBy("_lock")
+ private MseBlockWithStats _eos;
+ /// The current state of the queue.
+ ///
+ /// All changes to this field must be done by calling [#changeState(State, String)] in order to log the state
+ /// transitions.
+ @GuardedBy("_lock")
+ private State _state = State.FULL_OPEN;
+ /// The items in the queue.
+ ///
+ /// This is a circular array where [#_putIndex] is the index to add the next item and [#_takeIndex] is the index to
+ /// take the next item from. Only data blocks are stored in this array, the EOS block is stored in [#_eos].
+ ///
+ /// Like in normal blocking queues, elements are added when upstream threads call [#offer] and removed when the
+ /// downstream thread calls [#poll]. Unlike normal blocking queues, elements will be [removed][#drainDataBlocks()]
+ /// when transitioning to [State#WAITING_EOS] or [State#FULL_CLOSED].
+ @GuardedBy("_lock")
+ private final MseBlock.Data[] _dataBlocks;
+ @GuardedBy("_lock")
+ private int _takeIndex;
+ @GuardedBy("_lock")
+ private int _putIndex;
+ @GuardedBy("_lock")
+ private int _count;
+ /// Threads waiting to add more data to the queue.
+ ///
+ /// This is used to prevent the following situation:
+ /// 1. The queue is full.
+ /// 2. Thread A tries to add data. Thread A will be blocked waiting for space in the queue.
+ /// 3. Thread B adds an EOS block, which will transition the queue to [State#UPSTREAM_FINISHED].
+ /// 4. Thread C reads data from the queue in a loop, the scheduler doesn't give time to Thread A.
+ /// 5. Thread C consumes all data from the queue and then reads the EOS block.
+ /// 6. Finally Thread A is unblocked and adds data to the queue, even though the queue is already closed for write
+ ///
+ /// As a result the block from A will be lost. Instead, we use this counter to return null in [#poll] when the
+ /// queue is empty but there are still threads trying to add data to the queue.
+ @GuardedBy("_lock")
+ private int _pendingData;
+ private final ReentrantLock _lock = new ReentrantLock();
+ private final Condition _notFull = _lock.newCondition();
+
+ public CancellableBlockingQueue(String id, int capacity) {
+ _id = id;
+ _dataBlocks = new MseBlock.Data[capacity];
+ }
+
+ /// Notifies the downstream that there is data to read.
+ private void notifyReader() {
+ Reader reader = _reader;
+ if (reader != null) {
+ LOGGER.debug("Notifying reader");
+ reader.blockReadyToRead();
+ } else {
+ LOGGER.debug("No reader to notify");
+ }
+ }
+
+ /// Offers a successful or erroneous EOS block into the queue, returning the status of the operation.
+ ///
+ /// This method never blocks for long, as it doesn't need to wait for space in the queue.
Review Comment:
More documentation using /// format instead of the consistent /** */ format used in the rest of the class.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -315,4 +298,377 @@ public List<DataBuffer> getSerializedStats() {
return _serializedStats;
}
}
+
+ /// The state of the queue.
+ ///
+ /// ```
+ /// +-------------------+ offerEos +-------------------+
+ /// | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
+ /// +-------------------+ +-------------------+
+ /// | |
+ /// | earlyTerminate | poll -- when all pending data is read
+ /// v v
+ /// +-------------------+ offerEos +-------------------+
+ /// | WAITING_EOS | -----------> | FULL_CLOSED |
+ /// +-------------------+ +-------------------+
+ /// ```
+ private enum State {
+ /// The queue is open for both read and write.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue, or null if the queue is empty.
+ /// - [#offer] accepts both data and EOS blocks.
+ ///
+ /// Transitions to [State#UPSTREAM_FINISHED] when an EOS block is offered or to [State#WAITING_EOS] when
+ /// [#earlyTerminate()] is called.
+ FULL_OPEN,
+ /// The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
+ ///
+ /// - [#poll()] returns null.
+ /// - [#offer] rejects all data blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when an EOS block is offered.
+ WAITING_EOS,
+ /// The upstream has indicated that no more data will be sent.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue and then the EOS block.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when the EOS block is read by [#poll()].
+ UPSTREAM_FINISHED,
+ /// The queue is closed for both read and write.
+ ///
+ /// - [#poll()] always returns the EOS block, which is always not null.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// No transitions out of this state.
+ FULL_CLOSED
+ }
+
+ /// This is a special bounded blocking queue implementation similar to ArrayBlockingQueue, but:
+ /// - Only accepts a single reader (aka downstream).
+ /// - Only accepts a multiple concurrent writers (aka upstream)
+ /// - Can be [closed for write][#closeForWrite(MseBlock.Eos, List)].
+ /// - Can be [#earlyTerminate()]d.
+ ///
+ /// Read the [State] enum to understand the different states and their transitions.
+ ///
+ /// All methods of this class are thread-safe and may block, although only [#offer] should block for a long time.
+ @ThreadSafe
+ private static class CancellableBlockingQueue {
+ private final String _id;
+ @Nullable
+ private volatile Reader _reader;
+ /// This is set when the queue is in [State#FULL_CLOSED] or [State#UPSTREAM_FINISHED].
+ @Nullable
+ @GuardedBy("_lock")
+ private MseBlockWithStats _eos;
+ /// The current state of the queue.
+ ///
+ /// All changes to this field must be done by calling [#changeState(State, String)] in order to log the state
+ /// transitions.
+ @GuardedBy("_lock")
+ private State _state = State.FULL_OPEN;
+ /// The items in the queue.
+ ///
+ /// This is a circular array where [#_putIndex] is the index to add the next item and [#_takeIndex] is the index to
+ /// take the next item from. Only data blocks are stored in this array, the EOS block is stored in [#_eos].
+ ///
+ /// Like in normal blocking queues, elements are added when upstream threads call [#offer] and removed when the
+ /// downstream thread calls [#poll]. Unlike normal blocking queues, elements will be [removed][#drainDataBlocks()]
+ /// when transitioning to [State#WAITING_EOS] or [State#FULL_CLOSED].
+ @GuardedBy("_lock")
+ private final MseBlock.Data[] _dataBlocks;
+ @GuardedBy("_lock")
+ private int _takeIndex;
+ @GuardedBy("_lock")
+ private int _putIndex;
+ @GuardedBy("_lock")
+ private int _count;
+ /// Threads waiting to add more data to the queue.
+ ///
+ /// This is used to prevent the following situation:
+ /// 1. The queue is full.
+ /// 2. Thread A tries to add data. Thread A will be blocked waiting for space in the queue.
+ /// 3. Thread B adds an EOS block, which will transition the queue to [State#UPSTREAM_FINISHED].
+ /// 4. Thread C reads data from the queue in a loop, the scheduler doesn't give time to Thread A.
+ /// 5. Thread C consumes all data from the queue and then reads the EOS block.
+ /// 6. Finally Thread A is unblocked and adds data to the queue, even though the queue is already closed for write
+ ///
+ /// As a result the block from A will be lost. Instead, we use this counter to return null in [#poll] when the
+ /// queue is empty but there are still threads trying to add data to the queue.
+ @GuardedBy("_lock")
+ private int _pendingData;
+ private final ReentrantLock _lock = new ReentrantLock();
+ private final Condition _notFull = _lock.newCondition();
+
+ public CancellableBlockingQueue(String id, int capacity) {
+ _id = id;
+ _dataBlocks = new MseBlock.Data[capacity];
+ }
+
+ /// Notifies the downstream that there is data to read.
+ private void notifyReader() {
Review Comment:
Documentation style inconsistency continues with /// format instead of the standard /** */ format.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -204,59 +191,39 @@ private synchronized void updateWaitCpuTime() {
* Sets an error block into the mailbox. No more blocks are accepted after calling this method.
*/
public void setErrorBlock(ErrorMseBlock errorBlock, List<DataBuffer> serializedStats) {
- if (_errorBlock.compareAndSet(null, new MseBlockWithStats(errorBlock, serializedStats))) {
- _blocks.clear();
- notifyReader();
- }
+ _blocks.offerEos(errorBlock, serializedStats);
}
/**
- * Returns the first block from the mailbox, or {@code null} if there is no block received yet. Error block is
- * returned if exists.
+ * Returns the first block from the mailbox, or {@code null} if there is no block received yet.
*/
@Nullable
public MseBlockWithStats poll() {
- Preconditions.checkState(_reader != null, "A reader must be registered");
- MseBlockWithStats errorBlock = _errorBlock.get();
- return errorBlock != null ? errorBlock : _blocks.poll();
+ return _blocks.poll();
}
/**
- * Early terminate the mailbox, called when upstream doesn't expect any more data block.
+ * Early terminate the mailbox, called when upstream doesn't expect any more <em>data</em> block.
*/
public void earlyTerminate() {
- _isEarlyTerminated = true;
+ _blocks.earlyTerminate();
}
/**
- * Cancels the mailbox. No more blocks are accepted after calling this method. Should only be called by the receive
- * operator to clean up the remaining blocks.
+ * Cancels the mailbox. No more blocks are accepted after calling this method and {@link #poll()} will always return
+ * an error block.
*/
public void cancel() {
LOGGER.debug("Cancelling mailbox: {}", _id);
- if (_errorBlock.get() == null) {
- MseBlockWithStats errorBlock = new MseBlockWithStats(
- ErrorMseBlock.fromError(QueryErrorCode.EXECUTION_TIMEOUT, "Cancelled by receiver"),
- Collections.emptyList());
- if (_errorBlock.compareAndSet(null, errorBlock)) {
- _cancelledErrorBlock = errorBlock;
- _blocks.clear();
- }
- }
+ _blocks.offerEos(ErrorMseBlock.fromException(null), List.of());
}
+ /// Returns the number of pending **data** blocks in the mailbox.
+ ///
+ /// EOS blocks are not counted because they will be stored separately and once returned, following calls to [poll]
+ /// will always return the same EOS block.
Review Comment:
More documentation style inconsistency with /// format instead of /** */ format used elsewhere in the class.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2402343937
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java:
##########
@@ -393,7 +395,7 @@ protected ReceivingMailbox.MseBlockWithStats onTimeout() {
return onException(terminateException.getErrorCode(), terminateException.getMessage());
}
// TODO: Add the sender stage id to the error message
- String errMsg = "Timed out on stage " + _stageId + " waiting for data sent by a child stage";
+ String errMsg = "Timed out on stage " + _stageId + " waiting for data from child stage " + _senderStageId;
Review Comment:
This is wrong. StringBuilders won't make the code easier to read and String.format is orders of magnitude slower than string concatenation
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2425549831
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/exchange/BlockExchange.java:
##########
@@ -122,11 +116,10 @@ public boolean send(MseBlock.Data block)
* API to send a block to the destination mailboxes.
* @param eosBlock the block to be transferred
* @return true if all the mailboxes has been early terminated.
- * @throws IOException when sending stream unexpectedly closed.
- * @throws TimeoutException when sending stream timeout.
+ * @throws org.apache.pinot.spi.exception.QueryException if any mailbox fails to send the block, including on timeout.
*/
- public boolean send(MseBlock.Eos eosBlock, List<DataBuffer> serializedStats)
- throws IOException, TimeoutException {
+ // TODO: Remove throws, as they are never thrown
Review Comment:
Nop. Removed
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -60,44 +67,50 @@ public boolean isLocal() {
}
@Override
- public void send(MseBlock.Data data)
- throws IOException, TimeoutException {
+ public void send(MseBlock.Data data) {
sendPrivate(data, Collections.emptyList());
}
@Override
- public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
- throws IOException, TimeoutException {
+ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats) {
sendPrivate(block, serializedStats);
_isTerminated = true;
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ // We are not restoring the interrupt status because we are already throwing an exception
+ // Code that catches this exception must finish the work fast enough to comply the interrupt contract
+ // See https://github.com/apache/pinot/pull/16903#discussion_r2409003423
+ throw new QueryException(QueryErrorCode.INTERNAL, "Interrupted while sending data to mailbox: " + _id, e);
+ } catch (TimeoutException e) {
+ throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT, "Timed out adding block into mailbox: " + _id
+ + " with timeout: " + Duration.of(timeoutMs, ChronoUnit.MILLIS), e);
+ }
switch (status) {
case SUCCESS:
break;
- case CANCELLED:
- throw new QueryCancelledException(String.format("Mailbox: %s already cancelled from upstream", _id));
- case ERROR:
- throw new QueryException(QueryErrorCode.INTERNAL, String.format(
- "Mailbox: %s already errored out (received error block before)", _id));
- case TIMEOUT:
- throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
- String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
- case EARLY_TERMINATED:
+ case WAITING_EOS:
_isEarlyTerminated = true;
break;
+ case LAST_BLOCK:
+ _isTerminated = true;
+ break;
+ case ALREADY_TERMINATED:
+ LOGGER.warn("Trying to offer blocks to the already closed mailbox {}. This should not happen", _id);
Review Comment:
Done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2409860185
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -42,7 +42,15 @@ public class InMemorySendingMailbox implements SendingMailbox {
private final long _deadlineMs;
private ReceivingMailbox _receivingMailbox;
+ /**
Review Comment:
Done
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -71,31 +79,42 @@ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
sendPrivate(block, serializedStats);
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
Review Comment:
We can skip this if and only if we guarantee that the code that catches the query exception won't block. Otherwise, we would miss the interruption call and the thread will continue to do its stuff.
The alternative (set the flag) means that if the catching logic needs to block on any method that checks the flag, it will fail. This is safer, as we make sure the thread is actually interrupted.
I would say in the current code, there is no difference between restoring or not restoring the flag. I would feel more confident about the future code if we restore it, but I'll remove it if you prefer to not restore it.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -71,31 +79,42 @@ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
sendPrivate(block, serializedStats);
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new QueryException(QueryErrorCode.INTERNAL, "Interrupted while sending data to mailbox: " + _id, e);
+ } catch (TimeoutException e) {
+ throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
+ String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
+ }
+ _isEarlyTerminated = status != ReceivingMailbox.ReceivingMailboxStatus.SUCCESS;
switch (status) {
case SUCCESS:
+ case WAITING_EOS:
+ break;
+ case LAST_BLOCK:
+ _isTerminated = true;
break;
- case CANCELLED:
- throw new QueryCancelledException(String.format("Mailbox: %s already cancelled from upstream", _id));
- case ERROR:
- throw new QueryException(QueryErrorCode.INTERNAL, String.format(
- "Mailbox: %s already errored out (received error block before)", _id));
- case TIMEOUT:
- throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
- String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
- case EARLY_TERMINATED:
- _isEarlyTerminated = true;
+ case ALREADY_TERMINATED:
+ if (_isTerminated) {
+ LOGGER.debug("Local mailbox received a late message once the stream was closed. This can happen due to "
+ + "race condition between sending the last block and closing the stream on the sender side");
+ } else {
+ throw new QueryException(QueryErrorCode.INTERNAL, String.format(
Review Comment:
If we reach this code, the implementation is not correct. ALREADY_TERMINATED is only returned when the receiving mailbox is in FULL_CLOSED or UPSTREAM_FINISHED. This means it could only have happened if the same thread sent an EOS and then more blocks, which should never have happened. IIRC this is similar to the older ERROR case.
What do we do in this situation? Throwing an exception is the best idea. In fact, I would even remove the `if (_isTerminated)` branch, as that is the impossible case. If `is_terminated` is true, we don't even offer the block, as we have a check at the beginning of this method.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -315,4 +279,377 @@ public List<DataBuffer> getSerializedStats() {
return _serializedStats;
}
}
+
+ /// The state of the queue.
+ ///
+ /// ```
+ /// +-------------------+ offerEos +-------------------+
+ /// | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
+ /// +-------------------+ +-------------------+
+ /// | |
+ /// | earlyTerminate | poll -- when all pending data is read
+ /// v v
+ /// +-------------------+ offerEos +-------------------+
+ /// | WAITING_EOS | -----------> | FULL_CLOSED |
+ /// +-------------------+ +-------------------+
+ /// ```
+ private enum State {
+ /// The queue is open for both read and write.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue, or null if the queue is empty.
+ /// - [#offer] accepts both data and EOS blocks.
+ ///
+ /// Transitions to [State#UPSTREAM_FINISHED] when an EOS block is offered or to [State#WAITING_EOS] when
+ /// [#earlyTerminate()] is called.
+ FULL_OPEN,
+ /// The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
+ ///
+ /// - [#poll()] returns null.
+ /// - [#offer] rejects all data blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when an EOS block is offered.
+ WAITING_EOS,
+ /// The upstream has indicated that no more data will be sent.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue and then the EOS block.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when the EOS block is read by [#poll()].
+ UPSTREAM_FINISHED,
+ /// The queue is closed for both read and write.
+ ///
+ /// - [#poll()] always returns the EOS block, which is always not null.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// No transitions out of this state.
+ FULL_CLOSED
+ }
+
+ /// This is a special bounded blocking queue implementation similar to ArrayBlockingQueue, but:
+ /// - Only accepts a single reader (aka downstream).
+ /// - Only accepts a multiple concurrent writers (aka upstream)
+ /// - Can be [closed for write][#closeForWrite(MseBlock.Eos, List)].
+ /// - Can be [#earlyTerminate()]d.
+ ///
+ /// Read the [State] enum to understand the different states and their transitions.
+ ///
+ /// All methods of this class are thread-safe and may block, although only [#offer] should block for a long time.
+ @ThreadSafe
+ private static class CancellableBlockingQueue {
Review Comment:
I've added some tests, but it is difficult to test all situations without using brute force (ie running the test multiple times)
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/SendingMailbox.java:
##########
@@ -42,6 +42,7 @@ public interface SendingMailbox {
* and they should <b>not</b> acquire any resources when they are created. This method should throw if there was an
* error sending the data, since that would allow {@link BlockExchange} to exit early.
*/
+ // TODO: Remove throws as we don't throw these exceptions anymore
Review Comment:
I already applied that change but reverted it to reduce the number of changes in this PR and keep the focus on the queue change, which is already complex.
But if you are ok, I can add that change back.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java:
##########
@@ -55,11 +55,14 @@ public class MailboxService {
CacheBuilder.newBuilder().expireAfterAccess(DANGLING_RECEIVING_MAILBOX_EXPIRY_SECONDS, TimeUnit.SECONDS)
.removalListener((RemovalListener<String, ReceivingMailbox>) notification -> {
if (notification.wasEvicted()) {
- int numPendingBlocks = notification.getValue().getNumPendingBlocks();
+ ReceivingMailbox receivingMailbox = notification.getValue();
+ int numPendingBlocks = receivingMailbox.getNumPendingBlocks();
if (numPendingBlocks > 0) {
LOGGER.warn("Evicting dangling receiving mailbox: {} with {} pending blocks", notification.getKey(),
numPendingBlocks);
}
+ // In case there is a leak, we should cancel the mailbox to unblock any waiters and release resources.
+ receivingMailbox.cancel();
Review Comment:
Probably this won't ever be called now, given we also unblock waiters when the receiver operator finishes (either successfully or not), but it is safer to keep it just in case we missed something
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -268,7 +216,23 @@ public interface Reader {
}
public enum ReceivingMailboxStatus {
- SUCCESS, FIRST_ERROR, ERROR, TIMEOUT, CANCELLED, EARLY_TERMINATED
+ /// The block was successfully added to the mailbox.
+ ///
+ /// More blocks can be sent.
+ SUCCESS,
+ /// The block is rejected because downstream has early terminated and now is only waiting for EOS in order to
+ /// get the stats.
+ WAITING_EOS,
+ /// The received message is the last block the mailbox will ever read.
+ ///
+ /// This happens for example when an EOS block is added to the mailbox.
+ ///
+ /// No more blocks can be sent.
+ LAST_BLOCK,
+ /// The mailbox has been closed for write. Only EOS blocks are expected.
Review Comment:
No, the javadoc is incorrect here. I've just fixed it
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -71,31 +79,42 @@ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
sendPrivate(block, serializedStats);
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new QueryException(QueryErrorCode.INTERNAL, "Interrupted while sending data to mailbox: " + _id, e);
+ } catch (TimeoutException e) {
+ throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
+ String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
+ }
+ _isEarlyTerminated = status != ReceivingMailbox.ReceivingMailboxStatus.SUCCESS;
Review Comment:
That is the same discussion we had in https://github.com/apache/pinot/pull/16919#discussion_r2393173235. I think it is safe to mark early terminated in all cases but SUCCESS, but let's keep it as you suggest
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -315,4 +279,377 @@ public List<DataBuffer> getSerializedStats() {
return _serializedStats;
}
}
+
+ /// The state of the queue.
+ ///
+ /// ```
+ /// +-------------------+ offerEos +-------------------+
+ /// | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
+ /// +-------------------+ +-------------------+
+ /// | |
+ /// | earlyTerminate | poll -- when all pending data is read
+ /// v v
+ /// +-------------------+ offerEos +-------------------+
+ /// | WAITING_EOS | -----------> | FULL_CLOSED |
+ /// +-------------------+ +-------------------+
+ /// ```
+ private enum State {
+ /// The queue is open for both read and write.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue, or null if the queue is empty.
+ /// - [#offer] accepts both data and EOS blocks.
+ ///
+ /// Transitions to [State#UPSTREAM_FINISHED] when an EOS block is offered or to [State#WAITING_EOS] when
+ /// [#earlyTerminate()] is called.
+ FULL_OPEN,
+ /// The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
+ ///
+ /// - [#poll()] returns null.
+ /// - [#offer] rejects all data blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when an EOS block is offered.
+ WAITING_EOS,
+ /// The upstream has indicated that no more data will be sent.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue and then the EOS block.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when the EOS block is read by [#poll()].
+ UPSTREAM_FINISHED,
+ /// The queue is closed for both read and write.
+ ///
+ /// - [#poll()] always returns the EOS block, which is always not null.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// No transitions out of this state.
+ FULL_CLOSED
+ }
+
+ /// This is a special bounded blocking queue implementation similar to ArrayBlockingQueue, but:
+ /// - Only accepts a single reader (aka downstream).
+ /// - Only accepts a multiple concurrent writers (aka upstream)
+ /// - Can be [closed for write][#closeForWrite(MseBlock.Eos, List)].
+ /// - Can be [#earlyTerminate()]d.
+ ///
+ /// Read the [State] enum to understand the different states and their transitions.
+ ///
+ /// All methods of this class are thread-safe and may block, although only [#offer] should block for a long time.
+ @ThreadSafe
+ private static class CancellableBlockingQueue {
+ private final String _id;
+ @Nullable
+ private volatile Reader _reader;
+ /// This is set when the queue is in [State#FULL_CLOSED] or [State#UPSTREAM_FINISHED].
+ @Nullable
+ @GuardedBy("_lock")
+ private MseBlockWithStats _eos;
+ /// The current state of the queue.
+ ///
+ /// All changes to this field must be done by calling [#changeState(State, String)] in order to log the state
+ /// transitions.
+ @GuardedBy("_lock")
+ private State _state = State.FULL_OPEN;
+ /// The items in the queue.
+ ///
+ /// This is a circular array where [#_putIndex] is the index to add the next item and [#_takeIndex] is the index to
+ /// take the next item from. Only data blocks are stored in this array, the EOS block is stored in [#_eos].
+ ///
+ /// Like in normal blocking queues, elements are added when upstream threads call [#offer] and removed when the
+ /// downstream thread calls [#poll]. Unlike normal blocking queues, elements will be [removed][#drainDataBlocks()]
+ /// when transitioning to [State#WAITING_EOS] or [State#FULL_CLOSED].
+ @GuardedBy("_lock")
+ private final MseBlock.Data[] _dataBlocks;
+ @GuardedBy("_lock")
+ private int _takeIndex;
+ @GuardedBy("_lock")
+ private int _putIndex;
+ @GuardedBy("_lock")
+ private int _count;
+ /// Threads waiting to add more data to the queue.
+ ///
+ /// This is used to prevent the following situation:
+ /// 1. The queue is full.
+ /// 2. Thread A tries to add data. Thread A will be blocked waiting for space in the queue.
+ /// 3. Thread B adds an EOS block, which will transition the queue to [State#UPSTREAM_FINISHED].
+ /// 4. Thread C reads data from the queue in a loop, the scheduler doesn't give time to Thread A.
+ /// 5. Thread C consumes all data from the queue and then reads the EOS block.
+ /// 6. Finally Thread A is unblocked and adds data to the queue, even though the queue is already closed for write
+ ///
+ /// As a result the block from A will be lost. Instead, we use this counter to return null in [#poll] when the
+ /// queue is empty but there are still threads trying to add data to the queue.
+ @GuardedBy("_lock")
+ private int _pendingData;
+ private final ReentrantLock _lock = new ReentrantLock();
+ private final Condition _notFull = _lock.newCondition();
+
+ public CancellableBlockingQueue(String id, int capacity) {
+ _id = id;
+ _dataBlocks = new MseBlock.Data[capacity];
+ }
+
+ /// Notifies the downstream that there is data to read.
+ private void notifyReader() {
+ Reader reader = _reader;
+ if (reader != null) {
+ LOGGER.debug("Notifying reader");
+ reader.blockReadyToRead();
+ } else {
+ LOGGER.debug("No reader to notify");
+ }
+ }
+
+ /// Offers a successful or erroneous EOS block into the queue, returning the status of the operation.
+ ///
+ /// This method never blocks for long, as it doesn't need to wait for space in the queue.
+ public ReceivingMailboxStatus offerEos(MseBlock.Eos block, List<DataBuffer> stats) {
+ ReentrantLock lock = _lock;
+ lock.lock();
+ try {
+ switch (_state) {
+ case FULL_CLOSED:
+ case UPSTREAM_FINISHED:
+ // The queue is closed for write. Always reject the block.
+ LOGGER.debug("Mailbox: {} is already closed for write, ignoring the late {} block", _id, block);
+ return ReceivingMailboxStatus.ALREADY_TERMINATED;
+ case WAITING_EOS:
+ // We got the EOS block we expected. Close the queue for both read and write.
+ changeState(State.FULL_CLOSED, "received EOS block");
+ _eos = new MseBlockWithStats(block, stats);
+ notifyReader();
+ return ReceivingMailboxStatus.LAST_BLOCK;
+ case FULL_OPEN:
+ changeState(State.UPSTREAM_FINISHED, "received EOS block");
+ _eos = new MseBlockWithStats(block, stats);
+ notifyReader();
+ if (block.isError()) {
+ drainDataBlocks();
+ _notFull.signal();
+ }
+ return ReceivingMailboxStatus.LAST_BLOCK;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ /// Offers a data block into the queue within the timeout specified, returning the status of the operation.
+ public ReceivingMailboxStatus offerData(MseBlock.Data block, long timeout, TimeUnit timeUnit)
+ throws InterruptedException, TimeoutException {
+ ReentrantLock lock = _lock;
+ lock.lockInterruptibly();
+ try {
+ while (true) {
+ switch (_state) {
+ case FULL_CLOSED:
+ case UPSTREAM_FINISHED:
+ // The queue is closed for write. Always reject the block.
+ LOGGER.debug("Mailbox: {} is already closed for write, ignoring the late data block", _id);
+ return ReceivingMailboxStatus.ALREADY_TERMINATED;
+ case WAITING_EOS:
+ // The downstream is not interested in reading more data.
+ LOGGER.debug("Mailbox: {} is not interesting in late data block", _id);
+ return ReceivingMailboxStatus.WAITING_EOS;
+ case FULL_OPEN:
+ if (offerDataToBuffer(block, timeout, timeUnit)) {
+ notifyReader();
+ return ReceivingMailboxStatus.SUCCESS;
+ }
+ // otherwise transitioned to FULL_CLOSED or WAITING_EOS while waiting for space in the queue
+ // and we need to re-evaluate the state
+ break;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ /// Offers a data block into the queue within the timeout specified, returning true if the block was added
+ /// successfully.
+ ///
+ /// This method can only be called while the queue is in the FULL_OPEN state and the lock is held.
+ ///
+ /// This method can time out, in which case we automatically transition to the [State#FULL_CLOSED] state.
+ /// But instead of returning false, we throw a [TimeoutException]. This is because the caller may want to
+ /// distinguish between a timeout and other reasons for not being able to add the block to the queue in order to
+ /// report different error messages.
+ ///
+ /// @return true if the block was added successfully, false if the state changed while waiting.
+ /// @throws InterruptedException if the thread is interrupted while waiting for space in the queue.
+ /// @throws TimeoutException if the timeout specified elapsed before space was available in the queue.
+ @GuardedBy("_lock")
+ private boolean offerDataToBuffer(MseBlock.Data block, long timeout, TimeUnit timeUnit)
+ throws InterruptedException, TimeoutException {
+
+ assert _state == State.FULL_OPEN;
+
+ long nanos = timeUnit.toNanos(timeout);
+ MseBlock.Data[] items = _dataBlocks;
+ _pendingData++;
+ try {
+ while (_count == items.length && nanos > 0L) {
+ nanos = _notFull.awaitNanos(nanos);
+
+ switch (_state) {
+ case FULL_OPEN: // we are in the same state, continue waiting for space
+ break;
+ case FULL_CLOSED:
+ case WAITING_EOS:
+ // The queue is closed and the reader is not interested in reading more data.
+ return false;
+ case UPSTREAM_FINISHED:
+ // Another thread offered the EOS while we were waiting for space.
+ assert _eos != null;
+ if (_eos._block.isSuccess()) { // If closed with EOS, the reader is still interested in reading our block
+ continue;
+ }
+ // if closed with an error, the reader is not interested in reading our block
+ return false;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ }
+ if (nanos <= 0L) { // timed out
+ String errorMessage = "Timed out while waiting for receive operator to consume data from mailbox: " + _id;
+ ErrorMseBlock timeoutBlock = ErrorMseBlock.fromError(QueryErrorCode.EXECUTION_TIMEOUT, errorMessage);
+ changeState(State.FULL_CLOSED, "timed out while waiting to offer data block");
+ drainDataBlocks();
+ _eos = new MseBlockWithStats(timeoutBlock, List.of());
+ notifyReader();
+ throw new TimeoutException(errorMessage);
+ }
+ items[_putIndex] = block;
+ if (++_putIndex == items.length) {
+ _putIndex = 0;
+ }
+ _count++;
+ return true;
+ } finally {
+ _pendingData--;
+ }
+ }
+
+ /// Returns the first block from the queue, or `null` if there is no block in the queue. The returned block will be
+ /// an error block if the queue has been cancelled or has encountered an error.
+ ///
+ /// This method may block briefly while acquiring the lock, but it doesn't actually require waiting for data in the
+ /// queue.
+ @Nullable
+ public MseBlockWithStats poll() {
+ Preconditions.checkState(_reader != null, "A reader must be registered");
+ ReentrantLock lock = _lock;
+ lock.lock();
+ try {
+ switch (_state) {
+ case FULL_CLOSED:
+ // The queue is closed for both read and write. Always return the error block.
+ assert _eos != null;
+ return _eos;
+ case WAITING_EOS:
+ // The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
+ // Polls returns null and only EOS blocks are accepted by offer.
+ assert _eos == null;
+ return null;
+ case UPSTREAM_FINISHED:
+ // The upstream has indicated that no more data will be sent. Poll returns pending blocks and then the EOS
+ // block.
+ if (_count == 0) {
+ if (_pendingData > 0) {
+ // There are still threads trying to add data to the queue. We should wait for them to finish.
+ LOGGER.debug("Mailbox: {} has pending {} data blocks, waiting for them to finish", _id, _pendingData);
+ return null;
+ } else {
+ changeState(State.FULL_CLOSED, "read all data blocks");
+ return _eos;
+ }
+ }
+ break;
+ case FULL_OPEN:
+ if (_count == 0) {
+ assert _eos == null;
+ return null;
+ }
+ break;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ assert _count > 0 : "if we reach here, there must be data in the queue";
+ MseBlock.Data[] items = _dataBlocks;
+ MseBlock.Data block = items[_takeIndex];
+ assert block != null : "data block in the queue must not be null";
+ items[_takeIndex] = null;
+ if (++_takeIndex == items.length) {
+ _takeIndex = 0;
+ }
+ _count--;
+ _notFull.signal();
+
+ return new MseBlockWithStats(block, List.of());
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ @GuardedBy("_lock")
+ private void changeState(State newState, String desc) {
+ LOGGER.debug("Mailbox: {} {}, transitioning from {} to {}", _id, desc, _state, newState);
+ _state = newState;
+ }
+
+ @GuardedBy("_lock")
+ private void drainDataBlocks() {
+ Arrays.fill(_dataBlocks, null);
+ _count = 0;
Review Comment:
Yes, you are right. I'm moving the signal from the caller of this method to this method. There were some callers who (incorrectly) didn't signal.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2414725991
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java:
##########
@@ -55,11 +55,14 @@ public class MailboxService {
CacheBuilder.newBuilder().expireAfterAccess(DANGLING_RECEIVING_MAILBOX_EXPIRY_SECONDS, TimeUnit.SECONDS)
.removalListener((RemovalListener<String, ReceivingMailbox>) notification -> {
if (notification.wasEvicted()) {
- int numPendingBlocks = notification.getValue().getNumPendingBlocks();
+ ReceivingMailbox receivingMailbox = notification.getValue();
+ int numPendingBlocks = receivingMailbox.getNumPendingBlocks();
if (numPendingBlocks > 0) {
LOGGER.warn("Evicting dangling receiving mailbox: {} with {} pending blocks", notification.getKey(),
numPendingBlocks);
}
+ // In case there is a leak, we should cancel the mailbox to unblock any waiters and release resources.
+ receivingMailbox.cancel();
Review Comment:
What could happen (especially when query is killed) is that receiver is never set up, and this mailbox will be dangling without a receiver pulling from it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2402251567
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -315,4 +280,377 @@ public List<DataBuffer> getSerializedStats() {
return _serializedStats;
}
}
+
+ /// The state of the queue.
+ ///
+ /// ```
+ /// +-------------------+ offerEos +-------------------+
+ /// | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
+ /// +-------------------+ +-------------------+
+ /// | |
+ /// | earlyTerminate | poll -- when all pending data is read
+ /// v v
+ /// +-------------------+ offerEos +-------------------+
+ /// | WAITING_EOS | -----------> | FULL_CLOSED |
+ /// +-------------------+ +-------------------+
+ /// ```
+ private enum State {
+ /// The queue is open for both read and write.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue, or null if the queue is empty.
+ /// - [#offer] accepts both data and EOS blocks.
+ ///
+ /// Transitions to [State#UPSTREAM_FINISHED] when an EOS block is offered or to [State#WAITING_EOS] when
+ /// [#earlyTerminate()] is called.
+ FULL_OPEN,
+ /// The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
+ ///
+ /// - [#poll()] returns null.
+ /// - [#offer] rejects all data blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when an EOS block is offered.
+ WAITING_EOS,
+ /// The upstream has indicated that no more data will be sent.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue and then the EOS block.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when the EOS block is read by [#poll()].
+ UPSTREAM_FINISHED,
+ /// The queue is closed for both read and write.
+ ///
+ /// - [#poll()] always returns the EOS block, which is always not null.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// No transitions out of this state.
+ FULL_CLOSED
+ }
+
+ /// This is a special bounded blocking queue implementation similar to ArrayBlockingQueue, but:
+ /// - Only accepts a single reader (aka downstream).
+ /// - Only accepts a multiple concurrent writers (aka upstream)
+ /// - Can be [closed for write][#closeForWrite(MseBlock.Eos, List)].
+ /// - Can be [#earlyTerminate()]d.
+ ///
+ /// Read the [State] enum to understand the different states and their transitions.
+ ///
+ /// All methods of this class are thread-safe and may block, although only [#offer] should block for a long time.
+ @ThreadSafe
+ private static class CancellableBlockingQueue {
+ private final String _id;
+ @Nullable
+ private volatile Reader _reader;
+ /// This is set when the queue is in [State#FULL_CLOSED] or [State#UPSTREAM_FINISHED].
+ @Nullable
+ @GuardedBy("_lock")
+ private MseBlockWithStats _eos;
+ /// The current state of the queue.
+ ///
+ /// All changes to this field must be done by calling [#changeState(State, String)] in order to log the state
+ /// transitions.
+ @GuardedBy("_lock")
+ private State _state = State.FULL_OPEN;
+ /// The items in the queue.
+ ///
+ /// This is a circular array where [#_putIndex] is the index to add the next item and [#_takeIndex] is the index to
+ /// take the next item from. Only data blocks are stored in this array, the EOS block is stored in [#_eos].
+ ///
+ /// Like in normal blocking queues, elements are added when upstream threads call [#offer] and removed when the
+ /// downstream thread calls [#poll]. Unlike normal blocking queues, elements will be [removed][#drainDataBlocks()]
+ /// when transitioning to [State#WAITING_EOS] or [State#FULL_CLOSED].
+ @GuardedBy("_lock")
+ private final MseBlock.Data[] _dataBlocks;
+ @GuardedBy("_lock")
+ private int _takeIndex;
+ @GuardedBy("_lock")
+ private int _putIndex;
+ @GuardedBy("_lock")
+ private int _count;
+ /// Threads waiting to add more data to the queue.
+ ///
+ /// This is used to prevent the following situation:
+ /// 1. The queue is full.
+ /// 2. Thread A tries to add data. Thread A will be blocked waiting for space in the queue.
+ /// 3. Thread B adds an EOS block, which will transition the queue to [State#UPSTREAM_FINISHED].
+ /// 4. Thread C reads data from the queue in a loop, the scheduler doesn't give time to Thread A.
+ /// 5. Thread C consumes all data from the queue and then reads the EOS block.
+ /// 6. Finally Thread A is unblocked and adds data to the queue, even though the queue is already closed for write
+ ///
+ /// As a result the block from A will be lost. Instead, we use this counter to return null in [#poll] when the
+ /// queue is empty but there are still threads trying to add data to the queue.
+ @GuardedBy("_lock")
+ private int _pendingData;
+ private final ReentrantLock _lock = new ReentrantLock();
+ private final Condition _notFull = _lock.newCondition();
+
+ public CancellableBlockingQueue(String id, int capacity) {
+ _id = id;
+ _dataBlocks = new MseBlock.Data[capacity];
+ }
+
+ /// Notifies the downstream that there is data to read.
+ private void notifyReader() {
+ Reader reader = _reader;
+ if (reader != null) {
+ LOGGER.debug("Notifying reader");
+ reader.blockReadyToRead();
+ } else {
+ LOGGER.debug("No reader to notify");
+ }
+ }
+
+ /// Offers a successful or erroneous EOS block into the queue, returning the status of the operation.
+ ///
+ /// This method never blocks for long, as it doesn't need to wait for space in the queue.
+ public ReceivingMailboxStatus offerEos(MseBlock.Eos block, List<DataBuffer> stats) {
+ ReentrantLock lock = _lock;
+ lock.lock();
+ try {
+ switch (_state) {
+ case FULL_CLOSED:
+ case UPSTREAM_FINISHED:
+ // The queue is closed for write. Always reject the block.
+ LOGGER.debug("Mailbox: {} is already closed for write, ignoring the late {} block", _id, block);
+ return ReceivingMailboxStatus.ALREADY_TERMINATED;
+ case WAITING_EOS:
+ // We got the EOS block we expected. Close the queue for both read and write.
+ changeState(State.FULL_CLOSED, "received EOS block");
+ _eos = new MseBlockWithStats(block, stats);
+ notifyReader();
+ return ReceivingMailboxStatus.LAST_BLOCK;
+ case FULL_OPEN:
+ changeState(State.UPSTREAM_FINISHED, "received EOS block");
+ _eos = new MseBlockWithStats(block, stats);
+ notifyReader();
+ if (block.isError()) {
+ drainDataBlocks();
+ _notFull.signal();
+ }
+ return ReceivingMailboxStatus.LAST_BLOCK;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ /// Offers a data block into the queue within the timeout specified, returning the status of the operation.
+ public ReceivingMailboxStatus offerData(MseBlock.Data block, long timeout, TimeUnit timeUnit)
+ throws InterruptedException, TimeoutException {
+ ReentrantLock lock = _lock;
+ lock.lockInterruptibly();
+ try {
+ while (true) {
+ switch (_state) {
+ case FULL_CLOSED:
+ case UPSTREAM_FINISHED:
+ // The queue is closed for write. Always reject the block.
+ LOGGER.debug("Mailbox: {} is already closed for write, ignoring the late data block", _id);
+ return ReceivingMailboxStatus.ALREADY_TERMINATED;
+ case WAITING_EOS:
+ // The downstream is not interested in reading more data.
+ LOGGER.debug("Mailbox: {} is not interesting in late data block", _id);
+ return ReceivingMailboxStatus.WAITING_EOS;
+ case FULL_OPEN:
+ if (offerDataToBuffer(block, timeout, timeUnit)) {
+ notifyReader();
+ return ReceivingMailboxStatus.SUCCESS;
+ }
+ // otherwise transitioned to FULL_CLOSED or WAITING_EOS while waiting for space in the queue
+ // and we need to re-evaluate the state
+ break;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ /// Offers a data block into the queue within the timeout specified, returning true if the block was added
+ /// successfully.
+ ///
+ /// This method can only be called while the queue is in the FULL_OPEN state and the lock is held.
+ ///
+ /// This method can time out, in which case we automatically transition to the [State#FULL_CLOSED] state.
+ /// But instead of returning false, we throw a [TimeoutException]. This is because the caller may want to
+ /// distinguish between a timeout and other reasons for not being able to add the block to the queue in order to
+ /// report different error messages.
+ ///
+ /// @return true if the block was added successfully, false if the state changed while waiting.
+ /// @throws InterruptedException if the thread is interrupted while waiting for space in the queue.
+ /// @throws TimeoutException if the timeout specified elapsed before space was available in the queue.
+ @GuardedBy("_lock")
+ private synchronized boolean offerDataToBuffer(MseBlock.Data block, long timeout, TimeUnit timeUnit)
Review Comment:
The method is marked as @GuardedBy(\"_lock\") but also has 'synchronized' modifier. Since the calling method already acquires the lock, the synchronized modifier is redundant and should be removed.
```suggestion
private boolean offerDataToBuffer(MseBlock.Data block, long timeout, TimeUnit timeUnit)
```
##########
pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/ResourceBasedQueriesTest.java:
##########
@@ -60,6 +60,7 @@
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner;
import org.apache.pinot.spi.utils.JsonUtils;
import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+import org.assertj.core.api.Assertions;
Review Comment:
[nitpick] Consider using a static import for Assertions.assertThat() since it's only used once in the file. This would be more consistent with the existing static imports like Assert.assertTrue().
```suggestion
import static org.assertj.core.api.Assertions.assertThat;
```
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -200,63 +178,34 @@ private synchronized void updateWaitCpuTime() {
_lastArriveTime = now;
}
- /**
- * Sets an error block into the mailbox. No more blocks are accepted after calling this method.
- */
+ /// Sets an error block into the mailbox. No more blocks are accepted after calling this method.
public void setErrorBlock(ErrorMseBlock errorBlock, List<DataBuffer> serializedStats) {
- if (_errorBlock.compareAndSet(null, new MseBlockWithStats(errorBlock, serializedStats))) {
- _blocks.clear();
- notifyReader();
- }
+ _blocks.offerEos(errorBlock, serializedStats);
}
- /**
- * Returns the first block from the mailbox, or {@code null} if there is no block received yet. Error block is
- * returned if exists.
- */
+ /// Returns the first block from the mailbox, or {@code null} if there is no block received yet.
@Nullable
public MseBlockWithStats poll() {
- Preconditions.checkState(_reader != null, "A reader must be registered");
- MseBlockWithStats errorBlock = _errorBlock.get();
- return errorBlock != null ? errorBlock : _blocks.poll();
+ return _blocks.poll();
}
- /**
- * Early terminate the mailbox, called when upstream doesn't expect any more data block.
- */
+ /// Early terminate the mailbox, called when upstream doesn't expect any more *data* block.
public void earlyTerminate() {
- _isEarlyTerminated = true;
+ _blocks.earlyTerminate();
}
- /**
- * Cancels the mailbox. No more blocks are accepted after calling this method. Should only be called by the receive
- * operator to clean up the remaining blocks.
- */
+ /// Cancels the mailbox. No more blocks are accepted after calling this method and [#poll] will always return
+ /// an error block.
public void cancel() {
LOGGER.debug("Cancelling mailbox: {}", _id);
- if (_errorBlock.get() == null) {
- MseBlockWithStats errorBlock = new MseBlockWithStats(
- ErrorMseBlock.fromError(QueryErrorCode.EXECUTION_TIMEOUT, "Cancelled by receiver"),
- Collections.emptyList());
- if (_errorBlock.compareAndSet(null, errorBlock)) {
- _cancelledErrorBlock = errorBlock;
- _blocks.clear();
- }
- }
+ _blocks.offerEos(ErrorMseBlock.fromException(null), List.of());
}
+ /// Returns the number of pending **data** blocks in the mailbox.
+ ///
+ /// will always return the same EOS block.
Review Comment:
This comment fragment appears to be incomplete or misplaced. It doesn't relate to the getNumPendingBlocks() method below it and should be completed or removed.
```suggestion
```
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java:
##########
@@ -393,7 +395,7 @@ protected ReceivingMailbox.MseBlockWithStats onTimeout() {
return onException(terminateException.getErrorCode(), terminateException.getMessage());
}
// TODO: Add the sender stage id to the error message
- String errMsg = "Timed out on stage " + _stageId + " waiting for data sent by a child stage";
+ String errMsg = "Timed out on stage " + _stageId + " waiting for data from child stage " + _senderStageId;
Review Comment:
Use String.format() or StringBuilder for string concatenation involving multiple variables to improve readability and performance.
```suggestion
String errMsg = String.format("Timed out on stage %d waiting for data from child stage %d", _stageId, _senderStageId);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2418849530
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -71,31 +79,42 @@ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
sendPrivate(block, serializedStats);
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
Review Comment:
I've modified this class and GrpcReceivingMailbox to do not restore the interruption flag and added a comment explaining why
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2414716298
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -71,31 +79,42 @@ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
sendPrivate(block, serializedStats);
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
Review Comment:
Given we are already handling the exception (converting it into a `QueryException`), I feel we should not restore it. My concern is that if we restore it, the caller might override the `QueryException` into another different interrupted exception, which prevents the original cause to be captured
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz merged PR #16903:
URL: https://github.com/apache/pinot/pull/16903
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2402345529
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -315,4 +280,377 @@ public List<DataBuffer> getSerializedStats() {
return _serializedStats;
}
}
+
+ /// The state of the queue.
+ ///
+ /// ```
+ /// +-------------------+ offerEos +-------------------+
+ /// | FULL_OPEN | -----------> | UPSTREAM_FINISHED|
+ /// +-------------------+ +-------------------+
+ /// | |
+ /// | earlyTerminate | poll -- when all pending data is read
+ /// v v
+ /// +-------------------+ offerEos +-------------------+
+ /// | WAITING_EOS | -----------> | FULL_CLOSED |
+ /// +-------------------+ +-------------------+
+ /// ```
+ private enum State {
+ /// The queue is open for both read and write.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue, or null if the queue is empty.
+ /// - [#offer] accepts both data and EOS blocks.
+ ///
+ /// Transitions to [State#UPSTREAM_FINISHED] when an EOS block is offered or to [State#WAITING_EOS] when
+ /// [#earlyTerminate()] is called.
+ FULL_OPEN,
+ /// The downstream is not interested in reading more data but is waiting for an EOS block to get the stats.
+ ///
+ /// - [#poll()] returns null.
+ /// - [#offer] rejects all data blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when an EOS block is offered.
+ WAITING_EOS,
+ /// The upstream has indicated that no more data will be sent.
+ ///
+ /// - [#poll()] returns the pending blocks in the queue and then the EOS block.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// Transitions to [State#FULL_CLOSED] when the EOS block is read by [#poll()].
+ UPSTREAM_FINISHED,
+ /// The queue is closed for both read and write.
+ ///
+ /// - [#poll()] always returns the EOS block, which is always not null.
+ /// - [#offer] rejects all blocks.
+ ///
+ /// No transitions out of this state.
+ FULL_CLOSED
+ }
+
+ /// This is a special bounded blocking queue implementation similar to ArrayBlockingQueue, but:
+ /// - Only accepts a single reader (aka downstream).
+ /// - Only accepts a multiple concurrent writers (aka upstream)
+ /// - Can be [closed for write][#closeForWrite(MseBlock.Eos, List)].
+ /// - Can be [#earlyTerminate()]d.
+ ///
+ /// Read the [State] enum to understand the different states and their transitions.
+ ///
+ /// All methods of this class are thread-safe and may block, although only [#offer] should block for a long time.
+ @ThreadSafe
+ private static class CancellableBlockingQueue {
+ private final String _id;
+ @Nullable
+ private volatile Reader _reader;
+ /// This is set when the queue is in [State#FULL_CLOSED] or [State#UPSTREAM_FINISHED].
+ @Nullable
+ @GuardedBy("_lock")
+ private MseBlockWithStats _eos;
+ /// The current state of the queue.
+ ///
+ /// All changes to this field must be done by calling [#changeState(State, String)] in order to log the state
+ /// transitions.
+ @GuardedBy("_lock")
+ private State _state = State.FULL_OPEN;
+ /// The items in the queue.
+ ///
+ /// This is a circular array where [#_putIndex] is the index to add the next item and [#_takeIndex] is the index to
+ /// take the next item from. Only data blocks are stored in this array, the EOS block is stored in [#_eos].
+ ///
+ /// Like in normal blocking queues, elements are added when upstream threads call [#offer] and removed when the
+ /// downstream thread calls [#poll]. Unlike normal blocking queues, elements will be [removed][#drainDataBlocks()]
+ /// when transitioning to [State#WAITING_EOS] or [State#FULL_CLOSED].
+ @GuardedBy("_lock")
+ private final MseBlock.Data[] _dataBlocks;
+ @GuardedBy("_lock")
+ private int _takeIndex;
+ @GuardedBy("_lock")
+ private int _putIndex;
+ @GuardedBy("_lock")
+ private int _count;
+ /// Threads waiting to add more data to the queue.
+ ///
+ /// This is used to prevent the following situation:
+ /// 1. The queue is full.
+ /// 2. Thread A tries to add data. Thread A will be blocked waiting for space in the queue.
+ /// 3. Thread B adds an EOS block, which will transition the queue to [State#UPSTREAM_FINISHED].
+ /// 4. Thread C reads data from the queue in a loop, the scheduler doesn't give time to Thread A.
+ /// 5. Thread C consumes all data from the queue and then reads the EOS block.
+ /// 6. Finally Thread A is unblocked and adds data to the queue, even though the queue is already closed for write
+ ///
+ /// As a result the block from A will be lost. Instead, we use this counter to return null in [#poll] when the
+ /// queue is empty but there are still threads trying to add data to the queue.
+ @GuardedBy("_lock")
+ private int _pendingData;
+ private final ReentrantLock _lock = new ReentrantLock();
+ private final Condition _notFull = _lock.newCondition();
+
+ public CancellableBlockingQueue(String id, int capacity) {
+ _id = id;
+ _dataBlocks = new MseBlock.Data[capacity];
+ }
+
+ /// Notifies the downstream that there is data to read.
+ private void notifyReader() {
+ Reader reader = _reader;
+ if (reader != null) {
+ LOGGER.debug("Notifying reader");
+ reader.blockReadyToRead();
+ } else {
+ LOGGER.debug("No reader to notify");
+ }
+ }
+
+ /// Offers a successful or erroneous EOS block into the queue, returning the status of the operation.
+ ///
+ /// This method never blocks for long, as it doesn't need to wait for space in the queue.
+ public ReceivingMailboxStatus offerEos(MseBlock.Eos block, List<DataBuffer> stats) {
+ ReentrantLock lock = _lock;
+ lock.lock();
+ try {
+ switch (_state) {
+ case FULL_CLOSED:
+ case UPSTREAM_FINISHED:
+ // The queue is closed for write. Always reject the block.
+ LOGGER.debug("Mailbox: {} is already closed for write, ignoring the late {} block", _id, block);
+ return ReceivingMailboxStatus.ALREADY_TERMINATED;
+ case WAITING_EOS:
+ // We got the EOS block we expected. Close the queue for both read and write.
+ changeState(State.FULL_CLOSED, "received EOS block");
+ _eos = new MseBlockWithStats(block, stats);
+ notifyReader();
+ return ReceivingMailboxStatus.LAST_BLOCK;
+ case FULL_OPEN:
+ changeState(State.UPSTREAM_FINISHED, "received EOS block");
+ _eos = new MseBlockWithStats(block, stats);
+ notifyReader();
+ if (block.isError()) {
+ drainDataBlocks();
+ _notFull.signal();
+ }
+ return ReceivingMailboxStatus.LAST_BLOCK;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ /// Offers a data block into the queue within the timeout specified, returning the status of the operation.
+ public ReceivingMailboxStatus offerData(MseBlock.Data block, long timeout, TimeUnit timeUnit)
+ throws InterruptedException, TimeoutException {
+ ReentrantLock lock = _lock;
+ lock.lockInterruptibly();
+ try {
+ while (true) {
+ switch (_state) {
+ case FULL_CLOSED:
+ case UPSTREAM_FINISHED:
+ // The queue is closed for write. Always reject the block.
+ LOGGER.debug("Mailbox: {} is already closed for write, ignoring the late data block", _id);
+ return ReceivingMailboxStatus.ALREADY_TERMINATED;
+ case WAITING_EOS:
+ // The downstream is not interested in reading more data.
+ LOGGER.debug("Mailbox: {} is not interesting in late data block", _id);
+ return ReceivingMailboxStatus.WAITING_EOS;
+ case FULL_OPEN:
+ if (offerDataToBuffer(block, timeout, timeUnit)) {
+ notifyReader();
+ return ReceivingMailboxStatus.SUCCESS;
+ }
+ // otherwise transitioned to FULL_CLOSED or WAITING_EOS while waiting for space in the queue
+ // and we need to re-evaluate the state
+ break;
+ default:
+ throw new IllegalStateException("Unexpected state: " + _state);
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ /// Offers a data block into the queue within the timeout specified, returning true if the block was added
+ /// successfully.
+ ///
+ /// This method can only be called while the queue is in the FULL_OPEN state and the lock is held.
+ ///
+ /// This method can time out, in which case we automatically transition to the [State#FULL_CLOSED] state.
+ /// But instead of returning false, we throw a [TimeoutException]. This is because the caller may want to
+ /// distinguish between a timeout and other reasons for not being able to add the block to the queue in order to
+ /// report different error messages.
+ ///
+ /// @return true if the block was added successfully, false if the state changed while waiting.
+ /// @throws InterruptedException if the thread is interrupted while waiting for space in the queue.
+ /// @throws TimeoutException if the timeout specified elapsed before space was available in the queue.
+ @GuardedBy("_lock")
+ private synchronized boolean offerDataToBuffer(MseBlock.Data block, long timeout, TimeUnit timeUnit)
Review Comment:
Fixed. I missed that synchronize from a previous implementation
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java:
##########
@@ -200,63 +178,34 @@ private synchronized void updateWaitCpuTime() {
_lastArriveTime = now;
}
- /**
- * Sets an error block into the mailbox. No more blocks are accepted after calling this method.
- */
+ /// Sets an error block into the mailbox. No more blocks are accepted after calling this method.
public void setErrorBlock(ErrorMseBlock errorBlock, List<DataBuffer> serializedStats) {
- if (_errorBlock.compareAndSet(null, new MseBlockWithStats(errorBlock, serializedStats))) {
- _blocks.clear();
- notifyReader();
- }
+ _blocks.offerEos(errorBlock, serializedStats);
}
- /**
- * Returns the first block from the mailbox, or {@code null} if there is no block received yet. Error block is
- * returned if exists.
- */
+ /// Returns the first block from the mailbox, or {@code null} if there is no block received yet.
@Nullable
public MseBlockWithStats poll() {
- Preconditions.checkState(_reader != null, "A reader must be registered");
- MseBlockWithStats errorBlock = _errorBlock.get();
- return errorBlock != null ? errorBlock : _blocks.poll();
+ return _blocks.poll();
}
- /**
- * Early terminate the mailbox, called when upstream doesn't expect any more data block.
- */
+ /// Early terminate the mailbox, called when upstream doesn't expect any more *data* block.
public void earlyTerminate() {
- _isEarlyTerminated = true;
+ _blocks.earlyTerminate();
}
- /**
- * Cancels the mailbox. No more blocks are accepted after calling this method. Should only be called by the receive
- * operator to clean up the remaining blocks.
- */
+ /// Cancels the mailbox. No more blocks are accepted after calling this method and [#poll] will always return
+ /// an error block.
public void cancel() {
LOGGER.debug("Cancelling mailbox: {}", _id);
- if (_errorBlock.get() == null) {
- MseBlockWithStats errorBlock = new MseBlockWithStats(
- ErrorMseBlock.fromError(QueryErrorCode.EXECUTION_TIMEOUT, "Cancelled by receiver"),
- Collections.emptyList());
- if (_errorBlock.compareAndSet(null, errorBlock)) {
- _cancelledErrorBlock = errorBlock;
- _blocks.clear();
- }
- }
+ _blocks.offerEos(ErrorMseBlock.fromException(null), List.of());
}
+ /// Returns the number of pending **data** blocks in the mailbox.
+ ///
+ /// will always return the same EOS block.
Review Comment:
Fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2414722289
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -71,31 +79,42 @@ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
sendPrivate(block, serializedStats);
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new QueryException(QueryErrorCode.INTERNAL, "Interrupted while sending data to mailbox: " + _id, e);
+ } catch (TimeoutException e) {
+ throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
+ String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
+ }
+ _isEarlyTerminated = status != ReceivingMailbox.ReceivingMailboxStatus.SUCCESS;
switch (status) {
case SUCCESS:
+ case WAITING_EOS:
+ break;
+ case LAST_BLOCK:
+ _isTerminated = true;
break;
- case CANCELLED:
- throw new QueryCancelledException(String.format("Mailbox: %s already cancelled from upstream", _id));
- case ERROR:
- throw new QueryException(QueryErrorCode.INTERNAL, String.format(
- "Mailbox: %s already errored out (received error block before)", _id));
- case TIMEOUT:
- throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
- String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
- case EARLY_TERMINATED:
- _isEarlyTerminated = true;
+ case ALREADY_TERMINATED:
+ if (_isTerminated) {
+ LOGGER.debug("Local mailbox received a late message once the stream was closed. This can happen due to "
+ + "race condition between sending the last block and closing the stream on the sender side");
+ } else {
+ throw new QueryException(QueryErrorCode.INTERNAL, String.format(
Review Comment:
I see. Should we consider making the error more explicit about this shouldn't happen and indicate a bug in the code? Also probably log an error so that it can be discovered easier
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2422224901
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java:
##########
@@ -49,6 +50,7 @@ public class MailboxContentObserver implements StreamObserver<MailboxContent> {
private final MailboxService _mailboxService;
private final StreamObserver<MailboxStatus> _responseObserver;
private final List<ByteBuffer> _mailboxBuffers = Collections.synchronizedList(new ArrayList<>());
+ private boolean _closedStream = false;
Review Comment:
I think this needs to be `volatile`
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java:
##########
@@ -336,9 +335,6 @@ private void notifyErrorAfterSubmission(int stageId, ErrorMseBlock errorBlock, W
// current stage. We will need to fix this in future, but for now, we are sending the error block without
// the stats.
sendingMailbox.send(errorBlock, Collections.emptyList());
- } catch (TimeoutException e) {
- LOGGER.warn("Timed out sending error block to mailbox: {} for request: {}, stage: {}",
- routingInfo.getMailboxId(), requestId, stageId, e);
} catch (Exception e) {
Review Comment:
Do we want to check the response type to log error only with unexpected error?
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java:
##########
@@ -60,44 +67,50 @@ public boolean isLocal() {
}
@Override
- public void send(MseBlock.Data data)
- throws IOException, TimeoutException {
+ public void send(MseBlock.Data data) {
sendPrivate(data, Collections.emptyList());
}
@Override
- public void send(MseBlock.Eos block, List<DataBuffer> serializedStats)
- throws IOException, TimeoutException {
+ public void send(MseBlock.Eos block, List<DataBuffer> serializedStats) {
sendPrivate(block, serializedStats);
_isTerminated = true;
}
- private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats)
- throws TimeoutException {
+ private void sendPrivate(MseBlock block, List<DataBuffer> serializedStats) {
if (isTerminated() || (isEarlyTerminated() && block.isData())) {
+ LOGGER.debug("Mailbox {} already terminated, ignoring block {}", _id, block);
return;
}
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
_statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
- ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
-
+ ReceivingMailbox.ReceivingMailboxStatus status;
+ try {
+ status = _receivingMailbox.offer(block, serializedStats, timeoutMs);
+ } catch (InterruptedException e) {
+ // We are not restoring the interrupt status because we are already throwing an exception
+ // Code that catches this exception must finish the work fast enough to comply the interrupt contract
+ // See https://github.com/apache/pinot/pull/16903#discussion_r2409003423
+ throw new QueryException(QueryErrorCode.INTERNAL, "Interrupted while sending data to mailbox: " + _id, e);
+ } catch (TimeoutException e) {
+ throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT, "Timed out adding block into mailbox: " + _id
+ + " with timeout: " + Duration.of(timeoutMs, ChronoUnit.MILLIS), e);
+ }
switch (status) {
case SUCCESS:
break;
- case CANCELLED:
- throw new QueryCancelledException(String.format("Mailbox: %s already cancelled from upstream", _id));
- case ERROR:
- throw new QueryException(QueryErrorCode.INTERNAL, String.format(
- "Mailbox: %s already errored out (received error block before)", _id));
- case TIMEOUT:
- throw new QueryException(QueryErrorCode.EXECUTION_TIMEOUT,
- String.format("Timed out adding block into mailbox: %s with timeout: %dms", _id, timeoutMs));
- case EARLY_TERMINATED:
+ case WAITING_EOS:
_isEarlyTerminated = true;
break;
+ case LAST_BLOCK:
+ _isTerminated = true;
+ break;
+ case ALREADY_TERMINATED:
+ LOGGER.warn("Trying to offer blocks to the already closed mailbox {}. This should not happen", _id);
Review Comment:
Consider making it an `error` since it should not happen
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/exchange/BlockExchange.java:
##########
@@ -122,11 +116,10 @@ public boolean send(MseBlock.Data block)
* API to send a block to the destination mailboxes.
* @param eosBlock the block to be transferred
* @return true if all the mailboxes has been early terminated.
- * @throws IOException when sending stream unexpectedly closed.
- * @throws TimeoutException when sending stream timeout.
+ * @throws org.apache.pinot.spi.exception.QueryException if any mailbox fails to send the block, including on timeout.
*/
- public boolean send(MseBlock.Eos eosBlock, List<DataBuffer> serializedStats)
- throws IOException, TimeoutException {
+ // TODO: Remove throws, as they are never thrown
Review Comment:
Does this still apply?
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java:
##########
@@ -49,6 +50,7 @@ public class MailboxContentObserver implements StreamObserver<MailboxContent> {
private final MailboxService _mailboxService;
private final StreamObserver<MailboxStatus> _responseObserver;
private final List<ByteBuffer> _mailboxBuffers = Collections.synchronizedList(new ArrayList<>());
+ private boolean _closedStream = false;
Review Comment:
Do you find places where we close the stream multiple times?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Rewrite ReceivingMailbox to be able to unblock writers on cancellation or error [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16903:
URL: https://github.com/apache/pinot/pull/16903#discussion_r2425570239
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java:
##########
@@ -336,9 +335,6 @@ private void notifyErrorAfterSubmission(int stageId, ErrorMseBlock errorBlock, W
// current stage. We will need to fix this in future, but for now, we are sending the error block without
// the stats.
sendingMailbox.send(errorBlock, Collections.emptyList());
- } catch (TimeoutException e) {
- LOGGER.warn("Timed out sending error block to mailbox: {} for request: {}, stage: {}",
- routingInfo.getMailboxId(), requestId, stageId, e);
} catch (Exception e) {
Review Comment:
I've added some error handling for query exceptions. The main change is that EXECUTION_TIMEOUT is logged in the same way we used to log TimeoutExceptions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org