|
18 | 18 | */ |
19 | 19 | package org.apache.pinot.broker.routing.instanceselector; |
20 | 20 |
|
| 21 | +import java.time.Clock; |
| 22 | +import java.util.ArrayList; |
21 | 23 | import java.util.HashMap; |
22 | 24 | import java.util.List; |
23 | 25 | import java.util.Map; |
24 | 26 | import javax.annotation.Nullable; |
| 27 | +import org.apache.helix.store.zk.ZkHelixPropertyStore; |
| 28 | +import org.apache.helix.zookeeper.datamodel.ZNRecord; |
25 | 29 | import org.apache.pinot.broker.routing.adaptiveserverselector.AdaptiveServerSelector; |
26 | 30 | import org.apache.pinot.common.metrics.BrokerMetrics; |
27 | 31 | import org.apache.pinot.common.utils.HashUtil; |
|
37 | 41 | * Step1: Process seg1. Fetch server rankings. Pick the best server. |
38 | 42 | * Step2: Process seg2. Fetch server rankings (could have changed or not since Step 1). Pick the best server. |
39 | 43 | * Step3: Process seg3. Fetch server rankings (could have changed or not since Step 2). Pick the best server. |
40 | | -
|
| 44 | + * |
41 | 45 | * <p>If AdaptiveServerSelection is disabled, the selection algorithm will always evenly distribute the traffic to all |
42 | 46 | * replicas of each segment, and will try to select different replica id for each segment. The algorithm is very |
43 | 47 | * light-weight and will do best effort to balance the number of segments served by each selected server instance. |
44 | 48 | */ |
45 | 49 | public class BalancedInstanceSelector extends BaseInstanceSelector { |
46 | 50 |
|
47 | | - public BalancedInstanceSelector(String tableNameWithType, BrokerMetrics brokerMetrics, |
48 | | - @Nullable AdaptiveServerSelector adaptiveServerSelector) { |
49 | | - super(tableNameWithType, brokerMetrics, adaptiveServerSelector); |
| 51 | + public BalancedInstanceSelector(String tableNameWithType, ZkHelixPropertyStore<ZNRecord> propertyStore, |
| 52 | + BrokerMetrics brokerMetrics, @Nullable AdaptiveServerSelector adaptiveServerSelector, Clock clock) { |
| 53 | + super(tableNameWithType, propertyStore, brokerMetrics, adaptiveServerSelector, clock); |
50 | 54 | } |
51 | 55 |
|
52 | 56 | @Override |
53 | | - Map<String, String> select(List<String> segments, int requestId, |
54 | | - Map<String, List<String>> segmentToEnabledInstancesMap, Map<String, String> queryOptions) { |
| 57 | + Map<String, String> select(List<String> segments, int requestId, SegmentStates segmentStates, |
| 58 | + Map<String, String> queryOptions) { |
55 | 59 | Map<String, String> segmentToSelectedInstanceMap = new HashMap<>(HashUtil.getHashMapCapacity(segments.size())); |
56 | | - for (String segment : segments) { |
57 | | - List<String> enabledInstances = segmentToEnabledInstancesMap.get(segment); |
58 | | - // NOTE: enabledInstances can be null when there is no enabled instances for the segment, or the instance selector |
59 | | - // has not been updated (we update all components for routing in sequence) |
60 | | - if (enabledInstances == null) { |
61 | | - continue; |
| 60 | + if (_adaptiveServerSelector != null) { |
| 61 | + for (String segment : segments) { |
| 62 | + List<SegmentInstanceCandidate> candidates = segmentStates.getCandidates(segment); |
| 63 | + // NOTE: candidates can be null when there is no enabled instances for the segment, or the instance selector has |
| 64 | + // not been updated (we update all components for routing in sequence) |
| 65 | + if (candidates == null) { |
| 66 | + continue; |
| 67 | + } |
| 68 | + List<String> candidateInstances = new ArrayList<>(candidates.size()); |
| 69 | + for (SegmentInstanceCandidate candidate : candidates) { |
| 70 | + candidateInstances.add(candidate.getInstance()); |
| 71 | + } |
| 72 | + String selectedInstance = _adaptiveServerSelector.select(candidateInstances); |
| 73 | + if (candidates.get(candidateInstances.indexOf(selectedInstance)).isOnline()) { |
| 74 | + segmentToSelectedInstanceMap.put(segment, selectedInstance); |
| 75 | + } |
62 | 76 | } |
63 | | - |
64 | | - String selectedServer; |
65 | | - if (_adaptiveServerSelector != null) { |
66 | | - selectedServer = _adaptiveServerSelector.select(enabledInstances); |
67 | | - } else { |
68 | | - selectedServer = enabledInstances.get(requestId++ % enabledInstances.size()); |
| 77 | + } else { |
| 78 | + for (String segment : segments) { |
| 79 | + List<SegmentInstanceCandidate> candidates = segmentStates.getCandidates(segment); |
| 80 | + // NOTE: candidates can be null when there is no enabled instances for the segment, or the instance selector has |
| 81 | + // not been updated (we update all components for routing in sequence) |
| 82 | + if (candidates == null) { |
| 83 | + continue; |
| 84 | + } |
| 85 | + int selectedIdx = requestId++ % candidates.size(); |
| 86 | + SegmentInstanceCandidate selectedCandidate = candidates.get(selectedIdx); |
| 87 | + if (selectedCandidate.isOnline()) { |
| 88 | + segmentToSelectedInstanceMap.put(segment, selectedCandidate.getInstance()); |
| 89 | + } |
69 | 90 | } |
70 | | - |
71 | | - segmentToSelectedInstanceMap.put(segment, selectedServer); |
72 | 91 | } |
73 | | - |
74 | 92 | return segmentToSelectedInstanceMap; |
75 | 93 | } |
76 | 94 | } |
0 commit comments