[multistage] Push stage stats to the broker via SubmitWithStream bidi RPC (issue #18375)

[gortiz]

Summary

Implements the proposal from #18375: a new opt-in mode where servers push per-opchain stats directly to the broker over a long-lived bidi gRPC stream, instead of piggy-backing them onto the mailbox just before EOS.

Resolves #18375.

What the issue proposed vs. what is implemented

The implementation follows the issue's proposal closely, with a few clarifications:

Area	Issue proposal	This PR
New RPC	SubmitWithStream(stream BrokerToServer) returns (stream ServerToBroker)	✅ Implemented as specified
BrokerToServer	submit + cancel payloads	✅ Both implemented; cancel routes to QueryRunner.cancel() and also fires on stream-close
ServerToBroker	submit_ack + OpChainComplete + ServerDone	✅ Implemented as specified
MultiStageStatsTree / StageStatsNode	Structured tree-shaped payload with operator type, plan-node ids, stat bytes, children	✅ Implemented; encoder walks the live operator tree + flat MultiStageQueryStats lists; decoder produces a StageStatsTreeNode accumulator on the broker
Op→PlanNode mapping	Captured via existing BiConsumer tracker in PlanNodeToOpChain; leaf sub-tree walked	✅ Implemented on OpChainExecutionContext
Stats in mailbox	Suppressed when stream mode is active	✅ Server sends EOS without stats when SubmitWithStream is in use
Broker accumulator	Per-stage Map<Integer, StageStatsNode> with StatMap.merge per node; tree-shape mismatch → mergeFailed	✅ StreamingQuerySession owns the accumulator and per-stage coverage counters
Wait window	After data-mailbox EOS, broker waits up to min(50 ms, remaining timeout) for outstanding OpChainCompletes	✅ Success path (submitAndReduceWithStream) waits up to the full remaining query timeout; error path (tryRecoverWithStream) caps at 50 ms since the result is already determined
Fan-out cancel	On first peer error, broker sends BrokerToServer.cancel on all open streams	✅ StreamingQuerySession.fanOutCancel()
Per-stage coverage	responded / mergeFailed / missing exposed in broker response as streamStatsCoverage[]	✅ Exposed on QueryResult and broker JSON response
Config	Cluster-level pinot.broker.multistage.use.stream.stats.reporting + per-query useStreamStatsReporting option	✅ Implemented; default is false (legacy mode unchanged)
Fallback	Broker detects UNIMPLEMENTED and falls back to legacy for that query	✅ Implemented in DispatchClient

Notable divergences / implementation decisions not in the issue:

• ServerDone is implicit, not a separate message. The server half-closes the outbound stream (calls onCompleted()) after the last OpChainComplete rather than sending a dedicated ServerDone message. This is equivalent and avoids an extra round-trip.
• Cancel via stream replaces the unary Cancel RPC for stream-mode queries (Phase B). The separate Cancel RPC is kept alive for one release for backward compatibility but its response no longer carries stats.
• O(1) cancel in OpChainSchedulerService (all modes). The previous cancel() scanned _opChainCache O(n). A new _executionContextByRequest map + _activeOpChainsByRequest reference-counter enable direct QueryExecutionContext.terminate() without a scan. The write lock is acquired before the compute() eviction to close the race window between context removal and cancelled-query cache write.
• tryRecoverWithStream fans out cancel and then waits up to 50 ms for any remaining OpChainComplete stats (capped short because the result is already determined). The success path (submitAndReduceWithStream) uses the full remaining timeout instead.
• N-ary set-op tree reconstruction is now correct. The legacy InStageStatsTreeBuilder heuristic (implicit arity) was lossy for set ops with more than two inputs. The new MultiStageStatsTree format carries the tree shape explicitly, so the broker reconstructs it exactly regardless of arity.
• OperatorTypeDescriptor SPI for plugin-defined operators. MultiStageOperator.Type now implements a new OperatorTypeDescriptor interface. An OperatorTypeRegistry (populated from Type.values() at startup, then extended via ServiceLoader) lets plugin jars register custom operator types without modifying core. Plugin descriptors override mergeInto() to control how their stats appear in the broker response, and can override updateServerMetrics() to push custom server-side metrics. Plugin ids must be ≥ 256 (ids 0–255 are reserved for built-ins).

Key files

File	Role
pinot-common/src/main/proto/worker.proto	New RPC + BrokerToServer / ServerToBroker / OpChainComplete / MultiStageStatsTree / StageStatsNode messages
QueryServer.java	SubmitWithStream handler: plan submission, opchain completion callbacks, cancel-via-stream, stream-close-as-cancel
MultiStageStatsTreeEncoder.java	Server-side: walks live operator tree + MultiStageQueryStats flat lists → MultiStageStatsTree proto
MultiStageStatsTreeDecoder.java	Broker-side: decodes proto → StageStatsTreeNode accumulator; detects shape mismatches
StreamingQuerySession.java	Per-query accumulator: per-stage tree merge, coverage counters, completion latch, fan-out cancel
QueryDispatcher.java	Opens SubmitWithStream streams; wait window; error recovery; stats coverage in response
OpChainSchedulerService.java	O(1) cancel; completion listener registration for stream-mode stats push
OpChainExecutionContext.java	Op→PlanNode map captured at construction time
OperatorTypeDescriptor.java	SPI interface: getId(), name(), getStatKeyClass(), mergeInto(), updateServerMetrics()
OperatorTypeRegistry.java	Static registry: built-ins loaded from Type.values(), plugins discovered via ServiceLoader
StreamStatsReportingIntegrationTest.java	Integration tests: simple aggregation, join (≥3 stages), three-way UNION (N-ary set op regression), cluster-level config activation

Test plan

• StreamStatsReportingIntegrationTest — 4 tests covering simple aggregation, join, three-way UNION (N-ary set-op regression), cluster-level config; all pass
• MultiStageStatsTreeEncoderTest / MultiStageStatsTreeDecoderTest — round-trip, merge, shape-mismatch handling
• StreamingQuerySessionTest — coverage counters, fan-out cancel, latch semantics
• OpChainSchedulerServiceTest — O(1) cancel, multi-opchain context cleanup, completion listener lifecycle
• DispatchClientTest / StreamingDispatchObserverTest — SubmitWithStream client-side observer, UNIMPLEMENTED fallback
• OperatorTypeRegistryTest — all 16 built-in types registered, unknown id returns null, descriptor methods delegate to enum
• Legacy mode unchanged: all existing MSE integration tests unaffected (stream mode is opt-in, default off)

🤖 Generated with Claude Code

[codecov-commenter]

Codecov Report

❌ Patch coverage is 48.44633% with 365 lines in your changes missing coverage. Please review.
✅ Project coverage is 63.63%. Comparing base (bd76ff0) to head (886ab99).
⚠️ Report is 5 commits behind head on master.

Files with missing lines	Patch %	Lines
.../pinot/query/service/dispatch/QueryDispatcher.java	2.96%	130 Missing and 1 partial ⚠️
...apache/pinot/query/service/server/QueryServer.java	1.56%	126 Missing ⚠️
...requesthandler/MultiStageBrokerRequestHandler.java	4.16%	23 Missing ⚠️
...e/pinot/query/runtime/plan/StageStatsTreeNode.java	46.34%	21 Missing and 1 partial ⚠️
.../dispatch/streaming/StreamingDispatchObserver.java	70.00%	12 Missing and 3 partials ⚠️
...va/org/apache/pinot/query/runtime/QueryRunner.java	46.15%	6 Missing and 1 partial ⚠️
...vice/dispatch/streaming/StreamingQuerySession.java	93.20%	6 Missing and 1 partial ⚠️
...t/query/runtime/operator/OperatorTypeRegistry.java	53.84%	5 Missing and 1 partial ⚠️
...e/pinot/query/service/dispatch/DispatchClient.java	0.00%	6 Missing ⚠️
...query/runtime/plan/MultiStageStatsTreeEncoder.java	87.50%	4 Missing and 1 partial ⚠️
... and 8 more

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #18458      +/-   ##
============================================
- Coverage     63.68%   63.63%   -0.05%     
+ Complexity     1685     1684       -1     
============================================
  Files          3266     3273       +7     
  Lines        199821   200459     +638     
  Branches      31022    31108      +86     
============================================
+ Hits         127257   127564     +307     
- Misses        62425    62742     +317     
- Partials      10139    10153      +14     

Flag	Coverage Δ
custom-integration1	100.00% <ø> (ø)
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-21	63.63% <48.44%> (-0.05%)	⬇️
temurin	63.63% <48.44%> (-0.05%)	⬇️
unittests	63.63% <48.44%> (-0.05%)	⬇️
unittests1	55.81% <50.00%> (-0.02%)	⬇️
unittests2	34.83% <1.27%> (-0.12%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[xiangfu0]

See inline comment.

[yashmayya]

Review focused on performance/throughput overhead vs the existing legacy path (per the discussion on #18375). Architecture-wise this is sound — data path is unchanged, payload sizes are comparable, broker accumulator semantics look correct. But there are a handful of concrete things in this implementation that will cap throughput / regress latency at high QPS before legacy mode would, and that I'd want to see addressed before flipping _streamStats=true by default.

Headline concerns (each has an inline below):

• No 50ms wait window cap — submitAndReduceWithStream waits the full remaining query timeout for outstanding OpChainCompletes. A single stuck opchain blocks the client response until query deadline. The design doc on #18375 explicitly specified a min(50ms, remainingTimeout) hard cap.
• HTTP/2 stream limit — one ManagedChannel per server, no MAX_CONCURRENT_STREAMS bump. Long-lived streams now hold HTTP/2 slots for the full query lifetime → silent throughput cap at high QPS × many servers.
• Decode + merge on Netty event loops under a per-query lock — cross-query head-of-line blocking risk at high QPS.
• Error path also waits full timeout — strict regression vs legacy, which returns immediately on cancel.

Non-blocking but worth fixing before scale-up:

• Cancel fan-out is undebounced (transient errors amplify into cancel storms across N peers).
• Encoder failure on error path ⇒ broker has strictly less stats coverage than legacy on that opchain.
• Per-opchain IdentityHashMap allocations now run on the legacy hot path too.
• decodeNode has no depth guard.

Suggested validation before defaulting on: representative high-QPS soak (e.g. 5K QPS × 200 servers) measuring p50/p99/throughput in legacy vs stream; capture heap delta and Netty event-loop stall metrics. The success path of stream mode is comparable to legacy in theory — the threading + connection-cap details below need to be sorted for that to hold in practice. No correctness blockers — merge semantics, coverage accounting, O(1) cancel, and N-ary set-op reconstruction all look right.

[yashmayya]

No wait-window cap (latency regression). This uses the full remaining query timeout, but the design doc on #18375 specified min(50ms, remainingTimeout) as a hard cap. As written, a single stuck opchain on any server holds the broker response back until the entire query budget elapses — e.g. a 10s-timeout query that has data in hand at t=100ms but is missing one OpChainComplete will return to the client at t=10000ms instead of ~t=150ms.

For p99-sensitive workloads this is a material regression vs legacy mode. Suggest a configurable cap (min(STATS_WAIT_CAP_MS, remainingMs)) with a 50–100ms default. Same issue applies in tryRecoverWithStream below.

[yashmayya]

HTTP/2 stream cap will become the throughput ceiling. Each DispatchClient still holds a single ManagedChannel per server (the TODO at the top of the class notes this), and maxConcurrentCallsPerConnection isn't bumped on either side. Behavioral change vs unary submit: a stream slot used to be held for milliseconds, submitWithStream holds one for the full query lifetime.

With Q in-flight queries × N servers fanned out, each broker→server HTTP/2 connection needs ~Q concurrent streams. Netty/gRPC default MAX_CONCURRENT_STREAMS=100 means at high QPS new calls queue at the channel level — manifesting as broker slowness, not stream errors. Recommend before flipping _streamStats=true by default:

• Bump server-side NettyServerBuilder.maxConcurrentCallsPerConnection in QueryServer.buildGrpcServer (today's code doesn't set it).
• Implement the channel-pool TODO above this class, or document the per-server stream-count ceiling.
• Add numActiveStreamsByServer metric for capacity observability.

[gortiz]

This is incorrect. First, MAX_CONCURRENT_STREAMS is Integer.MAX_VALUE in Java (it seems it is 100 in Go). Anyway, the idea that new calls queue up in the stream make sense, but that should not be a problem.

Specifically, here we are not adding much noise. The new streams are not busy, they just send a small message when the opchains finish, which is data we don't send through the mailbox. In fact, the number of bytes being sent here is smaller, as the mailbox mechanism needs to send the stats of the deeper stages multiple times (as many as stages between them and the root stage).

That is why the suggested metric doesn't make much sense, as even the number of calls per connection is high, these should be silent calls.

Instead, I created two other metrics:

• MSE_STREAM_STATS_INCOMPLETE_COVERAGE: number of queries that returned partial stats
• MSE_STREAM_STATS_QUERIES: total queries using the new mechanism

[yashmayya]

Decode + merge under _lock on the Netty event loop. The class Javadoc says "doing it on the I/O thread is fine" — true at low QPS, but two real problems show up under load:

1. Cross-query head-of-line blocking. Netty client event loops are shared across all open streams (across queries). A slow decode here delays OpChainComplete delivery for other queries on the same loop.
2. Lock contention amplification. Two servers reporting near-simultaneously for the same query busy-spin the second I/O thread until the first releases.

MultiStageStatsTreeDecoder.decode is a recursive walk that allocates per node and deserializes StatMap bytes; mergeIntoAccumulatorLocked then walks again calling StatMap.merge per node. Suggest moving decode+merge to a small worker pool and keeping only the _completionLatch.countDown() + accumulator put under _lock. At minimum, please verify under a representative high-QPS soak before flipping the cluster default.

[gortiz]

This is a valid concern. I've modified the code to decode outside the lock, but we need the lock to protect the merge code. I think that, now that we have virtual threads, the best way to resolve this is to have a vthread handle decoding and merging the stats for each query. This vthread would read the stats in raw format from a queue, decode and merge (without having to use locks), similar to an actor. I added a comment suggesting that, but didn't want to introduce that to not make this PR more conflict

[yashmayya]

Cancel fan-out is undebounced. First success=false or transport-level error immediately broadcasts cancel to all N-1 peer servers. At scale this amplifies in two ways:

• A flaky server returning errors on a small fraction of queries now triggers N-1 cancel sends per affected query — at high QPS during a partial degradation this is a meaningful self-inflicted cancel storm.
• recordStreamError also triggers fan-out on any onError (broker GC pause beyond keep-alive, idle timeout, network reset). Premature-cancel probability grows with N × queryDuration.

Consider:

• Debounce (cancel only after K errors / K ms).
• Distinguish "server returned error" from "stream transport closed" — transport blips shouldn't kill peers.
• Emit a cancelFanoutsBroadcast{reason} metric so this is observable in prod.

[gortiz]

A flaky server returning errors on a small fraction of queries now triggers N-1 cancel sends per affected query — at high QPS during a partial degradation, this is a meaningful self-inflicted cancel storm.

The alternative is the system we have now, where opchains are not being stopped until they timeout or finish, even if there is no reader interested in their computations because they are early-terminated. The issue here is not so much in sending N-1 grpc messages but in how expensive it is to cancel each query, which was somehow expensive before because OpChaiScheduler needed to iterate over for (Map.Entry<OpChainId, Pair<MultiStageOperator, QueryExecutionContext>> entry : _opChainCache.asMap() looking for the opchains with the given id. Now the process is quite simpler.

Anyway, I think we can play it safe here and keep the original behavior unless a query option/broker config is used. I'm going to change that

[yashmayya]

Two concerns here:

1. Worse error-path coverage than legacy. When the encoder throws (most likely on treeSize != flatSize from an opchain that failed before emitting EOS), this emits success=false with empty stats. Stream mode also suppresses mailbox stats (see effectiveSendStats), so for this opchain the broker now has strictly less info than legacy mode would deliver via mailbox EOS or the cancel-RPC fallback. Worth either:
   • documenting this as a behavior change in the PR description, or
   • falling back to a best-effort flat encoder that emits whatever entries the flat list has, even when the tree shape doesn't align.
2. Encode + onNext under _streamLock on the opchain executor thread, with NettyServerBuilder.directExecutor() set at class init (line 217) — inbound handlers also run on the same loop. Wide trees stall the loop both directions. Suggest moving encode + send to a small bounded worker pool.

[yashmayya]

This synchronously kicks decode + merge on the gRPC client event loop via _session.recordOpChainComplete. See companion concern on StreamingQuerySession#recordOpChainComplete — at high QPS, expensive per-event work on the shared client event loop causes cross-query HOL blocking.

Aside: _opChainsReportedForThisServer is a plain int. The class Javadoc correctly notes gRPC serializes inbound callbacks per stream, but cancel() on this same instance is called from a different I/O thread (via fanOutCancel's iteration of _openStreams). Today that's safe because cancel() only touches _outbound (which is volatile), but the threading boundary is subtle — a short note on what runs on which thread would help future maintainers.

[gortiz]

This synchronously kicks decode + merge on the gRPC client event loop via _session.recordOpChainComplete. See companion concern on StreamingQuerySession#recordOpChainComplete — at high QPS, expensive per-event work on the shared client event loop causes cross-query HOL blocking.

This is the same as the agent pointed above, right? If that is the case, I've moved the decode phase outside the lock. Merging needs the lock and should be fast. In case it is not, we need to refactor the code a bit. I've added a comment indicating how to do that with a vhtread

[gortiz]

Aside: _opChainsReportedForThisServer is a plain int. The class Javadoc correctly notes gRPC serializes inbound callbacks per stream, but cancel() on this same instance is called from a different I/O thread (via fanOutCancel's iteration of _openStreams). Today that's safe because cancel() only touches _outbound (which is volatile), but the threading boundary is subtle — a short note on what runs on which thread would help future maintainers.

I'll improve the javadoc

[yashmayya]

Two concerns about this recursive walk:

1. Unbounded stack depth. Driven by the proto's children list. A pathological / malformed payload (or a future bug-injected loop) can overflow the Netty event loop stack. Add an explicit depth cap (e.g. MAX_OPERATOR_TREE_DEPTH = 64) and short-circuit to DecodeFailedException.
2. Per-query broker memory. Each opchain report keeps a fully-deserialized StageStatsTreeNode graph alive on the broker for the query's wait window. Aggregate broker footprint becomes ~O(numQueries × numServers × operatorsPerStage) — strictly larger than legacy's compact byte-merge buffer. Worth a heap-snapshot comparison vs legacy under representative load before defaulting on.

[yashmayya]

Allocates a fresh HashMap on every query whenever cluster default is on. At high QPS that's one extra alloc + put per query for what is effectively a constant. Cleaner: thread the cluster default through compileQuery / getQueryEnvConf so queryOptions carries it once, and drop this branch.

[yashmayya]

_operatorToPlanNodes and _planNodeIds are allocated and populated for every opchain via PlanNodeToOpChain.assignPlanNodeIds + record, regardless of whether stats reporting will use them. In legacy mode (which remains the default) these maps are never read — pure GC dead weight on the hot path.

At high QPS × deep stages this adds up. Suggest gating the puts in PlanNodeToOpChain behind context.isSendStats() == false (i.e. stream mode active per KEY_OF_STATS_REPORTING_MODE), or making these maps lazy.

[yashmayya]

Error path also waits full remaining timeout — same root issue as the success-path comment up the file. A query that hits an error at t=100ms with a 10s timeout doesn't return to the client until t=10000ms, vs legacy returning immediately on cancel.

For the error path specifically, blocking the client just to collect more stats is the wrong trade — you've already given up on the query. Use a much smaller cap here (e.g. 100–200ms) and represent the rest via streamStatsCoverage[].missing > 0.