gortiz commented on code in PR #18458: URL: https://github.com/apache/pinot/pull/18458#discussion_r3372060432
########## pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/streaming/StreamingQuerySession.java: ########## @@ -0,0 +1,336 @@ +/** + * 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.metrics.BrokerMeter; +import org.apache.pinot.common.metrics.BrokerMetrics; +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) { + boolean currentMerged = mergeIntoAccumulatorLocked(decoded.getCurrentStageId(), decoded.getCurrentStage()); + for (Map.Entry<Integer, StageStatsTreeNode> upstream : decoded.getUpstreamStages().entrySet()) { + mergeIntoAccumulatorLocked(upstream.getKey(), upstream.getValue()); + } + // Count this opchain as responded only when its current-stage merge succeeded. On a shape mismatch, + // mergeIntoAccumulatorLocked already recorded mergeFailed for the stage; counting it as responded too would + // make responded + mergeFailed exceed expected for that stage, so the coverage triple would not reconcile. + if (currentMerged) { + incrementLocked(_respondedByStage, stageId); + } + } else { + // Successful opchain with no stats tree (e.g. empty plan). Still counts as responded. + incrementLocked(_respondedByStage, stageId); + } + } finally { + _lock.unlock(); + } + _completionLatch.countDown(); + if (shouldFanOutCancel) { + fanOutCancel(); + } + } + + /** + * Merges {@code incoming} into the accumulator for {@code stageId}. Returns {@code true} if it was stored or merged + * cleanly, {@code false} if the merge hit a {@link StageStatsTreeNode.ShapeMismatchException} (in which case the + * stage is recorded as merge-failed). + */ + private boolean mergeIntoAccumulatorLocked(int stageId, StageStatsTreeNode incoming) { + StageStatsTreeNode existing = _stageAccumulator.get(stageId); + if (existing == null) { + _stageAccumulator.put(stageId, incoming); + return true; + } + try { + existing.merge(incoming); + return true; + } catch (StageStatsTreeNode.ShapeMismatchException e) { + // StageStatsTreeNode.merge mutates _statMap before recursing into children, so a ShapeMismatchException + // thrown during child recursion leaves the existing node in a partially-accumulated state. Remove it so + // subsequent opchains for this stage do not merge into corrupt state. + _stageAccumulator.remove(stageId); + LOGGER.warn("Shape mismatch merging stage {} on request {}: {}", stageId, _requestId, e.getMessage()); + incrementLocked(_mergeFailedByStage, stageId); + return false; + } + } + + private static void incrementLocked(Map<Integer, Integer> counter, int stageId) { + counter.merge(stageId, 1, Integer::sum); + } + + /** + * Records a transport-level error on one of the server streams (gRPC {@code onError}). Drains exactly + * {@code remainingExpected} entries from the latch (the number of opchains that will not now report) and triggers + * fan-out cancel if not already triggered. + * + * <p>The caller must pass the precise per-server remaining count rather than a fixed 1, because a single stream can + * carry multiple opchains (one per worker per stage). Passing an incorrect count either under-drains (causes + * {@link #awaitCompletion} to block until the query deadline) or over-drains (causes the latch to reach zero before + * all reports have arrived). + */ + public void recordStreamError(StreamingServerHandle stream, @Nullable Throwable error, int remainingExpected) { + boolean shouldFanOutCancel = false; + _lock.lock(); + try { + _openStreams.remove(stream); + if (!_peerErrorObserved) { + _peerErrorObserved = true; + shouldFanOutCancel = true; + } + } finally { + _lock.unlock(); + } + for (int i = 0; i < remainingExpected; i++) { + _completionLatch.countDown(); + } + LOGGER.warn("Stream error on request {} draining {} pending: {}", + _requestId, remainingExpected, error == null ? "<null>" : error.getMessage()); + if (shouldFanOutCancel) { + fanOutCancel(); + } + } + + /** + * Sends {@code BrokerToServer.cancel} on every open server stream. Used on the first peer error observed and when + * the broker's data mailbox reports a processing exception. Failures are swallowed — cancel is best-effort. + * Idempotent w.r.t. concurrent calls. + */ + public void fanOutCancel() { + Set<StreamingServerHandle> snapshot; + _lock.lock(); + try { + snapshot = new LinkedHashSet<>(_openStreams); + } finally { + _lock.unlock(); + } + if (!snapshot.isEmpty()) { + // Observable signal that a query broadcast cancel to its peers. fanOutCancel can fire more than once for a + // single query (e.g. a peer error then a recovery-path cancel), so this counts non-empty broadcasts, not + // affected queries — at high QPS during a partial degradation it can amplify, so it is worth watching in prod. + BrokerMetrics.get().addMeteredGlobalValue(BrokerMeter.MSE_STREAM_STATS_CANCEL_FANOUTS, 1L); + } + for (StreamingServerHandle stream : snapshot) { + try { + stream.cancel(_requestId); + } catch (Throwable t) { + LOGGER.warn("Failed to fan out cancel on request {}", _requestId, t); + } + } + } + + /** + * Blocks the calling thread until either every expected opchain has reported, or the timeout fires. + * + * <p>Returns {@code true} when all opchains reported before the timeout (early completion is the common case in + * stream mode). Returns {@code false} when the timeout fired first; the per-stage coverage exposed via + * {@link #snapshotCoverage()} indicates which stages are missing or had merge failures. + * + * <p>The dispatcher should only call this <strong>after</strong> the broker receiving mailbox has finished. That + * way a {@code true} return means both "data done" and "stats fully accounted for". + */ + public boolean awaitCompletion(long timeout, TimeUnit unit) + throws InterruptedException { + return _completionLatch.await(timeout, unit); + } + + /** + * Returns a snapshot of the per-stage coverage. Stage ids that received any responses (successful or + * merge-failed) appear in the map; missing stages are computed by the caller against the expected total. + */ + public Coverage snapshotCoverage() { + _lock.lock(); + try { + return new Coverage(new HashMap<>(_respondedByStage), new HashMap<>(_mergeFailedByStage), + new HashMap<>(_stageAccumulator)); Review Comment: Fixed in 64eddf626c — and thanks for the deterministic repro, that pinned it precisely. `snapshotCoverage()` now (a) deep-copies the per-stage trees via a new `StageStatsTreeNode.deepCopy()` (clones each node's `StatMap` with the existing copy ctor, recursively), so the returned snapshot is fully isolated from the live accumulator, and (b) sets a `_finalized` flag under `_lock` so any late `recordOpChainComplete` early-returns instead of merging. Both together: the broker flattens/serializes an independent copy, and the live accumulator is frozen the moment the broker stops waiting. Added `StreamingQuerySessionTest.testSnapshotIsolatedFromLateReports`, which mirrors your 5→12 scenario and fails on the old shallow-copy code. I went with deep-copy + freeze rather than cancelling the open streams on the success-partial path, to avoid cancelling slow-but-healthy opchains; the `_finalized` guard means their late reports are simply ignored (bounded by the deadline). ########## 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: Fixed in 64eddf626c. The response observer is now cast to `ServerCallStreamObserver` and the per-opchain `OpChainComplete` send is gated on `isReady()`: when the outbound buffer is backed up the (best-effort) stats for that opchain are dropped rather than buffered unboundedly on worker threads — the opchain is still counted, so `ServerDone` still fires and the broker reports that stage as `missing` in its coverage. I also set a `WRITE_BUFFER_WATER_MARK` (1MB/8MB) on the server builder so `isReady()` is a meaningful bound. This caps the server-side outbound memory the way the legacy `GrpcSendingMailbox` backpressure does, trading completeness-under-backpressure (acceptable for best-effort stats) for the memory bound. A buffer + `setOnReadyHandler` drain that preserves completeness is a possible future enhancement if needed. ########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java: ########## @@ -805,7 +814,18 @@ private BrokerResponse query(QueryEnvironment.CompiledQuery query, long requestI } } - fillOldBrokerResponseStats(brokerResponse, queryResults.getQueryStats(), dispatchableSubPlan); + fillOldBrokerResponseStats(brokerResponse, queryResults.getQueryStats(), dispatchableSubPlan, + queryResults.getStageCoverage()); + + if (QueryOptionsUtils.isStreamStats(query.getOptions(), _streamStatsDefault)) { + _brokerMetrics.addMeteredGlobalValue(BrokerMeter.MSE_STREAM_STATS_QUERIES, 1); + List<QueryDispatcher.QueryResult.StageCoverage> coverage = queryResults.getStageCoverage(); + if (coverage != null && coverage.stream() + .anyMatch(c -> c != null && (c.getMissing() > 0 || c.getMergeFailed() > 0))) { + _brokerMetrics.addMeteredGlobalValue(BrokerMeter.MSE_STREAM_STATS_INCOMPLETE_COVERAGE, 1); Review Comment: All four addressed in 64eddf626c: 1. **Meter on error path** — `MSE_STREAM_STATS_INCOMPLETE_COVERAGE` is now gated on `processingException == null`, so it no longer spikes on every failed/timed-out query (which returns partial coverage by definition) and stays meaningful as a "persistent stats gap on otherwise-healthy queries" signal. 2. **Broker inbound size** — `DispatchClient` now sets `maxInboundMessageSize(64MB)` on the dispatch channel (matching the server inbound), so a large `MultiStageStatsTree` no longer fails the query with `RESOURCE_EXHAUSTED`. 3. **Stage-id bound** — `mergeSessionStatsIntoResult` now sizes the result array from the broker's own stage count + `expectedByStage` only, never the server-supplied accumulator keys; out-of-bounds stage ids are skipped and logged, so a buggy server can't force a giant allocation. 4. **Pre-register cancel** — `QueryServer.handleCancel` now completes the stream (`cleanupListener` + `sendDoneAndComplete`, gated on the ack and a matching requestId) once the request is cancelled, so opchains cancelled before they register no longer make the broker wait out its full drain window. Combined with (1) this also removes the associated meter noise. -- 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]
