From ab39e5bd27e0ca5a2160cb6a1aeed70233ea62e9 Mon Sep 17 00:00:00 2001 From: Timothy Elgersma Date: Thu, 18 Jun 2026 16:43:45 -0400 Subject: [PATCH 1/8] Disable adaptive routing for Pinot tables using StrictReplicaGroup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../InstanceSelectorFactory.java | 11 ++ .../manager/BaseBrokerRoutingManager.java | 2 + .../InstanceSelectorTest.java | 167 +++++++++++------- 3 files changed, 116 insertions(+), 64 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorFactory.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorFactory.java index 49536b93a8fc..1696de7d1e82 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorFactory.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorFactory.java @@ -107,6 +107,17 @@ public static InstanceSelector getInstanceSelector(TableConfig tableConfig, case RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE: { LOGGER.info("Using StrictReplicaGroupInstanceSelector for table: {}", tableNameWithType); instanceSelector = new StrictReplicaGroupInstanceSelector(); + + // Adaptive Routing is incompatible with Strict Replica Group routing, because Adaptive Routing routes + // each segment independently to the lowest-latency server, violating the same-replica-group guarantee. + // TODO: rework StrictReplicaGroupInstanceSelector::select to do pool-based adaptive routing. + // See https://github.com/apache/pinot/issues/12507 + if (adaptiveServerSelector != null) { + LOGGER.info( + "Disabling adaptive routing for StrictReplicaGroupInstanceSelector table: {}", + tableNameWithType); + adaptiveServerSelector = null; + } break; } case RoutingConfig.MULTI_STAGE_REPLICA_GROUP_SELECTOR_TYPE: { diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java index aae1fa8ef7a5..0f49ff6a77b4 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java @@ -768,6 +768,8 @@ private void buildRoutingInternal(String tableNameWithType) { AdaptiveServerSelector adaptiveServerSelector = AdaptiveServerSelectorFactory.getAdaptiveServerSelector(_serverRoutingStatsManager, _pinotConfig); + // StrictReplicaGroupInstanceSelector is incompatible with adaptive routing; + // InstanceSelectorFactory will nullify the selector for such tables automatically. InstanceSelector instanceSelector = InstanceSelectorFactory.getInstanceSelector(tableConfig, _propertyStore, _brokerMetrics, adaptiveServerSelector, _pinotConfig, _routableServerInstanceMap.keySet(), _enabledServerInstanceMap, diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorTest.java index 03e35d251da4..f08059ca168f 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorTest.java @@ -52,6 +52,7 @@ import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.mockito.Mock; import org.mockito.MockitoAnnotations; import org.testng.annotations.AfterMethod; @@ -73,6 +74,9 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; @@ -1858,73 +1862,108 @@ public void testExcludeIdealStateOffline(String selectorType) { assertTrue(selectionResult.getUnavailableSegments().isEmpty()); } + // Shared topology for AR tests: 3 segments across 5 instances in two replica groups. + // segment2 intentionally has instance4 (unranked) so AR falls back to round-robin for that segment. + private static final String AR_INSTANCE0 = "instance0"; + private static final String AR_INSTANCE1 = "instance1"; + private static final String AR_INSTANCE2 = "instance2"; + private static final String AR_INSTANCE3 = "instance3"; + private static final String AR_INSTANCE4 = "instance4"; + private static final String AR_SEGMENT0 = "segment0"; + private static final String AR_SEGMENT1 = "segment1"; + private static final String AR_SEGMENT2 = "segment2"; + private static final List AR_SEGMENTS = + Arrays.asList(AR_SEGMENT0, AR_SEGMENT1, AR_SEGMENT2); + // Rankings: instance3 best → instance2 → instance1 → instance0 worst; instance4 absent (triggers AR fallback) + private static final List> AR_SERVER_RANKS = Arrays.asList( + new ImmutablePair<>(AR_INSTANCE3, 1.0), + new ImmutablePair<>(AR_INSTANCE2, 2.0), + new ImmutablePair<>(AR_INSTANCE1, 3.0), + new ImmutablePair<>(AR_INSTANCE0, 4.0) + ); + + private SegmentStates buildArSegmentStates() { + Map> candidatesMap = new HashMap<>(); + // segment0 → instance0, instance1 + candidatesMap.put(AR_SEGMENT0, Arrays.asList( + new SegmentInstanceCandidate(AR_INSTANCE0, true), + new SegmentInstanceCandidate(AR_INSTANCE1, true))); + // segment1 → instance2, instance3 + candidatesMap.put(AR_SEGMENT1, Arrays.asList( + new SegmentInstanceCandidate(AR_INSTANCE2, true), + new SegmentInstanceCandidate(AR_INSTANCE3, true))); + // segment2 → instance4 (unranked), instance3 + candidatesMap.put(AR_SEGMENT2, Arrays.asList( + new SegmentInstanceCandidate(AR_INSTANCE4, true), + new SegmentInstanceCandidate(AR_INSTANCE3, true))); + return new SegmentStates(candidatesMap, new HashSet<>(AR_SEGMENTS), null); + } + + private ReplicaGroupInstanceSelector buildArSelector(String selectorType, HybridSelector hybridSelector) { + RoutingConfig routingConfig = new RoutingConfig(null, null, selectorType, false); + TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testTable") + .setRoutingConfig(routingConfig).build(); + IdealState idealState = createIdealState(Map.of( + AR_SEGMENT0, List.of(Pair.of(AR_INSTANCE0, ONLINE), Pair.of(AR_INSTANCE1, ONLINE)), + AR_SEGMENT1, List.of(Pair.of(AR_INSTANCE2, ONLINE), Pair.of(AR_INSTANCE3, ONLINE)), + AR_SEGMENT2, List.of(Pair.of(AR_INSTANCE3, ONLINE), Pair.of(AR_INSTANCE4, ONLINE)))); + ExternalView externalView = createExternalView(Map.of( + AR_SEGMENT0, List.of(Pair.of(AR_INSTANCE0, ONLINE), Pair.of(AR_INSTANCE1, ONLINE)), + AR_SEGMENT1, List.of(Pair.of(AR_INSTANCE2, ONLINE), Pair.of(AR_INSTANCE3, ONLINE)), + AR_SEGMENT2, List.of(Pair.of(AR_INSTANCE3, ONLINE), Pair.of(AR_INSTANCE4, ONLINE)))); + return (ReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, + mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, + mock(PinotConfiguration.class), + Set.of(AR_INSTANCE0, AR_INSTANCE1, AR_INSTANCE2, AR_INSTANCE3, AR_INSTANCE4), + Map.of(), idealState, externalView, new HashSet<>(AR_SEGMENTS)); + } + @Test public void testReplicaGroupAdaptiveServerSelector() { - // Arrange - ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); - BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); HybridSelector hybridSelector = mock(HybridSelector.class); - ReplicaGroupInstanceSelector instanceSelector = new ReplicaGroupInstanceSelector(); + ReplicaGroupInstanceSelector instanceSelector = + buildArSelector(REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector); + + assertTrue(instanceSelector instanceof ReplicaGroupInstanceSelector); + assertFalse(instanceSelector instanceof StrictReplicaGroupInstanceSelector); + assertNotNull(instanceSelector._adaptiveServerSelector); + assertNotNull(instanceSelector._priorityPoolInstanceSelector); + + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(AR_SERVER_RANKS); + Pair, Map> result = + instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null); + + // AR prefers the better-ranked server when all candidates are ranked: + // segment0: instance1 (rank 3) over instance0 (rank 4) + // segment1: instance3 (rank 1) over instance2 (rank 2) + // segment2: instance4 unranked → AR falls back to round-robin → index 0 = instance4 + assertEquals(result.getLeft(), Map.of( + AR_SEGMENT0, AR_INSTANCE1, + AR_SEGMENT1, AR_INSTANCE3, + AR_SEGMENT2, AR_INSTANCE4)); + } - // Define instances and segments - String instance0 = "instance0"; - String instance1 = "instance1"; - String instance2 = "instance2"; - String instance3 = "instance3"; - String instance4 = "instance4"; - String segment0 = "segment0"; - String segment1 = "segment1"; - String segment2 = "segment2"; - List segments = Arrays.asList(segment0, segment1, segment2); - - // Define candidates for each segment - Map> instanceCandidatesMap = new HashMap<>(); - // segment0 -> instance0, instance1 - instanceCandidatesMap.put(segment0, - Arrays.asList(new SegmentInstanceCandidate(instance0, true), new SegmentInstanceCandidate(instance1, true))); - // segment1 -> instance2, instance3 - instanceCandidatesMap.put(segment1, - Arrays.asList(new SegmentInstanceCandidate(instance2, true), new SegmentInstanceCandidate(instance3, true))); - // segment2 -> instance3, instance4 // instance4 is not in the hybrid selector's server ranking - instanceCandidatesMap.put(segment2, - Arrays.asList(new SegmentInstanceCandidate(instance4, true), new SegmentInstanceCandidate(instance3, true))); - - IdealState idealState = createIdealState( - Map.of(segment0, List.of(Pair.of(instance0, ONLINE), Pair.of(instance1, ONLINE)), segment1, - List.of(Pair.of(instance2, ONLINE), Pair.of(instance3, ONLINE)), segment2, - List.of(Pair.of(instance3, ONLINE), Pair.of(instance4, ONLINE)))); - - ExternalView externalView = createExternalView( - Map.of(segment0, List.of(Pair.of(instance0, ONLINE), Pair.of(instance1, ONLINE)), segment1, - List.of(Pair.of(instance2, ONLINE), Pair.of(instance3, ONLINE)), segment2, - List.of(Pair.of(instance3, ONLINE), Pair.of(instance4, ONLINE)))); - - instanceSelector.init(_tableConfig, propertyStore, brokerMetrics, hybridSelector, Clock.systemUTC(), - INSTANCE_SELECTOR_CONFIG, Set.of(instance0, instance1, instance2, instance3, instance4), EMPTY_SERVER_MAP, - idealState, externalView, new HashSet<>(segments)); - - // Define the segment states - SegmentStates segmentStates = new SegmentStates(instanceCandidatesMap, new HashSet<>(segments), null); - - // Define server rankings - List> serverRanks = Arrays.asList( - new ImmutablePair<>(instance3, 1.0), - new ImmutablePair<>(instance2, 2.0), - new ImmutablePair<>(instance1, 3.0), - new ImmutablePair<>(instance0, 4.0) - ); - when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(serverRanks); - - // Act - Pair, Map> selectedResult = - instanceSelector.select(segments, 0, segmentStates, null); - - // Assert - Map expectedSelection = new HashMap<>(); - expectedSelection.put(segment0, instance1); - expectedSelection.put(segment1, instance3); - expectedSelection.put(segment2, instance4); - - assertEquals(selectedResult.getLeft(), expectedSelection); + @Test + public void testStrictReplicaGroupDoesNotUseAdaptiveServerSelector() { + HybridSelector hybridSelector = mock(HybridSelector.class); + ReplicaGroupInstanceSelector instanceSelector = + buildArSelector(STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector); + + assertTrue(instanceSelector instanceof StrictReplicaGroupInstanceSelector); + assertNull(instanceSelector._adaptiveServerSelector); + assertNull(instanceSelector._priorityPoolInstanceSelector); + + Pair, Map> result = + instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null); + + // hybridSelector never consulted during routing despite being passed to the factory + verify(hybridSelector, never()).fetchServerRankingsWithScores(any()); + + // Round-robin with requestId=0 picks candidate index 0 per segment — opposite of AR result above, + // which returns {segment0→instance1, segment1→instance3} for the same rankings. + assertEquals(result.getLeft(), Map.of( + AR_SEGMENT0, AR_INSTANCE0, + AR_SEGMENT1, AR_INSTANCE2, + AR_SEGMENT2, AR_INSTANCE4)); } } From d9ecee64472e93af5e7681ada3f717318a92c2a6 Mon Sep 17 00:00:00 2001 From: Timothy Elgersma Date: Mon, 6 Jul 2026 18:44:25 -0400 Subject: [PATCH 2/8] 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. Cleaner diffs, logical test separation Unit test only change. --- .../InstanceSelectorTest.java | 350 -------------- .../ReplicaGroupSelectorTest.java | 439 ++++++++++++++++++ 2 files changed, 439 insertions(+), 350 deletions(-) create mode 100644 pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorTest.java index f08059ca168f..1097556eeddc 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorTest.java @@ -34,13 +34,11 @@ import java.util.Set; import java.util.TreeMap; import java.util.concurrent.TimeUnit; -import org.apache.commons.lang3.tuple.ImmutablePair; import org.apache.commons.lang3.tuple.Pair; import org.apache.helix.model.ExternalView; import org.apache.helix.model.IdealState; import org.apache.helix.store.zk.ZkHelixPropertyStore; import org.apache.helix.zookeeper.datamodel.ZNRecord; -import org.apache.pinot.broker.routing.adaptiveserverselector.HybridSelector; import org.apache.pinot.common.metadata.ZKMetadataProvider; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.metrics.BrokerMetrics; @@ -52,7 +50,6 @@ import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.env.PinotConfiguration; -import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.mockito.Mock; import org.mockito.MockitoAnnotations; import org.testng.annotations.AfterMethod; @@ -74,9 +71,6 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertNotNull; -import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; @@ -107,13 +101,6 @@ public class InstanceSelectorTest { private static final String BALANCED_INSTANCE_SELECTOR = "balanced"; - private static List getSegments() { - return SEGMENTS; - } - - private final static List SEGMENTS = - Arrays.asList("segment0", "segment1", "segment2", "segment3", "segment4", "segment5", "segment6", "segment7", - "segment8", "segment9", "segment10", "segment11"); private static final long NEW_SEGMENT_EXPIRATION_SECONDS = 300; private static final long NEW_SEGMENT_EXPIRATION_MILLIS = TimeUnit.SECONDS.toMillis(NEW_SEGMENT_EXPIRATION_SECONDS); private final static InstanceSelectorConfig INSTANCE_SELECTOR_CONFIG = @@ -767,238 +754,6 @@ public void testInstanceSelector() { assertTrue(selectionResult.getUnavailableSegments().isEmpty()); } - @Test - public void testReplicaGroupInstanceSelectorNumReplicaGroupsToQuery() { - String offlineTableName = "testTable_OFFLINE"; - ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); - BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); - BrokerRequest brokerRequest = mock(BrokerRequest.class); - PinotQuery pinotQuery = mock(PinotQuery.class); - Map queryOptions = new HashMap<>(); - // numReplicas = 3, fanning the query to 2 replica groups - queryOptions.put("numReplicaGroupsToQuery", "2"); - when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery); - when(pinotQuery.getQueryOptions()).thenReturn(queryOptions); - - ReplicaGroupInstanceSelector replicaGroupInstanceSelector = new ReplicaGroupInstanceSelector(); - - Set enabledInstances = new HashSet<>(); - IdealState idealState = new IdealState(offlineTableName); - Map> idealStateSegmentAssignment = idealState.getRecord().getMapFields(); - ExternalView externalView = new ExternalView(offlineTableName); - Map> externalViewSegmentAssignment = externalView.getRecord().getMapFields(); - Set onlineSegments = new HashSet<>(); - - // 12 online segments with each segment having all 3 instances as online - // replicas are 3 - String instance0 = "instance0"; - String instance1 = "instance1"; - String instance2 = "instance2"; - enabledInstances.add(instance0); - enabledInstances.add(instance1); - enabledInstances.add(instance2); - - Map idealStateInstanceStateMap0 = new TreeMap<>(); - Map externalViewInstanceStateMap0 = new TreeMap<>(); - - for (String instance : enabledInstances) { - idealStateInstanceStateMap0.put(instance, ONLINE); - externalViewInstanceStateMap0.put(instance, ONLINE); - } - - List segments = getSegments(); - // add all segments to both idealStateSegmentAssignment and externalViewSegmentAssignment maps and also to online - // segments - for (String segment : segments) { - idealStateSegmentAssignment.put(segment, idealStateInstanceStateMap0); - externalViewSegmentAssignment.put(segment, externalViewInstanceStateMap0); - onlineSegments.add(segment); - } - - replicaGroupInstanceSelector.init(_tableConfig, propertyStore, brokerMetrics, null, Clock.systemUTC(), - INSTANCE_SELECTOR_CONFIG, enabledInstances, EMPTY_SERVER_MAP, idealState, externalView, onlineSegments); - // ReplicaGroupInstanceSelector - // segment0 -> instance0 - // segment2 -> instance0 - // segment4 -> instance0 - // segment6 -> instance0 - // segment8 -> instance0 - // segment10 -> instance0 - // segment1 -> instance1 - // segment3 -> instance1 - // segment5 -> instance1 - // segment7 -> instance1 - // segment9 -> instance1 - // segment11 -> instance1 - - Map expectedReplicaGroupInstanceSelectorResult = new HashMap<>(); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(0), instance0); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(1), instance1); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(2), instance0); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(3), instance1); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(4), instance0); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(5), instance1); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(6), instance0); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(7), instance1); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(8), instance0); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(9), instance1); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(10), instance0); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(11), instance1); - InstanceSelector.SelectionResult selectionResult = replicaGroupInstanceSelector.select(brokerRequest, segments, 0); - assertEquals(selectionResult.getSegmentToInstanceMap(), expectedReplicaGroupInstanceSelectorResult); - assertTrue(selectionResult.getUnavailableSegments().isEmpty()); - } - - @Test - public void testReplicaGroupInstanceSelectorNumReplicaGroupsToQueryGreaterThanReplicas() { - String offlineTableName = "testTable_OFFLINE"; - ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); - BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); - BrokerRequest brokerRequest = mock(BrokerRequest.class); - PinotQuery pinotQuery = mock(PinotQuery.class); - Map queryOptions = new HashMap<>(); - queryOptions.put("numReplicaGroupsToQuery", "4"); - - when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery); - when(pinotQuery.getQueryOptions()).thenReturn(queryOptions); - - ReplicaGroupInstanceSelector replicaGroupInstanceSelector = new ReplicaGroupInstanceSelector(); - - Set enabledInstances = new HashSet<>(); - IdealState idealState = new IdealState(offlineTableName); - Map> idealStateSegmentAssignment = idealState.getRecord().getMapFields(); - ExternalView externalView = new ExternalView(offlineTableName); - Map> externalViewSegmentAssignment = externalView.getRecord().getMapFields(); - Set onlineSegments = new HashSet<>(); - - // 12 online segments with each segment having all 3 instances as online - // replicas are 3 - String instance0 = "instance0"; - String instance1 = "instance1"; - String instance2 = "instance2"; - enabledInstances.add(instance0); - enabledInstances.add(instance1); - enabledInstances.add(instance2); - - List segments = getSegments(); - - Map idealStateInstanceStateMap0 = new TreeMap<>(); - Map externalViewInstanceStateMap0 = new TreeMap<>(); - - for (String instance : enabledInstances) { - idealStateInstanceStateMap0.put(instance, ONLINE); - externalViewInstanceStateMap0.put(instance, ONLINE); - } - - // add all segments to both idealStateSegmentAssignment and externalViewSegmentAssignment maps and also to online - // segments - for (String segment : segments) { - idealStateSegmentAssignment.put(segment, idealStateInstanceStateMap0); - externalViewSegmentAssignment.put(segment, externalViewInstanceStateMap0); - onlineSegments.add(segment); - } - - replicaGroupInstanceSelector.init(_tableConfig, propertyStore, brokerMetrics, null, Clock.systemUTC(), - INSTANCE_SELECTOR_CONFIG, enabledInstances, EMPTY_SERVER_MAP, idealState, externalView, onlineSegments); - // ReplicaGroupInstanceSelector - // segment0 -> instance0 - // segment3 -> instance0 - // segment6 -> instance0 - // segment9 -> instance0 - // segment1 -> instance1 - // segment4 -> instance1 - // segment7 -> instance1 - // segment10 -> instance1 - // segment2 -> instance2 - // segment5 -> instance2 - // segment8 -> instance2 - // segment11 -> instance2 - - Map expectedReplicaGroupInstanceSelectorResult = new HashMap<>(); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(0), instance0); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(1), instance1); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(2), instance2); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(3), instance0); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(4), instance1); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(5), instance2); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(6), instance0); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(7), instance1); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(8), instance2); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(9), instance0); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(10), instance1); - expectedReplicaGroupInstanceSelectorResult.put(segments.get(11), instance2); - InstanceSelector.SelectionResult selectionResult = replicaGroupInstanceSelector.select(brokerRequest, segments, 0); - assertEquals(selectionResult.getSegmentToInstanceMap(), expectedReplicaGroupInstanceSelectorResult); - assertTrue(selectionResult.getUnavailableSegments().isEmpty()); - } - - @Test - public void testReplicaGroupInstanceSelectorNumReplicaGroupsNotSet() { - String offlineTableName = "testTable_OFFLINE"; - ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); - BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); - BrokerRequest brokerRequest = mock(BrokerRequest.class); - PinotQuery pinotQuery = mock(PinotQuery.class); - Map queryOptions = new HashMap<>(); - - when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery); - when(pinotQuery.getQueryOptions()).thenReturn(queryOptions); - - ReplicaGroupInstanceSelector replicaGroupInstanceSelector = new ReplicaGroupInstanceSelector(); - - Set enabledInstances = new HashSet<>(); - IdealState idealState = new IdealState(offlineTableName); - Map> idealStateSegmentAssignment = idealState.getRecord().getMapFields(); - ExternalView externalView = new ExternalView(offlineTableName); - Map> externalViewSegmentAssignment = externalView.getRecord().getMapFields(); - Set onlineSegments = new HashSet<>(); - - // 12 online segments with each segment having all 3 instances as online - // replicas are 3 - String instance0 = "instance0"; - String instance1 = "instance1"; - String instance2 = "instance2"; - enabledInstances.add(instance0); - enabledInstances.add(instance1); - enabledInstances.add(instance2); - - List segments = getSegments(); - - Map idealStateInstanceStateMap0 = new TreeMap<>(); - Map externalViewInstanceStateMap0 = new TreeMap<>(); - - for (String instance : enabledInstances) { - idealStateInstanceStateMap0.put(instance, ONLINE); - externalViewInstanceStateMap0.put(instance, ONLINE); - } - - // add all segments to both idealStateSegmentAssignment and externalViewSegmentAssignment maps and also to online - // segments - for (String segment : segments) { - idealStateSegmentAssignment.put(segment, idealStateInstanceStateMap0); - externalViewSegmentAssignment.put(segment, externalViewInstanceStateMap0); - onlineSegments.add(segment); - } - - replicaGroupInstanceSelector.init(_tableConfig, propertyStore, brokerMetrics, null, Clock.systemUTC(), - INSTANCE_SELECTOR_CONFIG, enabledInstances, EMPTY_SERVER_MAP, idealState, externalView, onlineSegments); - // since numReplicaGroupsToQuery is not set, first query should go to first replica group, - // 2nd query should go to next replica group - - Map expectedReplicaGroupInstanceSelectorResult = new HashMap<>(); - for (String segment : segments) { - expectedReplicaGroupInstanceSelectorResult.put(segment, instance0); - } - InstanceSelector.SelectionResult selectionResult = replicaGroupInstanceSelector.select(brokerRequest, segments, 0); - assertEquals(selectionResult.getSegmentToInstanceMap(), expectedReplicaGroupInstanceSelectorResult); - - for (String segment : segments) { - expectedReplicaGroupInstanceSelectorResult.put(segment, instance1); - } - selectionResult = replicaGroupInstanceSelector.select(brokerRequest, segments, 1); - assertEquals(selectionResult.getSegmentToInstanceMap(), expectedReplicaGroupInstanceSelectorResult); - } - @Test public void testUnavailableSegments() { String offlineTableName = "testTable_OFFLINE"; @@ -1861,109 +1616,4 @@ public void testExcludeIdealStateOffline(String selectorType) { assertEquals(selectionResult.getSegmentToInstanceMap(), expectedBalancedInstanceSelectorResult); assertTrue(selectionResult.getUnavailableSegments().isEmpty()); } - - // Shared topology for AR tests: 3 segments across 5 instances in two replica groups. - // segment2 intentionally has instance4 (unranked) so AR falls back to round-robin for that segment. - private static final String AR_INSTANCE0 = "instance0"; - private static final String AR_INSTANCE1 = "instance1"; - private static final String AR_INSTANCE2 = "instance2"; - private static final String AR_INSTANCE3 = "instance3"; - private static final String AR_INSTANCE4 = "instance4"; - private static final String AR_SEGMENT0 = "segment0"; - private static final String AR_SEGMENT1 = "segment1"; - private static final String AR_SEGMENT2 = "segment2"; - private static final List AR_SEGMENTS = - Arrays.asList(AR_SEGMENT0, AR_SEGMENT1, AR_SEGMENT2); - // Rankings: instance3 best → instance2 → instance1 → instance0 worst; instance4 absent (triggers AR fallback) - private static final List> AR_SERVER_RANKS = Arrays.asList( - new ImmutablePair<>(AR_INSTANCE3, 1.0), - new ImmutablePair<>(AR_INSTANCE2, 2.0), - new ImmutablePair<>(AR_INSTANCE1, 3.0), - new ImmutablePair<>(AR_INSTANCE0, 4.0) - ); - - private SegmentStates buildArSegmentStates() { - Map> candidatesMap = new HashMap<>(); - // segment0 → instance0, instance1 - candidatesMap.put(AR_SEGMENT0, Arrays.asList( - new SegmentInstanceCandidate(AR_INSTANCE0, true), - new SegmentInstanceCandidate(AR_INSTANCE1, true))); - // segment1 → instance2, instance3 - candidatesMap.put(AR_SEGMENT1, Arrays.asList( - new SegmentInstanceCandidate(AR_INSTANCE2, true), - new SegmentInstanceCandidate(AR_INSTANCE3, true))); - // segment2 → instance4 (unranked), instance3 - candidatesMap.put(AR_SEGMENT2, Arrays.asList( - new SegmentInstanceCandidate(AR_INSTANCE4, true), - new SegmentInstanceCandidate(AR_INSTANCE3, true))); - return new SegmentStates(candidatesMap, new HashSet<>(AR_SEGMENTS), null); - } - - private ReplicaGroupInstanceSelector buildArSelector(String selectorType, HybridSelector hybridSelector) { - RoutingConfig routingConfig = new RoutingConfig(null, null, selectorType, false); - TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testTable") - .setRoutingConfig(routingConfig).build(); - IdealState idealState = createIdealState(Map.of( - AR_SEGMENT0, List.of(Pair.of(AR_INSTANCE0, ONLINE), Pair.of(AR_INSTANCE1, ONLINE)), - AR_SEGMENT1, List.of(Pair.of(AR_INSTANCE2, ONLINE), Pair.of(AR_INSTANCE3, ONLINE)), - AR_SEGMENT2, List.of(Pair.of(AR_INSTANCE3, ONLINE), Pair.of(AR_INSTANCE4, ONLINE)))); - ExternalView externalView = createExternalView(Map.of( - AR_SEGMENT0, List.of(Pair.of(AR_INSTANCE0, ONLINE), Pair.of(AR_INSTANCE1, ONLINE)), - AR_SEGMENT1, List.of(Pair.of(AR_INSTANCE2, ONLINE), Pair.of(AR_INSTANCE3, ONLINE)), - AR_SEGMENT2, List.of(Pair.of(AR_INSTANCE3, ONLINE), Pair.of(AR_INSTANCE4, ONLINE)))); - return (ReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, - mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, - mock(PinotConfiguration.class), - Set.of(AR_INSTANCE0, AR_INSTANCE1, AR_INSTANCE2, AR_INSTANCE3, AR_INSTANCE4), - Map.of(), idealState, externalView, new HashSet<>(AR_SEGMENTS)); - } - - @Test - public void testReplicaGroupAdaptiveServerSelector() { - HybridSelector hybridSelector = mock(HybridSelector.class); - ReplicaGroupInstanceSelector instanceSelector = - buildArSelector(REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector); - - assertTrue(instanceSelector instanceof ReplicaGroupInstanceSelector); - assertFalse(instanceSelector instanceof StrictReplicaGroupInstanceSelector); - assertNotNull(instanceSelector._adaptiveServerSelector); - assertNotNull(instanceSelector._priorityPoolInstanceSelector); - - when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(AR_SERVER_RANKS); - Pair, Map> result = - instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null); - - // AR prefers the better-ranked server when all candidates are ranked: - // segment0: instance1 (rank 3) over instance0 (rank 4) - // segment1: instance3 (rank 1) over instance2 (rank 2) - // segment2: instance4 unranked → AR falls back to round-robin → index 0 = instance4 - assertEquals(result.getLeft(), Map.of( - AR_SEGMENT0, AR_INSTANCE1, - AR_SEGMENT1, AR_INSTANCE3, - AR_SEGMENT2, AR_INSTANCE4)); - } - - @Test - public void testStrictReplicaGroupDoesNotUseAdaptiveServerSelector() { - HybridSelector hybridSelector = mock(HybridSelector.class); - ReplicaGroupInstanceSelector instanceSelector = - buildArSelector(STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector); - - assertTrue(instanceSelector instanceof StrictReplicaGroupInstanceSelector); - assertNull(instanceSelector._adaptiveServerSelector); - assertNull(instanceSelector._priorityPoolInstanceSelector); - - Pair, Map> result = - instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null); - - // hybridSelector never consulted during routing despite being passed to the factory - verify(hybridSelector, never()).fetchServerRankingsWithScores(any()); - - // Round-robin with requestId=0 picks candidate index 0 per segment — opposite of AR result above, - // which returns {segment0→instance1, segment1→instance3} for the same rankings. - assertEquals(result.getLeft(), Map.of( - AR_SEGMENT0, AR_INSTANCE0, - AR_SEGMENT1, AR_INSTANCE2, - AR_SEGMENT2, AR_INSTANCE4)); - } } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java new file mode 100644 index 000000000000..2810b5faa09f --- /dev/null +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java @@ -0,0 +1,439 @@ +/** + * 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.broker.routing.instanceselector; + +import java.time.Clock; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import org.apache.commons.lang3.tuple.ImmutablePair; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.helix.model.ExternalView; +import org.apache.helix.model.IdealState; +import org.apache.helix.store.zk.ZkHelixPropertyStore; +import org.apache.helix.zookeeper.datamodel.ZNRecord; +import org.apache.pinot.broker.routing.adaptiveserverselector.HybridSelector; +import org.apache.pinot.common.metrics.BrokerMetrics; +import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.core.transport.ServerInstance; +import org.apache.pinot.spi.config.table.RoutingConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.apache.pinot.spi.config.table.RoutingConfig.REPLICA_GROUP_INSTANCE_SELECTOR_TYPE; +import static org.apache.pinot.spi.config.table.RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE; +import static org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel.ONLINE; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +/** + * Tests for {@link ReplicaGroupInstanceSelector} and {@link StrictReplicaGroupInstanceSelector}, + * including adaptive server selection (pool-level routing for strict replica groups). + */ +@SuppressWarnings("unchecked") +public class ReplicaGroupSelectorTest { + private AutoCloseable _mocks; + + @Mock + private TableConfig _tableConfig; + + private static final String TABLE_NAME = "testTable_OFFLINE"; + private static final Map EMPTY_SERVER_MAP = Collections.EMPTY_MAP; + private static final InstanceSelectorConfig INSTANCE_SELECTOR_CONFIG = new InstanceSelectorConfig(false, 300, false); + private static final List SEGMENTS = + Arrays.asList("segment0", "segment1", "segment2", "segment3", "segment4", "segment5", "segment6", "segment7", + "segment8", "segment9", "segment10", "segment11"); + + @BeforeMethod + public void setUp() { + _mocks = MockitoAnnotations.openMocks(this); + when(_tableConfig.getTableName()).thenReturn(TABLE_NAME); + } + + @AfterMethod + public void tearDown() + throws Exception { + _mocks.close(); + } + + // --- Shared helpers --- + + static IdealState createIdealState(Map>> segmentState) { + IdealState idealState = new IdealState(TABLE_NAME); + Map> idealStateSegmentAssignment = idealState.getRecord().getMapFields(); + for (Map.Entry>> entry : segmentState.entrySet()) { + Map instanceStateMap = new TreeMap<>(); + for (Pair instanceState : entry.getValue()) { + instanceStateMap.put(instanceState.getLeft(), instanceState.getRight()); + } + idealStateSegmentAssignment.put(entry.getKey(), instanceStateMap); + } + return idealState; + } + + static ExternalView createExternalView(Map>> segmentState) { + ExternalView externalView = new ExternalView(TABLE_NAME); + Map> externalViewSegmentAssignment = externalView.getRecord().getMapFields(); + for (Map.Entry>> entry : segmentState.entrySet()) { + Map instanceStateMap = new TreeMap<>(); + for (Pair instanceState : entry.getValue()) { + instanceStateMap.put(instanceState.getLeft(), instanceState.getRight()); + } + externalViewSegmentAssignment.put(entry.getKey(), instanceStateMap); + } + return externalView; + } + + // --- ReplicaGroupInstanceSelector: numReplicaGroupsToQuery tests --- + + @Test + public void testReplicaGroupInstanceSelectorNumReplicaGroupsToQuery() { + String offlineTableName = "testTable_OFFLINE"; + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); + BrokerRequest brokerRequest = mock(BrokerRequest.class); + PinotQuery pinotQuery = mock(PinotQuery.class); + Map queryOptions = new HashMap<>(); + // numReplicas = 3, fanning the query to 2 replica groups + queryOptions.put("numReplicaGroupsToQuery", "2"); + when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery); + when(pinotQuery.getQueryOptions()).thenReturn(queryOptions); + + ReplicaGroupInstanceSelector replicaGroupInstanceSelector = new ReplicaGroupInstanceSelector(); + + Set enabledInstances = new HashSet<>(); + IdealState idealState = new IdealState(offlineTableName); + Map> idealStateSegmentAssignment = idealState.getRecord().getMapFields(); + ExternalView externalView = new ExternalView(offlineTableName); + Map> externalViewSegmentAssignment = externalView.getRecord().getMapFields(); + Set onlineSegments = new HashSet<>(); + + // 12 online segments with each segment having all 3 instances as online + // replicas are 3 + String instance0 = "instance0"; + String instance1 = "instance1"; + String instance2 = "instance2"; + enabledInstances.add(instance0); + enabledInstances.add(instance1); + enabledInstances.add(instance2); + + Map idealStateInstanceStateMap0 = new TreeMap<>(); + Map externalViewInstanceStateMap0 = new TreeMap<>(); + + for (String instance : enabledInstances) { + idealStateInstanceStateMap0.put(instance, ONLINE); + externalViewInstanceStateMap0.put(instance, ONLINE); + } + + List segments = SEGMENTS; + // add all segments to both idealStateSegmentAssignment and externalViewSegmentAssignment maps and also to online + // segments + for (String segment : segments) { + idealStateSegmentAssignment.put(segment, idealStateInstanceStateMap0); + externalViewSegmentAssignment.put(segment, externalViewInstanceStateMap0); + onlineSegments.add(segment); + } + + replicaGroupInstanceSelector.init(_tableConfig, propertyStore, brokerMetrics, null, Clock.systemUTC(), + INSTANCE_SELECTOR_CONFIG, enabledInstances, EMPTY_SERVER_MAP, idealState, externalView, onlineSegments); + + Map expectedReplicaGroupInstanceSelectorResult = new HashMap<>(); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(0), instance0); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(1), instance1); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(2), instance0); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(3), instance1); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(4), instance0); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(5), instance1); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(6), instance0); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(7), instance1); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(8), instance0); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(9), instance1); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(10), instance0); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(11), instance1); + InstanceSelector.SelectionResult selectionResult = replicaGroupInstanceSelector.select(brokerRequest, segments, 0); + assertEquals(selectionResult.getSegmentToInstanceMap(), expectedReplicaGroupInstanceSelectorResult); + assertTrue(selectionResult.getUnavailableSegments().isEmpty()); + } + + @Test + public void testReplicaGroupInstanceSelectorNumReplicaGroupsToQueryGreaterThanReplicas() { + String offlineTableName = "testTable_OFFLINE"; + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); + BrokerRequest brokerRequest = mock(BrokerRequest.class); + PinotQuery pinotQuery = mock(PinotQuery.class); + Map queryOptions = new HashMap<>(); + queryOptions.put("numReplicaGroupsToQuery", "4"); + + when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery); + when(pinotQuery.getQueryOptions()).thenReturn(queryOptions); + + ReplicaGroupInstanceSelector replicaGroupInstanceSelector = new ReplicaGroupInstanceSelector(); + + Set enabledInstances = new HashSet<>(); + IdealState idealState = new IdealState(offlineTableName); + Map> idealStateSegmentAssignment = idealState.getRecord().getMapFields(); + ExternalView externalView = new ExternalView(offlineTableName); + Map> externalViewSegmentAssignment = externalView.getRecord().getMapFields(); + Set onlineSegments = new HashSet<>(); + + // 12 online segments with each segment having all 3 instances as online + // replicas are 3 + String instance0 = "instance0"; + String instance1 = "instance1"; + String instance2 = "instance2"; + enabledInstances.add(instance0); + enabledInstances.add(instance1); + enabledInstances.add(instance2); + + List segments = SEGMENTS; + + Map idealStateInstanceStateMap0 = new TreeMap<>(); + Map externalViewInstanceStateMap0 = new TreeMap<>(); + + for (String instance : enabledInstances) { + idealStateInstanceStateMap0.put(instance, ONLINE); + externalViewInstanceStateMap0.put(instance, ONLINE); + } + + // add all segments to both idealStateSegmentAssignment and externalViewSegmentAssignment maps and also to online + // segments + for (String segment : segments) { + idealStateSegmentAssignment.put(segment, idealStateInstanceStateMap0); + externalViewSegmentAssignment.put(segment, externalViewInstanceStateMap0); + onlineSegments.add(segment); + } + + replicaGroupInstanceSelector.init(_tableConfig, propertyStore, brokerMetrics, null, Clock.systemUTC(), + INSTANCE_SELECTOR_CONFIG, enabledInstances, EMPTY_SERVER_MAP, idealState, externalView, onlineSegments); + + Map expectedReplicaGroupInstanceSelectorResult = new HashMap<>(); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(0), instance0); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(1), instance1); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(2), instance2); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(3), instance0); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(4), instance1); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(5), instance2); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(6), instance0); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(7), instance1); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(8), instance2); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(9), instance0); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(10), instance1); + expectedReplicaGroupInstanceSelectorResult.put(segments.get(11), instance2); + InstanceSelector.SelectionResult selectionResult = replicaGroupInstanceSelector.select(brokerRequest, segments, 0); + assertEquals(selectionResult.getSegmentToInstanceMap(), expectedReplicaGroupInstanceSelectorResult); + assertTrue(selectionResult.getUnavailableSegments().isEmpty()); + } + + @Test + public void testReplicaGroupInstanceSelectorNumReplicaGroupsNotSet() { + String offlineTableName = "testTable_OFFLINE"; + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); + BrokerRequest brokerRequest = mock(BrokerRequest.class); + PinotQuery pinotQuery = mock(PinotQuery.class); + Map queryOptions = new HashMap<>(); + + when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery); + when(pinotQuery.getQueryOptions()).thenReturn(queryOptions); + + ReplicaGroupInstanceSelector replicaGroupInstanceSelector = new ReplicaGroupInstanceSelector(); + + Set enabledInstances = new HashSet<>(); + IdealState idealState = new IdealState(offlineTableName); + Map> idealStateSegmentAssignment = idealState.getRecord().getMapFields(); + ExternalView externalView = new ExternalView(offlineTableName); + Map> externalViewSegmentAssignment = externalView.getRecord().getMapFields(); + Set onlineSegments = new HashSet<>(); + + // 12 online segments with each segment having all 3 instances as online + // replicas are 3 + String instance0 = "instance0"; + String instance1 = "instance1"; + String instance2 = "instance2"; + enabledInstances.add(instance0); + enabledInstances.add(instance1); + enabledInstances.add(instance2); + + List segments = SEGMENTS; + + Map idealStateInstanceStateMap0 = new TreeMap<>(); + Map externalViewInstanceStateMap0 = new TreeMap<>(); + + for (String instance : enabledInstances) { + idealStateInstanceStateMap0.put(instance, ONLINE); + externalViewInstanceStateMap0.put(instance, ONLINE); + } + + // add all segments to both idealStateSegmentAssignment and externalViewSegmentAssignment maps and also to online + // segments + for (String segment : segments) { + idealStateSegmentAssignment.put(segment, idealStateInstanceStateMap0); + externalViewSegmentAssignment.put(segment, externalViewInstanceStateMap0); + onlineSegments.add(segment); + } + + replicaGroupInstanceSelector.init(_tableConfig, propertyStore, brokerMetrics, null, Clock.systemUTC(), + INSTANCE_SELECTOR_CONFIG, enabledInstances, EMPTY_SERVER_MAP, idealState, externalView, onlineSegments); + // since numReplicaGroupsToQuery is not set, first query should go to first replica group, + // 2nd query should go to next replica group + + Map expectedReplicaGroupInstanceSelectorResult = new HashMap<>(); + for (String segment : segments) { + expectedReplicaGroupInstanceSelectorResult.put(segment, instance0); + } + InstanceSelector.SelectionResult selectionResult = replicaGroupInstanceSelector.select(brokerRequest, segments, 0); + assertEquals(selectionResult.getSegmentToInstanceMap(), expectedReplicaGroupInstanceSelectorResult); + + for (String segment : segments) { + expectedReplicaGroupInstanceSelectorResult.put(segment, instance1); + } + selectionResult = replicaGroupInstanceSelector.select(brokerRequest, segments, 1); + assertEquals(selectionResult.getSegmentToInstanceMap(), expectedReplicaGroupInstanceSelectorResult); + } + + // --- Adaptive server selection tests --- + + // Shared topology for AR tests: 3 segments across 5 instances in two replica groups. + // segment2 intentionally has instance4 (unranked) so AR falls back to round-robin for that segment. + private static final String AR_INSTANCE0 = "instance0"; + private static final String AR_INSTANCE1 = "instance1"; + private static final String AR_INSTANCE2 = "instance2"; + private static final String AR_INSTANCE3 = "instance3"; + private static final String AR_INSTANCE4 = "instance4"; + private static final String AR_SEGMENT0 = "segment0"; + private static final String AR_SEGMENT1 = "segment1"; + private static final String AR_SEGMENT2 = "segment2"; + private static final List AR_SEGMENTS = + Arrays.asList(AR_SEGMENT0, AR_SEGMENT1, AR_SEGMENT2); + // Rankings: instance3 best → instance2 → instance1 → instance0 worst; instance4 absent (triggers AR fallback) + private static final List> AR_SERVER_RANKS = Arrays.asList( + new ImmutablePair<>(AR_INSTANCE3, 1.0), + new ImmutablePair<>(AR_INSTANCE2, 2.0), + new ImmutablePair<>(AR_INSTANCE1, 3.0), + new ImmutablePair<>(AR_INSTANCE0, 4.0) + ); + + private SegmentStates buildArSegmentStates() { + Map> candidatesMap = new HashMap<>(); + // segment0 → instance0, instance1 + candidatesMap.put(AR_SEGMENT0, Arrays.asList( + new SegmentInstanceCandidate(AR_INSTANCE0, true), + new SegmentInstanceCandidate(AR_INSTANCE1, true))); + // segment1 → instance2, instance3 + candidatesMap.put(AR_SEGMENT1, Arrays.asList( + new SegmentInstanceCandidate(AR_INSTANCE2, true), + new SegmentInstanceCandidate(AR_INSTANCE3, true))); + // segment2 → instance4 (unranked), instance3 + candidatesMap.put(AR_SEGMENT2, Arrays.asList( + new SegmentInstanceCandidate(AR_INSTANCE4, true), + new SegmentInstanceCandidate(AR_INSTANCE3, true))); + return new SegmentStates(candidatesMap, new HashSet<>(AR_SEGMENTS), null); + } + + private ReplicaGroupInstanceSelector buildArSelector(String selectorType, HybridSelector hybridSelector) { + RoutingConfig routingConfig = new RoutingConfig(null, null, selectorType, false); + TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testTable") + .setRoutingConfig(routingConfig).build(); + IdealState idealState = createIdealState(Map.of( + AR_SEGMENT0, List.of(Pair.of(AR_INSTANCE0, ONLINE), Pair.of(AR_INSTANCE1, ONLINE)), + AR_SEGMENT1, List.of(Pair.of(AR_INSTANCE2, ONLINE), Pair.of(AR_INSTANCE3, ONLINE)), + AR_SEGMENT2, List.of(Pair.of(AR_INSTANCE3, ONLINE), Pair.of(AR_INSTANCE4, ONLINE)))); + ExternalView externalView = createExternalView(Map.of( + AR_SEGMENT0, List.of(Pair.of(AR_INSTANCE0, ONLINE), Pair.of(AR_INSTANCE1, ONLINE)), + AR_SEGMENT1, List.of(Pair.of(AR_INSTANCE2, ONLINE), Pair.of(AR_INSTANCE3, ONLINE)), + AR_SEGMENT2, List.of(Pair.of(AR_INSTANCE3, ONLINE), Pair.of(AR_INSTANCE4, ONLINE)))); + return (ReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, + mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, + mock(PinotConfiguration.class), + Set.of(AR_INSTANCE0, AR_INSTANCE1, AR_INSTANCE2, AR_INSTANCE3, AR_INSTANCE4), + Map.of(), idealState, externalView, new HashSet<>(AR_SEGMENTS)); + } + + @Test + public void testReplicaGroupAdaptiveServerSelector() { + HybridSelector hybridSelector = mock(HybridSelector.class); + ReplicaGroupInstanceSelector instanceSelector = + buildArSelector(REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector); + + assertTrue(instanceSelector instanceof ReplicaGroupInstanceSelector); + assertFalse(instanceSelector instanceof StrictReplicaGroupInstanceSelector); + assertNotNull(instanceSelector._adaptiveServerSelector); + assertNotNull(instanceSelector._priorityPoolInstanceSelector); + + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(AR_SERVER_RANKS); + Pair, Map> result = + instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null); + + // AR prefers the better-ranked server when all candidates are ranked: + // segment0: instance1 (rank 3) over instance0 (rank 4) + // segment1: instance3 (rank 1) over instance2 (rank 2) + // segment2: instance4 unranked → AR falls back to round-robin → index 0 = instance4 + assertEquals(result.getLeft(), Map.of( + AR_SEGMENT0, AR_INSTANCE1, + AR_SEGMENT1, AR_INSTANCE3, + AR_SEGMENT2, AR_INSTANCE4)); + } + + @Test + public void testStrictReplicaGroupDoesNotUseAdaptiveServerSelector() { + HybridSelector hybridSelector = mock(HybridSelector.class); + ReplicaGroupInstanceSelector instanceSelector = + buildArSelector(STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector); + + assertTrue(instanceSelector instanceof StrictReplicaGroupInstanceSelector); + assertNull(instanceSelector._adaptiveServerSelector); + assertNull(instanceSelector._priorityPoolInstanceSelector); + + Pair, Map> result = + instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null); + + // hybridSelector never consulted during routing despite being passed to the factory + verify(hybridSelector, never()).fetchServerRankingsWithScores(any()); + + // Round-robin with requestId=0 picks candidate index 0 per segment — opposite of AR result above, + // which returns {segment0→instance1, segment1→instance3} for the same rankings. + assertEquals(result.getLeft(), Map.of( + AR_SEGMENT0, AR_INSTANCE0, + AR_SEGMENT1, AR_INSTANCE2, + AR_SEGMENT2, AR_INSTANCE4)); + } +} From dc75cfd7f75c94ce78386ac4c49757c6f75963f6 Mon Sep 17 00:00:00 2001 From: Timothy Elgersma Date: Tue, 7 Jul 2026 17:24:35 -0400 Subject: [PATCH 3/8] Replace Pair return type on BaseInstanceSelector::select with a record MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The protected abstract `BaseInstanceSelector::select` method previously returned `Pair, Map>`, where the positional meaning of left vs. right was only conveyed by a comment. This PR replaces that with a nested `record InstanceMapping(Map segmentToInstanceMap, Map 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. --- .../BalancedInstanceSelector.java | 5 ++-- .../BaseInstanceSelector.java | 20 +++++++-------- .../instanceselector/InstanceSelector.java | 25 ++++++++++++++----- .../MultiStageReplicaGroupSelector.java | 14 +++++------ .../ReplicaGroupInstanceSelector.java | 7 +++--- .../manager/BaseBrokerRoutingManager.java | 3 +-- .../ReplicaGroupSelectorTest.java | 8 +++--- pom.xml | 4 ++- 8 files changed, 48 insertions(+), 38 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BalancedInstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BalancedInstanceSelector.java index 13c5c0668d9d..10526285acfd 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BalancedInstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BalancedInstanceSelector.java @@ -21,7 +21,6 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.broker.routing.adaptiveserverselector.ServerSelectionContext; import org.apache.pinot.common.metrics.BrokerMeter; import org.apache.pinot.common.metrics.BrokerMetrics; @@ -45,7 +44,7 @@ public class BalancedInstanceSelector extends BaseInstanceSelector { @Override - public Pair, Map> select(List segments, int requestId, + public InstanceMapping select(List segments, int requestId, SegmentStates segmentStates, Map queryOptions) { Map segmentToSelectedInstanceMap = new HashMap<>(HashUtil.getHashMapCapacity(segments.size())); // No need to adjust this map per total segment numbers, as optional segments should be empty most of the time. @@ -88,6 +87,6 @@ public Pair, Map> select(List segmen _brokerMetrics.addMeteredValue(BrokerMeter.POOL_SEG_QUERIES, entry.getValue(), BrokerMetrics.getTagForPreferredPool(queryOptions), String.valueOf(entry.getKey())); } - return Pair.of(segmentToSelectedInstanceMap, optionalSegmentToInstanceMap); + return new InstanceMapping(segmentToSelectedInstanceMap, optionalSegmentToInstanceMap); } } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java index 3108e7e6b840..96a9c987bb47 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java @@ -30,7 +30,6 @@ import java.util.TreeMap; import java.util.TreeSet; import javax.annotation.Nullable; -import org.apache.commons.lang3.tuple.Pair; import org.apache.helix.AccessOption; import org.apache.helix.model.ExternalView; import org.apache.helix.model.IdealState; @@ -465,12 +464,11 @@ public SelectionResult select(BrokerRequest brokerRequest, List segments // Copy the volatile reference so that segmentToInstanceMap and unavailableSegments can have a consistent view of // the state. SegmentStates segmentStates = _segmentStates; - Pair, Map> segmentToInstanceMap = - select(segments, requestIdInt, segmentStates, queryOptions); + InstanceMapping mapping = select(segments, requestIdInt, segmentStates, queryOptions); Set unavailableSegments = segmentStates.getUnavailableSegments(); if (unavailableSegments.isEmpty()) { - return new SelectionResult(segmentToInstanceMap, List.of(), 0); + return new SelectionResult(mapping, Collections.emptyList(), 0); } else { List unavailableSegmentsForRequest = new ArrayList<>(); for (String segment : segments) { @@ -478,7 +476,7 @@ public SelectionResult select(BrokerRequest brokerRequest, List segments unavailableSegmentsForRequest.add(segment); } } - return new SelectionResult(segmentToInstanceMap, unavailableSegmentsForRequest, 0); + return new SelectionResult(mapping, unavailableSegmentsForRequest, 0); } } @@ -501,11 +499,11 @@ int getPool(String instanceID) { } /** - * Selects the server instances for the given segments based on the request id and segment states. Returns two maps - * from segment to selected server instance hosting the segment. The 2nd map is for optional segments. The optional - * segments are used to get the new segments that are not online yet. Instead of simply skipping them by broker at - * routing time, we can send them to servers and let servers decide how to handle them. + * Selects the server instances for the given segments based on the request id and segment states. Returns an + * {@link InstanceSelector.InstanceMapping} containing two maps from segment to selected server instance. The optional + * map covers new segments that are not online yet; instead of simply skipping them at the broker, we can send them to + * servers and let servers decide how to handle them. */ - protected abstract Pair, Map/*optional segments*/> select(List segments, - int requestId, SegmentStates segmentStates, Map queryOptions); + protected abstract InstanceMapping select(List segments, int requestId, SegmentStates segmentStates, + Map queryOptions); } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java index 664ea8463681..75fc27fab920 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java @@ -19,11 +19,11 @@ package org.apache.pinot.broker.routing.instanceselector; import java.time.Clock; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Set; import javax.annotation.Nullable; -import org.apache.commons.lang3.tuple.Pair; import org.apache.helix.model.ExternalView; import org.apache.helix.model.IdealState; import org.apache.helix.store.zk.ZkHelixPropertyStore; @@ -83,23 +83,36 @@ void init(TableConfig tableConfig, ZkHelixPropertyStore propertyStore, */ Set getServingInstances(); + /** + * Holds the result of an instance selection: {@code segmentToInstanceMap} maps each segment to its selected server + * instance, and {@code optionalSegmentToInstanceMap} maps segments not yet fully online that the server may skip. + */ + record InstanceMapping(Map segmentToInstanceMap, + Map optionalSegmentToInstanceMap) { + static final InstanceMapping EMPTY = new InstanceMapping(Collections.emptyMap(), Collections.emptyMap()); + } + class SelectionResult { - private final Pair, Map/*optional segments*/> _segmentToInstanceMap; + private final InstanceMapping _instanceMapping; private final List _unavailableSegments; private int _numPrunedSegments; - public SelectionResult(Pair, Map> segmentToInstanceMap, + public SelectionResult(InstanceMapping instanceMapping, List unavailableSegments, int numPrunedSegments) { - _segmentToInstanceMap = segmentToInstanceMap; + _instanceMapping = instanceMapping; _unavailableSegments = unavailableSegments; _numPrunedSegments = numPrunedSegments; } + public static SelectionResult empty(int numPrunedSegments) { + return new SelectionResult(InstanceMapping.EMPTY, Collections.emptyList(), numPrunedSegments); + } + /** * Returns the map from segment to selected server instance hosting the segment. */ public Map getSegmentToInstanceMap() { - return _segmentToInstanceMap.getLeft(); + return _instanceMapping.segmentToInstanceMap(); } /** @@ -107,7 +120,7 @@ public Map getSegmentToInstanceMap() { * Optional segments can be skipped by broker or server upon any issue w/o failing the query. */ public Map getOptionalSegmentToInstanceMap() { - return _segmentToInstanceMap.getRight(); + return _instanceMapping.optionalSegmentToInstanceMap(); } /** diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/MultiStageReplicaGroupSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/MultiStageReplicaGroupSelector.java index 451b47dfe89b..6b0087656c7a 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/MultiStageReplicaGroupSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/MultiStageReplicaGroupSelector.java @@ -29,7 +29,6 @@ import java.util.Set; import javax.annotation.Nullable; import org.apache.commons.collections4.CollectionUtils; -import org.apache.commons.lang3.tuple.Pair; import org.apache.helix.model.ExternalView; import org.apache.helix.model.IdealState; import org.apache.helix.store.zk.ZkHelixPropertyStore; @@ -86,7 +85,7 @@ public void onAssignmentChange(IdealState idealState, ExternalView externalView, } @Override - public Pair, Map> select(List segments, int requestId, + public InstanceMapping select(List segments, int requestId, SegmentStates segmentStates, Map queryOptions) { // Create a copy of InstancePartitions to avoid race-condition with event-listeners above. InstancePartitions instancePartitions = _instancePartitions; @@ -110,8 +109,9 @@ public Pair, Map> select(List segmen * Returns a map from the segmentName to the corresponding server. It tries to select all servers from the * preferredReplicaGroup, but if it fails, it will try to select the relevant server from other instance partitions. * - * @return A pair of maps, where the first map contains the segments that are assigned to a server and the second - * map contains the segments that are optional (i.e., the server is not online to serve that segment). + * @return An {@link InstanceMapping} whose {@code segmentToInstanceMap} contains the segments assigned to a server + * and whose {@code optionalSegmentToInstanceMap} contains segments that are optional (i.e., the server is not + * online to serve that segment). * Example: * { * "required_segments": { @@ -124,7 +124,7 @@ public Pair, Map> select(List segmen * } * } */ - private Pair, Map> assign(Set segments, + private InstanceMapping assign(Set segments, SegmentStates segmentStates, InstancePartitions instancePartitions, int preferredReplicaId) { Map instanceToPartitionMap = instancePartitions.getInstanceToPartitionIdMap(); Map> instanceToSegmentsMap = new HashMap<>(); @@ -229,7 +229,7 @@ private void getSelectedInstancesForPartition(Map> instanceT * Based on whether the selected instance for the segment is online to serve that segment or not, * this method computes the segments that are optional and the segments that are not. */ - private Pair, Map> computeOptionalSegments( + private InstanceMapping computeOptionalSegments( Map segmentToSelectedInstanceMap, SegmentStates segmentStates) { Map segmentsToInstanceMap = new HashMap<>(); @@ -258,7 +258,7 @@ private Pair, Map> computeOptionalSegments( } } - return Pair.of(segmentsToInstanceMap, optionalSegmentToInstanceMap); + return new InstanceMapping(segmentsToInstanceMap, optionalSegmentToInstanceMap); } @VisibleForTesting diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupInstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupInstanceSelector.java index 8f35ba28b6af..75c844fc3af6 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupInstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupInstanceSelector.java @@ -27,7 +27,6 @@ import java.util.Set; import javax.annotation.Nullable; import org.apache.commons.collections.MapUtils; -import org.apache.commons.lang3.tuple.Pair; import org.apache.helix.model.ExternalView; import org.apache.helix.model.IdealState; import org.apache.pinot.broker.routing.adaptiveserverselector.ServerSelectionContext; @@ -70,7 +69,7 @@ public class ReplicaGroupInstanceSelector extends BaseInstanceSelector { private static final Logger LOGGER = LoggerFactory.getLogger(ReplicaGroupInstanceSelector.class); @Override - public Pair, Map> select(List segments, int requestId, + public InstanceMapping select(List segments, int requestId, SegmentStates segmentStates, Map queryOptions) { ServerSelectionContext ctx = new ServerSelectionContext(queryOptions, _config); if (_adaptiveServerSelector != null) { @@ -91,7 +90,7 @@ public Pair, Map> select(List segmen } } - private Pair, Map> selectServers(List segments, int requestId, + private InstanceMapping selectServers(List segments, int requestId, SegmentStates segmentStates, @Nullable Map serverRankMap, ServerSelectionContext ctx) { Map segmentToSelectedInstanceMap = new HashMap<>(HashUtil.getHashMapCapacity(segments.size())); @@ -147,7 +146,7 @@ private Pair, Map> selectServers(List fetchCandidateServersForQuery(List segments, diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java index 0f49ff6a77b4..4699bd15bfb1 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java @@ -1490,8 +1490,7 @@ InstanceSelector.SelectionResult calculateRouting(BrokerRequest brokerRequest, l selectionResult.setNumPrunedSegments(numPrunedSegments); return selectionResult; } else { - return new InstanceSelector.SelectionResult(Pair.of(Map.of(), Map.of()), - List.of(), numPrunedSegments); + return InstanceSelector.SelectionResult.empty(numPrunedSegments); } } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java index 2810b5faa09f..3f46bb3c1881 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java @@ -400,14 +400,14 @@ public void testReplicaGroupAdaptiveServerSelector() { assertNotNull(instanceSelector._priorityPoolInstanceSelector); when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(AR_SERVER_RANKS); - Pair, Map> result = + InstanceSelector.InstanceMapping result = instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null); // AR prefers the better-ranked server when all candidates are ranked: // segment0: instance1 (rank 3) over instance0 (rank 4) // segment1: instance3 (rank 1) over instance2 (rank 2) // segment2: instance4 unranked → AR falls back to round-robin → index 0 = instance4 - assertEquals(result.getLeft(), Map.of( + assertEquals(result.segmentToInstanceMap(), Map.of( AR_SEGMENT0, AR_INSTANCE1, AR_SEGMENT1, AR_INSTANCE3, AR_SEGMENT2, AR_INSTANCE4)); @@ -423,7 +423,7 @@ public void testStrictReplicaGroupDoesNotUseAdaptiveServerSelector() { assertNull(instanceSelector._adaptiveServerSelector); assertNull(instanceSelector._priorityPoolInstanceSelector); - Pair, Map> result = + InstanceSelector.InstanceMapping result = instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null); // hybridSelector never consulted during routing despite being passed to the factory @@ -431,7 +431,7 @@ public void testStrictReplicaGroupDoesNotUseAdaptiveServerSelector() { // Round-robin with requestId=0 picks candidate index 0 per segment — opposite of AR result above, // which returns {segment0→instance1, segment1→instance3} for the same rankings. - assertEquals(result.getLeft(), Map.of( + assertEquals(result.segmentToInstanceMap(), Map.of( AR_SEGMENT0, AR_INSTANCE0, AR_SEGMENT1, AR_INSTANCE2, AR_SEGMENT2, AR_INSTANCE4)); diff --git a/pom.xml b/pom.xml index d28b9fd02b6c..9fead38d5c89 100644 --- a/pom.xml +++ b/pom.xml @@ -2071,7 +2071,9 @@ ,\# - + + cleanthat-javaparser-unnecessaryimport + From 6dd33388e6d54ce646caf4bde92050f6eb52e4a7 Mon Sep 17 00:00:00 2001 From: Timothy Elgersma Date: Wed, 8 Jul 2026 14:28:33 -0400 Subject: [PATCH 4/8] Implement adaptive pool selection for StrictReplicaGroupInstanceSelector MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 https://github.com/apache/pinot/issues/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!) 00000000011111111112222222222333333333344444444445 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 Screenshot 2026-06-19 at 3 55 53 pm **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. Screenshot 2026-06-19 at 4 00 19 pm (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). Screenshot 2026-06-19 at 4 02 46 pm --- .../BaseInstanceSelector.java | 5 +- .../instanceselector/InstanceSelector.java | 13 +- .../InstanceSelectorFactory.java | 15 +- .../ReplicaGroupInstanceSelector.java | 4 +- .../StrictReplicaGroupInstanceSelector.java | 173 +++++++ .../manager/BaseBrokerRoutingManager.java | 4 +- .../ReplicaGroupSelectorTest.java | 465 +++++++++++++++++- .../pinot/spi/utils/CommonConstants.java | 7 + 8 files changed, 662 insertions(+), 24 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java index 96a9c987bb47..2224e538cba0 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java @@ -466,11 +466,12 @@ public SelectionResult select(BrokerRequest brokerRequest, List segments SegmentStates segmentStates = _segmentStates; InstanceMapping mapping = select(segments, requestIdInt, segmentStates, queryOptions); Set unavailableSegments = segmentStates.getUnavailableSegments(); + List mappingUnavailable = mapping.unavailableSegments(); - if (unavailableSegments.isEmpty()) { + if (unavailableSegments.isEmpty() && mappingUnavailable.isEmpty()) { return new SelectionResult(mapping, Collections.emptyList(), 0); } else { - List unavailableSegmentsForRequest = new ArrayList<>(); + List unavailableSegmentsForRequest = new ArrayList<>(mappingUnavailable); for (String segment : segments) { if (unavailableSegments.contains(segment)) { unavailableSegmentsForRequest.add(segment); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java index 75fc27fab920..f857f2d5ead8 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java @@ -85,11 +85,18 @@ void init(TableConfig tableConfig, ZkHelixPropertyStore propertyStore, /** * Holds the result of an instance selection: {@code segmentToInstanceMap} maps each segment to its selected server - * instance, and {@code optionalSegmentToInstanceMap} maps segments not yet fully online that the server may skip. + * instance, {@code optionalSegmentToInstanceMap} maps segments not yet fully online that the server may skip, and + * {@code unavailableSegments} lists segments that have candidates but could not be routed. */ record InstanceMapping(Map segmentToInstanceMap, - Map optionalSegmentToInstanceMap) { - static final InstanceMapping EMPTY = new InstanceMapping(Collections.emptyMap(), Collections.emptyMap()); + Map optionalSegmentToInstanceMap, + List unavailableSegments) { + static final InstanceMapping EMPTY = + new InstanceMapping(Collections.emptyMap(), Collections.emptyMap(), Collections.emptyList()); + + InstanceMapping(Map segmentToInstanceMap, Map optionalSegmentToInstanceMap) { + this(segmentToInstanceMap, optionalSegmentToInstanceMap, Collections.emptyList()); + } } class SelectionResult { diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorFactory.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorFactory.java index 1696de7d1e82..bd8923960a04 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorFactory.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorFactory.java @@ -106,18 +106,15 @@ public static InstanceSelector getInstanceSelector(TableConfig tableConfig, } case RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE: { LOGGER.info("Using StrictReplicaGroupInstanceSelector for table: {}", tableNameWithType); - instanceSelector = new StrictReplicaGroupInstanceSelector(); - - // Adaptive Routing is incompatible with Strict Replica Group routing, because Adaptive Routing routes - // each segment independently to the lowest-latency server, violating the same-replica-group guarantee. - // TODO: rework StrictReplicaGroupInstanceSelector::select to do pool-based adaptive routing. - // See https://github.com/apache/pinot/issues/12507 - if (adaptiveServerSelector != null) { - LOGGER.info( - "Disabling adaptive routing for StrictReplicaGroupInstanceSelector table: {}", + boolean enableStrictReplicaGroupAdaptiveRouting = brokerConfig.getProperty( + CommonConstants.Broker.AdaptiveServerSelector.CONFIG_OF_ENABLE_STRICT_REPLICA_GROUP, + CommonConstants.Broker.AdaptiveServerSelector.DEFAULT_ENABLE_STRICT_REPLICA_GROUP); + if (!enableStrictReplicaGroupAdaptiveRouting && adaptiveServerSelector != null) { + LOGGER.info("Adaptive routing disabled for StrictReplicaGroupInstanceSelector table: {}", tableNameWithType); adaptiveServerSelector = null; } + instanceSelector = new StrictReplicaGroupInstanceSelector(); break; } case RoutingConfig.MULTI_STAGE_REPLICA_GROUP_SELECTOR_TYPE: { diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupInstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupInstanceSelector.java index 75c844fc3af6..4726a7469f4a 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupInstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupInstanceSelector.java @@ -90,7 +90,7 @@ public InstanceMapping select(List segments, int requestId, } } - private InstanceMapping selectServers(List segments, int requestId, + protected InstanceMapping selectServers(List segments, int requestId, SegmentStates segmentStates, @Nullable Map serverRankMap, ServerSelectionContext ctx) { Map segmentToSelectedInstanceMap = new HashMap<>(HashUtil.getHashMapCapacity(segments.size())); @@ -120,7 +120,7 @@ private InstanceMapping selectServers(List segments, int requestId, } else if (MapUtils.isNotEmpty(serverRankMap)) { // Adaptive Server Selection is enabled. // Use the instance with the best rank if all servers have stats populated, else use the round-robin selected - // instance + // instance. As of 8 July 2026, this fallback is unreachable, but new implementations could require it. selectedInstance = candidates.stream() .anyMatch(candidate -> !serverRankMap.containsKey(candidate.getInstance())) ? selectedInstance diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java index 04e3b25b359d..3d8f2db1e66a 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java @@ -18,10 +18,25 @@ */ package org.apache.pinot.broker.routing.instanceselector; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import java.util.Set; +import java.util.stream.Collectors; import org.apache.helix.model.ExternalView; import org.apache.helix.model.IdealState; +import org.apache.pinot.broker.routing.adaptiveserverselector.ServerSelectionContext; +import org.apache.pinot.common.metrics.BrokerMeter; +import org.apache.pinot.common.metrics.BrokerMetrics; +import org.apache.pinot.common.utils.HashUtil; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + /** * Instance selector for strict replica-group routing strategy. @@ -50,14 +65,172 @@ * segments with the same assignment ([S1, S2, S3]) down on S1 to ensure that we always route the segments to the same * replica-group. * + * When adaptive server selection is enabled, this selector uses pool-level adaptive routing: it picks the best + * replica group for the entire query (using worst-case server rank within each group) and routes all segments to + * that group. This preserves the same-replica-group guarantee while benefiting from adaptive routing intelligence. + * * Note that new segments won't be used to exclude instances from serving when the segment is unavailable. * */ public class StrictReplicaGroupInstanceSelector extends ReplicaGroupInstanceSelector { + private static final Logger LOGGER = LoggerFactory.getLogger(StrictReplicaGroupInstanceSelector.class); @Override void updateSegmentMaps(IdealState idealState, ExternalView externalView, Set onlineSegments, Map newSegmentCreationTimeMap) { super.updateSegmentMapsForUpsertTable(idealState, externalView, onlineSegments, newSegmentCreationTimeMap); } + + @Override + public InstanceMapping select(List segments, int requestId, + SegmentStates segmentStates, Map queryOptions) { + + if (_adaptiveServerSelector == null || _priorityPoolInstanceSelector == null) { + ServerSelectionContext ctx = new ServerSelectionContext(queryOptions, _config); + return selectServers(segments, requestId, segmentStates, null, ctx); + } + + // Build a map: pool (replica group) -> distinct servers in that group that are candidates for this query. + // Keyed by instance name to deduplicate (a server hosting multiple segments appears once per group). + // Also track how many query segments each pool can serve (segment coverage) in a single pass. + // The guard defends against multiple candidates per pool per segment (e.g., during rebalancing). + Map> replicaGroupToQueryServers = new LinkedHashMap<>(); + Map replicaGroupSegmentCount = new HashMap<>(); + int totalSegmentsWithCandidates = 0; + + Set seenPools = new HashSet<>(); // allocate once and clear per segment to avoid O(n) allocations + for (String segment : segments) { + List candidates = segmentStates.getCandidates(segment); + if (candidates == null) { + continue; + } + + seenPools.clear(); + + totalSegmentsWithCandidates++; + for (SegmentInstanceCandidate candidate : candidates) { + int pool = candidate.getPool(); + replicaGroupToQueryServers + .computeIfAbsent(pool, k -> new LinkedHashMap<>()) + .putIfAbsent(candidate.getInstance(), candidate); + if (seenPools.add(pool)) { + replicaGroupSegmentCount.merge(pool, 1, Integer::sum); + } + } + } + + if (replicaGroupToQueryServers.isEmpty()) { + return new InstanceMapping(Collections.emptyMap(), Collections.emptyMap()); + } + + // Collect all distinct query-relevant candidates from replicaGroupToQueryServers, which already + // deduplicates by instance name. This avoids a second traversal of segments. + List allQueryCandidates = replicaGroupToQueryServers.values().stream() + .flatMap(m -> m.values().stream()) + .collect(Collectors.toList()); + ServerSelectionContext ctx = new ServerSelectionContext(queryOptions, _config); + int bestReplicaGroupId = chooseBestReplicaGroup( + replicaGroupToQueryServers, replicaGroupSegmentCount, totalSegmentsWithCandidates, + allQueryCandidates, ctx, requestId); + return selectServersForReplicaGroup(segments, bestReplicaGroupId, segmentStates, queryOptions); + } + + /** + * Scores and ranks replica groups using adaptive server selection stats. + * + * Each replica group is scored by the worst (maximum) rank among its query-relevant servers. + * Since scatter-gather query latency is bounded by the slowest server, we pick the group + * whose bottleneck server is the best (lowest worst-case rank). + * + * Groups that can serve ALL query segments (full coverage) are preferred over groups that + * would drop segments. Among full-coverage groups, the one with the best worst-case rank wins. + * If no group has full coverage, the best-ranked group is chosen regardless of coverage. + */ + private int chooseBestReplicaGroup( + Map> replicaGroupToQueryServers, + Map replicaGroupSegmentCount, + int totalSegmentsWithCandidates, + List allQueryCandidates, + ServerSelectionContext ctx, + int requestId) { + + List rankedServers = _priorityPoolInstanceSelector.rank(ctx, allQueryCandidates); + + // If no stats are available yet (empty ranking), fall back to pool based round-robin. + if (rankedServers.isEmpty()) { + List groupIds = new ArrayList<>(replicaGroupToQueryServers.keySet()); + return groupIds.get(Math.abs(requestId) % groupIds.size()); + } + + Map serverRankMap = new HashMap<>(HashUtil.getHashMapCapacity(rankedServers.size())); + for (int i = 0; i < rankedServers.size(); i++) { + serverRankMap.put(rankedServers.get(i), i); + } + + // Pick the group with full coverage first, then best worst-case rank. + // getOrDefault guards against AdaptiveServerSelector implementations that may not return every + // submitted server — unranked servers get rank -1 (best), matching HybridSelector's convention. + // As of 8 July 2026, this fallback is unreachable, but new implementations could require it. + return replicaGroupToQueryServers.entrySet().stream() + .min(Comparator.>>comparingInt( + e -> hasFullCoverage(e.getKey(), replicaGroupSegmentCount, totalSegmentsWithCandidates) ? 0 : 1) + .thenComparingInt(e -> e.getValue().values().stream() + .mapToInt(c -> serverRankMap.getOrDefault(c.getInstance(), -1)) + .max().orElse(Integer.MAX_VALUE))) + .map(Map.Entry::getKey) + .orElse(-1); + } + + private boolean hasFullCoverage(int pool, Map replicaGroupSegmentCount, + int totalSegmentsWithCandidates) { + return replicaGroupSegmentCount.getOrDefault(pool, 0) >= totalSegmentsWithCandidates; + } + + /** + * Routes all segments to the chosen replica group. For each segment, picks the candidate whose + * pool matches the selected replica group ID. If no candidate matches (transitional state during + * rebalancing), the segment is reported as unavailable. + */ + private InstanceMapping selectServersForReplicaGroup( + List segments, int replicaGroupId, SegmentStates segmentStates, Map queryOptions) { + + Map segmentToInstance = new HashMap<>(HashUtil.getHashMapCapacity(segments.size())); + Map optionalSegmentToInstance = new HashMap<>(); + List unavailableSegments = new ArrayList<>(); + + for (String segment : segments) { + List candidates = segmentStates.getCandidates(segment); + if (candidates == null) { + continue; + } + + SegmentInstanceCandidate selected = candidates.stream() + .filter(c -> c.getPool() == replicaGroupId) + .findFirst() + .orElse(null); + + // Skip segments with no candidate in the chosen replica group to preserve the strict same-replica-group + // invariant. We prefer to use groups that cover all queried segments, so this warning only occurs when we've + // selected a partial group. + if (selected == null) { + LOGGER.debug("No candidate found in replica group {} for segment {}; reporting as unavailable", + replicaGroupId, segment); + unavailableSegments.add(segment); + continue; + } + + if (selected.isOnline()) { + segmentToInstance.put(segment, selected.getInstance()); + } else { + optionalSegmentToInstance.put(segment, selected.getInstance()); + } + } + + // Emit POOL_SEG_QUERIES metric for observability. + _brokerMetrics.addMeteredValue(BrokerMeter.POOL_SEG_QUERIES, + segmentToInstance.size() + optionalSegmentToInstance.size(), + BrokerMetrics.getTagForPreferredPool(queryOptions), String.valueOf(replicaGroupId)); + + return new InstanceMapping(segmentToInstance, optionalSegmentToInstance, unavailableSegments); + } } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java index 4699bd15bfb1..2720fdfce9e0 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java @@ -768,8 +768,8 @@ private void buildRoutingInternal(String tableNameWithType) { AdaptiveServerSelector adaptiveServerSelector = AdaptiveServerSelectorFactory.getAdaptiveServerSelector(_serverRoutingStatsManager, _pinotConfig); - // StrictReplicaGroupInstanceSelector is incompatible with adaptive routing; - // InstanceSelectorFactory will nullify the selector for such tables automatically. + // For StrictReplicaGroupInstanceSelector tables, pool-level adaptive routing is used, + // preserving the same-replica-group guarantee while benefiting from adaptive server selection. InstanceSelector instanceSelector = InstanceSelectorFactory.getInstanceSelector(tableConfig, _propertyStore, _brokerMetrics, adaptiveServerSelector, _pinotConfig, _routableServerInstanceMap.keySet(), _enabledServerInstanceMap, diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java index 3f46bb3c1881..d56eb525f20b 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java @@ -34,6 +34,7 @@ import org.apache.helix.store.zk.ZkHelixPropertyStore; import org.apache.helix.zookeeper.datamodel.ZNRecord; import org.apache.pinot.broker.routing.adaptiveserverselector.HybridSelector; +import org.apache.pinot.common.metrics.BrokerMeter; import org.apache.pinot.common.metrics.BrokerMetrics; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.PinotQuery; @@ -42,7 +43,9 @@ import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.MockitoAnnotations; import org.testng.annotations.AfterMethod; @@ -53,6 +56,7 @@ import static org.apache.pinot.spi.config.table.RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE; import static org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel.ONLINE; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -370,6 +374,11 @@ private SegmentStates buildArSegmentStates() { } private ReplicaGroupInstanceSelector buildArSelector(String selectorType, HybridSelector hybridSelector) { + return buildArSelector(selectorType, hybridSelector, new PinotConfiguration(Collections.emptyMap())); + } + + private ReplicaGroupInstanceSelector buildArSelector(String selectorType, HybridSelector hybridSelector, + PinotConfiguration brokerConfig) { RoutingConfig routingConfig = new RoutingConfig(null, null, selectorType, false); TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testTable") .setRoutingConfig(routingConfig).build(); @@ -383,7 +392,7 @@ private ReplicaGroupInstanceSelector buildArSelector(String selectorType, Hybrid AR_SEGMENT2, List.of(Pair.of(AR_INSTANCE3, ONLINE), Pair.of(AR_INSTANCE4, ONLINE)))); return (ReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, - mock(PinotConfiguration.class), + brokerConfig, Set.of(AR_INSTANCE0, AR_INSTANCE1, AR_INSTANCE2, AR_INSTANCE3, AR_INSTANCE4), Map.of(), idealState, externalView, new HashSet<>(AR_SEGMENTS)); } @@ -414,26 +423,470 @@ public void testReplicaGroupAdaptiveServerSelector() { } @Test - public void testStrictReplicaGroupDoesNotUseAdaptiveServerSelector() { + public void testStrictReplicaGroupAdaptiveDisabledByFeatureFlag() { + // when the feature flag is disabled, the factory nulls out the adaptive selector so the + // selector falls back to round-robin. HybridSelector hybridSelector = mock(HybridSelector.class); + PinotConfiguration brokerConfig = new PinotConfiguration(Map.of( + CommonConstants.Broker.AdaptiveServerSelector.CONFIG_OF_ENABLE_STRICT_REPLICA_GROUP, "false")); ReplicaGroupInstanceSelector instanceSelector = - buildArSelector(STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector); + buildArSelector(STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector, brokerConfig); assertTrue(instanceSelector instanceof StrictReplicaGroupInstanceSelector); assertNull(instanceSelector._adaptiveServerSelector); assertNull(instanceSelector._priorityPoolInstanceSelector); - InstanceSelector.InstanceMapping result = + BaseInstanceSelector.InstanceMapping result = instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null); // hybridSelector never consulted during routing despite being passed to the factory verify(hybridSelector, never()).fetchServerRankingsWithScores(any()); - // Round-robin with requestId=0 picks candidate index 0 per segment — opposite of AR result above, - // which returns {segment0→instance1, segment1→instance3} for the same rankings. + // Round-robin with requestId=0 picks candidate index 0 per segment assertEquals(result.segmentToInstanceMap(), Map.of( AR_SEGMENT0, AR_INSTANCE0, AR_SEGMENT1, AR_INSTANCE2, AR_SEGMENT2, AR_INSTANCE4)); } + + @Test + public void testStrictReplicaGroupUsesAdaptiveServerSelector() { + HybridSelector hybridSelector = mock(HybridSelector.class); + ReplicaGroupInstanceSelector instanceSelector = + buildArSelector(STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector); + + assertTrue(instanceSelector instanceof StrictReplicaGroupInstanceSelector); + assertNotNull(instanceSelector._adaptiveServerSelector); + assertNotNull(instanceSelector._priorityPoolInstanceSelector); + } + + // --- Strict replica group pool-level adaptive routing tests --- + + // Topology: 2 replica groups (pool 0 and pool 1), 3 segments + // Pool 0: server_a (segment0, segment1), server_b (segment2) + // Pool 1: server_c (segment0, segment1), server_d (segment2) + private static final String SRG_SERVER_A = "server_a"; + private static final String SRG_SERVER_B = "server_b"; + private static final String SRG_SERVER_C = "server_c"; + private static final String SRG_SERVER_D = "server_d"; + private static final String SRG_SEGMENT0 = "seg0"; + private static final String SRG_SEGMENT1 = "seg1"; + private static final String SRG_SEGMENT2 = "seg2"; + private static final List SRG_SEGMENTS = Arrays.asList(SRG_SEGMENT0, SRG_SEGMENT1, SRG_SEGMENT2); + + private SegmentStates buildStrictReplicaGroupSegmentStates() { + Map> candidatesMap = new HashMap<>(); + // segment0 → server_a (pool 0), server_c (pool 1). + // Pool 0 is inserted into the LinkedHashMap before pool 1 because server_a (pool 0) appears + // first here. The round-robin fallback test depends on this insertion order. + candidatesMap.put(SRG_SEGMENT0, Arrays.asList( + new SegmentInstanceCandidate(SRG_SERVER_A, true, 0, 0), + new SegmentInstanceCandidate(SRG_SERVER_C, true, 1, 0))); + // segment1 → server_a (pool 0), server_c (pool 1) + candidatesMap.put(SRG_SEGMENT1, Arrays.asList( + new SegmentInstanceCandidate(SRG_SERVER_A, true, 0, 0), + new SegmentInstanceCandidate(SRG_SERVER_C, true, 1, 0))); + // segment2 → server_b (pool 0), server_d (pool 1) + candidatesMap.put(SRG_SEGMENT2, Arrays.asList( + new SegmentInstanceCandidate(SRG_SERVER_B, true, 0, 0), + new SegmentInstanceCandidate(SRG_SERVER_D, true, 1, 0))); + return new SegmentStates(candidatesMap, new HashSet<>(SRG_SEGMENTS), null); + } + + private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(HybridSelector hybridSelector) { + return buildStrictReplicaGroupArSelector(hybridSelector, mock(BrokerMetrics.class)); + } + + private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(HybridSelector hybridSelector, + BrokerMetrics brokerMetrics) { + return buildStrictReplicaGroupArSelector(hybridSelector, brokerMetrics, + new PinotConfiguration(Collections.emptyMap())); + } + + private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(HybridSelector hybridSelector, + BrokerMetrics brokerMetrics, PinotConfiguration brokerConfig) { + RoutingConfig routingConfig = new RoutingConfig(null, null, STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false); + TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testUpsertTable") + .setRoutingConfig(routingConfig).build(); + // Ideal state: mirrors the topology above + IdealState idealState = createIdealState(Map.of( + SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), + SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), + SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE), Pair.of(SRG_SERVER_D, ONLINE)))); + ExternalView externalView = createExternalView(Map.of( + SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), + SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), + SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE), Pair.of(SRG_SERVER_D, ONLINE)))); + // Provide server instances with pool assignments + ServerInstance serverA = mock(ServerInstance.class); + when(serverA.getPool()).thenReturn(0); + ServerInstance serverB = mock(ServerInstance.class); + when(serverB.getPool()).thenReturn(0); + ServerInstance serverC = mock(ServerInstance.class); + when(serverC.getPool()).thenReturn(1); + ServerInstance serverD = mock(ServerInstance.class); + when(serverD.getPool()).thenReturn(1); + Map serverMap = Map.of( + SRG_SERVER_A, serverA, SRG_SERVER_B, serverB, + SRG_SERVER_C, serverC, SRG_SERVER_D, serverD); + return (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, + mock(ZkHelixPropertyStore.class), brokerMetrics, hybridSelector, + brokerConfig, + Set.of(SRG_SERVER_A, SRG_SERVER_B, SRG_SERVER_C, SRG_SERVER_D), + serverMap, idealState, externalView, new HashSet<>(SRG_SEGMENTS)); + } + + @Test + public void testStrictReplicaGroupAdaptivePicksBestReplicaGroupByWorstCaseServer() { + // Pool 0 servers: server_a (rank 2), server_b (rank 3) → worst = 3 + // Pool 1 servers: server_c (rank 0), server_d (rank 1) → worst = 1 + // Pool 1 is better (lower worst-case rank), so all segments should route to pool 1 + HybridSelector hybridSelector = mock(HybridSelector.class); + BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); + StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector, + brokerMetrics); + + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( + new ImmutablePair<>(SRG_SERVER_C, 1.0), // rank 0 (best) + new ImmutablePair<>(SRG_SERVER_D, 2.0), // rank 1 + new ImmutablePair<>(SRG_SERVER_A, 3.0), // rank 2 + new ImmutablePair<>(SRG_SERVER_B, 4.0) // rank 3 (worst) + )); + + InstanceSelector.InstanceMapping result = + instanceSelector.select(SRG_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); + + // All segments routed to pool 1 + assertEquals(result.segmentToInstanceMap(), Map.of( + SRG_SEGMENT0, SRG_SERVER_C, + SRG_SEGMENT1, SRG_SERVER_C, + SRG_SEGMENT2, SRG_SERVER_D)); + verify(brokerMetrics).addMeteredValue(eq(BrokerMeter.POOL_SEG_QUERIES), eq(3L), + eq(BrokerMetrics.getTagForPreferredPool(null)), eq("1")); + } + + @Test + public void testStrictReplicaGroupAdaptivePicksPool0WhenItHasBetterWorstCase() { + // Pool 0 servers: server_a (rank 0), server_b (rank 1) → worst = 1 + // Pool 1 servers: server_c (rank 2), server_d (rank 3) → worst = 3 + // Pool 0 is better + HybridSelector hybridSelector = mock(HybridSelector.class); + StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector); + + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( + new ImmutablePair<>(SRG_SERVER_A, 1.0), + new ImmutablePair<>(SRG_SERVER_B, 2.0), + new ImmutablePair<>(SRG_SERVER_C, 3.0), + new ImmutablePair<>(SRG_SERVER_D, 4.0) + )); + + InstanceSelector.InstanceMapping result = + instanceSelector.select(SRG_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); + + // All segments routed to pool 0 + assertEquals(result.segmentToInstanceMap(), Map.of( + SRG_SEGMENT0, SRG_SERVER_A, + SRG_SEGMENT1, SRG_SERVER_A, + SRG_SEGMENT2, SRG_SERVER_B)); + } + + @Test + public void testStrictReplicaGroupAdaptiveNoStatsFallsBackToRoundRobin() { + // When fetchServerRankingsWithScores returns empty, fall back to round-robin + HybridSelector hybridSelector = mock(HybridSelector.class); + StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector); + + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Collections.emptyList()); + + // requestId=0 → picks group at index 0 (pool 0) + InstanceSelector.InstanceMapping result = + instanceSelector.select(SRG_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); + assertEquals(result.segmentToInstanceMap(), Map.of( + SRG_SEGMENT0, SRG_SERVER_A, + SRG_SEGMENT1, SRG_SERVER_A, + SRG_SEGMENT2, SRG_SERVER_B)); + + // requestId=1 → picks group at index 1 (pool 1) + result = instanceSelector.select(SRG_SEGMENTS, 1, buildStrictReplicaGroupSegmentStates(), null); + assertEquals(result.segmentToInstanceMap(), Map.of( + SRG_SEGMENT0, SRG_SERVER_C, + SRG_SEGMENT1, SRG_SERVER_C, + SRG_SEGMENT2, SRG_SERVER_D)); + } + + @Test + public void testStrictReplicaGroupAdaptiveConsistencyAllSegmentsSameReplicaGroup() { + // For any set of adaptive scores, all segments must route to the same replica group + HybridSelector hybridSelector = mock(HybridSelector.class); + StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector); + + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( + new ImmutablePair<>(SRG_SERVER_C, 1.0), + new ImmutablePair<>(SRG_SERVER_D, 2.0), + new ImmutablePair<>(SRG_SERVER_A, 3.0), + new ImmutablePair<>(SRG_SERVER_B, 4.0) + )); + + InstanceSelector.InstanceMapping result = + instanceSelector.select(SRG_SEGMENTS, 42, buildStrictReplicaGroupSegmentStates(), null); + + // Collect all pool IDs from selected instances + Map selections = result.segmentToInstanceMap(); + Set poolsUsed = new HashSet<>(); + SegmentStates states = buildStrictReplicaGroupSegmentStates(); + for (Map.Entry entry : selections.entrySet()) { + String segment = entry.getKey(); + String selectedInstance = entry.getValue(); + for (SegmentInstanceCandidate candidate : states.getCandidates(segment)) { + if (candidate.getInstance().equals(selectedInstance)) { + poolsUsed.add(candidate.getPool()); + break; + } + } + } + // All segments must be from the same pool + assertEquals(poolsUsed.size(), 1, "All segments should be routed to the same replica group"); + } + + @Test + public void testStrictReplicaGroupAdaptivePrefersFullCoverageGroup() { + // Topology: pool 0 has all 3 segments, pool 1 only has segment0 and segment1 (missing segment2). + // Even though pool 1 has better adaptive ranks, pool 0 should be chosen because it has full coverage. + // Pool 0: server_a (seg0, seg1), server_b (seg2) — full coverage + // Pool 1: server_c (seg0, seg1) — missing seg2, partial coverage + String serverC = "server_c"; + HybridSelector hybridSelector = mock(HybridSelector.class); + + RoutingConfig routingConfig = new RoutingConfig(null, null, STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false); + TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testUpsertTable") + .setRoutingConfig(routingConfig).build(); + // Pool 1 only has seg0 and seg1 (no seg2) + IdealState idealState = createIdealState(Map.of( + SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), + SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), + SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE)))); + ExternalView externalView = createExternalView(Map.of( + SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), + SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), + SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE)))); + ServerInstance serverA = mock(ServerInstance.class); + when(serverA.getPool()).thenReturn(0); + ServerInstance serverB = mock(ServerInstance.class); + when(serverB.getPool()).thenReturn(0); + ServerInstance serverCInstance = mock(ServerInstance.class); + when(serverCInstance.getPool()).thenReturn(1); + Map serverMap = Map.of( + SRG_SERVER_A, serverA, SRG_SERVER_B, serverB, serverC, serverCInstance); + StrictReplicaGroupInstanceSelector instanceSelector = + (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, + mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, + new PinotConfiguration(Collections.emptyMap()), + Set.of(SRG_SERVER_A, SRG_SERVER_B, serverC), + serverMap, idealState, externalView, new HashSet<>(SRG_SEGMENTS)); + + // Pool 1 (server_c) has rank 0 — best rank, but only covers 2 of 3 segments + // Pool 0 (server_a rank 2, server_b rank 3) — worse ranks, but full coverage + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( + new ImmutablePair<>(serverC, 1.0), // rank 0 (best) + new ImmutablePair<>(SRG_SERVER_A, 3.0), // rank 1 + new ImmutablePair<>(SRG_SERVER_B, 4.0) // rank 2 + )); + + // Build segment states where pool 1 is missing seg2 + Map> candidatesMap = new HashMap<>(); + candidatesMap.put(SRG_SEGMENT0, Arrays.asList( + new SegmentInstanceCandidate(SRG_SERVER_A, true, 0, 0), + new SegmentInstanceCandidate(serverC, true, 1, 0))); + candidatesMap.put(SRG_SEGMENT1, Arrays.asList( + new SegmentInstanceCandidate(SRG_SERVER_A, true, 0, 0), + new SegmentInstanceCandidate(serverC, true, 1, 0))); + candidatesMap.put(SRG_SEGMENT2, List.of( + new SegmentInstanceCandidate(SRG_SERVER_B, true, 0, 0))); + SegmentStates segmentStates = new SegmentStates(candidatesMap, new HashSet<>(SRG_SEGMENTS), null); + + InstanceSelector.InstanceMapping result = + instanceSelector.select(SRG_SEGMENTS, 0, segmentStates, null); + + // Pool 0 should be chosen (full coverage) even though pool 1 has better ranks + assertEquals(result.segmentToInstanceMap(), Map.of( + SRG_SEGMENT0, SRG_SERVER_A, + SRG_SEGMENT1, SRG_SERVER_A, + SRG_SEGMENT2, SRG_SERVER_B)); + } + + @Test + public void testStrictReplicaGroupAdaptiveNewSegmentNotFalselyUnavailable() { + // Topology: server_a (pool 0) hosts seg0 only. seg1 has no online instance (simulating a + // new/unassigned segment with null candidates in segmentStates). Pool 0 is chosen (only pool). + // seg1 should NOT appear in unavailable — it is a new segment, not a dropped one. + String serverA = "new_server_a"; + String seg0 = "new_seg0"; + String seg1 = "new_seg1"; + List segments = Arrays.asList(seg0, seg1); + HybridSelector hybridSelector = mock(HybridSelector.class); + + RoutingConfig routingConfig = new RoutingConfig(null, null, STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false); + TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testNewSegTable") + .setRoutingConfig(routingConfig).build(); + // seg1 intentionally absent from idealState/externalView → null candidates (new segment) + IdealState idealState = createIdealState(Map.of( + seg0, List.of(Pair.of(serverA, ONLINE)))); + ExternalView externalView = createExternalView(Map.of( + seg0, List.of(Pair.of(serverA, ONLINE)))); + ServerInstance serverAInstance = mock(ServerInstance.class); + when(serverAInstance.getPool()).thenReturn(0); + StrictReplicaGroupInstanceSelector instanceSelector = + (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, + mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, + new PinotConfiguration(Collections.emptyMap()), + Set.of(serverA), + Map.of(serverA, serverAInstance), + idealState, externalView, new HashSet<>(List.of(seg0))); + + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( + new ImmutablePair<>(serverA, 1.0))); + + BrokerRequest brokerRequest = mock(BrokerRequest.class); + PinotQuery pinotQuery = mock(PinotQuery.class); + when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery); + when(pinotQuery.getQueryOptions()).thenReturn(null); + + InstanceSelector.SelectionResult result = instanceSelector.select(brokerRequest, segments, 0); + assertEquals(result.getSegmentToInstanceMap(), Map.of(seg0, serverA)); + // seg1 has null candidates (not yet in the selector's state) → must NOT appear as unavailable + assertTrue(result.getUnavailableSegments().isEmpty()); + } + + @Test + public void testStrictReplicaGroupAdaptiveDroppedSegmentReportedAsUnavailable() { + // Topology: no replica group has full coverage. + // Pool 0: server_a hosts seg0, seg1 (missing seg2) + // Pool 1: server_c hosts seg1, seg2 (missing seg0) + // Pool 1 has better ranks -> chosen. seg0 has no candidate in pool 1 -> reported unavailable. + String serverA = "partial_server_a"; + String serverC = "partial_server_c"; + String seg0 = "partial_seg0"; + String seg1 = "partial_seg1"; + String seg2 = "partial_seg2"; + List segments = Arrays.asList(seg0, seg1, seg2); + HybridSelector hybridSelector = mock(HybridSelector.class); + + RoutingConfig routingConfig = new RoutingConfig(null, null, STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false); + TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testPartialCoverage") + .setRoutingConfig(routingConfig).build(); + IdealState idealState = createIdealState(Map.of( + seg0, List.of(Pair.of(serverA, ONLINE)), + seg1, List.of(Pair.of(serverA, ONLINE), Pair.of(serverC, ONLINE)), + seg2, List.of(Pair.of(serverC, ONLINE)))); + ExternalView externalView = createExternalView(Map.of( + seg0, List.of(Pair.of(serverA, ONLINE)), + seg1, List.of(Pair.of(serverA, ONLINE), Pair.of(serverC, ONLINE)), + seg2, List.of(Pair.of(serverC, ONLINE)))); + ServerInstance serverAInstance = mock(ServerInstance.class); + when(serverAInstance.getPool()).thenReturn(0); + ServerInstance serverCInstance = mock(ServerInstance.class); + when(serverCInstance.getPool()).thenReturn(1); + Map serverMap = Map.of(serverA, serverAInstance, serverC, serverCInstance); + BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); + StrictReplicaGroupInstanceSelector instanceSelector = + (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, + mock(ZkHelixPropertyStore.class), brokerMetrics, hybridSelector, + new PinotConfiguration(Collections.emptyMap()), + Set.of(serverA, serverC), + serverMap, idealState, externalView, new HashSet<>(segments)); + + // Pool 1 (server_c) ranked better -> pool 1 chosen + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( + new ImmutablePair<>(serverC, 1.0), + new ImmutablePair<>(serverA, 3.0) + )); + + BrokerRequest brokerRequest = mock(BrokerRequest.class); + PinotQuery pinotQuery = mock(PinotQuery.class); + when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery); + when(pinotQuery.getQueryOptions()).thenReturn(null); + + InstanceSelector.SelectionResult result = instanceSelector.select(brokerRequest, segments, 0); + + // seg1 and seg2 routed to pool 1 + assertEquals(result.getSegmentToInstanceMap(), Map.of(seg1, serverC, seg2, serverC)); + // seg0 has candidates (pool 0 only) but no match in chosen pool 1 -> reported unavailable + assertEquals(result.getUnavailableSegments(), List.of(seg0)); + } + + @Test + public void testStrictReplicaGroupAdaptiveOnlyQueryRelevantServersScored() { + // Topology: pool 1 has an extra server (server_e) that hosts a segment NOT in the query. + // server_e has a very high rank but should not affect pool 1's score since it hosts no query segments. + // Pool 0: server_a (rank 2), server_b (rank 3) → worst = 3 + // Pool 1: server_c (rank 0), server_d (rank 1) → worst = 1 + // (server_e at rank 99 hosts segment "seg_extra" which is not in the query) + // Pool 1 should win because only query-relevant servers matter (server_e is ignored). + String serverE = "server_e"; + String segExtra = "seg_extra"; + HybridSelector hybridSelector = mock(HybridSelector.class); + + // Build selector with server_e in pool 1, hosting seg_extra + RoutingConfig routingConfig = new RoutingConfig(null, null, STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false); + TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testUpsertTable") + .setRoutingConfig(routingConfig).build(); + IdealState idealState = createIdealState(Map.of( + SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), + SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), + SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE), Pair.of(SRG_SERVER_D, ONLINE)), + segExtra, List.of(Pair.of(serverE, ONLINE)))); + ExternalView externalView = createExternalView(Map.of( + SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), + SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), + SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE), Pair.of(SRG_SERVER_D, ONLINE)), + segExtra, List.of(Pair.of(serverE, ONLINE)))); + ServerInstance serverA = mock(ServerInstance.class); + when(serverA.getPool()).thenReturn(0); + ServerInstance serverB = mock(ServerInstance.class); + when(serverB.getPool()).thenReturn(0); + ServerInstance serverC = mock(ServerInstance.class); + when(serverC.getPool()).thenReturn(1); + ServerInstance serverD = mock(ServerInstance.class); + when(serverD.getPool()).thenReturn(1); + ServerInstance serverEInstance = mock(ServerInstance.class); + when(serverEInstance.getPool()).thenReturn(1); + Map serverMap = Map.of( + SRG_SERVER_A, serverA, SRG_SERVER_B, serverB, + SRG_SERVER_C, serverC, SRG_SERVER_D, serverD, + serverE, serverEInstance); + Set allSegments = new HashSet<>(SRG_SEGMENTS); + allSegments.add(segExtra); + StrictReplicaGroupInstanceSelector instanceSelector = + (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, + mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, + new PinotConfiguration(Collections.emptyMap()), + Set.of(SRG_SERVER_A, SRG_SERVER_B, SRG_SERVER_C, SRG_SERVER_D, serverE), + serverMap, idealState, externalView, allSegments); + + // server_e has rank 99 but hosts no query segments — should be irrelevant + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( + new ImmutablePair<>(SRG_SERVER_C, 1.0), // rank 0 + new ImmutablePair<>(SRG_SERVER_D, 2.0), // rank 1 + new ImmutablePair<>(SRG_SERVER_A, 3.0), // rank 2 + new ImmutablePair<>(SRG_SERVER_B, 4.0), // rank 3 + new ImmutablePair<>(serverE, 99.0) // rank 4 (high score, not query-relevant) + )); + + // Query only SRG_SEGMENTS (not seg_extra), so server_e is not query-relevant + InstanceSelector.InstanceMapping result = + instanceSelector.select(SRG_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); + + // Pool 1 wins (worst rank among query-relevant servers = 1) over pool 0 (worst = 3) + // server_e's rank 4 is not considered because it hosts no query segments + assertEquals(result.segmentToInstanceMap(), Map.of( + SRG_SEGMENT0, SRG_SERVER_C, + SRG_SEGMENT1, SRG_SERVER_C, + SRG_SEGMENT2, SRG_SERVER_D)); + + ArgumentCaptor> rankedCandidatesCaptor = ArgumentCaptor.forClass(List.class); + verify(hybridSelector).fetchServerRankingsWithScores(rankedCandidatesCaptor.capture()); + assertEquals(new HashSet<>(rankedCandidatesCaptor.getValue()), + Set.of(SRG_SERVER_A, SRG_SERVER_B, SRG_SERVER_C, SRG_SERVER_D)); + assertEquals(rankedCandidatesCaptor.getValue().size(), 4); + } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index ee25435988a6..3e7c547d9f50 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -1271,6 +1271,13 @@ public enum Type { public static final String CONFIG_OF_STATS_METRIC_EXPORT_INTERVAL_MS = CONFIG_PREFIX + ".stats.metric.export.interval.ms"; public static final long DEFAULT_STATS_METRIC_EXPORT_INTERVAL_MS = 10 * 1000; + + // Controls whether pool-level adaptive routing is enabled for StrictReplicaGroupInstanceSelector. + // When false, StrictReplicaGroupInstanceSelector falls back to round-robin even if adaptive server + // selection is configured. + public static final String CONFIG_OF_ENABLE_STRICT_REPLICA_GROUP = + CONFIG_PREFIX + ".enable.strict.replica.group"; + public static final boolean DEFAULT_ENABLE_STRICT_REPLICA_GROUP = true; } public static class Grpc { From f45af915d6a4dd488bd5f6a2a980f4b4de9a8f8c Mon Sep 17 00:00:00 2001 From: Tim Elgersma Date: Thu, 9 Jul 2026 12:24:35 -0400 Subject: [PATCH 5/8] remove Collections.* uses --- .../instanceselector/BaseInstanceSelector.java | 2 +- .../routing/instanceselector/InstanceSelector.java | 7 +++---- .../StrictReplicaGroupInstanceSelector.java | 3 +-- .../routing/manager/BaseBrokerRoutingManager.java | 1 - .../instanceselector/ReplicaGroupSelectorTest.java | 14 +++++++------- 5 files changed, 12 insertions(+), 15 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java index 2224e538cba0..7382ff4290fb 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java @@ -469,7 +469,7 @@ public SelectionResult select(BrokerRequest brokerRequest, List segments List mappingUnavailable = mapping.unavailableSegments(); if (unavailableSegments.isEmpty() && mappingUnavailable.isEmpty()) { - return new SelectionResult(mapping, Collections.emptyList(), 0); + return new SelectionResult(mapping, List.of(), 0); } else { List unavailableSegmentsForRequest = new ArrayList<>(mappingUnavailable); for (String segment : segments) { diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java index f857f2d5ead8..5622007d2b08 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java @@ -19,7 +19,6 @@ package org.apache.pinot.broker.routing.instanceselector; import java.time.Clock; -import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Set; @@ -92,10 +91,10 @@ record InstanceMapping(Map segmentToInstanceMap, Map optionalSegmentToInstanceMap, List unavailableSegments) { static final InstanceMapping EMPTY = - new InstanceMapping(Collections.emptyMap(), Collections.emptyMap(), Collections.emptyList()); + new InstanceMapping(Map.of(), Map.of(), List.of()); InstanceMapping(Map segmentToInstanceMap, Map optionalSegmentToInstanceMap) { - this(segmentToInstanceMap, optionalSegmentToInstanceMap, Collections.emptyList()); + this(segmentToInstanceMap, optionalSegmentToInstanceMap, List.of()); } } @@ -112,7 +111,7 @@ public SelectionResult(InstanceMapping instanceMapping, } public static SelectionResult empty(int numPrunedSegments) { - return new SelectionResult(InstanceMapping.EMPTY, Collections.emptyList(), numPrunedSegments); + return new SelectionResult(InstanceMapping.EMPTY, List.of(), numPrunedSegments); } /** diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java index 3d8f2db1e66a..10ce6432a4bf 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java @@ -19,7 +19,6 @@ package org.apache.pinot.broker.routing.instanceselector; import java.util.ArrayList; -import java.util.Collections; import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; @@ -120,7 +119,7 @@ public InstanceMapping select(List segments, int requestId, } if (replicaGroupToQueryServers.isEmpty()) { - return new InstanceMapping(Collections.emptyMap(), Collections.emptyMap()); + return new InstanceMapping(Map.of(), Map.of()); } // Collect all distinct query-relevant candidates from replicaGroupToQueryServers, which already diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java index 2720fdfce9e0..fb9908c4ec64 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java @@ -44,7 +44,6 @@ import javax.annotation.concurrent.GuardedBy; import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.lang3.StringUtils; -import org.apache.commons.lang3.tuple.Pair; import org.apache.helix.AccessOption; import org.apache.helix.BaseDataAccessor; import org.apache.helix.HelixConstants.ChangeType; diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java index d56eb525f20b..34d733505754 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java @@ -374,7 +374,7 @@ private SegmentStates buildArSegmentStates() { } private ReplicaGroupInstanceSelector buildArSelector(String selectorType, HybridSelector hybridSelector) { - return buildArSelector(selectorType, hybridSelector, new PinotConfiguration(Collections.emptyMap())); + return buildArSelector(selectorType, hybridSelector, new PinotConfiguration(Map.of())); } private ReplicaGroupInstanceSelector buildArSelector(String selectorType, HybridSelector hybridSelector, @@ -500,7 +500,7 @@ private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(Hyb private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(HybridSelector hybridSelector, BrokerMetrics brokerMetrics) { return buildStrictReplicaGroupArSelector(hybridSelector, brokerMetrics, - new PinotConfiguration(Collections.emptyMap())); + new PinotConfiguration(Map.of())); } private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(HybridSelector hybridSelector, @@ -596,7 +596,7 @@ public void testStrictReplicaGroupAdaptiveNoStatsFallsBackToRoundRobin() { HybridSelector hybridSelector = mock(HybridSelector.class); StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector); - when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Collections.emptyList()); + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(List.of()); // requestId=0 → picks group at index 0 (pool 0) InstanceSelector.InstanceMapping result = @@ -680,7 +680,7 @@ public void testStrictReplicaGroupAdaptivePrefersFullCoverageGroup() { StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, - new PinotConfiguration(Collections.emptyMap()), + new PinotConfiguration(Map.of()), Set.of(SRG_SERVER_A, SRG_SERVER_B, serverC), serverMap, idealState, externalView, new HashSet<>(SRG_SEGMENTS)); @@ -738,7 +738,7 @@ public void testStrictReplicaGroupAdaptiveNewSegmentNotFalselyUnavailable() { StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, - new PinotConfiguration(Collections.emptyMap()), + new PinotConfiguration(Map.of()), Set.of(serverA), Map.of(serverA, serverAInstance), idealState, externalView, new HashSet<>(List.of(seg0))); @@ -791,7 +791,7 @@ public void testStrictReplicaGroupAdaptiveDroppedSegmentReportedAsUnavailable() StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), brokerMetrics, hybridSelector, - new PinotConfiguration(Collections.emptyMap()), + new PinotConfiguration(Map.of()), Set.of(serverA, serverC), serverMap, idealState, externalView, new HashSet<>(segments)); @@ -859,7 +859,7 @@ public void testStrictReplicaGroupAdaptiveOnlyQueryRelevantServersScored() { StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, - new PinotConfiguration(Collections.emptyMap()), + new PinotConfiguration(Map.of()), Set.of(SRG_SERVER_A, SRG_SERVER_B, SRG_SERVER_C, SRG_SERVER_D, serverE), serverMap, idealState, externalView, allSegments); From 093aa8585bd6022359957bc4d6f8689949cdc5b8 Mon Sep 17 00:00:00 2001 From: Timothy Elgersma Date: Thu, 9 Jul 2026 20:44:11 +0000 Subject: [PATCH 6/8] Group by idealStateReplicaId not pool 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 --- .../StrictReplicaGroupInstanceSelector.java | 37 ++++++++++--------- .../ReplicaGroupSelectorTest.java | 12 +++--- 2 files changed, 26 insertions(+), 23 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java index 10ce6432a4bf..249618d93435 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java @@ -64,7 +64,7 @@ * segments with the same assignment ([S1, S2, S3]) down on S1 to ensure that we always route the segments to the same * replica-group. * - * When adaptive server selection is enabled, this selector uses pool-level adaptive routing: it picks the best + * When adaptive server selection is enabled, this selector uses replica-group-level adaptive routing: it picks the best * replica group for the entire query (using worst-case server rank within each group) and routes all segments to * that group. This preserves the same-replica-group guarantee while benefiting from adaptive routing intelligence. * @@ -89,31 +89,31 @@ public InstanceMapping select(List segments, int requestId, return selectServers(segments, requestId, segmentStates, null, ctx); } - // Build a map: pool (replica group) -> distinct servers in that group that are candidates for this query. - // Keyed by instance name to deduplicate (a server hosting multiple segments appears once per group). - // Also track how many query segments each pool can serve (segment coverage) in a single pass. - // The guard defends against multiple candidates per pool per segment (e.g., during rebalancing). + // Build a map: idealStateReplicaId (replica group) -> distinct servers in that group that are candidates for this + // query. Keyed by instance name to deduplicate (a server hosting multiple segments appears once per group). + // Also track how many query segments each replica group can serve (segment coverage) in a single pass. + // The guard defends against multiple candidates per replica group per segment (e.g., during rebalancing). Map> replicaGroupToQueryServers = new LinkedHashMap<>(); Map replicaGroupSegmentCount = new HashMap<>(); int totalSegmentsWithCandidates = 0; - Set seenPools = new HashSet<>(); // allocate once and clear per segment to avoid O(n) allocations + Set seenReplicaIds = new HashSet<>(); // allocate once and clear per segment to avoid O(n) allocations for (String segment : segments) { List candidates = segmentStates.getCandidates(segment); if (candidates == null) { continue; } - seenPools.clear(); + seenReplicaIds.clear(); totalSegmentsWithCandidates++; for (SegmentInstanceCandidate candidate : candidates) { - int pool = candidate.getPool(); + int replicaId = candidate.getIdealStateReplicaId(); replicaGroupToQueryServers - .computeIfAbsent(pool, k -> new LinkedHashMap<>()) + .computeIfAbsent(replicaId, k -> new LinkedHashMap<>()) .putIfAbsent(candidate.getInstance(), candidate); - if (seenPools.add(pool)) { - replicaGroupSegmentCount.merge(pool, 1, Integer::sum); + if (seenReplicaIds.add(replicaId)) { + replicaGroupSegmentCount.merge(replicaId, 1, Integer::sum); } } } @@ -187,7 +187,7 @@ private boolean hasFullCoverage(int pool, Map replicaGroupSegm /** * Routes all segments to the chosen replica group. For each segment, picks the candidate whose - * pool matches the selected replica group ID. If no candidate matches (transitional state during + * idealStateReplicaId matches the selected replica group ID. If no candidate matches (transitional state during * rebalancing), the segment is reported as unavailable. */ private InstanceMapping selectServersForReplicaGroup( @@ -196,6 +196,7 @@ private InstanceMapping selectServersForReplicaGroup( Map segmentToInstance = new HashMap<>(HashUtil.getHashMapCapacity(segments.size())); Map optionalSegmentToInstance = new HashMap<>(); List unavailableSegments = new ArrayList<>(); + Map poolToSegmentCount = new HashMap<>(); for (String segment : segments) { List candidates = segmentStates.getCandidates(segment); @@ -204,7 +205,7 @@ private InstanceMapping selectServersForReplicaGroup( } SegmentInstanceCandidate selected = candidates.stream() - .filter(c -> c.getPool() == replicaGroupId) + .filter(c -> c.getIdealStateReplicaId() == replicaGroupId) .findFirst() .orElse(null); @@ -223,12 +224,14 @@ private InstanceMapping selectServersForReplicaGroup( } else { optionalSegmentToInstance.put(segment, selected.getInstance()); } + poolToSegmentCount.merge(selected.getPool(), 1, Integer::sum); } - // Emit POOL_SEG_QUERIES metric for observability. - _brokerMetrics.addMeteredValue(BrokerMeter.POOL_SEG_QUERIES, - segmentToInstance.size() + optionalSegmentToInstance.size(), - BrokerMetrics.getTagForPreferredPool(queryOptions), String.valueOf(replicaGroupId)); + // Emit per-pool POOL_SEG_QUERIES metric for observability, matching the parent's pattern. + for (Map.Entry entry : poolToSegmentCount.entrySet()) { + _brokerMetrics.addMeteredValue(BrokerMeter.POOL_SEG_QUERIES, entry.getValue(), + BrokerMetrics.getTagForPreferredPool(queryOptions), String.valueOf(entry.getKey())); + } return new InstanceMapping(segmentToInstance, optionalSegmentToInstance, unavailableSegments); } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java index 34d733505754..52ee3f6cc58b 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java @@ -500,7 +500,7 @@ private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(Hyb private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(HybridSelector hybridSelector, BrokerMetrics brokerMetrics) { return buildStrictReplicaGroupArSelector(hybridSelector, brokerMetrics, - new PinotConfiguration(Map.of())); + new PinotConfiguration(Collections.emptyMap())); } private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(HybridSelector hybridSelector, @@ -596,7 +596,7 @@ public void testStrictReplicaGroupAdaptiveNoStatsFallsBackToRoundRobin() { HybridSelector hybridSelector = mock(HybridSelector.class); StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector); - when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(List.of()); + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Collections.emptyList()); // requestId=0 → picks group at index 0 (pool 0) InstanceSelector.InstanceMapping result = @@ -680,7 +680,7 @@ public void testStrictReplicaGroupAdaptivePrefersFullCoverageGroup() { StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, - new PinotConfiguration(Map.of()), + new PinotConfiguration(Collections.emptyMap()), Set.of(SRG_SERVER_A, SRG_SERVER_B, serverC), serverMap, idealState, externalView, new HashSet<>(SRG_SEGMENTS)); @@ -738,7 +738,7 @@ public void testStrictReplicaGroupAdaptiveNewSegmentNotFalselyUnavailable() { StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, - new PinotConfiguration(Map.of()), + new PinotConfiguration(Collections.emptyMap()), Set.of(serverA), Map.of(serverA, serverAInstance), idealState, externalView, new HashSet<>(List.of(seg0))); @@ -791,7 +791,7 @@ public void testStrictReplicaGroupAdaptiveDroppedSegmentReportedAsUnavailable() StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), brokerMetrics, hybridSelector, - new PinotConfiguration(Map.of()), + new PinotConfiguration(Collections.emptyMap()), Set.of(serverA, serverC), serverMap, idealState, externalView, new HashSet<>(segments)); @@ -859,7 +859,7 @@ public void testStrictReplicaGroupAdaptiveOnlyQueryRelevantServersScored() { StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, - new PinotConfiguration(Map.of()), + new PinotConfiguration(Collections.emptyMap()), Set.of(SRG_SERVER_A, SRG_SERVER_B, SRG_SERVER_C, SRG_SERVER_D, serverE), serverMap, idealState, externalView, allSegments); From 597273ee9889c0ab749342bdced8c60c78f4bb00 Mon Sep 17 00:00:00 2001 From: Tim Elgersma Date: Thu, 9 Jul 2026 17:40:34 -0400 Subject: [PATCH 7/8] remove more `Collections.emptyMap()` uses --- .../instanceselector/ReplicaGroupSelectorTest.java | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java index 52ee3f6cc58b..34d733505754 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java @@ -500,7 +500,7 @@ private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(Hyb private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(HybridSelector hybridSelector, BrokerMetrics brokerMetrics) { return buildStrictReplicaGroupArSelector(hybridSelector, brokerMetrics, - new PinotConfiguration(Collections.emptyMap())); + new PinotConfiguration(Map.of())); } private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(HybridSelector hybridSelector, @@ -596,7 +596,7 @@ public void testStrictReplicaGroupAdaptiveNoStatsFallsBackToRoundRobin() { HybridSelector hybridSelector = mock(HybridSelector.class); StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector); - when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Collections.emptyList()); + when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(List.of()); // requestId=0 → picks group at index 0 (pool 0) InstanceSelector.InstanceMapping result = @@ -680,7 +680,7 @@ public void testStrictReplicaGroupAdaptivePrefersFullCoverageGroup() { StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, - new PinotConfiguration(Collections.emptyMap()), + new PinotConfiguration(Map.of()), Set.of(SRG_SERVER_A, SRG_SERVER_B, serverC), serverMap, idealState, externalView, new HashSet<>(SRG_SEGMENTS)); @@ -738,7 +738,7 @@ public void testStrictReplicaGroupAdaptiveNewSegmentNotFalselyUnavailable() { StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, - new PinotConfiguration(Collections.emptyMap()), + new PinotConfiguration(Map.of()), Set.of(serverA), Map.of(serverA, serverAInstance), idealState, externalView, new HashSet<>(List.of(seg0))); @@ -791,7 +791,7 @@ public void testStrictReplicaGroupAdaptiveDroppedSegmentReportedAsUnavailable() StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), brokerMetrics, hybridSelector, - new PinotConfiguration(Collections.emptyMap()), + new PinotConfiguration(Map.of()), Set.of(serverA, serverC), serverMap, idealState, externalView, new HashSet<>(segments)); @@ -859,7 +859,7 @@ public void testStrictReplicaGroupAdaptiveOnlyQueryRelevantServersScored() { StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, - new PinotConfiguration(Collections.emptyMap()), + new PinotConfiguration(Map.of()), Set.of(SRG_SERVER_A, SRG_SERVER_B, SRG_SERVER_C, SRG_SERVER_D, serverE), serverMap, idealState, externalView, allSegments); From 5708837ab651c404c557dc4c7e97e109c4544e96 Mon Sep 17 00:00:00 2001 From: Timothy Elgersma Date: Mon, 13 Jul 2026 10:41:49 -0400 Subject: [PATCH 8/8] update test for replica-group change. include RG / pool number in the test names to make test behaviour clearer --- .../ReplicaGroupSelectorTest.java | 573 +++++++++--------- 1 file changed, 297 insertions(+), 276 deletions(-) diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java index 34d733505754..68cf4b33a161 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java @@ -54,6 +54,7 @@ import static org.apache.pinot.spi.config.table.RoutingConfig.REPLICA_GROUP_INSTANCE_SELECTOR_TYPE; import static org.apache.pinot.spi.config.table.RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE; +import static org.apache.pinot.spi.utils.CommonConstants.Broker.FALLBACK_POOL_ID; import static org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel.ONLINE; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; @@ -70,7 +71,7 @@ /** * Tests for {@link ReplicaGroupInstanceSelector} and {@link StrictReplicaGroupInstanceSelector}, - * including adaptive server selection (pool-level routing for strict replica groups). + * including adaptive server selection (replica-group-level routing for strict replica groups). */ @SuppressWarnings("unchecked") public class ReplicaGroupSelectorTest { @@ -336,40 +337,40 @@ public void testReplicaGroupInstanceSelectorNumReplicaGroupsNotSet() { // --- Adaptive server selection tests --- - // Shared topology for AR tests: 3 segments across 5 instances in two replica groups. - // segment2 intentionally has instance4 (unranked) so AR falls back to round-robin for that segment. - private static final String AR_INSTANCE0 = "instance0"; - private static final String AR_INSTANCE1 = "instance1"; - private static final String AR_INSTANCE2 = "instance2"; - private static final String AR_INSTANCE3 = "instance3"; - private static final String AR_INSTANCE4 = "instance4"; + // Shared topology for AR tests: 3 segments across 5 instances in two pools / replica groups. + // segment2 intentionally has AR_P0_RG0_SERVER_E (unranked) so AR falls back to round-robin for that segment. + private static final String AR_P0_RG0_SERVER_A = "ar_p0_rg0_server_a"; + private static final String AR_P1_RG1_SERVER_B = "ar_p1_rg1_server_b"; + private static final String AR_P0_RG0_SERVER_C = "ar_p0_rg0_server_c"; + private static final String AR_P1_RG1_SERVER_D = "ar_p1_rg1_server_d"; + private static final String AR_P0_RG0_SERVER_E = "ar_p0_rg0_server_e"; private static final String AR_SEGMENT0 = "segment0"; private static final String AR_SEGMENT1 = "segment1"; private static final String AR_SEGMENT2 = "segment2"; private static final List AR_SEGMENTS = Arrays.asList(AR_SEGMENT0, AR_SEGMENT1, AR_SEGMENT2); - // Rankings: instance3 best → instance2 → instance1 → instance0 worst; instance4 absent (triggers AR fallback) + // Rankings: D best → C → B → A worst; E absent (triggers AR fallback) private static final List> AR_SERVER_RANKS = Arrays.asList( - new ImmutablePair<>(AR_INSTANCE3, 1.0), - new ImmutablePair<>(AR_INSTANCE2, 2.0), - new ImmutablePair<>(AR_INSTANCE1, 3.0), - new ImmutablePair<>(AR_INSTANCE0, 4.0) + new ImmutablePair<>(AR_P1_RG1_SERVER_D, 1.0), + new ImmutablePair<>(AR_P0_RG0_SERVER_C, 2.0), + new ImmutablePair<>(AR_P1_RG1_SERVER_B, 3.0), + new ImmutablePair<>(AR_P0_RG0_SERVER_A, 4.0) ); private SegmentStates buildArSegmentStates() { Map> candidatesMap = new HashMap<>(); - // segment0 → instance0, instance1 + // segment0 → pool 0 / replica group 0 server A, pool 1 / replica group 1 server B candidatesMap.put(AR_SEGMENT0, Arrays.asList( - new SegmentInstanceCandidate(AR_INSTANCE0, true), - new SegmentInstanceCandidate(AR_INSTANCE1, true))); - // segment1 → instance2, instance3 + new SegmentInstanceCandidate(AR_P0_RG0_SERVER_A, true, 0, 0), + new SegmentInstanceCandidate(AR_P1_RG1_SERVER_B, true, 1, 1))); + // segment1 → pool 0 / replica group 0 server C, pool 1 / replica group 1 server D candidatesMap.put(AR_SEGMENT1, Arrays.asList( - new SegmentInstanceCandidate(AR_INSTANCE2, true), - new SegmentInstanceCandidate(AR_INSTANCE3, true))); - // segment2 → instance4 (unranked), instance3 + new SegmentInstanceCandidate(AR_P0_RG0_SERVER_C, true, 0, 0), + new SegmentInstanceCandidate(AR_P1_RG1_SERVER_D, true, 1, 1))); + // segment2 → pool 0 / replica group 0 server E (unranked), pool 1 / replica group 1 server D candidatesMap.put(AR_SEGMENT2, Arrays.asList( - new SegmentInstanceCandidate(AR_INSTANCE4, true), - new SegmentInstanceCandidate(AR_INSTANCE3, true))); + new SegmentInstanceCandidate(AR_P0_RG0_SERVER_E, true, 0, 0), + new SegmentInstanceCandidate(AR_P1_RG1_SERVER_D, true, 1, 1))); return new SegmentStates(candidatesMap, new HashSet<>(AR_SEGMENTS), null); } @@ -383,18 +384,34 @@ private ReplicaGroupInstanceSelector buildArSelector(String selectorType, Hybrid TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testTable") .setRoutingConfig(routingConfig).build(); IdealState idealState = createIdealState(Map.of( - AR_SEGMENT0, List.of(Pair.of(AR_INSTANCE0, ONLINE), Pair.of(AR_INSTANCE1, ONLINE)), - AR_SEGMENT1, List.of(Pair.of(AR_INSTANCE2, ONLINE), Pair.of(AR_INSTANCE3, ONLINE)), - AR_SEGMENT2, List.of(Pair.of(AR_INSTANCE3, ONLINE), Pair.of(AR_INSTANCE4, ONLINE)))); + AR_SEGMENT0, List.of(Pair.of(AR_P0_RG0_SERVER_A, ONLINE), Pair.of(AR_P1_RG1_SERVER_B, ONLINE)), + AR_SEGMENT1, List.of(Pair.of(AR_P0_RG0_SERVER_C, ONLINE), Pair.of(AR_P1_RG1_SERVER_D, ONLINE)), + AR_SEGMENT2, List.of(Pair.of(AR_P1_RG1_SERVER_D, ONLINE), Pair.of(AR_P0_RG0_SERVER_E, ONLINE)))); ExternalView externalView = createExternalView(Map.of( - AR_SEGMENT0, List.of(Pair.of(AR_INSTANCE0, ONLINE), Pair.of(AR_INSTANCE1, ONLINE)), - AR_SEGMENT1, List.of(Pair.of(AR_INSTANCE2, ONLINE), Pair.of(AR_INSTANCE3, ONLINE)), - AR_SEGMENT2, List.of(Pair.of(AR_INSTANCE3, ONLINE), Pair.of(AR_INSTANCE4, ONLINE)))); + AR_SEGMENT0, List.of(Pair.of(AR_P0_RG0_SERVER_A, ONLINE), Pair.of(AR_P1_RG1_SERVER_B, ONLINE)), + AR_SEGMENT1, List.of(Pair.of(AR_P0_RG0_SERVER_C, ONLINE), Pair.of(AR_P1_RG1_SERVER_D, ONLINE)), + AR_SEGMENT2, List.of(Pair.of(AR_P1_RG1_SERVER_D, ONLINE), Pair.of(AR_P0_RG0_SERVER_E, ONLINE)))); + ServerInstance serverA = mock(ServerInstance.class); + when(serverA.getPool()).thenReturn(0); + ServerInstance serverB = mock(ServerInstance.class); + when(serverB.getPool()).thenReturn(1); + ServerInstance serverC = mock(ServerInstance.class); + when(serverC.getPool()).thenReturn(0); + ServerInstance serverD = mock(ServerInstance.class); + when(serverD.getPool()).thenReturn(1); + ServerInstance serverE = mock(ServerInstance.class); + when(serverE.getPool()).thenReturn(0); + Map serverMap = Map.of( + AR_P0_RG0_SERVER_A, serverA, + AR_P1_RG1_SERVER_B, serverB, + AR_P0_RG0_SERVER_C, serverC, + AR_P1_RG1_SERVER_D, serverD, + AR_P0_RG0_SERVER_E, serverE); return (ReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, brokerConfig, - Set.of(AR_INSTANCE0, AR_INSTANCE1, AR_INSTANCE2, AR_INSTANCE3, AR_INSTANCE4), - Map.of(), idealState, externalView, new HashSet<>(AR_SEGMENTS)); + Set.of(AR_P0_RG0_SERVER_A, AR_P1_RG1_SERVER_B, AR_P0_RG0_SERVER_C, AR_P1_RG1_SERVER_D, AR_P0_RG0_SERVER_E), + serverMap, idealState, externalView, new HashSet<>(AR_SEGMENTS)); } @Test @@ -413,44 +430,47 @@ public void testReplicaGroupAdaptiveServerSelector() { instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null); // AR prefers the better-ranked server when all candidates are ranked: - // segment0: instance1 (rank 3) over instance0 (rank 4) - // segment1: instance3 (rank 1) over instance2 (rank 2) - // segment2: instance4 unranked → AR falls back to round-robin → index 0 = instance4 + // segment0: B (rank 3) over A (rank 4) + // segment1: D (rank 1) over C (rank 2) + // segment2: E unranked → AR falls back to round-robin → index 0 = E assertEquals(result.segmentToInstanceMap(), Map.of( - AR_SEGMENT0, AR_INSTANCE1, - AR_SEGMENT1, AR_INSTANCE3, - AR_SEGMENT2, AR_INSTANCE4)); + AR_SEGMENT0, AR_P1_RG1_SERVER_B, + AR_SEGMENT1, AR_P1_RG1_SERVER_D, + AR_SEGMENT2, AR_P0_RG0_SERVER_E)); } @Test public void testStrictReplicaGroupAdaptiveDisabledByFeatureFlag() { - // when the feature flag is disabled, the factory nulls out the adaptive selector so the - // selector falls back to round-robin. + // When the feature flag is disabled, the factory nulls out the adaptive selector so the selector falls back to the + // strict selector's non-adaptive, round-robin-by-replica-index path. HybridSelector hybridSelector = mock(HybridSelector.class); + BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); PinotConfiguration brokerConfig = new PinotConfiguration(Map.of( CommonConstants.Broker.AdaptiveServerSelector.CONFIG_OF_ENABLE_STRICT_REPLICA_GROUP, "false")); - ReplicaGroupInstanceSelector instanceSelector = - buildArSelector(STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector, brokerConfig); + StrictReplicaGroupInstanceSelector instanceSelector = + buildStrictReplicaGroupArSelector(hybridSelector, brokerMetrics, brokerConfig); - assertTrue(instanceSelector instanceof StrictReplicaGroupInstanceSelector); assertNull(instanceSelector._adaptiveServerSelector); assertNull(instanceSelector._priorityPoolInstanceSelector); BaseInstanceSelector.InstanceMapping result = - instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null); + instanceSelector.select(STRICT_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); - // hybridSelector never consulted during routing despite being passed to the factory verify(hybridSelector, never()).fetchServerRankingsWithScores(any()); + assertEquals(result.segmentToInstanceMap(), Map.of( + STRICT_SEGMENT0, STRICT_RG0_SERVER_A, + STRICT_SEGMENT1, STRICT_RG0_SERVER_A, + STRICT_SEGMENT2, STRICT_RG0_SERVER_B)); - // Round-robin with requestId=0 picks candidate index 0 per segment + result = instanceSelector.select(STRICT_SEGMENTS, 1, buildStrictReplicaGroupSegmentStates(), null); assertEquals(result.segmentToInstanceMap(), Map.of( - AR_SEGMENT0, AR_INSTANCE0, - AR_SEGMENT1, AR_INSTANCE2, - AR_SEGMENT2, AR_INSTANCE4)); + STRICT_SEGMENT0, STRICT_RG1_SERVER_C, + STRICT_SEGMENT1, STRICT_RG1_SERVER_C, + STRICT_SEGMENT2, STRICT_RG1_SERVER_D)); } @Test - public void testStrictReplicaGroupUsesAdaptiveServerSelector() { + public void testStrictReplicaGroupFactoryEnablesAdaptiveRouting() { HybridSelector hybridSelector = mock(HybridSelector.class); ReplicaGroupInstanceSelector instanceSelector = buildArSelector(STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector); @@ -460,37 +480,39 @@ public void testStrictReplicaGroupUsesAdaptiveServerSelector() { assertNotNull(instanceSelector._priorityPoolInstanceSelector); } - // --- Strict replica group pool-level adaptive routing tests --- - - // Topology: 2 replica groups (pool 0 and pool 1), 3 segments - // Pool 0: server_a (segment0, segment1), server_b (segment2) - // Pool 1: server_c (segment0, segment1), server_d (segment2) - private static final String SRG_SERVER_A = "server_a"; - private static final String SRG_SERVER_B = "server_b"; - private static final String SRG_SERVER_C = "server_c"; - private static final String SRG_SERVER_D = "server_d"; - private static final String SRG_SEGMENT0 = "seg0"; - private static final String SRG_SEGMENT1 = "seg1"; - private static final String SRG_SEGMENT2 = "seg2"; - private static final List SRG_SEGMENTS = Arrays.asList(SRG_SEGMENT0, SRG_SEGMENT1, SRG_SEGMENT2); + // --- Strict replica group adaptive routing tests --- + + // Topology: 2 replica groups, 3 segments + // Replica group 0: server_a (segment0, segment1), server_b (segment2) + // Replica group 1: server_c (segment0, segment1), server_d (segment2) + // If a strict-fixture server name omits `P#`, its pool is FALLBACK_POOL_ID. + private static final String STRICT_RG0_SERVER_A = "strict_rg0_server_a"; + private static final String STRICT_RG0_SERVER_B = "strict_rg0_server_b"; + private static final String STRICT_RG1_SERVER_C = "strict_rg1_server_c"; + private static final String STRICT_RG1_SERVER_D = "strict_rg1_server_d"; + private static final String STRICT_SEGMENT0 = "seg0"; + private static final String STRICT_SEGMENT1 = "seg1"; + private static final String STRICT_SEGMENT2 = "seg2"; + private static final List STRICT_SEGMENTS = Arrays.asList(STRICT_SEGMENT0, STRICT_SEGMENT1, STRICT_SEGMENT2); + + private List buildStrictReplicaGroupCandidates(String replicaGroup0Instance, + String replicaGroup1Instance) { + return Arrays.asList( + new SegmentInstanceCandidate(replicaGroup0Instance, true, FALLBACK_POOL_ID, 0), + new SegmentInstanceCandidate(replicaGroup1Instance, true, FALLBACK_POOL_ID, 1)); + } private SegmentStates buildStrictReplicaGroupSegmentStates() { Map> candidatesMap = new HashMap<>(); - // segment0 → server_a (pool 0), server_c (pool 1). - // Pool 0 is inserted into the LinkedHashMap before pool 1 because server_a (pool 0) appears - // first here. The round-robin fallback test depends on this insertion order. - candidatesMap.put(SRG_SEGMENT0, Arrays.asList( - new SegmentInstanceCandidate(SRG_SERVER_A, true, 0, 0), - new SegmentInstanceCandidate(SRG_SERVER_C, true, 1, 0))); - // segment1 → server_a (pool 0), server_c (pool 1) - candidatesMap.put(SRG_SEGMENT1, Arrays.asList( - new SegmentInstanceCandidate(SRG_SERVER_A, true, 0, 0), - new SegmentInstanceCandidate(SRG_SERVER_C, true, 1, 0))); - // segment2 → server_b (pool 0), server_d (pool 1) - candidatesMap.put(SRG_SEGMENT2, Arrays.asList( - new SegmentInstanceCandidate(SRG_SERVER_B, true, 0, 0), - new SegmentInstanceCandidate(SRG_SERVER_D, true, 1, 0))); - return new SegmentStates(candidatesMap, new HashSet<>(SRG_SEGMENTS), null); + // All strict-fixture servers intentionally share FALLBACK_POOL_ID so these tests fail if strict routing regresses + // back to grouping or filtering by pool instead of idealStateReplicaId. + candidatesMap.put(STRICT_SEGMENT0, + buildStrictReplicaGroupCandidates(STRICT_RG0_SERVER_A, STRICT_RG1_SERVER_C)); + candidatesMap.put(STRICT_SEGMENT1, + buildStrictReplicaGroupCandidates(STRICT_RG0_SERVER_A, STRICT_RG1_SERVER_C)); + candidatesMap.put(STRICT_SEGMENT2, + buildStrictReplicaGroupCandidates(STRICT_RG0_SERVER_B, STRICT_RG1_SERVER_D)); + return new SegmentStates(candidatesMap, new HashSet<>(STRICT_SEGMENTS), null); } private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(HybridSelector hybridSelector) { @@ -510,214 +532,217 @@ private StrictReplicaGroupInstanceSelector buildStrictReplicaGroupArSelector(Hyb .setRoutingConfig(routingConfig).build(); // Ideal state: mirrors the topology above IdealState idealState = createIdealState(Map.of( - SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), - SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), - SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE), Pair.of(SRG_SERVER_D, ONLINE)))); + STRICT_SEGMENT0, List.of( + Pair.of(STRICT_RG0_SERVER_A, ONLINE), + Pair.of(STRICT_RG1_SERVER_C, ONLINE)), + STRICT_SEGMENT1, List.of( + Pair.of(STRICT_RG0_SERVER_A, ONLINE), + Pair.of(STRICT_RG1_SERVER_C, ONLINE)), + STRICT_SEGMENT2, List.of( + Pair.of(STRICT_RG0_SERVER_B, ONLINE), + Pair.of(STRICT_RG1_SERVER_D, ONLINE)))); ExternalView externalView = createExternalView(Map.of( - SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), - SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), - SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE), Pair.of(SRG_SERVER_D, ONLINE)))); - // Provide server instances with pool assignments + STRICT_SEGMENT0, List.of( + Pair.of(STRICT_RG0_SERVER_A, ONLINE), + Pair.of(STRICT_RG1_SERVER_C, ONLINE)), + STRICT_SEGMENT1, List.of( + Pair.of(STRICT_RG0_SERVER_A, ONLINE), + Pair.of(STRICT_RG1_SERVER_C, ONLINE)), + STRICT_SEGMENT2, List.of( + Pair.of(STRICT_RG0_SERVER_B, ONLINE), + Pair.of(STRICT_RG1_SERVER_D, ONLINE)))); + // All servers use the fallback pool to prove strict routing keys off idealStateReplicaId rather than pool. ServerInstance serverA = mock(ServerInstance.class); - when(serverA.getPool()).thenReturn(0); + when(serverA.getPool()).thenReturn(FALLBACK_POOL_ID); ServerInstance serverB = mock(ServerInstance.class); - when(serverB.getPool()).thenReturn(0); + when(serverB.getPool()).thenReturn(FALLBACK_POOL_ID); ServerInstance serverC = mock(ServerInstance.class); - when(serverC.getPool()).thenReturn(1); + when(serverC.getPool()).thenReturn(FALLBACK_POOL_ID); ServerInstance serverD = mock(ServerInstance.class); - when(serverD.getPool()).thenReturn(1); + when(serverD.getPool()).thenReturn(FALLBACK_POOL_ID); Map serverMap = Map.of( - SRG_SERVER_A, serverA, SRG_SERVER_B, serverB, - SRG_SERVER_C, serverC, SRG_SERVER_D, serverD); + STRICT_RG0_SERVER_A, serverA, STRICT_RG0_SERVER_B, serverB, + STRICT_RG1_SERVER_C, serverC, STRICT_RG1_SERVER_D, serverD); return (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), brokerMetrics, hybridSelector, brokerConfig, - Set.of(SRG_SERVER_A, SRG_SERVER_B, SRG_SERVER_C, SRG_SERVER_D), - serverMap, idealState, externalView, new HashSet<>(SRG_SEGMENTS)); + Set.of(STRICT_RG0_SERVER_A, STRICT_RG0_SERVER_B, + STRICT_RG1_SERVER_C, STRICT_RG1_SERVER_D), + serverMap, idealState, externalView, new HashSet<>(STRICT_SEGMENTS)); } @Test public void testStrictReplicaGroupAdaptivePicksBestReplicaGroupByWorstCaseServer() { - // Pool 0 servers: server_a (rank 2), server_b (rank 3) → worst = 3 - // Pool 1 servers: server_c (rank 0), server_d (rank 1) → worst = 1 - // Pool 1 is better (lower worst-case rank), so all segments should route to pool 1 + // Replica group 0 servers: server_a (rank 2), server_b (rank 3) → worst = 3 + // Replica group 1 servers: server_c (rank 0), server_d (rank 1) → worst = 1 + // Replica group 1 is better (lower worst-case rank), so all segments should route there HybridSelector hybridSelector = mock(HybridSelector.class); BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector, brokerMetrics); when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( - new ImmutablePair<>(SRG_SERVER_C, 1.0), // rank 0 (best) - new ImmutablePair<>(SRG_SERVER_D, 2.0), // rank 1 - new ImmutablePair<>(SRG_SERVER_A, 3.0), // rank 2 - new ImmutablePair<>(SRG_SERVER_B, 4.0) // rank 3 (worst) + new ImmutablePair<>(STRICT_RG1_SERVER_C, 1.0), // rank 0 (best) + new ImmutablePair<>(STRICT_RG1_SERVER_D, 2.0), // rank 1 + new ImmutablePair<>(STRICT_RG0_SERVER_A, 3.0), // rank 2 + new ImmutablePair<>(STRICT_RG0_SERVER_B, 4.0) // rank 3 (worst) )); InstanceSelector.InstanceMapping result = - instanceSelector.select(SRG_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); + instanceSelector.select(STRICT_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); - // All segments routed to pool 1 + // All segments routed to replica group 1 assertEquals(result.segmentToInstanceMap(), Map.of( - SRG_SEGMENT0, SRG_SERVER_C, - SRG_SEGMENT1, SRG_SERVER_C, - SRG_SEGMENT2, SRG_SERVER_D)); + STRICT_SEGMENT0, STRICT_RG1_SERVER_C, + STRICT_SEGMENT1, STRICT_RG1_SERVER_C, + STRICT_SEGMENT2, STRICT_RG1_SERVER_D)); + // Metrics are still reported per actual pool, even though routing is chosen by replica group. verify(brokerMetrics).addMeteredValue(eq(BrokerMeter.POOL_SEG_QUERIES), eq(3L), - eq(BrokerMetrics.getTagForPreferredPool(null)), eq("1")); + eq(BrokerMetrics.getTagForPreferredPool(null)), eq(String.valueOf(FALLBACK_POOL_ID))); } @Test - public void testStrictReplicaGroupAdaptivePicksPool0WhenItHasBetterWorstCase() { - // Pool 0 servers: server_a (rank 0), server_b (rank 1) → worst = 1 - // Pool 1 servers: server_c (rank 2), server_d (rank 3) → worst = 3 - // Pool 0 is better + public void testStrictReplicaGroupAdaptivePicksGroup0WhenItHasBetterWorstCase() { + // Replica group 0 servers: server_a (rank 0), server_b (rank 1) → worst = 1 + // Replica group 1 servers: server_c (rank 2), server_d (rank 3) → worst = 3 + // Replica group 0 is better HybridSelector hybridSelector = mock(HybridSelector.class); StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector); when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( - new ImmutablePair<>(SRG_SERVER_A, 1.0), - new ImmutablePair<>(SRG_SERVER_B, 2.0), - new ImmutablePair<>(SRG_SERVER_C, 3.0), - new ImmutablePair<>(SRG_SERVER_D, 4.0) + new ImmutablePair<>(STRICT_RG0_SERVER_A, 1.0), + new ImmutablePair<>(STRICT_RG0_SERVER_B, 2.0), + new ImmutablePair<>(STRICT_RG1_SERVER_C, 3.0), + new ImmutablePair<>(STRICT_RG1_SERVER_D, 4.0) )); InstanceSelector.InstanceMapping result = - instanceSelector.select(SRG_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); + instanceSelector.select(STRICT_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); - // All segments routed to pool 0 + // All segments routed to replica group 0 assertEquals(result.segmentToInstanceMap(), Map.of( - SRG_SEGMENT0, SRG_SERVER_A, - SRG_SEGMENT1, SRG_SERVER_A, - SRG_SEGMENT2, SRG_SERVER_B)); + STRICT_SEGMENT0, STRICT_RG0_SERVER_A, + STRICT_SEGMENT1, STRICT_RG0_SERVER_A, + STRICT_SEGMENT2, STRICT_RG0_SERVER_B)); } @Test - public void testStrictReplicaGroupAdaptiveNoStatsFallsBackToRoundRobin() { - // When fetchServerRankingsWithScores returns empty, fall back to round-robin + public void testStrictReplicaGroupAdaptiveEmptyRankingsFallBackToRoundRobin() { + // When adaptive ranking returns no servers, fall back to round-robin by replica-group index. HybridSelector hybridSelector = mock(HybridSelector.class); StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector); when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(List.of()); - // requestId=0 → picks group at index 0 (pool 0) + // requestId=0 → picks group at index 0 (replica group 0) InstanceSelector.InstanceMapping result = - instanceSelector.select(SRG_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); + instanceSelector.select(STRICT_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); assertEquals(result.segmentToInstanceMap(), Map.of( - SRG_SEGMENT0, SRG_SERVER_A, - SRG_SEGMENT1, SRG_SERVER_A, - SRG_SEGMENT2, SRG_SERVER_B)); + STRICT_SEGMENT0, STRICT_RG0_SERVER_A, + STRICT_SEGMENT1, STRICT_RG0_SERVER_A, + STRICT_SEGMENT2, STRICT_RG0_SERVER_B)); - // requestId=1 → picks group at index 1 (pool 1) - result = instanceSelector.select(SRG_SEGMENTS, 1, buildStrictReplicaGroupSegmentStates(), null); + // requestId=1 → picks group at index 1 (replica group 1) + result = instanceSelector.select(STRICT_SEGMENTS, 1, buildStrictReplicaGroupSegmentStates(), null); assertEquals(result.segmentToInstanceMap(), Map.of( - SRG_SEGMENT0, SRG_SERVER_C, - SRG_SEGMENT1, SRG_SERVER_C, - SRG_SEGMENT2, SRG_SERVER_D)); + STRICT_SEGMENT0, STRICT_RG1_SERVER_C, + STRICT_SEGMENT1, STRICT_RG1_SERVER_C, + STRICT_SEGMENT2, STRICT_RG1_SERVER_D)); } @Test public void testStrictReplicaGroupAdaptiveConsistencyAllSegmentsSameReplicaGroup() { - // For any set of adaptive scores, all segments must route to the same replica group + // When adaptive scores favor replica group 1, all segments must route to that same replica group. HybridSelector hybridSelector = mock(HybridSelector.class); StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector); when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( - new ImmutablePair<>(SRG_SERVER_C, 1.0), - new ImmutablePair<>(SRG_SERVER_D, 2.0), - new ImmutablePair<>(SRG_SERVER_A, 3.0), - new ImmutablePair<>(SRG_SERVER_B, 4.0) + new ImmutablePair<>(STRICT_RG1_SERVER_C, 1.0), + new ImmutablePair<>(STRICT_RG1_SERVER_D, 2.0), + new ImmutablePair<>(STRICT_RG0_SERVER_A, 3.0), + new ImmutablePair<>(STRICT_RG0_SERVER_B, 4.0) )); InstanceSelector.InstanceMapping result = - instanceSelector.select(SRG_SEGMENTS, 42, buildStrictReplicaGroupSegmentStates(), null); - - // Collect all pool IDs from selected instances - Map selections = result.segmentToInstanceMap(); - Set poolsUsed = new HashSet<>(); - SegmentStates states = buildStrictReplicaGroupSegmentStates(); - for (Map.Entry entry : selections.entrySet()) { - String segment = entry.getKey(); - String selectedInstance = entry.getValue(); - for (SegmentInstanceCandidate candidate : states.getCandidates(segment)) { - if (candidate.getInstance().equals(selectedInstance)) { - poolsUsed.add(candidate.getPool()); - break; - } - } - } - // All segments must be from the same pool - assertEquals(poolsUsed.size(), 1, "All segments should be routed to the same replica group"); + instanceSelector.select(STRICT_SEGMENTS, 42, buildStrictReplicaGroupSegmentStates(), null); + + assertEquals(result.segmentToInstanceMap(), Map.of( + STRICT_SEGMENT0, STRICT_RG1_SERVER_C, + STRICT_SEGMENT1, STRICT_RG1_SERVER_C, + STRICT_SEGMENT2, STRICT_RG1_SERVER_D)); } @Test public void testStrictReplicaGroupAdaptivePrefersFullCoverageGroup() { - // Topology: pool 0 has all 3 segments, pool 1 only has segment0 and segment1 (missing segment2). - // Even though pool 1 has better adaptive ranks, pool 0 should be chosen because it has full coverage. - // Pool 0: server_a (seg0, seg1), server_b (seg2) — full coverage - // Pool 1: server_c (seg0, seg1) — missing seg2, partial coverage + // Topology: replica group 0 has all 3 segments, replica group 1 only has segment0 and segment1. + // Even though replica group 1 has better adaptive ranks, replica group 0 should be chosen (full coverage). + // Replica group 0: server_a (seg0, seg1), server_b (seg2) — full coverage + // Replica group 1: server_c (seg0, seg1) — missing seg2, partial coverage + // Pools intentionally do not match replica groups: server_a and server_c share pool 0, while server_b is pool 1. String serverC = "server_c"; HybridSelector hybridSelector = mock(HybridSelector.class); RoutingConfig routingConfig = new RoutingConfig(null, null, STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false); TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testUpsertTable") .setRoutingConfig(routingConfig).build(); - // Pool 1 only has seg0 and seg1 (no seg2) + // Replica group 1 only has seg0 and seg1 (no seg2) IdealState idealState = createIdealState(Map.of( - SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), - SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), - SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE)))); + STRICT_SEGMENT0, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), + STRICT_SEGMENT1, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), + STRICT_SEGMENT2, List.of(Pair.of(STRICT_RG0_SERVER_B, ONLINE)))); ExternalView externalView = createExternalView(Map.of( - SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), - SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), - SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE)))); + STRICT_SEGMENT0, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), + STRICT_SEGMENT1, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE), Pair.of(serverC, ONLINE)), + STRICT_SEGMENT2, List.of(Pair.of(STRICT_RG0_SERVER_B, ONLINE)))); ServerInstance serverA = mock(ServerInstance.class); when(serverA.getPool()).thenReturn(0); ServerInstance serverB = mock(ServerInstance.class); - when(serverB.getPool()).thenReturn(0); + when(serverB.getPool()).thenReturn(1); ServerInstance serverCInstance = mock(ServerInstance.class); - when(serverCInstance.getPool()).thenReturn(1); + when(serverCInstance.getPool()).thenReturn(0); Map serverMap = Map.of( - SRG_SERVER_A, serverA, SRG_SERVER_B, serverB, serverC, serverCInstance); + STRICT_RG0_SERVER_A, serverA, STRICT_RG0_SERVER_B, serverB, serverC, serverCInstance); StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, new PinotConfiguration(Map.of()), - Set.of(SRG_SERVER_A, SRG_SERVER_B, serverC), - serverMap, idealState, externalView, new HashSet<>(SRG_SEGMENTS)); + Set.of(STRICT_RG0_SERVER_A, STRICT_RG0_SERVER_B, serverC), + serverMap, idealState, externalView, new HashSet<>(STRICT_SEGMENTS)); - // Pool 1 (server_c) has rank 0 — best rank, but only covers 2 of 3 segments - // Pool 0 (server_a rank 2, server_b rank 3) — worse ranks, but full coverage + // Replica group 1 (server_c) has rank 0 — best rank, but only covers 2 of 3 segments + // Replica group 0 (server_a rank 1, server_b rank 2) — worse ranks, but full coverage when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( new ImmutablePair<>(serverC, 1.0), // rank 0 (best) - new ImmutablePair<>(SRG_SERVER_A, 3.0), // rank 1 - new ImmutablePair<>(SRG_SERVER_B, 4.0) // rank 2 + new ImmutablePair<>(STRICT_RG0_SERVER_A, 3.0), // rank 1 + new ImmutablePair<>(STRICT_RG0_SERVER_B, 4.0) // rank 2 )); - // Build segment states where pool 1 is missing seg2 + // Build segment states where replica group 1 is missing seg2 Map> candidatesMap = new HashMap<>(); - candidatesMap.put(SRG_SEGMENT0, Arrays.asList( - new SegmentInstanceCandidate(SRG_SERVER_A, true, 0, 0), - new SegmentInstanceCandidate(serverC, true, 1, 0))); - candidatesMap.put(SRG_SEGMENT1, Arrays.asList( - new SegmentInstanceCandidate(SRG_SERVER_A, true, 0, 0), - new SegmentInstanceCandidate(serverC, true, 1, 0))); - candidatesMap.put(SRG_SEGMENT2, List.of( - new SegmentInstanceCandidate(SRG_SERVER_B, true, 0, 0))); - SegmentStates segmentStates = new SegmentStates(candidatesMap, new HashSet<>(SRG_SEGMENTS), null); + candidatesMap.put(STRICT_SEGMENT0, Arrays.asList( + new SegmentInstanceCandidate(STRICT_RG0_SERVER_A, true, 0, 0), + new SegmentInstanceCandidate(serverC, true, 0, 1))); + candidatesMap.put(STRICT_SEGMENT1, Arrays.asList( + new SegmentInstanceCandidate(STRICT_RG0_SERVER_A, true, 0, 0), + new SegmentInstanceCandidate(serverC, true, 0, 1))); + candidatesMap.put(STRICT_SEGMENT2, List.of( + new SegmentInstanceCandidate(STRICT_RG0_SERVER_B, true, 1, 0))); + SegmentStates segmentStates = new SegmentStates(candidatesMap, new HashSet<>(STRICT_SEGMENTS), null); InstanceSelector.InstanceMapping result = - instanceSelector.select(SRG_SEGMENTS, 0, segmentStates, null); + instanceSelector.select(STRICT_SEGMENTS, 0, segmentStates, null); - // Pool 0 should be chosen (full coverage) even though pool 1 has better ranks + // Replica group 0 should be chosen (full coverage) even though replica group 1 has better ranks assertEquals(result.segmentToInstanceMap(), Map.of( - SRG_SEGMENT0, SRG_SERVER_A, - SRG_SEGMENT1, SRG_SERVER_A, - SRG_SEGMENT2, SRG_SERVER_B)); + STRICT_SEGMENT0, STRICT_RG0_SERVER_A, + STRICT_SEGMENT1, STRICT_RG0_SERVER_A, + STRICT_SEGMENT2, STRICT_RG0_SERVER_B)); } @Test public void testStrictReplicaGroupAdaptiveNewSegmentNotFalselyUnavailable() { - // Topology: server_a (pool 0) hosts seg0 only. seg1 has no online instance (simulating a - // new/unassigned segment with null candidates in segmentStates). Pool 0 is chosen (only pool). + // Topology: server_a (replica group 0) hosts seg0 only. seg1 has no online instance (simulating a + // new/unassigned segment with null candidates in segmentStates). Replica group 0 is chosen (only group). // seg1 should NOT appear in unavailable — it is a new segment, not a dropped one. String serverA = "new_server_a"; String seg0 = "new_seg0"; @@ -734,7 +759,7 @@ public void testStrictReplicaGroupAdaptiveNewSegmentNotFalselyUnavailable() { ExternalView externalView = createExternalView(Map.of( seg0, List.of(Pair.of(serverA, ONLINE)))); ServerInstance serverAInstance = mock(ServerInstance.class); - when(serverAInstance.getPool()).thenReturn(0); + when(serverAInstance.getPool()).thenReturn(FALLBACK_POOL_ID); StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, @@ -759,134 +784,130 @@ public void testStrictReplicaGroupAdaptiveNewSegmentNotFalselyUnavailable() { @Test public void testStrictReplicaGroupAdaptiveDroppedSegmentReportedAsUnavailable() { - // Topology: no replica group has full coverage. - // Pool 0: server_a hosts seg0, seg1 (missing seg2) - // Pool 1: server_c hosts seg1, seg2 (missing seg0) - // Pool 1 has better ranks -> chosen. seg0 has no candidate in pool 1 -> reported unavailable. - String serverA = "partial_server_a"; - String serverC = "partial_server_c"; - String seg0 = "partial_seg0"; - String seg1 = "partial_seg1"; - String seg2 = "partial_seg2"; - List segments = Arrays.asList(seg0, seg1, seg2); + // Topology: no replica group (idealStateReplicaId) has full coverage. + // ReplicaId 0 (server_a): covers seg0, seg1 (missing seg2) + // ReplicaId 1 (server_c, server_d): covers seg1, seg2 (missing seg0) + // ReplicaId 1 has better rank -> chosen. seg0 has no replicaId-1 candidate -> reported unavailable. HybridSelector hybridSelector = mock(HybridSelector.class); + StrictReplicaGroupInstanceSelector instanceSelector = buildStrictReplicaGroupArSelector(hybridSelector); - RoutingConfig routingConfig = new RoutingConfig(null, null, STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false); - TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testPartialCoverage") - .setRoutingConfig(routingConfig).build(); - IdealState idealState = createIdealState(Map.of( - seg0, List.of(Pair.of(serverA, ONLINE)), - seg1, List.of(Pair.of(serverA, ONLINE), Pair.of(serverC, ONLINE)), - seg2, List.of(Pair.of(serverC, ONLINE)))); - ExternalView externalView = createExternalView(Map.of( - seg0, List.of(Pair.of(serverA, ONLINE)), - seg1, List.of(Pair.of(serverA, ONLINE), Pair.of(serverC, ONLINE)), - seg2, List.of(Pair.of(serverC, ONLINE)))); - ServerInstance serverAInstance = mock(ServerInstance.class); - when(serverAInstance.getPool()).thenReturn(0); - ServerInstance serverCInstance = mock(ServerInstance.class); - when(serverCInstance.getPool()).thenReturn(1); - Map serverMap = Map.of(serverA, serverAInstance, serverC, serverCInstance); - BrokerMetrics brokerMetrics = mock(BrokerMetrics.class); - StrictReplicaGroupInstanceSelector instanceSelector = - (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, - mock(ZkHelixPropertyStore.class), brokerMetrics, hybridSelector, - new PinotConfiguration(Map.of()), - Set.of(serverA, serverC), - serverMap, idealState, externalView, new HashSet<>(segments)); - - // Pool 1 (server_c) ranked better -> pool 1 chosen when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( - new ImmutablePair<>(serverC, 1.0), - new ImmutablePair<>(serverA, 3.0) + new ImmutablePair<>(STRICT_RG1_SERVER_C, 1.0), + new ImmutablePair<>(STRICT_RG0_SERVER_A, 3.0) )); - BrokerRequest brokerRequest = mock(BrokerRequest.class); - PinotQuery pinotQuery = mock(PinotQuery.class); - when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery); - when(pinotQuery.getQueryOptions()).thenReturn(null); + // Craft segment states with partial coverage per replica group: + // seg0: only in replicaId 0 (server_a) + // seg1: in both replicaId 0 (server_a) and replicaId 1 (server_c) + // seg2: only in replicaId 1 (server_d) + Map> candidatesMap = new HashMap<>(); + candidatesMap.put(STRICT_SEGMENT0, List.of( + new SegmentInstanceCandidate(STRICT_RG0_SERVER_A, true, FALLBACK_POOL_ID, 0))); + candidatesMap.put(STRICT_SEGMENT1, Arrays.asList( + new SegmentInstanceCandidate(STRICT_RG0_SERVER_A, true, FALLBACK_POOL_ID, 0), + new SegmentInstanceCandidate(STRICT_RG1_SERVER_C, true, FALLBACK_POOL_ID, 1))); + candidatesMap.put(STRICT_SEGMENT2, List.of( + new SegmentInstanceCandidate(STRICT_RG1_SERVER_D, true, FALLBACK_POOL_ID, 1))); + SegmentStates segmentStates = new SegmentStates(candidatesMap, new HashSet<>(STRICT_SEGMENTS), null); - InstanceSelector.SelectionResult result = instanceSelector.select(brokerRequest, segments, 0); + InstanceSelector.InstanceMapping result = + instanceSelector.select(STRICT_SEGMENTS, 0, segmentStates, null); - // seg1 and seg2 routed to pool 1 - assertEquals(result.getSegmentToInstanceMap(), Map.of(seg1, serverC, seg2, serverC)); - // seg0 has candidates (pool 0 only) but no match in chosen pool 1 -> reported unavailable - assertEquals(result.getUnavailableSegments(), List.of(seg0)); + // seg1 and seg2 routed to replicaId 1 + assertEquals(result.segmentToInstanceMap(), Map.of( + STRICT_SEGMENT1, STRICT_RG1_SERVER_C, + STRICT_SEGMENT2, STRICT_RG1_SERVER_D)); + // seg0 has candidates (replicaId 0 only) but no match in chosen replicaId 1 -> reported unavailable + assertEquals(result.unavailableSegments(), List.of(STRICT_SEGMENT0)); } @Test public void testStrictReplicaGroupAdaptiveOnlyQueryRelevantServersScored() { - // Topology: pool 1 has an extra server (server_e) that hosts a segment NOT in the query. - // server_e has a very high rank but should not affect pool 1's score since it hosts no query segments. - // Pool 0: server_a (rank 2), server_b (rank 3) → worst = 3 - // Pool 1: server_c (rank 0), server_d (rank 1) → worst = 1 - // (server_e at rank 99 hosts segment "seg_extra" which is not in the query) - // Pool 1 should win because only query-relevant servers matter (server_e is ignored). - String serverE = "server_e"; + // Replica group 1 has an extra server (server_e) that only hosts seg_extra, which is not in the query. + // server_e should never be scored for a query over STRICT_SEGMENTS. + // Replica group 0: server_a (rank 2), server_b (rank 3) → worst = 3 + // Replica group 1: server_c (rank 0), server_d (rank 1) → worst = 1 + // Replica group 1 should win because only query-relevant servers matter. + String rg1ServerE = "server_e"; String segExtra = "seg_extra"; HybridSelector hybridSelector = mock(HybridSelector.class); - // Build selector with server_e in pool 1, hosting seg_extra + // Build selector with server_e in replica group 1, hosting seg_extra. RoutingConfig routingConfig = new RoutingConfig(null, null, STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false); TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testUpsertTable") .setRoutingConfig(routingConfig).build(); + // server_e's pool intentionally differs from the queried servers to show pool is irrelevant. + // segExtra places server_d at position 0 and server_e at position 1 (idealStateReplicaId=1). IdealState idealState = createIdealState(Map.of( - SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), - SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), - SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE), Pair.of(SRG_SERVER_D, ONLINE)), - segExtra, List.of(Pair.of(serverE, ONLINE)))); + STRICT_SEGMENT0, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE), Pair.of(STRICT_RG1_SERVER_C, ONLINE)), + STRICT_SEGMENT1, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE), Pair.of(STRICT_RG1_SERVER_C, ONLINE)), + STRICT_SEGMENT2, List.of(Pair.of(STRICT_RG0_SERVER_B, ONLINE), Pair.of(STRICT_RG1_SERVER_D, ONLINE)), + segExtra, List.of(Pair.of(STRICT_RG1_SERVER_D, ONLINE), Pair.of(rg1ServerE, ONLINE)))); ExternalView externalView = createExternalView(Map.of( - SRG_SEGMENT0, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), - SRG_SEGMENT1, List.of(Pair.of(SRG_SERVER_A, ONLINE), Pair.of(SRG_SERVER_C, ONLINE)), - SRG_SEGMENT2, List.of(Pair.of(SRG_SERVER_B, ONLINE), Pair.of(SRG_SERVER_D, ONLINE)), - segExtra, List.of(Pair.of(serverE, ONLINE)))); + STRICT_SEGMENT0, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE), Pair.of(STRICT_RG1_SERVER_C, ONLINE)), + STRICT_SEGMENT1, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE), Pair.of(STRICT_RG1_SERVER_C, ONLINE)), + STRICT_SEGMENT2, List.of(Pair.of(STRICT_RG0_SERVER_B, ONLINE), Pair.of(STRICT_RG1_SERVER_D, ONLINE)), + segExtra, List.of(Pair.of(STRICT_RG1_SERVER_D, ONLINE), Pair.of(rg1ServerE, ONLINE)))); ServerInstance serverA = mock(ServerInstance.class); - when(serverA.getPool()).thenReturn(0); + when(serverA.getPool()).thenReturn(FALLBACK_POOL_ID); ServerInstance serverB = mock(ServerInstance.class); - when(serverB.getPool()).thenReturn(0); + when(serverB.getPool()).thenReturn(FALLBACK_POOL_ID); ServerInstance serverC = mock(ServerInstance.class); - when(serverC.getPool()).thenReturn(1); + when(serverC.getPool()).thenReturn(FALLBACK_POOL_ID); ServerInstance serverD = mock(ServerInstance.class); - when(serverD.getPool()).thenReturn(1); + when(serverD.getPool()).thenReturn(FALLBACK_POOL_ID); ServerInstance serverEInstance = mock(ServerInstance.class); - when(serverEInstance.getPool()).thenReturn(1); + when(serverEInstance.getPool()).thenReturn(99); Map serverMap = Map.of( - SRG_SERVER_A, serverA, SRG_SERVER_B, serverB, - SRG_SERVER_C, serverC, SRG_SERVER_D, serverD, - serverE, serverEInstance); - Set allSegments = new HashSet<>(SRG_SEGMENTS); + STRICT_RG0_SERVER_A, serverA, STRICT_RG0_SERVER_B, serverB, + STRICT_RG1_SERVER_C, serverC, STRICT_RG1_SERVER_D, serverD, + rg1ServerE, serverEInstance); + Set allSegments = new HashSet<>(STRICT_SEGMENTS); allSegments.add(segExtra); StrictReplicaGroupInstanceSelector instanceSelector = (StrictReplicaGroupInstanceSelector) InstanceSelectorFactory.getInstanceSelector(tableConfig, mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class), hybridSelector, new PinotConfiguration(Map.of()), - Set.of(SRG_SERVER_A, SRG_SERVER_B, SRG_SERVER_C, SRG_SERVER_D, serverE), + Set.of(STRICT_RG0_SERVER_A, STRICT_RG0_SERVER_B, + STRICT_RG1_SERVER_C, STRICT_RG1_SERVER_D, rg1ServerE), serverMap, idealState, externalView, allSegments); - // server_e has rank 99 but hosts no query segments — should be irrelevant + // server_e has a much worse mocked score but hosts no query segments, so it should never be ranked. when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList( - new ImmutablePair<>(SRG_SERVER_C, 1.0), // rank 0 - new ImmutablePair<>(SRG_SERVER_D, 2.0), // rank 1 - new ImmutablePair<>(SRG_SERVER_A, 3.0), // rank 2 - new ImmutablePair<>(SRG_SERVER_B, 4.0), // rank 3 - new ImmutablePair<>(serverE, 99.0) // rank 4 (high score, not query-relevant) + new ImmutablePair<>(STRICT_RG1_SERVER_C, 1.0), // rank 0 + new ImmutablePair<>(STRICT_RG1_SERVER_D, 2.0), // rank 1 + new ImmutablePair<>(STRICT_RG0_SERVER_A, 3.0), // rank 2 + new ImmutablePair<>(STRICT_RG0_SERVER_B, 4.0), // rank 3 + new ImmutablePair<>(rg1ServerE, 99.0) // rank 4 (not query-relevant) )); - // Query only SRG_SEGMENTS (not seg_extra), so server_e is not query-relevant + Map> candidatesMap = new HashMap<>(); + candidatesMap.put(STRICT_SEGMENT0, + buildStrictReplicaGroupCandidates(STRICT_RG0_SERVER_A, STRICT_RG1_SERVER_C)); + candidatesMap.put(STRICT_SEGMENT1, + buildStrictReplicaGroupCandidates(STRICT_RG0_SERVER_A, STRICT_RG1_SERVER_C)); + candidatesMap.put(STRICT_SEGMENT2, + buildStrictReplicaGroupCandidates(STRICT_RG0_SERVER_B, STRICT_RG1_SERVER_D)); + candidatesMap.put(segExtra, Arrays.asList( + new SegmentInstanceCandidate(STRICT_RG1_SERVER_D, true, FALLBACK_POOL_ID, 0), + new SegmentInstanceCandidate(rg1ServerE, true, 99, 1))); + SegmentStates segmentStates = new SegmentStates(candidatesMap, allSegments, null); + InstanceSelector.InstanceMapping result = - instanceSelector.select(SRG_SEGMENTS, 0, buildStrictReplicaGroupSegmentStates(), null); + instanceSelector.select(STRICT_SEGMENTS, 0, segmentStates, null); - // Pool 1 wins (worst rank among query-relevant servers = 1) over pool 0 (worst = 3) - // server_e's rank 4 is not considered because it hosts no query segments + // Replica group 1 wins (worst rank among query-relevant servers = 1) over replica group 0 (worst = 3) + // server_e is not considered because it hosts no query segments assertEquals(result.segmentToInstanceMap(), Map.of( - SRG_SEGMENT0, SRG_SERVER_C, - SRG_SEGMENT1, SRG_SERVER_C, - SRG_SEGMENT2, SRG_SERVER_D)); + STRICT_SEGMENT0, STRICT_RG1_SERVER_C, + STRICT_SEGMENT1, STRICT_RG1_SERVER_C, + STRICT_SEGMENT2, STRICT_RG1_SERVER_D)); ArgumentCaptor> rankedCandidatesCaptor = ArgumentCaptor.forClass(List.class); verify(hybridSelector).fetchServerRankingsWithScores(rankedCandidatesCaptor.capture()); assertEquals(new HashSet<>(rankedCandidatesCaptor.getValue()), - Set.of(SRG_SERVER_A, SRG_SERVER_B, SRG_SERVER_C, SRG_SERVER_D)); + Set.of(STRICT_RG0_SERVER_A, STRICT_RG0_SERVER_B, + STRICT_RG1_SERVER_C, STRICT_RG1_SERVER_D)); assertEquals(rankedCandidatesCaptor.getValue().size(), 4); } }