yashmayya commented on code in PR #18458:
URL: https://github.com/apache/pinot/pull/18458#discussion_r3359228673
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java:
##########
@@ -450,6 +454,296 @@ private void
submitTimeSeriesInternal(Worker.TimeSeriesQueryRequest request,
}
}
+ /// Stream-mode submission handler. The broker keeps the stream open for the
query lifetime; the server replies
+ /// with a {@code submit_ack} as the first message and (in subsequent
commits) per-opchain
+ /// {@link Worker.OpChainComplete} messages followed by a final {@link
Worker.ServerDone}.
+ ///
+ /// This skeleton wires up the gRPC mechanics + plan submission via the
existing submission path. It does NOT yet
+ /// emit OpChainComplete / ServerDone — those need a per-opchain completion
hook on
+ /// {@link org.apache.pinot.query.runtime.executor.OpChainSchedulerService},
which is layered on next.
+ /// Cancel still routes through the existing unary {@link
#cancel(Worker.CancelRequest, StreamObserver)} RPC; broker
+ /// stream-close also triggers a cancel here.
+ @Override
+ public StreamObserver<Worker.BrokerToServer> submitWithStream(
+ StreamObserver<Worker.ServerToBroker> responseObserver) {
+ return new SubmitWithStreamObserver(responseObserver);
+ }
+
+ /// Per-query state for an open {@code SubmitWithStream} call. Owns the
response stream and serialises every
+ /// {@code onNext} call on it via a {@code synchronized} block — gRPC
requires {@code StreamObserver.onNext} to be
+ /// called serially.
+ ///
+ /// Tracks the expected number of opchains for the request (sum of
WorkerMetadata across all stages). An
+ /// {@link OpChainCompletionListener} registered with {@link
QueryRunner#registerOpChainCompletionListener}
+ /// fires once per opchain finishing, encodes its stats via {@link
MultiStageStatsTreeEncoder}, and emits an
+ /// {@link Worker.OpChainComplete} on the response stream. When the
per-request completed-count reaches the
+ /// expected total, {@link Worker.ServerDone} is emitted and the stream is
closed.
+ ///
+ /// All blocking work (plan deserialization, opchain construction) runs on
+ /// {@link QueryServer#_submissionExecutorService}.
+ private final class SubmitWithStreamObserver implements
StreamObserver<Worker.BrokerToServer> {
+ private final StreamObserver<Worker.ServerToBroker> _responseObserver;
+ /// Serialises onNext calls on the response stream and guards mutable
session state.
+ private final Object _streamLock = new Object();
+ /// True once we've received the first {@code submit} and dispatched it.
+ private final AtomicBoolean _submitted = new AtomicBoolean(false);
+ /// True once we've completed the response stream (success or error).
Idempotent guard.
+ private final AtomicBoolean _completed = new AtomicBoolean(false);
+ /// Number of opchains we expect to report for this request — set after we
deserialize the plan.
+ private final AtomicInteger _expectedOpChains = new AtomicInteger(-1);
+ /// Number of opchains that have reported so far via the completion
listener.
+ private final AtomicInteger _completedOpChains = new AtomicInteger(0);
+ /// Set once we successfully parse the request id from the submit
metadata. Used by cancel-via-stream.
+ private volatile long _requestId = -1;
+ /// Completed once {@link #sendSubmitAck} has been called (success or
error path). Guards the ack/done race:
+ /// if a trivial opchain finishes before the {@code whenComplete} callback
fires, {@code onOpChainComplete}
+ /// waits for this future via {@code thenRun} instead of calling {@link
#sendDoneAndComplete} immediately,
+ /// ensuring the broker always receives the {@code submit_ack} before
{@code ServerDone}.
+ private final CompletableFuture<Void> _ackSentFuture = new
CompletableFuture<>();
+
+ SubmitWithStreamObserver(StreamObserver<Worker.ServerToBroker>
responseObserver) {
+ _responseObserver = responseObserver;
+ }
+
+ @Override
+ public void onNext(Worker.BrokerToServer message) {
+ switch (message.getPayloadCase()) {
+ case SUBMIT:
+ handleSubmit(message.getSubmit());
+ break;
+ case CANCEL:
+ handleCancel(message.getCancel());
+ break;
+ case PAYLOAD_NOT_SET:
+ default:
+ sendErrorAndComplete("Unexpected BrokerToServer payload: " +
message.getPayloadCase());
+ break;
+ }
+ }
+
+ @Override
+ public void onError(Throwable t) {
+ // Broker-side stream error / disconnect. Treat like a cancel and clean
up; do not reply on the response stream
+ // (the underlying transport is gone). Use compareAndSet to be
consistent with sendDoneAndComplete and avoid
+ // double-cancelling if onOpChainComplete already completed the stream
first.
+ LOGGER.warn("SubmitWithStream stream error for request {}: {}",
_requestId, t.getMessage());
+ if (_completed.compareAndSet(false, true)) {
+ cleanupListener();
+ cancelIfSubmitted();
+ }
+ }
+
+ @Override
+ public void onCompleted() {
+ // Broker has half-closed (no more inbound messages). The server stream
stays open until all opchains have
+ // reported via the completion listener — it's the listener's job to
emit ServerDone and complete the stream.
+ // If the broker half-closes before the server is done, that's OK; we
keep emitting on the response stream
+ // until our own completion criterion is met.
+ }
+
+ private void handleSubmit(Worker.QueryRequest request) {
+ if (!_submitted.compareAndSet(false, true)) {
+ sendErrorAndComplete("Multiple submit messages on the same stream are
not allowed");
+ return;
+ }
+ ServerMetrics.get().addMeteredGlobalValue(ServerMeter.MSE_QUERIES, 1L);
+ Map<String, String> deserializedMetadata;
+ try {
+ deserializedMetadata =
QueryPlanSerDeUtils.fromProtoProperties(request.getMetadata());
+ } catch (Exception e) {
+ LOGGER.error("Caught exception while deserializing request metadata",
e);
+ sendErrorAndComplete("Caught exception while deserializing request
metadata: " + e.getMessage());
+ return;
+ }
+ // Override the cluster-level _sendStats decision for this request:
stats travel out-of-band on the bidi stream
+ // (via the OpChainCompletionListener), so we suppress the mailbox-side
stats path. The override is read by
+ // QueryRunner.effectiveSendStats(...).
+ Map<String, String> reqMetadata = new HashMap<>(deserializedMetadata);
+
reqMetadata.put(CommonConstants.MultiStageQueryRunner.KEY_OF_STATS_REPORTING_MODE,
+ CommonConstants.MultiStageQueryRunner.STATS_REPORTING_MODE_STREAM);
+ try {
+ _requestId = Long.parseLong(reqMetadata.get(MetadataKeys.REQUEST_ID));
+ } catch (Exception ignored) {
+ // _requestId stays at -1; cancel-on-stream-close will just be a no-op.
+ }
+ // Count how many opchains will run on this server: sum of
WorkerMetadata across all stage plans.
+ int opChainCount = 0;
+ for (Worker.StagePlan stagePlan : request.getStagePlanList()) {
+ opChainCount += stagePlan.getStageMetadata().getWorkerMetadataCount();
Review Comment:
Re-reviewed the latest commits — default-off looks clean, the master merge
is fine, and the other fixes (hot-path gating, leak backout, fail-loud, drain
config, coverage reconcile) all check out. One thing still open on this one.
The commit message says the graft is "correct even for leaves with multiple
build sides," but I don't think it is yet. The graft is unconditional while the
fold isn't, and they diverge in two reachable cases. Both end the same way:
`encode` throws `treeSize != flatSize`, the catch keeps `success=true` with
empty stats, and the stage reports as responded/missing=0 — the same silent
loss this fix targets.
**1. Multiple build sides — default config.** `SELECT ... FROM a WHERE x IN
(SELECT ...) AND y IN (SELECT ...)` produces one `PipelineBreakerOperator` with
two `MAILBOX_RECEIVE` children, so the graft walks 5 ops (SEND, LEAF, PB, RECV,
RECV). But the folded flat list has only 4:
`PipelineBreakerOperator.calculateUpstreamStats` reduces the receives via
`mergeUpstream`, and since both are in the same (leaf) stage, `currentDiff ==
0` skips carrying over the second receive's current-stage entry
(`MultiStageQueryStats` ~L265). So flatSize=4, treeSize=5 → throw. Verified
against the planner (the two-semi-join plan in `PinotHintablePlans.json` shows
two `PIPELINE_BREAKER` exchanges feeding one `a` scan) and with two unit tests
I added locally to `MultiStageStatsTreeEncoderTest`: one shows `mergeUpstream`
collapsing two same-stage receives to a single entry, the other shows the
2-receive graft throwing `IllegalStateException`. No special config needed.
**2. `pinot.query.mse.skip.pipeline.breaker.stats=true`.** Here the leaf
doesn't fold the PB at all (the fold at `PlanNodeToOpChain` L161 is gated on
`isKeepPipelineBreakerStats()`), but the graft still adds it → treeSize(4) !=
flatSize(2) → same drop. So the "its stats are not lost" comment in
`onOpChainComplete` doesn't hold under that (non-default) flag.
Root cause: the graft is driven structurally by the live PB tree (always all
N receives, regardless of keepPB), while the fold is conditional (keepPB) and
collapsing (`mergeUpstream` same-stage). Could the graft be made to track the
fold — gate it on `isKeepPipelineBreakerStats()`, and for the multi-receive
case either graft only the receive(s) the flat list actually kept, or stop
collapsing same-stage receives in `mergeUpstream`? A two-semi-join test plus a
keepPB=false test would lock it down — note the keepPB=false drop currently
passes `assertFullCoverage`, so only a per-stage tree-content assertion catches
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
