gortiz commented on code in PR #18458:
URL: https://github.com/apache/pinot/pull/18458#discussion_r3373172931
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java:
##########
@@ -450,6 +455,329 @@ 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);
+ /// Pipeline-breaker root operators awaiting their stage's main (leaf)
opchain, keyed by {@link OpChainId}. A
+ /// pipeline-breaker opchain shares the same OpChainId as its leaf opchain
and is not reported as a separate
+ /// stage; its operators are folded into the leaf's flat stats but are
absent from the leaf's live operator tree,
+ /// so we stash the PB root here and graft it onto the leaf when the leaf
opchain reports (see
+ /// {@link #onOpChainComplete} and {@link MultiStageStatsTreeEncoder}).
Keyed by OpChainId to support multiple
+ /// PB-bearing leaf opchains per request (multiple leaf stages / workers).
The PB opchain completes strictly
+ /// before its leaf — the leaf consumes the PB results, gated on the PB's
CountDownLatch in
+ /// {@code PipelineBreakerExecutor}, which establishes a happens-before
from this {@code put} to the leaf's
+ /// {@code remove}; {@link ConcurrentHashMap} also makes the
cross-executor-thread read safe directly.
+ private final Map<OpChainId, MultiStageOperator> _pipelineBreakerRoots =
new ConcurrentHashMap<>();
+ /// 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;
+ }
+ // Count stream-path (SubmitWithStream) queries the same way the unary
submit() path does, via the MseMetrics
+ // abstraction master introduced (MseMeter.QUERIES forwards to
ServerMeter.MSE_QUERIES in SERVER mode).
+ MseMetrics.get().addMeteredGlobalValue(MseMeter.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();
+ }
+ final int expected = opChainCount;
+ // Must set _expectedOpChains BEFORE registerOpChainCompletionListener.
The AtomicInteger.set is a volatile
+ // write; ConcurrentHashMap.put (inside
registerOpChainCompletionListener) provides a subsequent happens-before
+ // edge, so any thread that reads via ConcurrentHashMap.get (the
listener callback) is guaranteed to observe
+ // the _expectedOpChains value set here. Reordering these two lines
would break the JMM guarantee.
+ _expectedOpChains.set(expected);
+
+ // Register the per-request completion listener BEFORE submitting.
Otherwise short opchains could finish before
+ // we've registered and we'd miss their events.
+ if (_requestId >= 0 && expected > 0) {
+ _queryRunner.registerOpChainCompletionListener(_requestId,
this::onOpChainComplete);
+ }
+
+ long timeoutMs =
Long.parseLong(reqMetadata.get(QueryOptionKey.TIMEOUT_MS));
+ CompletableFuture.runAsync(() -> submitInternal(request, reqMetadata),
_submissionExecutorService)
+ .orTimeout(timeoutMs, TimeUnit.MILLISECONDS)
+ .whenComplete((result, error) -> {
+ if (error != null) {
+ LOGGER.error("Caught exception while submitting request: {}",
_requestId, error);
+ sendSubmitAck(buildErrorResponse("Caught exception while
submitting request: " + error.getMessage()));
+ // Submission failed — no opchains will run, so emit ServerDone
immediately and clean up.
+ cleanupListener();
+ sendDoneAndComplete();
+ } else {
+ sendSubmitAck(buildOkResponse());
+ // If for some reason expected was 0 (empty plan), close the
stream now.
+ if (expected == 0) {
+ cleanupListener();
+ sendDoneAndComplete();
+ }
+ }
+ // Signal that the ack has been sent (regardless of
success/error). Any onOpChainComplete invocation
+ // that raced ahead of this whenComplete callback will be
unblocked via thenRun.
+ _ackSentFuture.complete(null);
+ });
+ }
+
+ /**
+ * Fires once per opchain on this server completing. Encodes the stats
into a {@link Worker.MultiStageStatsTree},
+ * emits an {@link Worker.OpChainComplete}, and emits {@link
Worker.ServerDone} once all expected opchains have
+ * reported.
+ */
+ private void onOpChainComplete(OpChainId opChainId, MultiStageOperator
rootOperator,
+ @Nullable MultiStageQueryStats stats, OpChainExecutionContext context,
@Nullable Throwable error) {
+ // A pipeline-breaker opchain (dynamic-broadcast semi-join / lookup-join
build side) is built from the SAME
+ // OpChainExecutionContext as its stage's main (leaf) opchain, so it
carries an identical OpChainId and fires
+ // this listener too. It must NOT be reported as a separate stage:
+ // - _expectedOpChains counts only the main per-worker opchains, so
counting the PB here would make
+ // _completedOpChains reach the expected total one early,
prematurely firing ServerDone and causing the
+ // real opchain's OpChainComplete to be dropped by the
!_completed.get() guard at the send site.
+ // - its stats are not lost: the PB's operators are already folded
into the leaf opchain's flat stats (via
+ // LeafOperator.calculateUpstreamStats, exactly as in the legacy
mailbox path), but they are absent from
+ // the leaf's live operator tree, which is what the encoder walks.
So we stash the PB root here and graft
+ // it onto the leaf below, letting the stage report once — when its
leaf opchain completes — with one
+ // coherent tree.
+ if (rootOperator.getOperatorType() ==
MultiStageOperator.Type.PIPELINE_BREAKER) {
+ _pipelineBreakerRoots.put(opChainId, rootOperator);
+ return;
+ }
+ Worker.OpChainComplete.Builder builder =
Worker.OpChainComplete.newBuilder()
+ .setStageId(opChainId.getStageId())
+ .setWorkerId(opChainId.getVirtualServerId())
+ .setSuccess(error == null);
+ if (error != null) {
+ builder.setErrorMsg(error.getMessage() == null ?
error.getClass().getSimpleName() : error.getMessage());
+ }
+ // Claim any stashed pipeline-breaker root for this stage (remove() so
the stash self-cleans, even on the
+ // leaf-error path below where it is unused). The PB always completed
before this leaf opchain (see
+ // _pipelineBreakerRoots), so it is already present when the stage
folded one.
+ MultiStageOperator pipelineBreakerRoot =
_pipelineBreakerRoots.remove(opChainId);
+ if (stats != null) {
+ // If this stage folded a pipeline breaker, graft its stashed operator
subtree onto the leaf so the encoder's
+ // tree walk realigns with the leaf opchain's folded flat stats.
+ try {
+ builder.setStats(MultiStageStatsTreeEncoder.encode(rootOperator,
stats, context, pipelineBreakerRoot));
+ } catch (Throwable t) {
+ // Encoding failed — no partial stats can be recovered. The encoder
is all-or-nothing by design:
+ // 1. The upfront treeSize != flatSize check throws
IllegalStateException before any proto node is built.
+ // 2. An IOException from serializeStatMap mid-walk leaves partial
StageStatsNode builders only on the
+ // Java call stack; they are discarded when the exception
unwinds. No partially-built
+ // MultiStageStatsTree is ever returned.
+ // We deliberately do NOT set success=false here. The opchain
computation itself succeeded (error==null
+ // at the top of this method), so we preserve success=true and send
empty stats instead. Setting
+ // success=false would cause the broker to treat the opchain as a
peer error and fire fanOutCancel,
+ // cancelling a completely healthy query just because stats
serialization failed.
+ LOGGER.warn("Failed to encode stats tree for opchain {}", opChainId,
t);
+ builder.setErrorMsg(builder.getErrorMsg().isEmpty() ? "stats encode
failed: " + t.getMessage()
+ : builder.getErrorMsg() + "; stats encode failed: " +
t.getMessage());
+ }
+ }
+ Worker.ServerToBroker message =
Worker.ServerToBroker.newBuilder().setOpchain(builder).build();
+ try {
+ synchronized (_streamLock) {
+ if (!_completed.get()) {
+ _responseObserver.onNext(message);
Review Comment:
Correction / follow-up: I reverted the `isReady()`-drop approach in
bd10711bcb. `ServerCallStreamObserver.isReady()` is transiently `false` at
stream start (before the first HTTP/2 window update), so dropping the
OpChainComplete then loses the leaf opchain's stats on the **success** path —
it failed `StreamStatsReportingIntegrationTest.testSemiJoinPipelineBreaker` in
CI (stage rendered `EMPTY_MAILBOX_SEND`). Dropping best-effort stats on
transient not-ready is a correctness regression, so that was the wrong fix.
The correct bound is a small per-stream buffer drained via
`setOnReadyHandler` (loses nothing on transient not-ready; drops only under
*sustained* overflow, and keeps `ServerDone` ordered after the drain). I'm
deferring that to a focused follow-up rather than rushing it here, since it's a
non-trivial change to the response path and this item is opt-in +
deadline-bounded. The unbounded-buffering concern therefore remains open for
now — tracked as a follow-up. Sorry for the churn on this one.
--
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]
