gortiz commented on code in PR #18458:
URL: https://github.com/apache/pinot/pull/18458#discussion_r3356939222
##########
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:
Fixed in 1476e498ba. Wrapped `_executorService.submit(task)` in a try/catch
on `RejectedExecutionException`: the increment of the two maps always runs
before `submit()`, so on rejection (the task never runs, the directExecutor
callback never fires) the catch calls `decrementActiveOpChains(requestId)` to
back out the per-request entry — including the pinned `QueryExecutionContext` —
then rethrows so the caller propagates it as a stage error. No
double-decrement: submit either enqueues (callback decrements once later) or
throws (catch decrements once).
##########
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:
Fixed in 1476e498ba. The gate is now
`OpChainExecutionContext.isStreamStatsReporting()` (reads
`KEY_OF_STATS_REPORTING_MODE == stream`) instead of `isSendStats()`. As you
noted, `isSendStats()` is true in legacy SAFE (the default) where these ids are
never read — and, conversely, it is *false* in stream mode (where they ARE
needed), so the old gate both wasted work on the default path and skipped the
walk where it mattered. The two `IdentityHashMap`s are now lazily allocated
too, so legacy queries allocate neither.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainExecutionContext.java:
##########
@@ -63,6 +67,25 @@ public class OpChainExecutionContext {
private ServerPlanRequestContext _leafStageContext;
private final boolean _sendStats;
private final boolean _keepPipelineBreakerStats;
+ /**
+ * Map of MultiStageOperator -> PlanNodes that compile down to that
operator. Populated by
+ * {@link org.apache.pinot.query.runtime.plan.PlanNodeToOpChain} during
opchain construction. Cardinality is
+ * one-to-many: an intermediate operator maps to a single PlanNode, but the
leaf operator maps to the whole sub-tree
+ * of v1 plan nodes below the leaf-stage boundary. Used by the stream-mode
stats reporting path
+ * ({@code MultiStageStatsTreeEncoder}) to attach plan-node identifiers to
each operator's stats.
+ * <p>
+ * Identity-based (IdentityHashMap) because PlanNode equality is structural
and two distinct nodes can compare equal.
+ */
+ private final Map<MultiStageOperator, List<PlanNode>> _operatorToPlanNodes =
new IdentityHashMap<>();
Review Comment:
Addressed in 1476e498ba. The `_operatorToPlanNodes` / `_planNodeIds` maps
are now lazily allocated (null until first put), and the pre-walk that
populates them is gated on `isStreamStatsReporting()` rather than
`isSendStats()`. So in legacy mode (the default) neither map is allocated or
populated — no GC dead weight on the hot path. Writes happen single-threaded at
opchain construction and are read at stats-report time after the opchain has
run, so the lifecycle gives the happens-before; no volatile needed.
##########
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:
Fixed in 1476e498ba. `mergeIntoAccumulatorLocked` now returns whether the
merge succeeded, and an opchain is counted as `responded` only when its
current-stage merge succeeds. On a shape mismatch it is counted as
`mergeFailed` only, so `responded + mergeFailed` reconciles against the
expected count for the stage. Updated `testShapeMismatchDoesNotAbortQuery`
accordingly.
##########
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:
Fixed in 1476e498ba. Made the drain window configurable via
`pinot.broker.mse.stream.stats.drain.ms` (default 50ms; `_statsDrainMs`
threaded through the dispatcher), and fixed the stale comment that referenced
the non-existent `STATS_DRAIN_ON_SUCCESS_MS`. (1) and the ~50ms healthy-query
wait you noted are now tunable by operators.
##########
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:
Fixed in 1476e498ba — corrected the comment: these are NOT
`PlanNode.getNodeId()`, they are stage-scoped sequential ids assigned by
`PlanNodeToOpChain`'s pre-order walk; all workers of a stage walk the same
plan, and the ids are serialized on the wire.
--
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]
