yashmayya commented on code in PR #18458:
URL: https://github.com/apache/pinot/pull/18458#discussion_r3352990417
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/DispatchClient.java:
##########
@@ -130,6 +133,38 @@ public void submit(Worker.QueryRequest request,
QueryServerInstance virtualServe
_dispatchStub.withDeadline(deadline).submit(request, new
LastValueDispatchObserver<>(virtualServer, callback));
}
+ /**
+ * Opens a {@code SubmitWithStream} bidi RPC for one server, sends the
initial {@code submit}, and registers the
+ * resulting {@link StreamingDispatchObserver} with {@code session} for
cancel fan-out and {@code OpChainComplete}
+ * accumulation.
+ *
+ * <p>The submit-ack callback is invoked exactly once: with the {@link
Worker.QueryResponse} on the first
+ * {@code submit_ack} from the server, or with a non-null {@link Throwable}
if the stream errors before the ack
+ * arrives.
+ *
+ * @param request the plan submission
+ * @param virtualServer server identity (used in callbacks for
routing decisions on failure)
+ * @param deadline gRPC deadline for the call
+ * @param session broker-side streaming session — the returned
observer registers itself here
+ * @param expectedOpChainsForThisServer number of opchains this server is
expected to report; used to drain the
+ * session latch correctly when the stream
errors before all opchains have responded
+ * @param ackCallback receives the submit-ack response or a
failure throwable
+ * @return the observer, also exposed as
+ * {@link
org.apache.pinot.query.service.dispatch.streaming.StreamingServerHandle} on the
session for
+ * cancel fan-out
+ */
+ public StreamingDispatchObserver submitWithStream(Worker.QueryRequest
request, QueryServerInstance virtualServer,
+ Deadline deadline, StreamingQuerySession session, int
expectedOpChainsForThisServer,
+ BiConsumer<Worker.QueryResponse, Throwable> ackCallback) {
+ StreamingDispatchObserver observer = new
StreamingDispatchObserver(virtualServer, session,
+ expectedOpChainsForThisServer, ackCallback);
+ StreamObserver<Worker.BrokerToServer> outbound =
_dispatchStub.withDeadline(deadline).submitWithStream(observer);
Review Comment:
There's no `UNIMPLEMENTED` fallback here. The PR description's compat table
says it's "Implemented in DispatchClient", but this just opens the stream and
sends — the accurate doc is the one on `submitAndReduceWithStream` ("No
automatic fallback … will cause query failures").
So if stream mode is enabled while any server is still on an older version
(rolling upgrade, or a single rollback), every MSE query fails. It's behind the
default-off flag so that's tolerable, but the description should match the
code, and a loud warning at the call site (or an actual fallback) would help
stop someone flipping this mid-upgrade.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java:
##########
@@ -243,19 +459,52 @@ private QueryResult tryRecover(long requestId,
Set<QueryServerInstance> servers,
} else if (ex instanceof QueryException) {
errorCode = ((QueryException) ex).getErrorCode();
} else {
- // in case of unknown exceptions, the exception will be rethrown, so we
don't need stats
cancel(requestId, servers);
throw ex;
}
- // in case of known exceptions (timeout or query exception), we need can
build here the erroneous QueryResult
- // that include the stats.
- LOGGER.warn("Query failed with a known exception. Trying to cancel the
other opchains");
- MultiStageQueryStats stats = cancelWithStats(requestId, servers);
- if (stats == null) {
+ LOGGER.warn("Query failed with a known exception. Cancelling remaining
opchains.");
+ cancel(requestId, servers);
+ QueryProcessingException processingException = new
QueryProcessingException(errorCode, ex.getMessage());
+ return new QueryResult(processingException,
MultiStageQueryStats.emptyStats(0), 0L);
Review Comment:
Worth noting this changes legacy behavior too: on master this path returned
partial stats via `cancelWithStats(...)`; now it returns empty stats. It only
bites when dispatch throws a known exception (a normal broker timeout still
returns an error block *with* stats), and dropping the per-failure cancel
fan-out is a reasonable trade — but it is a default-path change, so a note in
the description/release notes would be good.
Minor: the `QueryRunner.cancel` Javadoc saying the consumer "was reverted at
the same time" reads as if it was already gone — master's `tryRecover` did
consume it; this PR removes it.
##########
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:
`_expectedOpChains` counts only main per-worker opchains, but
pipeline-breaker opchains also fire the same completion listener — they go
through `scheduler.register` in `PipelineBreakerExecutor` with the same
requestId, and they carry the *same* `OpChainId` (stageId, workerId) as the
stage's main opchain. So a stage with a pipeline breaker fires the listener
twice but is only counted once: `_completedOpChains` hits the expected total
early, `ServerDone` fires, and the real opchain's `OpChainComplete` then gets
dropped by the `!_completed.get()` guard at the send site.
Depending on completion order the broker either records the
pipeline-breaker's tree as that stage's stats (with `missing=0`, so nothing
signals it's wrong), or hits a shape mismatch and marks some other stage
`missing`/`mergeFailed`. Repro: `... WHERE key IN (SELECT key FROM small)` —
the default planner turns that into a dynamic-broadcast semi-join with a
`PIPELINE_BREAKER` on a dispatched stage.
The current tests miss it: `testJoinQuery` is a plain inner join (no PB),
and `assertFullCoverage` only checks the counters, never the per-stage tree.
Simplest fix is probably to not fire the stream listener for PB opchains (skip
when the root is `Type.PIPELINE_BREAKER` — they shouldn't be reported as a
stage anyway). Would be good to add a semi-join regression test that asserts
the actual tree per stage.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/streaming/StreamingQuerySession.java:
##########
@@ -0,0 +1,316 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.service.dispatch.streaming;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.ReentrantLock;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.proto.Worker;
+import org.apache.pinot.query.runtime.plan.MultiStageStatsTreeDecoder;
+import org.apache.pinot.query.runtime.plan.StageStatsTreeNode;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Broker-side per-query state for the {@code SubmitWithStream} dispatch path.
Owns the per-stage tree accumulator,
+ * the outstanding-opchain count, the per-stage coverage counters, and the set
of open server streams.
+ *
+ * <p>Concurrency model — all mutating methods acquire the per-session lock,
so the accumulator and counters need no
+ * additional internal synchronization. gRPC client {@code onNext} callbacks
land on I/O threads and call into this
+ * session directly. Stat decoding (proto → {@link StageStatsTreeNode}) is
done <em>outside</em> the lock to minimise
+ * lock hold time; only the map mutations that update the accumulator are
performed under the lock.
+ *
+ * <p>Completion semantics — {@link #awaitCompletion(long, TimeUnit)} returns
{@code true} as soon as every expected
+ * opchain has reported (early completion), and {@code false} if the timeout
fires first. The dispatcher should call
+ * it <strong>only after</strong> the broker receiving mailbox has finished,
so that a successful return means both
+ * "data done" and "stats fully accounted for". When it returns {@code false}
the per-stage coverage exposes which
+ * stages are missing.
+ */
+public class StreamingQuerySession {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(StreamingQuerySession.class);
+
+ private final long _requestId;
+ private final int _expectedOpChains;
+ private final CountDownLatch _completionLatch;
+ /**
+ * Guards {@link #_stageAccumulator}, {@link #_respondedByStage}, {@link
#_mergeFailedByStage},
+ * {@link #_openStreams}, and {@link #_peerErrorObserved}. Lock hold time is
proportional to the merge work (a few
+ * map operations), not to proto decode; see {@link #recordOpChainComplete}
for why decode is done outside.
+ *
+ * <p>If lock contention becomes a bottleneck at high QPS, a virtual-thread
actor (one VT per query draining from
+ * an {@code ArrayBlockingQueue}, with gRPC I/O threads simply enqueuing)
would eliminate the lock entirely and
+ * avoid any contention between concurrent inbound callbacks.
+ */
+ private final ReentrantLock _lock = new ReentrantLock();
+
+ /** Per-stage merged accumulator. Mutated under {@link #_lock}. */
+ private final Map<Integer, StageStatsTreeNode> _stageAccumulator = new
HashMap<>();
+ /** Per-stage count of opchains that have responded successfully and merged
cleanly. */
+ private final Map<Integer, Integer> _respondedByStage = new HashMap<>();
+ /** Per-stage count of opchains that responded but the broker couldn't merge
their payload. */
+ private final Map<Integer, Integer> _mergeFailedByStage = new HashMap<>();
+
+ /** Set of open server streams. Iteration order is insertion order so cancel
fan-out is deterministic. */
+ private final Set<StreamingServerHandle> _openStreams = new
LinkedHashSet<>();
+
+ /** True after the first peer error (success=false OpChainComplete or stream
onError). Used to trigger fan-out
+ * cancel idempotently. */
+ private boolean _peerErrorObserved = false;
+
+ public StreamingQuerySession(long requestId, int expectedOpChains) {
+ _requestId = requestId;
+ _expectedOpChains = expectedOpChains;
+ _completionLatch = new CountDownLatch(expectedOpChains);
+ }
+
+ public long getRequestId() {
+ return _requestId;
+ }
+
+ public int getExpectedOpChains() {
+ return _expectedOpChains;
+ }
+
+ /**
+ * Registers an open server stream so the session can iterate them later for
fan-out cancel. Must be called by the
+ * dispatcher when the {@code SubmitWithStream} call is opened.
+ */
+ public void registerStream(StreamingServerHandle stream) {
+ _lock.lock();
+ try {
+ _openStreams.add(stream);
+ } finally {
+ _lock.unlock();
+ }
+ }
+
+ /**
+ * Removes a stream from the open-streams set. Called when the server emits
{@code ServerDone} (clean close) or the
+ * stream errors. Idempotent.
+ */
+ public void unregisterStream(StreamingServerHandle stream) {
+ _lock.lock();
+ try {
+ _openStreams.remove(stream);
+ } finally {
+ _lock.unlock();
+ }
+ }
+
+ /**
+ * Records an {@link Worker.OpChainComplete} message decoded from a server
stream. Decrements the outstanding count
+ * and merges the contained tree into the per-stage accumulator (or marks
the stage {@code mergeFailed} on a shape
+ * mismatch / decode failure). Also records {@code success=false} reports as
peer errors so fan-out cancel can fire.
+ *
+ * <p>Decoding (proto → {@link StageStatsTreeNode}) is performed
<em>before</em> acquiring {@link #_lock} because
+ * the input proto is immutable and {@link
MultiStageStatsTreeDecoder.Decoded} is a fresh allocation. Only the map
+ * mutations are done under the lock, which keeps lock hold time
proportional to the (small) merge work rather than
+ * the full recursive decode.
+ */
+ public void recordOpChainComplete(Worker.OpChainComplete message) {
+ int stageId = message.getStageId();
+ boolean isSuccess = message.getSuccess();
+ Worker.MultiStageStatsTree statsTree = message.getStats();
+
+ // Decode outside the lock — proto is immutable, Decoded is a fresh
allocation with no shared state.
+ MultiStageStatsTreeDecoder.Decoded decoded = null;
+ MultiStageStatsTreeDecoder.DecodeFailedException decodeError = null;
+ if (statsTree.hasCurrentStage()) {
+ try {
+ decoded = MultiStageStatsTreeDecoder.decode(statsTree);
+ } catch (MultiStageStatsTreeDecoder.DecodeFailedException e) {
+ decodeError = e;
+ }
+ }
+
+ boolean shouldFanOutCancel = false;
+ _lock.lock();
+ try {
+ if (!isSuccess) {
+ if (!_peerErrorObserved) {
+ _peerErrorObserved = true;
+ shouldFanOutCancel = true;
+ }
+ }
+ if (decodeError != null) {
+ LOGGER.warn("Decode failed for opchain stage={} worker={} on request
{}: {}",
+ stageId, message.getWorkerId(), _requestId,
decodeError.getMessage());
+ incrementLocked(_mergeFailedByStage, stageId);
+ } else if (decoded != null) {
+ mergeIntoAccumulatorLocked(decoded.getCurrentStageId(),
decoded.getCurrentStage());
+ for (Map.Entry<Integer, StageStatsTreeNode> upstream :
decoded.getUpstreamStages().entrySet()) {
+ mergeIntoAccumulatorLocked(upstream.getKey(), upstream.getValue());
+ }
+ incrementLocked(_respondedByStage, stageId);
Review Comment:
Minor: on a shape mismatch the merge increments `_mergeFailedByStage` (in
the helper) and then this line also increments `_respondedByStage`, so
`responded + mergeFailed` can exceed `expected` for the stage (the `missing =
max(0, …)` in the dispatcher hides the negative). Results are fine, but the
coverage triple won't reconcile in exactly the version-skew case someone would
be reading it for.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanNodeToOpChain.java:
##########
@@ -96,16 +96,56 @@ public static OpChain convert(PlanNode node,
OpChainExecutionContext context) {
*/
public static OpChain convert(PlanNode node, OpChainExecutionContext context,
BiConsumer<PlanNode, MultiStageOperator> tracker) {
- MyVisitor visitor = new MyVisitor(tracker);
+ // Assign deterministic stage-scoped ids to every PlanNode reachable from
the root before constructing operators,
+ // so the encoder can attach plan_node_ids to each StageStatsNode without
needing to mutate or re-walk the plan.
+ // Both broker and server perform this same pre-walk over the same plan
structure, producing matching ids.
+ // Skip the walk (and the per-opchain map population) when stats are not
going to be sent — this is the common case
+ // in legacy mode and avoids O(depth) allocations on the hot path.
+ if (context.isSendStats()) {
Review Comment:
This gate is `isSendStats()`, which is true by default (SAFE mode) — not
"stream mode on". So the id pre-walk + the leaf `collectPlanNodeSubTree` walk +
the two `IdentityHashMap`s all run on every legacy query, even though only the
stream-mode encoder ever reads those maps. The comment below ("skip … common
case in legacy mode") doesn't hold.
It's cheap per stage, but it's avoidable work added to the default path —
gating on stream mode instead would keep legacy untouched.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java:
##########
@@ -225,10 +247,204 @@ public QueryResult submitAndReduce(RequestContext
context, DispatchableSubPlan d
}
}
- /// Tries to recover from an exception thrown during query dispatching.
+ /// Streaming variant of {@link #submitAndReduce}: opens one {@code
SubmitWithStream} bidi RPC per server, runs the
+ /// broker's stage 0 reducer, and once the receiving mailbox finishes awaits
the per-stage stats with early
+ /// completion (returns as soon as every expected opchain has reported, or
when the wait window fires — whichever
+ /// happens first). Stats from the session accumulator are then merged into
the broker's local stage 0 stats to
+ /// build the final {@link QueryResult}.
+ ///
+ /// The wait window is bounded by the query's remaining timeout: if {@code
submitWithStream + runReducer} consumed
+ /// most of the budget, the per-stage stats may end up partial (visible via
the per-stage {@code mergeFailed} /
+ /// {@code missing} counts the session exposes).
///
- /// [QueryException] and [TimeoutException] are handled by returning a
[QueryResult] with the error code and stats,
- /// while other exceptions are not known, so they are directly rethrown.
+ /// Cancel is handled via {@link StreamingQuerySession#fanOutCancel()} — no
unary Cancel RPCs are issued for this
+ /// query path. On any error, fan-out cancel is broadcast over the open
streams, then the broker waits for remaining
+ /// stats before building the final result.
+ ///
+ /// <b>Mixed-version policy.</b> No automatic fallback to the unary {@link
#submit} path. Enabling
+ /// {@link CommonConstants.Broker.Request.QueryOptionKey#STREAM_STATS}
requires every server in the
+ /// cluster to implement {@code SubmitWithStream}; if any server returns
{@code UNIMPLEMENTED} or any other
+ /// transport error during dispatch, {@link #submitWithStream} surfaces the
throwable through the ack queue,
+ /// {@link #processResults} throws, and this method fans out cancel via the
session before propagating the failure.
+ private QueryResult submitAndReduceWithStream(RequestContext context,
DispatchableSubPlan dispatchableSubPlan,
+ long timeoutMs, Map<String, String> queryOptions)
+ throws Exception {
+ long requestId = context.getRequestId();
+ long deadlineMs = System.currentTimeMillis() + timeoutMs;
+ Set<QueryServerInstance> servers = new HashSet<>();
+
+ // The session's expected-opchain count must equal the total number of
opchains across every (server, non-root
+ // stage) pair — that's how many OpChainComplete messages we expect to
receive.
+ Set<DispatchablePlanFragment> stagePlansWithoutRoot =
dispatchableSubPlan.getQueryStagesWithoutRoot();
+ int totalExpected = 0;
+ Map<Integer, Integer> expectedByStage = new HashMap<>();
+ for (DispatchablePlanFragment stagePlan : stagePlansWithoutRoot) {
+ int stageId = stagePlan.getPlanFragment().getFragmentId();
+ int stageCount = 0;
+ for (List<Integer> workerIds :
stagePlan.getServerInstanceToWorkerIdMap().values()) {
+ stageCount += workerIds.size();
+ }
+ totalExpected += stageCount;
+ expectedByStage.put(stageId, stageCount);
+ }
+ StreamingQuerySession session = new StreamingQuerySession(requestId,
totalExpected);
+
+ try {
+ submitWithStream(requestId, dispatchableSubPlan, timeoutMs, servers,
queryOptions, session);
+ QueryResult brokerResult = runReducer(dispatchableSubPlan, queryOptions,
_mailboxService);
+
+ // Receiving mailbox finished — data is ready. Wait for stats on a
best-effort basis; cap at
+ // STATS_DRAIN_ON_SUCCESS_MS so a single slow opchain cannot delay the
client response.
+ long statsWaitMs = Math.min(STATS_DRAIN_MS, Math.max(0, deadlineMs -
System.currentTimeMillis()));
Review Comment:
Two small things here:
1. The description says the success path waits the full remaining timeout,
but this correctly caps at `min(50ms, remaining)` — code's fine, description is
stale.
2. Healthy queries still pay this wait: the `OpChainComplete`s are emitted
*after* the data EOS that `runReducer` already returned on, so the latch is
rarely at zero yet and a successful stream-mode query waits up to ~50ms it
wouldn't in legacy. Bounded and fine, but might be worth making
`STATS_DRAIN_MS` configurable.
(The comment just above refers to `STATS_DRAIN_ON_SUCCESS_MS`, which doesn't
exist — it's `STATS_DRAIN_MS`.)
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java:
##########
@@ -69,6 +72,25 @@ public class OpChainSchedulerService {
private final ReadWriteLock[] _queryLocks;
private final Cache<Long, Boolean> _cancelledQueryCache;
private final Metrics _metrics = new Metrics();
+ /// Per-request opchain completion listeners used by the stream-mode
stats-reporting path. A listener is registered
+ /// before the broker submits the query and fires once for every opchain
that runs on this server for that request.
+ /// Listeners are responsible for unregistering themselves once they've
consumed all expected events.
+ private final ConcurrentMap<Long, OpChainCompletionListener>
_completionListeners = new ConcurrentHashMap<>();
+ /// Maps requestId → QueryExecutionContext for O(1) cancel. An entry is
added when the first opchain for a request
+ /// registers and removed when cancel() fires or when the last opchain for
that request completes. The counter map
+ /// below tracks how many opchains are still active per request so we know
when to clean up.
+ ///
+ /// Consistency invariant: all mutations of BOTH maps for the same requestId
must occur inside
+ /// `_activeOpChainsByRequest.compute(requestId, …)`. The compute() bin-lock
for the requestId key serializes
+ /// all register and decrement operations, keeping the two maps coherent.
cancel() acquires the query write lock
+ /// BEFORE calling compute() so that register()'s read lock cannot overlap
with the eviction window.
+ ///
+ /// NOTE: `_opChainCache.put()` in registerInternal must stay inside the
read lock. cancel()'s cache-invalidation
+ /// forEach runs OUTSIDE the write lock (after the cancelled-query cache is
written). It relies on the read lock
+ /// exclusion to guarantee that no register() call is mid-flight between the
cache.put and the counter increment
+ /// when the forEach observes the cache entry.
+ private final ConcurrentMap<Long, QueryExecutionContext>
_executionContextByRequest = new ConcurrentHashMap<>();
+ private final ConcurrentMap<Long, AtomicInteger> _activeOpChainsByRequest =
new ConcurrentHashMap<>();
Review Comment:
These two maps have no TTL and are cleaned up only by the completion
callback or `cancel()`. But `_executorService.submit(task)` below can throw —
the MSE executor is wrapped in `HardLimitExecutor` + heap-throttle — and on
rejection the `directExecutor` callback never fires, so the entry (and the
pinned `QueryExecutionContext`) leaks until a later cancel. The old code only
left a TTL'd `_opChainCache` entry. A try/catch around the submit that backs
these out on failure would close the gap.
##########
pinot-common/src/main/proto/worker.proto:
##########
@@ -100,3 +109,81 @@ message MailboxInfo {
message Properties {
map<string, string> property = 1;
}
+
+//
=============================================================================
+// SubmitWithStream — stream-mode stats reporting messages
+//
=============================================================================
+
+// Messages flowing from broker to server on the SubmitWithStream stream.
+// First message MUST be a `submit`. Subsequent messages may carry cancel.
+message BrokerToServer {
+ oneof payload {
+ // Plan submission. Same shape as today's unary Submit request.
+ QueryRequest submit = 1;
+ // Optional cancel message. In Phase A the broker still uses the unary
+ // Cancel RPC; this field is reserved and processed by the server but only
+ // emitted by the broker in Phase B.
+ CancelRequest cancel = 2;
+ }
+}
+
+// Messages flowing from server to broker on the SubmitWithStream stream.
+// The first server message is a `submit_ack`. Each opchain that runs on this
+// server emits exactly one `opchain` message when it finishes (success or
error).
+// After the last opchain has reported, the server emits `done` and
half-closes.
+message ServerToBroker {
+ oneof payload {
+ // Synchronous-style ack for the submission. Same shape as today's unary
+ // QueryResponse. Carries early errors (plan parsing, malformed metadata).
+ QueryResponse submit_ack = 1;
+ // Per-stage stats from one opchain that ran on this server.
+ OpChainComplete opchain = 2;
+ // Final marker — the server has no more messages for this query.
+ ServerDone done = 3;
+ }
+}
+
+// Sent once per opchain that ran on this server, regardless of
success/failure.
+message OpChainComplete {
+ int32 stage_id = 1;
+ int32 worker_id = 2;
+ bool success = 3;
+ string error_msg = 4; // populated when success = false
+ MultiStageStatsTree stats = 5; // structured, tree-shaped stats payload
+}
+
+// Final marker on the server-to-broker stream.
+message ServerDone {
+}
+
+// Multi-stage stats produced by one opchain. Carries the current stage as an
+// explicit operator tree, plus any upstream-stage trees this opchain
+// accumulated (today these reach an opchain via pipeline-breaker stats and via
+// mailbox-receive merges; in stream mode each upstream stage is also reported
+// by its own server, so the broker merges duplicates per stage).
+message MultiStageStatsTree {
+ int32 current_stage_id = 1;
+ StageStatsNode current_stage = 2;
+ // Sparse map of upstream stage id -> tree.
+ map<int32, StageStatsNode> upstream_stages = 3;
+}
+
+// One node of the operator tree for a single stage.
+message StageStatsNode {
+ // The MultiStageOperator.Type id, same byte we serialize today via
+ // MultiStageOperator.Type.getId().
+ int32 operator_type_id = 1;
+ // Plan node ids that compile down to this operator. One-to-many; the leaf
+ // operator typically owns the whole sub-tree of v1 plan nodes below the leaf
+ // boundary. Stable per request via PlanNode.getNodeId().
Review Comment:
Doc nit: these ids aren't `PlanNode.getNodeId()` — they're a stage-scoped
sequential counter from `assignPlanNodeIds`. Same with the "broker and server
perform the same walk" note in `OpChainExecutionContext`/`PlanNodeToOpChain`:
the ids are serialized on the wire, so it's really "all workers of a stage walk
the same plan".
--
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]
