Add Adaptive Routing MSE inflight reqs stats

[timothy-e]

Add separate adaptive routing per-server stats for MSE and SSE, preventing one engine's infra issues from poisoning the other's adaptive routing decisions. (We saw a case on a Pinot cluster where MSE latency was impacted but SSE traffic was completely fine. We don't understand this scenario, but without splitting, the happy SSE queries would dilute the signal from MSE, and the MSE signals would cause unnecessry routing changes to SSE queries).

To start, we record per-server in-flight request counts for MSE queries. Latency (recordStatsUponResponseArrival with real values) and per-sender arrival granularity are left to a follow-up PR once per-sender EOS interception is in place (see TODO in submitAndReduce).

• MultiStageBrokerRequestHandler: accept optional ServerRoutingStatsManager via new constructor overload; pass through to QueryDispatcher
• BaseBrokerStarter: pass _serverRoutingStatsManager to the handler
• QueryDispatcher.submitAndReduce: when statsManager is present, call recordStatsForQuerySubmission for each dispatched server after submit(), and decrement inflight for all servers in the finally block via recordStatsUponResponseArrival(..., -1)

Gaps

1. Submit could fail part way through so we never get to the incrementing. To handle this, we either have to:

• split up calculating the server list from the actual dispatch.
• send the stats manager (or a hook) into the dispatch
• be okay with sometimes under-reporting the number of inflight ops for a few servers. This approach is the simplest. Under-reporting is unlikely to result in a systemic bias towards a set of servers, and the stats are only used for adaptive routing.

2. This isn't as accurate as MSE. Servers are marked as inflight from the beginning of the query until the end, which is not equal to the amount of time they were actually working. Resolving this would require passing live stats through the entire execution tree during query execution. Since this is only used for adaptive routing, perfect accuracy isn't required.

• If a cluster generally handles queries that fan out to all/most servers, then most servers will have the same number of inflight requests, so this metric will have neglible impact.
• If a cluster generally handles queries that target a smaller number of servers: every query that involves a degraded server will result in higher inflight requests for each server involved in the query. But every query that does not involve the degraded server will have shorter inflight requests. Eventually, this should average out to identifying the degraded server as the one with the most inflight requests.
• Ultimately, we'll get more signal from the latency - it will just take a few seconds longer to appear.

Testing

With this deployed, we see MSE inflight_requests reported in the metrics.

When we ran MSE queries, then stopped, then ran SSE queries, we saw the inflight requests for each reflect the traffic.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 98.03922% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 64.29%. Comparing base (d88887e) to head (780f0ea).
⚠️ Report is 44 commits behind head on master.

Files with missing lines	Patch %	Lines
...requesthandler/MultiStageBrokerRequestHandler.java	66.66%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #18553      +/-   ##
============================================
+ Coverage     63.68%   64.29%   +0.60%     
+ Complexity     1684     1126     -558     
============================================
  Files          3262     3311      +49     
  Lines        199835   203856    +4021     
  Branches      31034    31726     +692     
============================================
+ Hits         127266   131064    +3798     
+ Misses        62416    62282     -134     
- Partials      10153    10510     +357     

Flag	Coverage Δ
custom-integration1	100.00% <ø> (ø)
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-21	64.29% <98.03%> (+0.60%)	⬆️
temurin	64.29% <98.03%> (+0.60%)	⬆️
unittests	64.28% <98.03%> (+0.60%)	⬆️
unittests1	56.74% <100.00%> (+0.96%)	⬆️
unittests2	35.52% <39.21%> (+0.57%)	⬆️

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.

[timothy-e]

@yashmayya @Jackie-Jiang Can you TAL at this PR that adds initial support for adaptive routing stats? I have a follow up PR that will enable latency tracking. Both are working well on our QA clusters.

[yashmayya]

The overall direction of separating MSE/SSE stats looks right. Six MAJOR concerns below, posted inline:

• A silent fall-through in getStatsMap() lumps TSE queries into the SSE bucket and silently routes non-query-thread callers to SSE.
• The /debug/serverRoutingStats endpoint and getServerRoutingStatsStr() still expose only the SSE map — operators triaging MSE adaptive routing via the debug surface are blind after this PR.
• The async _executorService is a 2-thread pool with FIFO submission but unordered execution — the decrement task can land before the increment task, producing a transient negative inflight count and a wrong EMA sample. Preexisting for SSE, but amplified for MSE because the increment→decrement window shrinks to runReducer duration.
• recordStatsUponResponseArrival(..., -1) skips the latency update — so for MSE, _latencyMsEMA stays at _avgInitializationVal forever, auto-decay never fires, and the new ADAPTIVE_SERVER_MSE_LATENCY_EMA / ADAPTIVE_SERVER_MSE_HYBRID_SCORE gauges export inert values that nonetheless feed HYBRID/LATENCY MSE routing decisions via the standard instance selectors.
• BaseBrokerStarter.createMultiStageBrokerRequestHandler is a protected extension point whose signature changed — downstream broker starters that overrode it will silently stop being invoked. The constructor on MultiStageBrokerRequestHandler itself preserved its old form; this factory didn't.
• QueryDispatcherTest.testStatsManagerRecordsSubmissionAndArrivalForDispatchedServers uses a Mockito.mock and only checks "method called at least once" — it doesn't verify ordering, doesn't exercise the partial-submit()-failure path the inline comments claim to defend, and doesn't cover the null-statsManager overload.

[yashmayya]

MAJOR — debug surface silently omits the new MSE map.

The new _mseServerQueryStatsMap is not exposed by any of the public read APIs of this class:

• getServerRoutingStats() (line 236) still returns only _serverQueryStatsMap, so the /debug/serverRoutingStats REST endpoint (PinotBrokerDebug.java:280) returns SSE-only data after this PR.
• getServerRoutingStatsStr() (line 243) still iterates only _serverQueryStatsMap.

After this PR, the only programmatic view into MSE adaptive-routing state is the Prometheus gauges. Operators triaging routing decisions via the debug REST endpoint will be silently blind for MSE.

Suggest either growing the response payload to Map<QueryType, Map<String, ServerRoutingStatsEntry>> or adding a parallel debug endpoint.

[yashmayya]

MAJOR — two silent fall-throughs in the new dispatch.

1. QueryExecutionContext.QueryType has three values: SSE, MSE, TSE. This ternary lumps TSE into the SSE bucket. Today TimeSeriesRequestHandler doesn't call recordStats*, so this is dormant — but the moment anyone wires TSE adaptive routing, TSE traffic will silently poison SSE stats, which is the exact engine cross-talk the PR exists to prevent.
2. The no-arg getStatsMap() below defaults to the SSE map when QueryThreadContext.getIfAvailable() is null. Any caller invoked from a non-query thread (debug REST handlers, periodic tasks, future SPI consumers) silently reads/writes the SSE map with no warning.

Suggest: switch exhaustively on the enum and either throw or log on an unknown/unset query type, or add a third TSE-map slot up front. At minimum, document on the method that the caller is expected to run with a QueryThreadContext set.

[yashmayya]

MAJOR — increment/decrement can be reordered on the stats-manager executor, causing a transient negative inflight count and a wrong EMA sample.

recordStatsForQuerySubmission (here) and recordStatsUponResponseArrival (line 219 in the finally) both enqueue tasks on ServerRoutingStatsManager._executorService, which is a newFixedThreadPool(2) by default. The FIFO submission queue does NOT guarantee execution order across the two worker threads — for any single server X, the decrement task can acquire the per-server write lock before the increment task. When that happens:

• _numInFlightRequests momentarily becomes -1, then 0 after the increment lands.
• The transient -1 is read by any concurrent fetchNumInFlightRequestsForServer because _numInFlightRequests is volatile.
• The EMA is sampled only inside updateNumInFlightRequestsForQuerySubmission (ServerRoutingStatsEntry.java:99), so it gets sampled at 0 instead of the expected +1.

This race is preexisting for SSE (AsyncQueryResponse uses the same executor), but the increment→decrement window in MSE is just runReducer duration (often milliseconds) rather than a server-side roundtrip — meaningfully more likely to reorder.

Options: serialize per-server (per-server work channel keyed by serverInstanceId), single-thread the executor, or apply the decrement synchronously while submitting the EMA update asynchronously.

[yashmayya]

MAJOR — latency EMA for MSE entries will be stuck at the init value forever, and the new MSE gauges will export that stuck value.

Passing latency = -1 is handled safely by updateStatsUponResponseArrival (it skips updateLatency), but the consequences go beyond "no latency yet":

1. _latencyMsEMA._lastUpdatedTimeMs stays at 0 for every MSE entry, so the EMA's auto-decay scheduled task never fires.
2. _latencyMsEMA.getAverage() returns _avgInitializationVal (default 1.0) for the lifetime of the broker.
3. ADAPTIVE_SERVER_MSE_LATENCY_EMA therefore exports a flat constant — operators wiring dashboards will assume the gauge is broken and stop trusting it.
4. ADAPTIVE_SERVER_MSE_HYBRID_SCORE is inflight^exp * latencyEma, which collapses to inflight^exp * initVal. If _avgInitializationVal == 0.0, the MSE hybrid score is identically 0 for all servers.

This isn't just a v0 inert-data concern: BalancedInstanceSelector and ReplicaGroupInstanceSelector (the common MSE leaf selectors via BaseInstanceSelector._priorityPoolInstanceSelector) DO invoke _adaptiveServerSelector — and via the new getStatsMap(), those reads now hit the MSE map, where these values are stuck. With HYBRID or LATENCY selector configured, MSE routing decisions will be derived from stuck data.

Until the TODO above lands, suggest either: gate ADAPTIVE_SERVER_MSE_LATENCY_EMA / ADAPTIVE_SERVER_MSE_HYBRID_SCORE from being exported (so dashboards don't show misleading flat lines), or fall back to the SSE entry's latency EMA when the MSE entry has never received a real latency sample, or document loudly in the metric Javadoc that these gauges are inert until follow-up.

[yashmayya]

MAJOR — this test does not actually verify the contract its production code claims to defend.

Mockito.verify(statsManager).recordStatsForQuerySubmission(...) only checks the method was called at least once per server. It does not catch any of the bugs this PR's production code comments are explicitly guarding against:

1. Ordering: a future refactor that moves the decrement out of the finally (or breaks the incrementedServers set bookkeeping) still passes this test.
2. Partial-failure path: the inline comment in submitAndReduce says "guarding against a partial failure in submit()". But the only error path exercised here is runReducer throwing NPE due to the mocked MailboxService — i.e., the path where submit() succeeded. The actual partial-submit()-failure case isn't covered.
3. Legacy submitAndReduce(... null) overload: no test asserts the null-statsManager overload doesn't throw or change behavior.
4. Real-executor behavior: the manager is a Mockito mock, so the async executor and per-server lock ordering aren't exercised at all (see related comment about the executor race on QueryDispatcher.java:201).

Suggested additions:

• A test where submit() throws after partial dispatch and asserts incrementedServers is empty / recordStatsUponResponseArrival is NOT called for any server.
• A test exercising the legacy 4-arg submitAndReduce overload.
• An end-to-end test using a real ServerRoutingStatsManager (not a mock) that asserts numInFlight == 0 for every server after submitAndReduce returns successfully — which is what the inline comments actually claim.