[#12507] Implement adaptive pool selection for StrictReplicaGroupInstanceSelector#18947
Open
timothy-e wants to merge 8 commits into
Open
[#12507] Implement adaptive pool selection for StrictReplicaGroupInstanceSelector#18947timothy-e wants to merge 8 commits into
timothy-e wants to merge 8 commits into
Conversation
Adaptive routing and strict replica group routing are incompatible.
Strict replica group routing guarantees that all segments for a given query are served by servers within the same replica group. This is a correctness requirement for upsert tables: because the upsert metadata map (primary key -> latest document) is maintained per-server, every segment for a given primary key must always land on the same server.
StrictReplicaGroupInstanceSelector.select()` calls `ReplicaGroupInstanceSelector.select()`. When adaptive routing is enabled, it `calls `PriorityPoolInstanceSelector.rank()` with the union of all candidate servers across all segments and all replica groups. It then builds a global rank map ordered by latency score. Each segment independently picks the lowest-ranked server from its own candidate list — `min(rank)` across all replicas. Because the ranking is global and each segment selects independently, different segments of the same query can be routed to servers in different replica groups, breaking upserts.
The pool-awareness in `PriorityPoolInstanceSelector` does not save us here. Pool-consistent ranking only activates when the query explicitly includes an `orderedPreferredPools` option. Without it, `rank()` returns a flat latency-ordered list with no pool or replica-group structure, and the cross-replica-group fan-out occurs on every query.
Automatically disables adaptive server selection for StrictReplicaGroup tables in `InstanceSelectorFactory` (upsert tables are required to use this instance selector type). When the broker has adaptive routing enabled globally (e.g. `pinot.broker.adaptive.server.selector.type = HYBRID`), upsert and dedup tables will now automatically fall back to non-adaptive routing without requiring any table config changes.
The change nullifies the `AdaptiveServerSelector` passed to `instanceSelector.init()` when `tableConfig.isUpsertEnabled()` or `tableConfig.isDedupEnabled()` is true. A clarifying comment is also added in `BaseBrokerRoutingManager` at the call site.
We want to allow adaptive routing as much as possible.
**Unit tests**
I reworked `testReplicaGroupAdaptiveServerSelector` to:
1. construct the selector using `InstanceSelectorFactory` to mirror the real code better.
2. extract constants into helper functons
Then I basically duplicated the test for StrictReplicaGroup and asserted that we use servers from the same replica group.
**Manual tests**
Deployed to an internal QA cluster. Created an upsert table, and at table creation, it logged:
```
INFO [InstanceSelectorFactory] [HelixTaskExecutor-message_handle_thread_23:219] Disabling adaptive server selection for StrictReplicaGroupInstanceSelector table: tim_upsert_test_REALTIME
```
(this should happen any time a broker restarts or a table is created)
Then, to validate that the log does what it says it does, I used this script:
```sh
BROKER="localhost:8099"
TABLE="${1:-tim_upsert_test_REALTIME}"
CALLS="${2:-10}"
passed=0
for i in $(seq 1 "$CALLS"); do
response=$(curl -sf "http://$BROKER/debug/routingTable/sql?query=SELECT%20*%20FROM%20$TABLE")
p1=$(echo "$response" | grep -o '__p1' | wc -l)
p2=$(echo "$response" | grep -o '__p2' | wc -l)
p3=$(echo "$response" | grep -o '__p3' | wc -l)
pools_used=0
[[ $p1 -gt 0 ]] && ((pools_used++))
[[ $p2 -gt 0 ]] && ((pools_used++))
[[ $p3 -gt 0 ]] && ((pools_used++))
if [[ $pools_used -eq 1 ]]; then
((passed++))
elif [[ $pools_used -eq 0 ]]; then
echo "[Call $i] ✗ ERROR (no pools found — curl may have failed)"
else
echo "[Call $i] ✗ FAIL (mixed pools: p1=$p1 p2=$p2 p3=$p3)"
fi
sleep 1
done
echo ""
[[ "$passed" -eq "$CALLS" ]] && echo "✓ All $CALLS calls OK — adaptive routing appears disabled" \
|| echo "✗ $((CALLS - passed))/$CALLS calls had mixed pools — adaptive routing may be active"
```
And I ran it for a regular table:
```
./check.sh billing_analytics_customer_change_events_v1_testing_REALTIME 1000
...
✗ 54/1000 calls had mixed pools — adaptive routing may be active
```
(as expected, the non-upsert table sometimes routes across pools).
and for my new upsert table:
```
./check.sh tim_upsert_test_REALTIME 1000
✓ All 1000 calls OK — adaptive routing appears disabled
```
This change affects broker-level query routing at table initialization time. The `LOGGER.info` message ("Disabling adaptive server selection for upsert/dedup table: ...") will be emitted once per upsert/dedup table when routing is initialized, which can be used to confirm the change is active. Revert by reverting this commit.
…t.java Because there's a lot of Adaptive Routing specific logic that will be added to the ReplicaGroupSelector tests, move it to a new file to isolate the behaviour better. Moving this before changing any logic allows the diff in the follow-up commits to be cleaner. Cleaner diffs, logical test separation Unit test only change.
The protected abstract `BaseInstanceSelector::select` method previously returned `Pair<Map<String, String>, Map<String, String>>`, where the positional meaning of left vs. right was only conveyed by a comment. This PR replaces that with a nested `record InstanceMapping(Map<String, String> segmentToInstanceMap, Map<String, String> optionalSegmentToInstanceMap)` defined inside `BaseInstanceSelector`, making the two maps self-documenting at call sites. `SelectionResult` in `InstanceSelector` is also updated to hold the two maps as named fields directly rather than wrapping a `Pair`, removing the last `Pair` usages from the instance-selector package. All three concrete selectors (`BalancedInstanceSelector`, `ReplicaGroupInstanceSelector`, `MultiStageReplicaGroupSelector`). We also needed a small change to pom.xml to change the engine used for removing unused imports because it gets tripped up by java records (unused in Pinot so far). Without this change, lint fails with ``` [ERROR] Step 'removeUnusedImports' found problem in 'BaseInstanceSelector.java': 513:112: error: ';' expected com.google.googlejavaformat.java.FormatterException: 513:112: error: ';' expected ``` or ``` /src/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java: com.google.googlejavaformat.java.FormatterException: 99:5: error: illegal start of expression ``` Reduces reliance on positional `Pair` semantics in a hot routing code path; named record accessors (`segmentToInstanceMap()`, `optionalSegmentToInstanceMap()`) are clearer than `.getLeft()` / `.getRight()`. This simplifies the follow up PR. No runtime impact — purely a refactor of return types within the instance-selector package. Safe to revert.
A previous commit disabled adaptive routing for `StrictReplicaGroupInstanceSelector`. This commit re-enables adaptive routing for `StrictReplicaGroupInstanceSelector` by implementing pool-level (replica-group-level) server selection. Instead of the parent class's per-segment independent adaptive routing (which violates the same-replica-group guarantee), this implementation: 1. Groups all query-relevant candidates by pool (replica group) 2. Scores each replica group by its **worst-case server rank** — since scatter-gather query latency is bounded by the slowest server, we want the group whose bottleneck is least bad 3. Picks the group with the lowest worst-case rank 4. Routes **all** segments to that single replica group Falls back to round-robin through replica groups when no adaptive stats are available yet. Also adds a flag `pinot.broker.adaptive.server.selector.enable.strict.replica.group` (default true) to control this behaviour, in case we discover any bugs with it - we can revert back to the simple "disable AR" A previous commit disabled adaptive routing entirely for strict replica group tables because per-segment AR violates the same-replica-group invariant required for correctness on upsert tables. This PR restores the latency benefits of adaptive routing while preserving correctness, by making the adaptive decision at the replica-group level rather than the segment level. See apache#12507 We deployed this an internal QA cluster, where we have an upsert table. Using a script [grapher.sh](https://git.corp.stripe.com/gist/timothye/077369ba3ad42f762dad02fc4a977f79) that prints the pool that each query is routed to, we see: **Before disabling AR + upsert tables**: (e.g. before this PR): we often hit both pools in a single query (which could lead to correctness results for upsert tables!) <img width="470" alt="00000000011111111112222222222333333333344444444445" src="https://git.corp.stripe.com/user-attachments/assets/3808f61e-eef0-44e2-8ad7-d27e0a314509" /> Correctness issues are still unlikely, because we'd need an upsert to span a segment boundary, and then the old and new segments would need to be returned by two different servers. Say: > * Servers p1-1, p2-1 (each different pools) contain the same segments A, B. > * key abc123 was in segment A, but we sealed A so now it's stored in segment B. > * There was an upsert across segment boundaries - so we want to disregard the value for abc123 from A, and use the value of it from B. > * But server p2-1's segment B is slightly out of date and doesn't have the new value of abc123 yet. > > Then, if I ask for segments A, B and I'm willing to go cross replicas, I could get: > > * segment A from p1-1 (won't give me abc123 because it's stale) > * segment B from p2-1 (won't give me abc123 because it doesn't have it yet) > > so we don't return abc123. > > or > * segment A from p2-1 (will give me abc123 because it doesn't know it's stale yet) > * segment B from p1-1 (will give me abc123 because it knows this is the new value) > > so we double-return abc123. **Steady State with this PR**: we hit p1 and p2 roughly equally <img width="1685" alt="Screenshot 2026-06-19 at 3 55 53 pm" src="https://git.corp.stripe.com/user-attachments/assets/60486140-2e52-41fc-890e-88b97e5591b0" /> **When a server in pool p2 is degraded with this PR**: we stop routing immediately to p2 and start routing to p1. When the server is healthy again, we can resume routing to p2. <img width="1674" alt="Screenshot 2026-06-19 at 4 00 19 pm" src="https://git.corp.stripe.com/user-attachments/assets/d4b30271-6d3a-442b-b3f4-cd641b6e0c3f" /> (We can see in the adaptive routing stats that the score for that server spiked, so it was adaptive routing that drove the move away from p2). <img width="726" alt="Screenshot 2026-06-19 at 4 02 46 pm" src="https://git.corp.stripe.com/user-attachments/assets/0faccda7-c8af-425d-b23d-7719753e9314" />
This was referenced Jul 9, 2026
Contributor
Author
|
@MeihanLi @ankitsultana @vvivekiyer I have a fix for the correctness bug identified in #12507. Can you TAL? |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #18947 +/- ##
============================================
+ Coverage 64.93% 65.05% +0.11%
- Complexity 1347 1403 +56
============================================
Files 3395 3399 +4
Lines 212533 212891 +358
Branches 33515 33581 +66
============================================
+ Hits 138012 138496 +484
+ Misses 63379 63250 -129
- Partials 11142 11145 +3
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
xiangfu0
reviewed
Jul 9, 2026
xiangfu0
left a comment
xiangfu0
left a comment
Contributor
There was a problem hiding this comment.
Found one high-signal issue; see inline comment.
| } | ||
|
|
||
| SegmentInstanceCandidate selected = candidates.stream() | ||
| .filter(c -> c.getPool() == replicaGroupId) |
Contributor
There was a problem hiding this comment.
This treats server pool as the strict replica-group id, but Pinot assignment can place multiple replica groups in one pool when numReplicaGroups > numPools. In that topology this filter can pick different replica groups for different segments that share the same pool, breaking the strict same-replica-group guarantee required by upsert/dedup routing and risking stale or duplicate primary-key results. Please select by the actual replica-group/replica slot, or explicitly enforce one replica group per pool and add a regression test for multiple replica groups sharing a pool.
Contributor
Author
There was a problem hiding this comment.
Good catch, I've updated it to use getIdealStateReplicaId instead.
Pool (infrastructure rack/zone tag) and replica group (logical segment-mirroring position) are orthogonal. When numReplicaGroups > numPools or all servers have FALLBACK_POOL_ID=-1, pool-based grouping merges distinct replica groups into one, allowing mixed selections that violate the strict same-replica-group invariant required for upsert/dedup correctness (stale or duplicate primary-key results). - Replace candidate.getPool() with candidate.getIdealStateReplicaId() for grouping in select() and filtering in selectServersForReplicaGroup() - Emit POOL_SEG_QUERIES metric per actual pool (from selected candidate) rather than per replicaGroupId, matching ReplicaGroupInstanceSelector.selectServers() - Fix test data: set correct idealStateReplicaId per sorted ideal state position - Add regression test with all servers at pool=-1 proving replicaId-based routing
… test names to make test behaviour clearer
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR is split into multiple logical commits for easier review. The bulk of the logic is in "Implement adaptive pool selection for StrictReplicaGroupInstanceSelector" commit.
Summary
Adaptive routing violates the same-replica-group invariant required for correctness on upsert tables. This PR restores the latency benefits of adaptive routing while preserving correctness, by making the routing decision at the replica-group level rather than the server level.
Instead of the parent class's per-segment independent adaptive routing (which violates the same-replica-group guarantee), this implementation:
Replica groups with no stats are preferred.
Also adds a flag
pinot.broker.adaptive.server.selector.enable.strict.replica.group(default true) to control this behaviour, in case we discover any bugs with the implementation - we can switch to a simple "disable adaptive routing for Strict Replica Groups"Fixes #12507
Docs changes:
Commit breakdown
1. Disable adaptive routing for Pinot tables using StrictReplicaGroup
Automatically disables adaptive server selection for StrictReplicaGroup tables in
InstanceSelectorFactory. When the broker has adaptive routing enabled globally (e.g.pinot.broker.adaptive.server.selector.type = HYBRID), upsert and dedup tables will now automatically fall back to non-adaptive routing without requiring any table config changes.2. Move ReplicaGroup selector tests to dedicated ReplicaGroupSelectorTest.java
Because there's a lot of Adaptive Routing specific logic that will be added to the ReplicaGroupSelector tests, move it to a new file to isolate the behaviour better. Moving this before changing any logic allows the diff in the follow-up commits to be cleaner.
3. Replace Pair return type on BaseInstanceSelector::select with a record
The protected abstract
BaseInstanceSelector::selectmethod previously returnedPair<Map<String, String>, Map<String, String>>, where the positional meaning of left vs. right was only conveyed by a comment. This replaces that with a self-documentingrecord.records aren't used yet in the code base, so we needed a small change to pom.xml because the linter would incorrectly error withcom.google.googlejavaformat.java.FormatterException: 99:5: error: illegal start of expression.4. Implement adaptive pool selection for StrictReplicaGroupInstanceSelector
The main feature commit (described in Summary above).
Test Plan
On a broker, with adaptive routing enabled and an upsert table, we used the script
pinot-pool-distribution-grapher.shto show:Correctness issues are still unlikely, because we'd need an upsert to span a segment boundary, and then the old and new segments would need to be returned by two different servers. Say:
Steady State with this PR: we hit p1 and p2 roughly equally
Steady state with this PR and
pinot.broker.adaptive.server.selector.enable.strict.replica.groupdisabled: we perform the same as above.When a server in pool p2 is degraded with this PR: we stop routing immediately to p2 and start routing to p1. When the server is healthy again, we can resume routing to p2.
