This is an automated email from the ASF dual-hosted git repository.
gortiz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 88ea694bae5 [#12507] Implement adaptive pool selection for
StrictReplicaGroupInstanceSelector (#18947)
88ea694bae5 is described below
commit 88ea694bae5f98760fb8c61f737fe6d3bc713427
Author: Timothy Elgersma <[email protected]>
AuthorDate: Tue Sep 1 02:54:58 2026 -0400
[#12507] Implement adaptive pool selection for
StrictReplicaGroupInstanceSelector (#18947)
---
.../instanceselector/BalancedInstanceSelector.java | 5 +-
.../instanceselector/BaseInstanceSelector.java | 24 +-
.../routing/instanceselector/InstanceSelector.java | 22 +-
.../instanceselector/InstanceSelectorFactory.java | 23 +-
.../MultiStageReplicaGroupSelector.java | 14 +-
.../ReplicaGroupInstanceSelector.java | 39 +-
.../instanceselector/SegmentInstanceCandidate.java | 2 +-
.../StrictReplicaGroupInstanceSelector.java | 62 +-
.../routing/manager/BaseBrokerRoutingManager.java | 6 +-
.../instanceselector/InstanceSelectorTest.java | 90 +-
.../instanceselector/ReplicaGroupSelectorTest.java | 1063 ++++++++++++++++++++
.../apache/pinot/spi/utils/CommonConstants.java | 7 +
pom.xml | 2 +-
13 files changed, 1214 insertions(+), 145 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 7707366ee22..a1a7ce09bf2 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 @@ package org.apache.pinot.broker.routing.instanceselector;
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;
@@ -44,7 +43,7 @@ import org.apache.pinot.common.utils.HashUtil;
public class BalancedInstanceSelector extends BaseInstanceSelector {
@Override
- public Pair<Map<String, String>, Map<String, String>> select(List<String>
segments, int requestId,
+ public InstanceMapping select(List<String> segments, int requestId,
SegmentStates segmentStates, Map<String, String> queryOptions) {
Map<String, String> 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.
@@ -87,6 +86,6 @@ public class BalancedInstanceSelector extends
BaseInstanceSelector {
_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 f322325519d..6c8d96cbf6d 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.SortedMap;
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;
@@ -517,15 +516,14 @@ public abstract class BaseInstanceSelector implements
InstanceSelector {
(brokerRequest.getPinotQuery() != null &&
brokerRequest.getPinotQuery().getQueryOptions() != null)
? brokerRequest.getPinotQuery().getQueryOptions() : Map.of();
int requestIdInt = (int) (requestId % MAX_REQUEST_ID);
- // Copy the volatile reference so that segmentToInstanceMap and
unavailableSegments can have a consistent view of
- // the state.
+ // Copy the volatile reference so that the segment-to-instance map and
unavailable segments have a consistent view
+ // of the state.
SegmentStates segmentStates = _segmentStates;
- Pair<Map<String, String>, Map<String, String>> segmentToInstanceMap =
- select(segments, requestIdInt, segmentStates, queryOptions);
+ InstanceMapping mapping = select(segments, requestIdInt, segmentStates,
queryOptions);
Set<String> unavailableSegments = segmentStates.getUnavailableSegments();
if (unavailableSegments.isEmpty()) {
- return new SelectionResult(segmentToInstanceMap, List.of(), 0);
+ return new SelectionResult(mapping, List.of(), 0);
} else {
List<String> unavailableSegmentsForRequest = new ArrayList<>();
for (String segment : segments) {
@@ -533,7 +531,7 @@ public abstract class BaseInstanceSelector implements
InstanceSelector {
unavailableSegmentsForRequest.add(segment);
}
}
- return new SelectionResult(segmentToInstanceMap,
unavailableSegmentsForRequest, 0);
+ return new SelectionResult(mapping, unavailableSegmentsForRequest, 0);
}
}
@@ -555,10 +553,10 @@ public abstract class BaseInstanceSelector implements
InstanceSelector {
return pool;
}
- /// 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.
- protected abstract Pair<Map<String, String>, Map<String, String>/*optional
segments*/> select(List<String> segments,
- int requestId, SegmentStates segmentStates, Map<String, String>
queryOptions);
+ /// Selects the server instances for the given segments based on the request
id and segment states. Returns an
+ /// [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 InstanceMapping select(List<String> segments, int
requestId, SegmentStates segmentStates,
+ Map<String, String> 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 06491ebd90d..5dd2d4f429e 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
@@ -23,7 +23,6 @@ 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;
@@ -88,27 +87,38 @@ public interface InstanceSelector {
return null;
}
+ /// Holds the result of an instance selection: `segmentToInstanceMap` maps
each segment to its selected server
+ /// instance, and `optionalSegmentToInstanceMap` maps segments not yet fully
online that the server may skip.
+ record InstanceMapping(Map<String, String> segmentToInstanceMap,
+ Map<String, String> optionalSegmentToInstanceMap) {
+ static final InstanceMapping EMPTY = new InstanceMapping(Map.of(),
Map.of());
+ }
+
class SelectionResult {
- private final Pair<Map<String, String>, Map<String, String>/*optional
segments*/> _segmentToInstanceMap;
+ private final InstanceMapping _instanceMapping;
private final List<String> _unavailableSegments;
private int _numPrunedSegments;
- public SelectionResult(Pair<Map<String, String>, Map<String, String>>
segmentToInstanceMap,
+ public SelectionResult(InstanceMapping instanceMapping,
List<String> unavailableSegments, int numPrunedSegments) {
- _segmentToInstanceMap = segmentToInstanceMap;
+ _instanceMapping = instanceMapping;
_unavailableSegments = unavailableSegments;
_numPrunedSegments = numPrunedSegments;
}
+ public static SelectionResult empty(int numPrunedSegments) {
+ return new SelectionResult(InstanceMapping.EMPTY, List.of(),
numPrunedSegments);
+ }
+
/// Returns the map from segment to selected server instance hosting the
segment.
public Map<String, String> getSegmentToInstanceMap() {
- return _segmentToInstanceMap.getLeft();
+ return _instanceMapping.segmentToInstanceMap();
}
/// Returns the map from optional segment to selected server instance
hosting the optional segment.
/// Optional segments can be skipped by broker or server upon any issue
w/o failing the query.
public Map<String, String> getOptionalSegmentToInstanceMap() {
- return _segmentToInstanceMap.getRight();
+ return _instanceMapping.optionalSegmentToInstanceMap();
}
/// Returns the unavailable segments (no enabled instance or all enabled
instances are in ERROR state).
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 49536b93a8f..79361106aba 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
@@ -74,6 +74,7 @@ public class InstanceSelectorFactory {
ExternalView externalView, Set<String> onlineSegments) {
String tableNameWithType = tableConfig.getTableName();
RoutingConfig routingConfig = tableConfig.getRoutingConfig();
+ AdaptiveServerSelector effectiveAdaptiveServerSelector =
adaptiveServerSelector;
boolean useFixedReplica =
brokerConfig.getProperty(CommonConstants.Broker.CONFIG_OF_USE_FIXED_REPLICA,
CommonConstants.Broker.DEFAULT_USE_FIXED_REPLICA);
if (routingConfig != null && routingConfig.getUseFixedReplica() != null) {
@@ -106,6 +107,19 @@ public class InstanceSelectorFactory {
}
case RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE: {
LOGGER.info("Using StrictReplicaGroupInstanceSelector for table:
{}", tableNameWithType);
+ boolean enableStrictReplicaGroupAdaptiveRouting =
brokerConfig.getProperty(
+
CommonConstants.Broker.AdaptiveServerSelector.CONFIG_OF_STRICT_REPLICA_GROUP_ENABLED,
+
CommonConstants.Broker.AdaptiveServerSelector.DEFAULT_STRICT_REPLICA_GROUP_ENABLED);
+ if (!enableStrictReplicaGroupAdaptiveRouting &&
effectiveAdaptiveServerSelector != null) {
+ LOGGER.info("Adaptive routing disabled for
StrictReplicaGroupInstanceSelector table: {}",
+ tableNameWithType);
+ effectiveAdaptiveServerSelector = null;
+ }
+ if (useFixedReplica && effectiveAdaptiveServerSelector != null) {
+ LOGGER.warn("Disabling adaptive routing for
StrictReplicaGroupInstanceSelector table {} using fixed "
+ + "replica routing", tableNameWithType);
+ effectiveAdaptiveServerSelector = null;
+ }
instanceSelector = new StrictReplicaGroupInstanceSelector();
break;
}
@@ -136,8 +150,15 @@ public class InstanceSelectorFactory {
if (instanceSelector == null) {
instanceSelector = new BalancedInstanceSelector();
}
+ if (instanceSelector.getClass() == ReplicaGroupInstanceSelector.class
+ && (tableConfig.isUpsertEnabled() || tableConfig.isDedupEnabled())
+ && effectiveAdaptiveServerSelector != null) {
+ LOGGER.warn("Disabling adaptive routing for legacy upsert/dedup table {}
using ReplicaGroupInstanceSelector",
+ tableNameWithType);
+ effectiveAdaptiveServerSelector = null;
+ }
- instanceSelector.init(tableConfig, propertyStore, brokerMetrics,
adaptiveServerSelector, clock,
+ instanceSelector.init(tableConfig, propertyStore, brokerMetrics,
effectiveAdaptiveServerSelector, clock,
config, enabledInstances, enabledServerMap, idealState, externalView,
onlineSegments);
return instanceSelector;
}
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 06dac8d64bf..fd120eed374 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.Map;
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;
@@ -84,7 +83,7 @@ public class MultiStageReplicaGroupSelector extends
BaseInstanceSelector {
}
@Override
- public Pair<Map<String, String>, Map<String, String>> select(List<String>
segments, int requestId,
+ public InstanceMapping select(List<String> segments, int requestId,
SegmentStates segmentStates, Map<String, String> queryOptions) {
// Create a copy of InstancePartitions to avoid race-condition with
event-listeners above.
InstancePartitions instancePartitions = _instancePartitions;
@@ -107,8 +106,9 @@ public class MultiStageReplicaGroupSelector extends
BaseInstanceSelector {
/// 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 [InstanceMapping] whose `segmentToInstanceMap` contains the
segments assigned to a server and whose
+ /// `optionalSegmentToInstanceMap` contains segments that are optional
(i.e., the server is not online to serve that
+ /// segment).
/// Example:
/// {
/// "required_segments": {
@@ -120,7 +120,7 @@ public class MultiStageReplicaGroupSelector extends
BaseInstanceSelector {
/// "segment4": "server4"
/// }
/// }
- private Pair<Map<String, String>, Map<String, String>> assign(Set<String>
segments,
+ private InstanceMapping assign(Set<String> segments,
SegmentStates segmentStates, InstancePartitions instancePartitions, int
preferredReplicaId) {
Map<String, Integer> instanceToPartitionMap =
instancePartitions.getInstanceToPartitionIdMap();
Map<String, Set<String>> instanceToSegmentsMap = new HashMap<>();
@@ -221,7 +221,7 @@ public class MultiStageReplicaGroupSelector extends
BaseInstanceSelector {
/// 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<String, String>, Map<String, String>>
computeOptionalSegments(
+ private InstanceMapping computeOptionalSegments(
Map<String, String> segmentToSelectedInstanceMap, SegmentStates
segmentStates) {
Map<String, String> segmentsToInstanceMap = new HashMap<>();
@@ -250,7 +250,7 @@ public class MultiStageReplicaGroupSelector extends
BaseInstanceSelector {
}
}
- 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 debc04a1912..493e9555f82 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
@@ -19,7 +19,6 @@
package org.apache.pinot.broker.routing.instanceselector;
import java.util.ArrayList;
-import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
@@ -27,7 +26,6 @@ import java.util.Map;
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;
@@ -73,9 +71,13 @@ public class ReplicaGroupInstanceSelector extends
BaseInstanceSelector {
private static final Logger LOGGER =
LoggerFactory.getLogger(ReplicaGroupInstanceSelector.class);
@Override
- public Pair<Map<String, String>, Map<String, String>> select(List<String>
segments, int requestId,
+ public InstanceMapping select(List<String> segments, int requestId,
SegmentStates segmentStates, Map<String, String> queryOptions) {
- ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ return selectWithContext(segments, requestId, segmentStates, new
ServerSelectionContext(queryOptions, _config));
+ }
+
+ protected InstanceMapping selectWithContext(List<String> segments, int
requestId,
+ SegmentStates segmentStates, ServerSelectionContext ctx) {
if (_adaptiveServerSelector != null) {
// Adaptive Server Selection is enabled.
List<SegmentInstanceCandidate> candidateServers =
fetchCandidateServersForQuery(segments, segmentStates);
@@ -94,7 +96,7 @@ public class ReplicaGroupInstanceSelector extends
BaseInstanceSelector {
}
}
- private Pair<Map<String, String>, Map<String, String>>
selectServers(List<String> segments, int requestId,
+ protected InstanceMapping selectServers(List<String> segments, int requestId,
SegmentStates segmentStates, @Nullable Map<String, Integer>
serverRankMap, ServerSelectionContext ctx) {
Map<String, String> segmentToSelectedInstanceMap = new
HashMap<>(HashUtil.getHashMapCapacity(segments.size()));
@@ -115,7 +117,7 @@ public class ReplicaGroupInstanceSelector extends
BaseInstanceSelector {
// Round-robin selection (default behavior)
int numCandidates = candidates.size();
- int instanceIdx = (requestId + replicaOffset) % numCandidates;
+ int instanceIdx = Math.floorMod(requestId + replicaOffset,
numCandidates);
SegmentInstanceCandidate selectedInstance = candidates.get(instanceIdx);
if (useFixedReplica) {
// Adaptive Server Selection cannot be used with fixed replica routing.
@@ -124,13 +126,19 @@ public class ReplicaGroupInstanceSelector extends
BaseInstanceSelector {
} 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
- selectedInstance = candidates.stream()
- .anyMatch(candidate ->
!serverRankMap.containsKey(candidate.getInstance()))
- ? selectedInstance
- : candidates.stream()
- .min(Comparator.comparingInt(candidate ->
serverRankMap.get(candidate.getInstance())))
- .orElse(selectedInstance);
+ // instance. As of 8 July 2026, this fallback is unreachable, but new
implementations could require it.
+ int bestRank = Integer.MAX_VALUE;
+ for (SegmentInstanceCandidate candidate : candidates) {
+ Integer rank = serverRankMap.get(candidate.getInstance());
+ if (rank == null) {
+ selectedInstance = candidates.get(instanceIdx);
+ break;
+ }
+ if (rank < bestRank) {
+ bestRank = rank;
+ selectedInstance = candidate;
+ }
+ }
}
poolToSegmentCount.merge(selectedInstance.getPool(), 1, Integer::sum);
@@ -150,7 +158,7 @@ public class ReplicaGroupInstanceSelector extends
BaseInstanceSelector {
_brokerMetrics.addMeteredValue(BrokerMeter.POOL_SEG_QUERIES,
entry.getValue(),
BrokerMetrics.getTagForPreferredPool(ctx.getQueryOptions()),
String.valueOf(entry.getKey()));
}
- return Pair.of(segmentToSelectedInstanceMap, optionalSegmentToInstanceMap);
+ return new InstanceMapping(segmentToSelectedInstanceMap,
optionalSegmentToInstanceMap);
}
private List<SegmentInstanceCandidate>
fetchCandidateServersForQuery(List<String> segments,
@@ -194,7 +202,6 @@ public class ReplicaGroupInstanceSelector extends
BaseInstanceSelector {
_oldSegmentExpectedReplicasMap.clear();
int newSegmentMapCapacity =
HashUtil.getHashMapCapacity(newSegmentCreationTimeMap.size());
_newSegmentStateMap = new HashMap<>(newSegmentMapCapacity);
-
Map<String, Map<String, String>> idealStateAssignment =
idealState.getRecord().getMapFields();
Map<String, Map<String, String>> externalViewAssignment =
externalView.getRecord().getMapFields();
@@ -247,7 +254,6 @@ public class ReplicaGroupInstanceSelector extends
BaseInstanceSelector {
// NOTE: onlineInstances is either a TreeSet or an EmptySet (sorted)
Set<String> onlineInstances = entry.getValue();
Map<String, String> idealStateInstanceStateMap =
idealStateAssignment.get(segment);
-
Set<String> unavailableInstances =
unavailableInstancesMap.get(idealStateInstanceStateMap.keySet());
List<SegmentInstanceCandidate> candidates = new
ArrayList<>(onlineInstances.size());
int idealStateReplicaId = 0;
@@ -267,7 +273,6 @@ public class ReplicaGroupInstanceSelector extends
BaseInstanceSelector {
Set<String> onlineInstances = entry.getValue();
Map<String, String> idealStateInstanceStateMap =
idealStateAssignment.get(segment);
Map<String, String> sortedIdealStateInstanceStateMap =
convertToSortedMap(idealStateInstanceStateMap);
-
Set<String> unavailableInstances =
unavailableInstancesMap.getOrDefault(idealStateInstanceStateMap.keySet(),
Set.of());
List<SegmentInstanceCandidate> candidates = new
ArrayList<>(idealStateInstanceStateMap.size());
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/SegmentInstanceCandidate.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/SegmentInstanceCandidate.java
index f501ca8dfab..631d58853f7 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/SegmentInstanceCandidate.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/SegmentInstanceCandidate.java
@@ -70,6 +70,6 @@ public class SegmentInstanceCandidate {
@Override
public String toString() {
return "SegmentInstanceCandidate{" + "_instance='" + _instance + '\'' + ",
_online=" + _online + ", _pool=" + _pool
- + '}';
+ + ", _idealStateReplicaId=" + _idealStateReplicaId + '}';
}
}
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 b7c6675033d..1f296a27c52 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,39 +18,35 @@
*/
package org.apache.pinot.broker.routing.instanceselector;
+import java.util.List;
import java.util.Map;
import java.util.Set;
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.utils.config.QueryOptionsUtils;
+
/// Instance selector for strict replica-group routing strategy.
///
-/// ```
-/// The strict replica-group routing strategy always routes the query to the
instances within the same replica-group.
-/// (Note that the replica-group information is derived from the ideal state
of the table, where the instances are
-/// sorted alphabetically in the instance state map, so the replica-groups in
the instance selector might not match the
-/// replica-groups in the instance partitions). The goal of this algorithm is
to ensure that segments from the same
-/// partition are never served from multiple different instances. The
instances in a replica-group should have all the
-/// online segments (segments with ONLINE/CONSUMING instances in the ideal
state and selected by the pre-selector)
-/// available (ONLINE/CONSUMING in the external view) in order to serve
queries. If any segment is unavailable in the
-/// replica-group, we mark the whole replica-group down and not serve queries
with this replica-group.
+/// The strict replica-group routing strategy always routes same-partition
segments to the same instance. During
+/// routing state construction, [#updateSegmentMapsForUpsertTable(IdealState,
ExternalView, Set, Map)] removes from
+/// every segment in a partition any replica that is unavailable for any old
segment in that partition. Consequently,
+/// all same-partition segments have identical, ordered candidate identities.
///
-/// The selection algorithm is the same as {@link
ReplicaGroupInstanceSelector}, and will always evenly distribute the
-/// traffic to all replica-groups that have all online segments available.
+/// Adaptive routing relies on two parent-selector contracts:
+/// - [#selectWithContext(List, int, SegmentStates, ServerSelectionContext)]
obtains one ranking snapshot before
+/// iterating the query's segments; it must not re-rank per segment.
+/// - [#selectServers(List, int, SegmentStates, Map, ServerSelectionContext)]
resolves equal-ranked candidates using
+/// its strictly-less-than (`rank < bestRank`) comparison. It consequently
retains the first candidate in the shared
+/// ordered candidate list.
///
-/// The algorithm relies on the mirror segment assignment from replica-group
segment assignment strategy. With mirror
-/// segment assignment, any server in one replica-group will always have a
corresponding server in other replica-groups
-/// that have the same segments assigned. For example, if S1 is a server in
replica-group 1, and it has mirror server S2
-/// in replica-group 2 and S3 in replica-group 3. All segments assigned to S1
will also be assigned to S2 and S3. In
-/// stable scenario (external view matches ideal state), all segments assigned
to S1 will have the same enabled
-/// instances of [S1, S2, S3] sorted (in alphabetical order). If we always
pick the same index of enabled instances for
-/// all segments, only one of S1, S2, S3 will be picked, and all the segments
are processed by the same server. In
-/// transitioning/error scenario (external view does not match ideal state),
if a segment is down on S1, we mark all
-/// segments with the same assignment ([S1, S2, S3]) down on S1 to ensure that
we always route the segments to the same
-/// replica-group.
+/// Since same-partition old segments have identical ordered candidates, these
contracts ensure they select the same
+/// instance. Changing either contract requires revisiting strict
replica-group routing. Different partitions may
+/// independently choose different replicas.
///
-/// Note that new segments won't be used to exclude instances from serving
when the segment is unavailable.
-/// ```
+/// New segments do not exclude a candidate when that segment is unavailable;
they remain optional so that the broker or
+/// server can skip them if necessary.
public class StrictReplicaGroupInstanceSelector extends
ReplicaGroupInstanceSelector {
@Override
@@ -58,4 +54,24 @@ public class StrictReplicaGroupInstanceSelector extends
ReplicaGroupInstanceSele
Map<String, Long> newSegmentCreationTimeMap) {
super.updateSegmentMapsForUpsertTable(idealState, externalView,
onlineSegments, newSegmentCreationTimeMap);
}
+
+ @Override
+ public InstanceMapping select(List<String> segments, int requestId,
+ SegmentStates segmentStates, Map<String, String> queryOptions) {
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ if (_adaptiveServerSelector != null && _priorityPoolInstanceSelector !=
null) {
+ if (ctx.isUseFixedReplica()) {
+ // Fixed-replica routing configured on the broker or table disables
adaptive routing in InstanceSelectorFactory.
+ // This fixed replica config can still be set per query, in which case
we'll prefer that over adaptive routing.
+ return selectServers(segments, requestId, segmentStates, null, ctx);
+ }
+ if (QueryOptionsUtils.getNumReplicaGroupsToQuery(ctx.getQueryOptions())
!= null) {
+ // This option intentionally fans segments across replica groups, so
preserve the non-adaptive behavior.
+ // This contradicts the behaviour specified by StrictReplicaGroup, so
we should eventually reject / ignore
+ // this combination.
+ return selectServers(segments, requestId, segmentStates, null, ctx);
+ }
+ }
+ return selectWithContext(segments, requestId, segmentStates, ctx);
+ }
}
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 269243cb5b6..dfde3e9c950 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
@@ -45,7 +45,6 @@ import javax.annotation.Nullable;
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;
@@ -815,6 +814,8 @@ public abstract class BaseBrokerRoutingManager implements
RoutingManager, Cluste
AdaptiveServerSelector adaptiveServerSelector =
AdaptiveServerSelectorFactory.getAdaptiveServerSelector(_serverRoutingStatsManager,
_pinotConfig);
+ // 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,
@@ -1551,8 +1552,7 @@ public abstract class BaseBrokerRoutingManager implements
RoutingManager, Cluste
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/InstanceSelectorTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorTest.java
index 796ee992b1d..a7fe5b32e10 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
@@ -78,6 +78,7 @@ 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.assertNull;
import static org.testng.Assert.assertTrue;
@@ -252,6 +253,25 @@ public class InstanceSelectorTest {
Set.of()) instanceof ReplicaGroupInstanceSelector);
}
+ @Test
+ public void
testFactoryDisablesAdaptiveRoutingForLegacyUpsertReplicaGroupSelector() {
+ TableConfig tableConfig = mock(TableConfig.class);
+ RoutingConfig routingConfig = mock(RoutingConfig.class);
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ when(tableConfig.getTableName()).thenReturn("legacyUpsert_REALTIME");
+ when(tableConfig.getRoutingConfig()).thenReturn(routingConfig);
+ when(tableConfig.isUpsertEnabled()).thenReturn(true);
+
when(routingConfig.getInstanceSelectorType()).thenReturn(REPLICA_GROUP_INSTANCE_SELECTOR_TYPE);
+
+ ReplicaGroupInstanceSelector instanceSelector =
+ (ReplicaGroupInstanceSelector)
InstanceSelectorFactory.getInstanceSelector(tableConfig, _propertyStore,
+ _brokerMetrics, hybridSelector, new PinotConfiguration(),
Set.of(), Map.of(),
+ new IdealState("legacyUpsert_REALTIME"), new
ExternalView("legacyUpsert_REALTIME"), Set.of());
+
+ assertNull(instanceSelector._adaptiveServerSelector);
+ assertNull(instanceSelector._priorityPoolInstanceSelector);
+ }
+
@Test
public void testInstanceSelector() {
String offlineTableName = "testTable_OFFLINE";
@@ -1863,76 +1883,6 @@ public class InstanceSelectorTest {
assertTrue(selectionResult.getUnavailableSegments().isEmpty());
}
- @Test
- public void testReplicaGroupAdaptiveServerSelector() {
- // Arrange
- ZkHelixPropertyStore<ZNRecord> propertyStore =
mock(ZkHelixPropertyStore.class);
- BrokerMetrics brokerMetrics = mock(BrokerMetrics.class);
- HybridSelector hybridSelector = mock(HybridSelector.class);
- ReplicaGroupInstanceSelector instanceSelector = new
ReplicaGroupInstanceSelector();
-
- // 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<String> segments = Arrays.asList(segment0, segment1, segment2);
-
- // Define candidates for each segment
- Map<String, List<SegmentInstanceCandidate>> 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<Pair<String, Double>> 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<String, String>, Map<String, String>> selectedResult =
- instanceSelector.select(segments, 0, segmentStates, null);
-
- // Assert
- Map<String, String> expectedSelection = new HashMap<>();
- expectedSelection.put(segment0, instance1);
- expectedSelection.put(segment1, instance3);
- expectedSelection.put(segment2, instance4);
-
- assertEquals(selectedResult.getLeft(), expectedSelection);
- }
-
// Replica health metrics
//
// The scenarios below all use the same three instances and assert on the
TableReplicaHealth the
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 00000000000..1442ef17672
--- /dev/null
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java
@@ -0,0 +1,1063 @@
+/**
+ * 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.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.BrokerMeter;
+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.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;
+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.Broker.FALLBACK_POOL_ID;
+import static
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel.OFFLINE;
+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;
+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 [ReplicaGroupInstanceSelector] and
[StrictReplicaGroupInstanceSelector], including adaptive server
+/// selection 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<String, ServerInstance> EMPTY_SERVER_MAP = Map.of();
+ private static final InstanceSelectorConfig INSTANCE_SELECTOR_CONFIG = new
InstanceSelectorConfig(false, 300, false);
+ private static final List<String> 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<String, List<Pair<String, String>>>
segmentState) {
+ IdealState idealState = new IdealState(TABLE_NAME);
+ Map<String, Map<String, String>> idealStateSegmentAssignment =
idealState.getRecord().getMapFields();
+ for (Map.Entry<String, List<Pair<String, String>>> entry :
segmentState.entrySet()) {
+ Map<String, String> instanceStateMap = new TreeMap<>();
+ for (Pair<String, String> instanceState : entry.getValue()) {
+ instanceStateMap.put(instanceState.getLeft(),
instanceState.getRight());
+ }
+ idealStateSegmentAssignment.put(entry.getKey(), instanceStateMap);
+ }
+ return idealState;
+ }
+
+ static ExternalView createExternalView(Map<String, List<Pair<String,
String>>> segmentState) {
+ ExternalView externalView = new ExternalView(TABLE_NAME);
+ Map<String, Map<String, String>> externalViewSegmentAssignment =
externalView.getRecord().getMapFields();
+ for (Map.Entry<String, List<Pair<String, String>>> entry :
segmentState.entrySet()) {
+ Map<String, String> instanceStateMap = new TreeMap<>();
+ for (Pair<String, String> 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<ZNRecord> propertyStore =
mock(ZkHelixPropertyStore.class);
+ BrokerMetrics brokerMetrics = mock(BrokerMetrics.class);
+ BrokerRequest brokerRequest = mock(BrokerRequest.class);
+ PinotQuery pinotQuery = mock(PinotQuery.class);
+ Map<String, String> 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<String> enabledInstances = new HashSet<>();
+ IdealState idealState = new IdealState(offlineTableName);
+ Map<String, Map<String, String>> idealStateSegmentAssignment =
idealState.getRecord().getMapFields();
+ ExternalView externalView = new ExternalView(offlineTableName);
+ Map<String, Map<String, String>> externalViewSegmentAssignment =
externalView.getRecord().getMapFields();
+ Set<String> 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<String, String> idealStateInstanceStateMap0 = new TreeMap<>();
+ Map<String, String> externalViewInstanceStateMap0 = new TreeMap<>();
+
+ for (String instance : enabledInstances) {
+ idealStateInstanceStateMap0.put(instance, ONLINE);
+ externalViewInstanceStateMap0.put(instance, ONLINE);
+ }
+
+ List<String> 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<String, String> 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<ZNRecord> propertyStore =
mock(ZkHelixPropertyStore.class);
+ BrokerMetrics brokerMetrics = mock(BrokerMetrics.class);
+ BrokerRequest brokerRequest = mock(BrokerRequest.class);
+ PinotQuery pinotQuery = mock(PinotQuery.class);
+ Map<String, String> queryOptions = new HashMap<>();
+ queryOptions.put("numReplicaGroupsToQuery", "4");
+
+ when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery);
+ when(pinotQuery.getQueryOptions()).thenReturn(queryOptions);
+
+ ReplicaGroupInstanceSelector replicaGroupInstanceSelector = new
ReplicaGroupInstanceSelector();
+
+ Set<String> enabledInstances = new HashSet<>();
+ IdealState idealState = new IdealState(offlineTableName);
+ Map<String, Map<String, String>> idealStateSegmentAssignment =
idealState.getRecord().getMapFields();
+ ExternalView externalView = new ExternalView(offlineTableName);
+ Map<String, Map<String, String>> externalViewSegmentAssignment =
externalView.getRecord().getMapFields();
+ Set<String> 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<String> segments = SEGMENTS;
+
+ Map<String, String> idealStateInstanceStateMap0 = new TreeMap<>();
+ Map<String, String> 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<String, String> 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<ZNRecord> propertyStore =
mock(ZkHelixPropertyStore.class);
+ BrokerMetrics brokerMetrics = mock(BrokerMetrics.class);
+ BrokerRequest brokerRequest = mock(BrokerRequest.class);
+ PinotQuery pinotQuery = mock(PinotQuery.class);
+ Map<String, String> queryOptions = new HashMap<>();
+
+ when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery);
+ when(pinotQuery.getQueryOptions()).thenReturn(queryOptions);
+
+ ReplicaGroupInstanceSelector replicaGroupInstanceSelector = new
ReplicaGroupInstanceSelector();
+
+ Set<String> enabledInstances = new HashSet<>();
+ IdealState idealState = new IdealState(offlineTableName);
+ Map<String, Map<String, String>> idealStateSegmentAssignment =
idealState.getRecord().getMapFields();
+ ExternalView externalView = new ExternalView(offlineTableName);
+ Map<String, Map<String, String>> externalViewSegmentAssignment =
externalView.getRecord().getMapFields();
+ Set<String> 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<String> segments = SEGMENTS;
+
+ Map<String, String> idealStateInstanceStateMap0 = new TreeMap<>();
+ Map<String, String> 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<String, String> 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 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<String> AR_SEGMENTS =
+ Arrays.asList(AR_SEGMENT0, AR_SEGMENT1, AR_SEGMENT2);
+ // Rankings: D best → C → B → A worst; E absent (triggers AR fallback)
+ private static final List<Pair<String, Double>> AR_SERVER_RANKS =
Arrays.asList(
+ 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<String, List<SegmentInstanceCandidate>> candidatesMap = new
HashMap<>();
+ // segment0 → pool 0 / replica group 0 server A, pool 1 / replica group 1
server B
+ candidatesMap.put(AR_SEGMENT0, Arrays.asList(
+ 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_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_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);
+ }
+
+ private ReplicaGroupInstanceSelector buildArSelector(String selectorType,
HybridSelector hybridSelector) {
+ return buildArSelector(selectorType, hybridSelector, new
PinotConfiguration(Map.of()));
+ }
+
+ 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();
+ IdealState idealState = createIdealState(Map.of(
+ 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_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<String, ServerInstance> 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_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
+ 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);
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null);
+
+ // AR prefers the better-ranked server when all candidates are ranked:
+ // 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_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 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_STRICT_REPLICA_GROUP_ENABLED,
"false"));
+ StrictReplicaGroupInstanceSelector instanceSelector =
+ buildStrictReplicaGroupArSelector(hybridSelector, brokerMetrics,
brokerConfig);
+
+ assertNull(instanceSelector._adaptiveServerSelector);
+ assertNull(instanceSelector._priorityPoolInstanceSelector);
+
+ BaseInstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), null);
+
+ 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));
+
+ result = instanceSelector.select(STRICT_SEGMENTS, 1,
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 testStrictReplicaGroupFactoryEnablesAdaptiveRouting() {
+ 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);
+ }
+
+ @Test
+ public void
testStrictReplicaGroupAdaptiveDisabledByFixedReplicaBrokerConfig() {
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ PinotConfiguration brokerConfig = new PinotConfiguration(Map.of(
+ CommonConstants.Broker.CONFIG_OF_USE_FIXED_REPLICA, "true"));
+ StrictReplicaGroupInstanceSelector instanceSelector =
+ buildStrictReplicaGroupArSelector(hybridSelector,
mock(BrokerMetrics.class), brokerConfig);
+
+ assertNull(instanceSelector._adaptiveServerSelector);
+ assertNull(instanceSelector._priorityPoolInstanceSelector);
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), null);
+ verify(hybridSelector, never()).fetchServerRankingsWithScores(any());
+ assertEquals(result.segmentToInstanceMap().get(STRICT_SEGMENT0),
+ result.segmentToInstanceMap().get(STRICT_SEGMENT1));
+ }
+
+ @Test
+ public void
testStrictReplicaGroupAdaptiveDisabledByFixedReplicaTableConfig() {
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector,
+ mock(BrokerMetrics.class), new PinotConfiguration(Map.of()), true);
+
+ assertNull(instanceSelector._adaptiveServerSelector);
+ assertNull(instanceSelector._priorityPoolInstanceSelector);
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), null);
+ verify(hybridSelector, never()).fetchServerRankingsWithScores(any());
+ assertEquals(result.segmentToInstanceMap().get(STRICT_SEGMENT0),
+ result.segmentToInstanceMap().get(STRICT_SEGMENT1));
+ }
+
+ // --- 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<String> STRICT_SEGMENTS =
Arrays.asList(STRICT_SEGMENT0, STRICT_SEGMENT1, STRICT_SEGMENT2);
+
+ private List<SegmentInstanceCandidate>
buildStrictReplicaGroupCandidates(String replicaGroup0Instance,
+ String replicaGroup1Instance) {
+ return buildStrictReplicaGroupCandidates(replicaGroup0Instance,
replicaGroup1Instance, FALLBACK_POOL_ID,
+ FALLBACK_POOL_ID);
+ }
+
+ private List<SegmentInstanceCandidate>
buildStrictReplicaGroupCandidates(String replicaGroup0Instance,
+ String replicaGroup1Instance, int replicaGroup0Pool, int
replicaGroup1Pool) {
+ return Arrays.asList(
+ new SegmentInstanceCandidate(replicaGroup0Instance, true,
replicaGroup0Pool, 0),
+ new SegmentInstanceCandidate(replicaGroup1Instance, true,
replicaGroup1Pool, 1));
+ }
+
+ private SegmentStates buildStrictReplicaGroupSegmentStates() {
+ return buildStrictReplicaGroupSegmentStates(FALLBACK_POOL_ID,
FALLBACK_POOL_ID);
+ }
+
+ private SegmentStates buildStrictReplicaGroupSegmentStates(int
replicaGroup0Pool, int replicaGroup1Pool) {
+ Map<String, List<SegmentInstanceCandidate>> candidatesMap = new
HashMap<>();
+ // 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, replicaGroup0Pool,
+ replicaGroup1Pool));
+ candidatesMap.put(STRICT_SEGMENT1,
+ buildStrictReplicaGroupCandidates(STRICT_RG0_SERVER_A,
STRICT_RG1_SERVER_C, replicaGroup0Pool,
+ replicaGroup1Pool));
+ candidatesMap.put(STRICT_SEGMENT2,
+ buildStrictReplicaGroupCandidates(STRICT_RG0_SERVER_B,
STRICT_RG1_SERVER_D, replicaGroup0Pool,
+ replicaGroup1Pool));
+ return new SegmentStates(candidatesMap, new HashSet<>(STRICT_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(Map.of()));
+ }
+
+ private StrictReplicaGroupInstanceSelector
buildStrictReplicaGroupArSelector(HybridSelector hybridSelector,
+ BrokerMetrics brokerMetrics, PinotConfiguration brokerConfig) {
+ return buildStrictReplicaGroupArSelector(hybridSelector, brokerMetrics,
brokerConfig, null);
+ }
+
+ private StrictReplicaGroupInstanceSelector
buildStrictReplicaGroupArSelector(HybridSelector hybridSelector,
+ BrokerMetrics brokerMetrics, PinotConfiguration brokerConfig, Boolean
useFixedReplica) {
+ RoutingConfig routingConfig = new RoutingConfig(null, null,
STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE,
+ useFixedReplica);
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.REALTIME).setTableName("testUpsertTable")
+ .setRoutingConfig(routingConfig).build();
+ // Ideal state: mirrors the topology above
+ IdealState idealState = createIdealState(Map.of(
+ 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(
+ 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(FALLBACK_POOL_ID);
+ ServerInstance serverB = mock(ServerInstance.class);
+ when(serverB.getPool()).thenReturn(FALLBACK_POOL_ID);
+ ServerInstance serverC = mock(ServerInstance.class);
+ when(serverC.getPool()).thenReturn(FALLBACK_POOL_ID);
+ ServerInstance serverD = mock(ServerInstance.class);
+ when(serverD.getPool()).thenReturn(FALLBACK_POOL_ID);
+ Map<String, ServerInstance> serverMap = Map.of(
+ 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(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
testStrictReplicaGroupAdaptivePicksHealthiestCandidatePerPartition() {
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ BrokerMetrics brokerMetrics = mock(BrokerMetrics.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector,
+ brokerMetrics);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList(
+ 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(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), null);
+
+ // Each partition independently prefers its healthier replica. Segments 0
and 1 share a candidate list, so they
+ // make the same choice.
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT1, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT2, STRICT_RG1_SERVER_D));
+ // Metrics are reported per selected pool.
+ verify(brokerMetrics).addMeteredValue(eq(BrokerMeter.POOL_SEG_QUERIES),
eq(3L),
+ eq(BrokerMetrics.getTagForPreferredPool(null)),
eq(String.valueOf(FALLBACK_POOL_ID)));
+ }
+
+ @Test
+ public void testStrictReplicaGroupAdaptivePicksOtherHealthyCandidates() {
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList(
+ 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(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), null);
+
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT1, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT2, STRICT_RG0_SERVER_B));
+ }
+
+ @Test
+ public void testStrictReplicaGroupAdaptivePrefersPoolBeforeAdaptiveHealth() {
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList(
+ new ImmutablePair<>(STRICT_RG1_SERVER_C, 1.0),
+ new ImmutablePair<>(STRICT_RG1_SERVER_D, 2.0),
+ new ImmutablePair<>(STRICT_RG0_SERVER_A, 90.0),
+ new ImmutablePair<>(STRICT_RG0_SERVER_B, 100.0)
+ ));
+
+ Map<String, String> queryOptions = Map.of(
+ CommonConstants.Broker.Request.QueryOptionKey.ORDERED_PREFERRED_POOLS,
"0");
+ InstanceSelector.InstanceMapping result =
instanceSelector.select(STRICT_SEGMENTS, 0,
+ buildStrictReplicaGroupSegmentStates(0, 1), queryOptions);
+
+ // Preferred pools are the primary selection criterion, even when another
pool is healthier.
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT1, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT2, STRICT_RG0_SERVER_B));
+ }
+
+ @Test
+ public void testStrictReplicaGroupAdaptiveUsesHealthWithinPreferredPool() {
+ String nonPreferredServerE = "strict_rg2_server_e";
+ String nonPreferredServerF = "strict_rg2_server_f";
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList(
+ new ImmutablePair<>(nonPreferredServerE, 0.0),
+ new ImmutablePair<>(nonPreferredServerF, 0.0),
+ new ImmutablePair<>(STRICT_RG1_SERVER_C, 1.0),
+ new ImmutablePair<>(STRICT_RG1_SERVER_D, 2.0),
+ new ImmutablePair<>(STRICT_RG0_SERVER_A, 20.0),
+ new ImmutablePair<>(STRICT_RG0_SERVER_B, 30.0)
+ ));
+
+ Map<String, List<SegmentInstanceCandidate>> candidatesMap = new
HashMap<>();
+ List<SegmentInstanceCandidate> partition0Candidates = Arrays.asList(
+ new SegmentInstanceCandidate(STRICT_RG0_SERVER_A, true, 0, 0),
+ new SegmentInstanceCandidate(STRICT_RG1_SERVER_C, true, 0, 1),
+ new SegmentInstanceCandidate(nonPreferredServerE, true, 1, 2));
+ candidatesMap.put(STRICT_SEGMENT0, partition0Candidates);
+ candidatesMap.put(STRICT_SEGMENT1, partition0Candidates);
+ candidatesMap.put(STRICT_SEGMENT2, Arrays.asList(
+ new SegmentInstanceCandidate(STRICT_RG0_SERVER_B, true, 0, 0),
+ new SegmentInstanceCandidate(STRICT_RG1_SERVER_D, true, 0, 1),
+ new SegmentInstanceCandidate(nonPreferredServerF, true, 1, 2)));
+ SegmentStates segmentStates = new SegmentStates(candidatesMap, new
HashSet<>(STRICT_SEGMENTS), null);
+
+ Map<String, String> queryOptions = Map.of(
+ CommonConstants.Broker.Request.QueryOptionKey.ORDERED_PREFERRED_POOLS,
"0");
+ InstanceSelector.InstanceMapping result =
instanceSelector.select(STRICT_SEGMENTS, 1, segmentStates,
+ queryOptions);
+
+ // The healthiest candidate within pool 0 wins, while the healthier
candidates in pool 1 are ignored.
+ 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
testStrictReplicaGroupAdaptiveEmptyRankingsFallBackToRoundRobin() {
+ // When adaptive ranking returns no servers, fall back to round-robin by
candidate index.
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(List.of());
+
+ // requestId=0 picks candidate index 0.
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), null);
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT1, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT2, STRICT_RG0_SERVER_B));
+
+ // requestId=1 picks candidate index 1.
+ result = instanceSelector.select(STRICT_SEGMENTS, 1,
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 testStrictReplicaGroupAdaptiveSupportsFixedReplicaQueryOption() {
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
+ Map<String, String> queryOptions = Map.of(
+ CommonConstants.Broker.Request.QueryOptionKey.USE_FIXED_REPLICA,
"true");
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), queryOptions);
+ verify(hybridSelector, never()).fetchServerRankingsWithScores(any());
+ assertEquals(result.segmentToInstanceMap().get(STRICT_SEGMENT0),
+ result.segmentToInstanceMap().get(STRICT_SEGMENT1));
+ }
+
+ @Test
+ public void testStrictReplicaGroupAdaptivePreservesNumReplicaGroupsToQuery()
{
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
+ Map<String, String> queryOptions = Map.of(
+
CommonConstants.Broker.Request.QueryOptionKey.NUM_REPLICA_GROUPS_TO_QUERY, "2");
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), queryOptions);
+
+ verify(hybridSelector, never()).fetchServerRankingsWithScores(any());
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT1, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT2, STRICT_RG0_SERVER_B));
+ }
+
+ @Test
+ public void testStrictReplicaGroupAdaptiveIsConsistentWithinEachPartition() {
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList(
+ new ImmutablePair<>(STRICT_RG1_SERVER_C, 1.0),
+ new ImmutablePair<>(STRICT_RG0_SERVER_B, 2.0),
+ new ImmutablePair<>(STRICT_RG0_SERVER_A, 3.0),
+ new ImmutablePair<>(STRICT_RG1_SERVER_D, 4.0)
+ ));
+
+ InstanceSelector.InstanceMapping result =
+ 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_RG0_SERVER_B));
+ }
+
+ @Test
+ public void
testStrictReplicaGroupAdaptiveUsesFilteredCandidatesPerPartition() {
+ String partition0Segment0 = "partition0_segment0";
+ String partition0Segment1 = "partition0_segment1";
+ String partition1Segment0 = "partition1_segment0";
+ String partition1Segment1 = "partition1_segment1";
+ List<String> segments = List.of(partition0Segment0, partition0Segment1,
partition1Segment0,
+ partition1Segment1);
+ String partition0Replica0 = "partition0_replica0";
+ String partition0Replica1 = "partition0_replica1";
+ String partition1Replica0 = "partition1_replica0";
+ String partition1Replica1 = "partition1_replica1";
+ 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("partitionedTable")
+ .setRoutingConfig(routingConfig).build();
+ IdealState idealState = createIdealState(Map.of(
+ partition0Segment0, List.of(Pair.of(partition0Replica0, ONLINE),
Pair.of(partition0Replica1, ONLINE)),
+ partition0Segment1, List.of(Pair.of(partition0Replica0, ONLINE),
Pair.of(partition0Replica1, ONLINE)),
+ partition1Segment0, List.of(Pair.of(partition1Replica0, ONLINE),
Pair.of(partition1Replica1, ONLINE)),
+ partition1Segment1, List.of(Pair.of(partition1Replica0, ONLINE),
Pair.of(partition1Replica1, ONLINE))));
+ // Degrade opposite replica positions in the two partitions. Strict state
construction must remove each degraded
+ // replica from every old segment with the same IdealState assignment.
+ ExternalView externalView = createExternalView(Map.of(
+ partition0Segment0, List.of(Pair.of(partition0Replica0, ONLINE),
Pair.of(partition0Replica1, ONLINE)),
+ partition0Segment1, List.of(Pair.of(partition0Replica0, OFFLINE),
Pair.of(partition0Replica1, ONLINE)),
+ partition1Segment0, List.of(Pair.of(partition1Replica0, ONLINE),
Pair.of(partition1Replica1, ONLINE)),
+ partition1Segment1, List.of(Pair.of(partition1Replica0, ONLINE),
Pair.of(partition1Replica1, OFFLINE))));
+
+ ServerInstance server0 = mock(ServerInstance.class);
+ when(server0.getPool()).thenReturn(FALLBACK_POOL_ID);
+ ServerInstance server1 = mock(ServerInstance.class);
+ when(server1.getPool()).thenReturn(FALLBACK_POOL_ID);
+ ServerInstance server2 = mock(ServerInstance.class);
+ when(server2.getPool()).thenReturn(FALLBACK_POOL_ID);
+ ServerInstance server3 = mock(ServerInstance.class);
+ when(server3.getPool()).thenReturn(FALLBACK_POOL_ID);
+ Map<String, ServerInstance> serverMap = Map.of(partition0Replica0,
server0, partition0Replica1, server1,
+ partition1Replica0, server2, partition1Replica1, server3);
+ StrictReplicaGroupInstanceSelector instanceSelector =
+ (StrictReplicaGroupInstanceSelector)
InstanceSelectorFactory.getInstanceSelector(tableConfig,
+ mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class),
hybridSelector,
+ new PinotConfiguration(Map.of()),
+ Set.of(partition0Replica0, partition0Replica1, partition1Replica0,
partition1Replica1), serverMap,
+ idealState, externalView, new HashSet<>(segments));
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(List.of(
+ new ImmutablePair<>(partition1Replica0, 1.0), new
ImmutablePair<>(partition0Replica1, 2.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(
+ partition0Segment0, partition0Replica1,
+ partition0Segment1, partition0Replica1,
+ partition1Segment0, partition1Replica0,
+ partition1Segment1, partition1Replica0));
+ assertTrue(result.getUnavailableSegments().isEmpty());
+ }
+
+ @Test
+ public void testStrictReplicaGroupAdaptiveNewSegmentNotFalselyUnavailable() {
+ // 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";
+ String seg1 = "new_seg1";
+ List<String> 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(FALLBACK_POOL_ID);
+ StrictReplicaGroupInstanceSelector instanceSelector =
+ (StrictReplicaGroupInstanceSelector)
InstanceSelectorFactory.getInstanceSelector(tableConfig,
+ mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class),
hybridSelector,
+ new PinotConfiguration(Map.of()),
+ 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
testStrictReplicaGroupAdaptiveNewSegmentWithCandidateRemainsOptional() {
+ String newSegment = "new_seg_with_candidate";
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList(
+ new ImmutablePair<>(STRICT_RG0_SERVER_A, 1.0),
+ new ImmutablePair<>(STRICT_RG1_SERVER_D, 2.0)
+ ));
+
+ Map<String, List<SegmentInstanceCandidate>> candidatesMap = new
HashMap<>();
+ candidatesMap.put(STRICT_SEGMENT0, List.of(
+ new SegmentInstanceCandidate(STRICT_RG0_SERVER_A, true,
FALLBACK_POOL_ID, 0)));
+ candidatesMap.put(newSegment, List.of(
+ new SegmentInstanceCandidate(STRICT_RG1_SERVER_D, false,
FALLBACK_POOL_ID, 1)));
+ SegmentStates segmentStates = new SegmentStates(candidatesMap,
+ Set.of(STRICT_RG0_SERVER_A, STRICT_RG1_SERVER_D), Set.of());
+
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(List.of(STRICT_SEGMENT0, newSegment), 0,
segmentStates, Map.of());
+
+ assertEquals(result.segmentToInstanceMap(), Map.of(STRICT_SEGMENT0,
STRICT_RG0_SERVER_A));
+ assertEquals(result.optionalSegmentToInstanceMap(), Map.of(newSegment,
STRICT_RG1_SERVER_D));
+ }
+
+ @Test
+ public void testStrictReplicaGroupAdaptiveOnlyQueryRelevantServersScored() {
+ // The selector 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.
+ String rg1ServerE = "server_e";
+ String segExtra = "seg_extra";
+ 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();
+ // 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(
+ 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(
+ 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(FALLBACK_POOL_ID);
+ ServerInstance serverB = mock(ServerInstance.class);
+ when(serverB.getPool()).thenReturn(FALLBACK_POOL_ID);
+ ServerInstance serverC = mock(ServerInstance.class);
+ when(serverC.getPool()).thenReturn(FALLBACK_POOL_ID);
+ ServerInstance serverD = mock(ServerInstance.class);
+ when(serverD.getPool()).thenReturn(FALLBACK_POOL_ID);
+ ServerInstance serverEInstance = mock(ServerInstance.class);
+ when(serverEInstance.getPool()).thenReturn(99);
+ Map<String, ServerInstance> serverMap = Map.of(
+ STRICT_RG0_SERVER_A, serverA, STRICT_RG0_SERVER_B, serverB,
+ STRICT_RG1_SERVER_C, serverC, STRICT_RG1_SERVER_D, serverD,
+ rg1ServerE, serverEInstance);
+ Set<String> 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(STRICT_RG0_SERVER_A, STRICT_RG0_SERVER_B,
+ STRICT_RG1_SERVER_C, STRICT_RG1_SERVER_D, rg1ServerE),
+ serverMap, idealState, externalView, allSegments);
+
+ // 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<>(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)
+ ));
+
+ Map<String, List<SegmentInstanceCandidate>> 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(STRICT_SEGMENTS, 0, segmentStates, null);
+
+ // Each query partition picks its healthiest candidate; server_e is not
considered because it hosts no query
+ // segments.
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT1, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT2, STRICT_RG1_SERVER_D));
+
+ ArgumentCaptor<List<String>> rankedCandidatesCaptor =
ArgumentCaptor.forClass(List.class);
+
verify(hybridSelector).fetchServerRankingsWithScores(rankedCandidatesCaptor.capture());
+ assertEquals(new HashSet<>(rankedCandidatesCaptor.getValue()),
+ Set.of(STRICT_RG0_SERVER_A, STRICT_RG0_SERVER_B,
+ STRICT_RG1_SERVER_C, STRICT_RG1_SERVER_D));
+ assertEquals(rankedCandidatesCaptor.getValue().size(), 4);
+ }
+
+ @Test
+ public void testStrictReplicaGroupAdaptiveHandlesMinimumRequestId() {
+ String replicaGroup2ServerE = "strict_rg2_server_e";
+ String replicaGroup2ServerF = "strict_rg2_server_f";
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(List.of());
+
+ Map<String, List<SegmentInstanceCandidate>> candidatesMap = new
HashMap<>();
+ List<SegmentInstanceCandidate> partition0Candidates = Arrays.asList(
+ new SegmentInstanceCandidate(STRICT_RG0_SERVER_A, true,
FALLBACK_POOL_ID, 0),
+ new SegmentInstanceCandidate(STRICT_RG1_SERVER_C, true,
FALLBACK_POOL_ID, 1),
+ new SegmentInstanceCandidate(replicaGroup2ServerE, true,
FALLBACK_POOL_ID, 2));
+ candidatesMap.put(STRICT_SEGMENT0, partition0Candidates);
+ candidatesMap.put(STRICT_SEGMENT1, partition0Candidates);
+ candidatesMap.put(STRICT_SEGMENT2, Arrays.asList(
+ new SegmentInstanceCandidate(STRICT_RG0_SERVER_B, true,
FALLBACK_POOL_ID, 0),
+ new SegmentInstanceCandidate(STRICT_RG1_SERVER_D, true,
FALLBACK_POOL_ID, 1),
+ new SegmentInstanceCandidate(replicaGroup2ServerF, true,
FALLBACK_POOL_ID, 2)));
+ SegmentStates segmentStates = new SegmentStates(candidatesMap, new
HashSet<>(STRICT_SEGMENTS), null);
+
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, Integer.MIN_VALUE,
segmentStates, null);
+
+ // floorMod(Integer.MIN_VALUE, 3) is 1; the previous '%' indexing returned
-2 and threw.
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT1, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT2, STRICT_RG1_SERVER_D));
+ }
+}
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 42f66b844c0..71298ee7cb3 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
@@ -1254,6 +1254,13 @@ public class CommonConstants {
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 replica-group-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_STRICT_REPLICA_GROUP_ENABLED =
+ CONFIG_PREFIX + ".strict.replica.group.enabled";
+ public static final boolean DEFAULT_STRICT_REPLICA_GROUP_ENABLED = true;
}
public static class Grpc {
diff --git a/pom.xml b/pom.xml
index c2cf6b1acfc..682494ff0d9 100644
--- a/pom.xml
+++ b/pom.xml
@@ -2103,7 +2103,7 @@
<plugin>
<groupId>com.diffplug.spotless</groupId>
<artifactId>spotless-maven-plugin</artifactId>
- <version>2.46.1</version>
+ <version>3.0.0</version>
<executions>
<execution>
<goals>
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
