xiangfu0 commented on code in PR #19916:
URL: https://github.com/apache/pinot/pull/19916#discussion_r4332399450
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -1622,174 +1656,280 @@ private static void handleErrorInstance(String
segmentName, String instanceName,
}
}
+ /// Returns the target assignment for the ideal state changed while waiting
for the external view to converge, when
+ /// the segments to be moved are unchanged: the current assignment with the
target of the segments to be moved. For
+ /// a partition-pinned table (strict replica-group routing), a partition
added while waiting keeps its current
+ /// instances, and a segment added while waiting, which is assigned to the
current instances of its partition, is
+ /// moved to the target instances of its partition, unless it is a CONSUMING
segment and CONSUMING segments are not
+ /// included in the rebalance. Returns `null` when the target assignment
must be recalculated instead: when the
+ /// partition id of such a segment is not in its name, or when CONSUMING
segments are not included and an ONLINE
+ /// segment was added to a pauseless table, because it can still be
committing (which only its ZK metadata tells),
+ /// in which case it is not moved either.
+ @Nullable
+ @VisibleForTesting
+ static TableAssignment getTargetAssignmentForNewIdealState(Map<String,
Map<String, String>> currentAssignment,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment, List<String>
segmentsToMove,
+ boolean enableStrictReplicaGroup, boolean includeConsuming, boolean
pauseless, Logger tableRebalanceLogger) {
+ // Other instance assignment code returns a TreeMap to keep it sorted,
doing the same here
+ Map<String, Map<String, String>> newTargetAssignment = new
TreeMap<>(currentAssignment);
+ for (String segment : segmentsToMove) {
+ newTargetAssignment.put(segment, targetAssignment.get(segment));
+ }
+ if (!enableStrictReplicaGroup) {
+ return new TableAssignment(newTargetAssignment,
targetPartitionAssignment);
+ }
+ Map<Integer, List<String>> newTargetPartitionAssignment = new
TreeMap<>(targetPartitionAssignment);
+ for (Map.Entry<Integer, List<String>> entry :
currentPartitionAssignment.entrySet()) {
+ newTargetPartitionAssignment.putIfAbsent(entry.getKey(),
entry.getValue());
+ }
+ for (Map.Entry<String, Map<String, String>> entry :
currentAssignment.entrySet()) {
+ String segment = entry.getKey();
+ if (targetAssignment.containsKey(segment)) {
+ continue;
+ }
+ Map<String, String> currentInstanceStateMap = entry.getValue();
+ String state;
+ if (currentInstanceStateMap.containsValue(SegmentStateModel.ONLINE)) {
+ if (!includeConsuming && pauseless) {
+ tableRebalanceLogger.info("Found new ONLINE segment: {} of pauseless
table while CONSUMING segments are not "
+ + "included, re-calculating the target assignment based on the
new IdealState", segment);
+ return null;
+ }
+ state = SegmentStateModel.ONLINE;
+ } else if
(currentInstanceStateMap.containsValue(SegmentStateModel.CONSUMING)) {
+ if (!includeConsuming) {
+ continue;
+ }
+ state = SegmentStateModel.CONSUMING;
+ } else {
+ // All instances are OFFLINE, and OFFLINE segments are not rebalanced
+ continue;
+ }
+ Integer partitionId =
SegmentUtils.getPartitionIdFromSegmentName(segment);
+ if (partitionId == null) {
+ tableRebalanceLogger.info("Found new segment: {} without partition id
in its name, re-calculating the target "
+ + "assignment based on the new IdealState", segment);
+ return null;
+ }
+ List<String> currentInstances =
currentPartitionAssignment.get(partitionId);
+ if (currentInstances == null || currentInstances.size() !=
currentInstanceStateMap.size()
+ || !currentInstanceStateMap.keySet().containsAll(currentInstances)) {
+ // Not assigned to the instances of its partition (e.g. relocated to
the COMPLETED instances), keep it as is
+ continue;
+ }
+ newTargetAssignment.put(segment,
Review Comment:
**A new ONLINE segment sitting on its partition entry is retargeted to the
partition's CONSUMING-placement target even when the table relocates COMPLETED
segments**, which disagrees with what `rebalanceTable` computes (COMPLETED
strategy). Example: non-upsert strictReplicaGroup table with COMPLETED instance
partitions, rebalanced with includeConsuming=false; during one EV wait
partition P commits twice, so segment S created and committed within that wait
is ONLINE on P's intermediate entry and absent from the target; this line sends
it to P's consuming target, leaving a completed segment on consuming servers
(extra moves, fixed only by a later relocation).
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -1622,174 +1656,280 @@ private static void handleErrorInstance(String
segmentName, String instanceName,
}
}
+ /// Returns the target assignment for the ideal state changed while waiting
for the external view to converge, when
+ /// the segments to be moved are unchanged: the current assignment with the
target of the segments to be moved. For
+ /// a partition-pinned table (strict replica-group routing), a partition
added while waiting keeps its current
+ /// instances, and a segment added while waiting, which is assigned to the
current instances of its partition, is
+ /// moved to the target instances of its partition, unless it is a CONSUMING
segment and CONSUMING segments are not
+ /// included in the rebalance. Returns `null` when the target assignment
must be recalculated instead: when the
+ /// partition id of such a segment is not in its name, or when CONSUMING
segments are not included and an ONLINE
+ /// segment was added to a pauseless table, because it can still be
committing (which only its ZK metadata tells),
+ /// in which case it is not moved either.
+ @Nullable
+ @VisibleForTesting
+ static TableAssignment getTargetAssignmentForNewIdealState(Map<String,
Map<String, String>> currentAssignment,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment, List<String>
segmentsToMove,
+ boolean enableStrictReplicaGroup, boolean includeConsuming, boolean
pauseless, Logger tableRebalanceLogger) {
+ // Other instance assignment code returns a TreeMap to keep it sorted,
doing the same here
+ Map<String, Map<String, String>> newTargetAssignment = new
TreeMap<>(currentAssignment);
+ for (String segment : segmentsToMove) {
+ newTargetAssignment.put(segment, targetAssignment.get(segment));
+ }
+ if (!enableStrictReplicaGroup) {
+ return new TableAssignment(newTargetAssignment,
targetPartitionAssignment);
+ }
+ Map<Integer, List<String>> newTargetPartitionAssignment = new
TreeMap<>(targetPartitionAssignment);
+ for (Map.Entry<Integer, List<String>> entry :
currentPartitionAssignment.entrySet()) {
+ newTargetPartitionAssignment.putIfAbsent(entry.getKey(),
entry.getValue());
+ }
+ for (Map.Entry<String, Map<String, String>> entry :
currentAssignment.entrySet()) {
+ String segment = entry.getKey();
+ if (targetAssignment.containsKey(segment)) {
+ continue;
+ }
+ Map<String, String> currentInstanceStateMap = entry.getValue();
+ String state;
+ if (currentInstanceStateMap.containsValue(SegmentStateModel.ONLINE)) {
+ if (!includeConsuming && pauseless) {
+ tableRebalanceLogger.info("Found new ONLINE segment: {} of pauseless
table while CONSUMING segments are not "
+ + "included, re-calculating the target assignment based on the
new IdealState", segment);
+ return null;
+ }
+ state = SegmentStateModel.ONLINE;
+ } else if
(currentInstanceStateMap.containsValue(SegmentStateModel.CONSUMING)) {
+ if (!includeConsuming) {
+ continue;
+ }
+ state = SegmentStateModel.CONSUMING;
+ } else {
+ // All instances are OFFLINE, and OFFLINE segments are not rebalanced
+ continue;
+ }
+ Integer partitionId =
SegmentUtils.getPartitionIdFromSegmentName(segment);
+ if (partitionId == null) {
+ tableRebalanceLogger.info("Found new segment: {} without partition id
in its name, re-calculating the target "
+ + "assignment based on the new IdealState", segment);
+ return null;
+ }
+ List<String> currentInstances =
currentPartitionAssignment.get(partitionId);
+ if (currentInstances == null || currentInstances.size() !=
currentInstanceStateMap.size()
+ || !currentInstanceStateMap.keySet().containsAll(currentInstances)) {
+ // Not assigned to the instances of its partition (e.g. relocated to
the COMPLETED instances), keep it as is
+ continue;
+ }
+ newTargetAssignment.put(segment,
+
SegmentAssignmentUtils.getInstanceStateMap(newTargetPartitionAssignment.get(partitionId),
state));
+ }
+ return new TableAssignment(newTargetAssignment,
newTargetPartitionAssignment);
+ }
+
/// Uses the default LOGGER
@VisibleForTesting
static Map<String, Map<String, String>> getNextAssignment(Map<String,
Map<String, String>> currentAssignment,
Map<String, Map<String, String>> targetAssignment, int
minAvailableReplicas, boolean enableStrictReplicaGroup,
- boolean lowDiskMode, int batchSizePerServer,
Object2IntOpenHashMap<String> segmentPartitionIdMap,
- PartitionIdFetcher partitionIdFetcher, DataLossRiskAssessor
dataLossRiskAssessor) {
- return getNextAssignment(currentAssignment, targetAssignment,
minAvailableReplicas, enableStrictReplicaGroup,
- lowDiskMode, batchSizePerServer, segmentPartitionIdMap,
partitionIdFetcher, dataLossRiskAssessor, LOGGER);
+ boolean lowDiskMode, int batchSizePerServer, DataLossRiskAssessor
dataLossRiskAssessor) {
+ return getNextAssignment(currentAssignment, targetAssignment, Map.of(),
Map.of(), minAvailableReplicas,
+ enableStrictReplicaGroup, lowDiskMode, batchSizePerServer,
dataLossRiskAssessor).getSegmentAssignment();
+ }
+
+ /// Uses the default LOGGER
+ @VisibleForTesting
+ static TableAssignment getNextAssignment(Map<String, Map<String, String>>
currentAssignment,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment, int
minAvailableReplicas, boolean enableStrictReplicaGroup,
+ boolean lowDiskMode, int batchSizePerServer, DataLossRiskAssessor
dataLossRiskAssessor) {
+ return getNextAssignment(currentAssignment, targetAssignment,
currentPartitionAssignment,
+ targetPartitionAssignment, minAvailableReplicas,
enableStrictReplicaGroup, lowDiskMode, batchSizePerServer,
+ dataLossRiskAssessor, LOGGER);
}
/// Returns the next assignment for the table based on the current
assignment and the target assignment with regard to
/// the minimum available replicas requirement. For strict replica-group
mode, track the available instances for all
/// the segments with the same instances in the next assignment, and ensure
the minimum available replicas requirement
/// is met. If adding the assignment for a segment breaks the requirement,
use the current assignment for the segment.
///
- /// For strict replica group routing only (where the segment assignment is
not StrictRealtimeSegmentAssignment)
- /// if batching is enabled, the instances assigned for the same partitionId
can be different for different segments.
- /// For strict replica group routing with StrictRealtimeSegmentAssignment on
the other hand, the assignment for a
- /// given
- /// partitionId will be the same across all segments. We can treat both
cases similarly by creating a mapping from
- /// partitionId -> unique set of instance assignments -> currentAssignment.
With StrictRealtimeSegmentAssignment,
- /// this map will have a single entry for 'unique set of instance
assignments'.
- ///
- /// TODO: Ideally if strict replica group routing is enabled then
StrictRealtimeSegmentAssignment should be used, but
- /// this is not enforced in the code today. Once enforcement is added,
then the routing side and assignment side
- /// will be equivalent and all segments belonging to a given
partitionId will be assigned to the same set of
- /// instances. Special handling to check each group of assigned
instances can be removed in that case. The
- /// strict replica group routing can also be utilized for OFFLINE
tables, thus StrictRealtimeSegmentAssignment
- /// also needs to be made more generic for the OFFLINE case.
- private static Map<String, Map<String, String>>
getNextAssignment(Map<String, Map<String, String>> currentAssignment,
- Map<String, Map<String, String>> targetAssignment, int
minAvailableReplicas, boolean enableStrictReplicaGroup,
- boolean lowDiskMode, int batchSizePerServer,
Object2IntOpenHashMap<String> segmentPartitionIdMap,
- PartitionIdFetcher partitionIdFetcher, DataLossRiskAssessor
dataLossRiskAssessor, Logger tableRebalanceLogger) {
+ /// The partition assignment (the instances of each partition of a
partition-pinned table, see
+ /// [PartitionAssignmentUtils]) is moved along with the segments: a
partition takes the same step as a segment with
+ /// the same current and target instances, so that a segment added to the
partition during the rebalance is assigned
+ /// to the same instances as the other segments of the partition at that
point. In strict replica-group mode, a
+ /// partition moves exactly when the segments with its current and target
instances move. A partition of the target
+ /// without entry in the current partition assignment is left without entry,
as its current instances are unknown:
+ /// the caller sets it once the segments reach the target. A partition of
the current partition assignment absent
+ /// from the target is dropped.
+ private static TableAssignment getNextAssignment(Map<String, Map<String,
String>> currentAssignment,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment, int
minAvailableReplicas, boolean enableStrictReplicaGroup,
+ boolean lowDiskMode, int batchSizePerServer, DataLossRiskAssessor
dataLossRiskAssessor,
+ Logger tableRebalanceLogger) {
return enableStrictReplicaGroup ?
getNextStrictReplicaGroupAssignment(currentAssignment, targetAssignment,
- minAvailableReplicas, lowDiskMode, batchSizePerServer,
segmentPartitionIdMap, partitionIdFetcher,
+ currentPartitionAssignment, targetPartitionAssignment,
minAvailableReplicas, lowDiskMode, batchSizePerServer,
dataLossRiskAssessor, tableRebalanceLogger)
- : getNextNonStrictReplicaGroupAssignment(currentAssignment,
targetAssignment, minAvailableReplicas, lowDiskMode,
- batchSizePerServer, dataLossRiskAssessor);
+ : getNextNonStrictReplicaGroupAssignment(currentAssignment,
targetAssignment, currentPartitionAssignment,
+ targetPartitionAssignment, minAvailableReplicas, lowDiskMode,
batchSizePerServer, dataLossRiskAssessor);
}
- private static Map<String, Map<String, String>>
getNextStrictReplicaGroupAssignment(
+ private static TableAssignment getNextStrictReplicaGroupAssignment(
Map<String, Map<String, String>> currentAssignment, Map<String,
Map<String, String>> targetAssignment,
- int minAvailableReplicas, boolean lowDiskMode, int batchSizePerServer,
- Object2IntOpenHashMap<String> segmentPartitionIdMap, PartitionIdFetcher
partitionIdFetcher,
- DataLossRiskAssessor dataLossRiskAssessor, Logger tableRebalanceLogger) {
+ Map<Integer, List<String>> currentPartitionAssignment, Map<Integer,
List<String>> targetPartitionAssignment,
+ int minAvailableReplicas, boolean lowDiskMode, int batchSizePerServer,
DataLossRiskAssessor dataLossRiskAssessor,
+ Logger tableRebalanceLogger) {
Map<String, Map<String, String>> nextAssignment = new TreeMap<>();
+ Map<Integer, List<String>> nextPartitionAssignment = new TreeMap<>();
Map<String, Integer> numSegmentsToOffloadMap =
getNumSegmentsToOffloadMap(currentAssignment, targetAssignment);
Map<Pair<Set<String>, Set<String>>, Set<String>> assignmentMap = new
HashMap<>();
Map<Set<String>, Set<String>> availableInstancesMap = new HashMap<>();
Map<String, Integer> serverToNumSegmentsAddedSoFar = new HashMap<>();
if (batchSizePerServer == RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER) {
// Directly update the nextAssignment with anyServerExhaustedBatchSize =
false and return if batching is disabled
- updateNextAssignmentForPartitionIdStrictReplicaGroup(currentAssignment,
targetAssignment, nextAssignment, false,
- minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap,
assignmentMap, availableInstancesMap,
- serverToNumSegmentsAddedSoFar, dataLossRiskAssessor);
- return nextAssignment;
- }
-
- // Batching is enabled, calculate the Pair(current instances, target
instances) -> partitionId -> currentAssignment
- Map<Pair<Set<String>, Set<String>>, Map<Integer, Map<String, Map<String,
String>>>>
- currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap =
-
getCurrentAndTargetInstancesToPartitionIdToCurrentAssignmentMap(currentAssignment,
targetAssignment,
- segmentPartitionIdMap, partitionIdFetcher);
-
- // Iterating over the unique pairs of current and target instances
- for (Map<Integer, Map<String, Map<String, String>>>
partitionIdToCurrentAssignment
- :
currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap.values()) {
- // Check if the servers of the first assignment for each unique
partition has any space left to move the
- // segments assigned to the partition and unique assigned instances as a
whole. If so, let's mark the partitions
- // as to be moved, otherwise we mark the partition as a whole as not
moveable.
- // Iterating over the partitionIds with the same unique pair of current
and assigned instances
- for (Map<String, Map<String, String>> curAssignment :
partitionIdToCurrentAssignment.values()) {
- Map.Entry<String, Map<String, String>> firstEntry =
curAssignment.entrySet().iterator().next();
- // It is enough to check for whether any server for one segment is
above the limit or not since all segments
- // in curAssignment will have the same current and target instances
and same partitionId
- Map<String, String> firstEntryInstanceStateMap = firstEntry.getValue();
- SingleSegmentAssignment firstAssignment =
- getNextSingleSegmentAssignment(firstEntryInstanceStateMap,
targetAssignment.get(firstEntry.getKey()),
- minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap,
assignmentMap);
- Set<String> serversAdded =
-
getServersAddedInSingleSegmentAssignment(firstEntryInstanceStateMap,
firstAssignment._instanceStateMap);
- boolean anyServerExhaustedBatchSize = false;
- for (String server : serversAdded) {
- int segmentsAddedToServerSoFar =
serverToNumSegmentsAddedSoFar.getOrDefault(server, 0);
- // Case I: We already exceeded the batchSizePerServer for this
server, cannot add any more segments
- // Case II: We have not yet exceeded the batchSizePerServer for this
server, but we don't have sufficient
- // space to host the segments for this assignment on the server, and
we have allocated some partitions so
- // far. If the batchSizePerServer is less than the number of
segments in a given partitionId, we must host
- // at least 1 partition and exceed the batchSizePerServer to ensure
progress is made. Thus, performing this
- // check only if segmentsAddedToServerSoFar > 0 is necessary.
- if ((segmentsAddedToServerSoFar >= batchSizePerServer) ||
(segmentsAddedToServerSoFar > 0
- && (segmentsAddedToServerSoFar + curAssignment.size()) >
batchSizePerServer)) {
- anyServerExhaustedBatchSize = true;
- break;
- }
+ StrictReplicaGroup group = new StrictReplicaGroup();
+ group._segments.putAll(currentAssignment);
+ for (int partitionId : currentPartitionAssignment.keySet()) {
+ if (targetPartitionAssignment.containsKey(partitionId)) {
+ group._partitionIds.add(partitionId);
}
- updateNextAssignmentForPartitionIdStrictReplicaGroup(curAssignment,
targetAssignment, nextAssignment,
- anyServerExhaustedBatchSize, minAvailableReplicas, lowDiskMode,
numSegmentsToOffloadMap, assignmentMap,
- availableInstancesMap, serverToNumSegmentsAddedSoFar,
dataLossRiskAssessor);
}
+ moveStrictReplicaGroup(group, false, targetAssignment,
currentPartitionAssignment, targetPartitionAssignment,
+ nextAssignment, nextPartitionAssignment, minAvailableReplicas,
lowDiskMode, numSegmentsToOffloadMap,
+ assignmentMap, availableInstancesMap, serverToNumSegmentsAddedSoFar,
dataLossRiskAssessor);
+ return new TableAssignment(nextAssignment, nextPartitionAssignment);
+ }
+
+ // Batching is enabled, move the segments and the partitions group by group
+ for (StrictReplicaGroup group : getStrictReplicaGroups(currentAssignment,
targetAssignment,
+ currentPartitionAssignment, targetPartitionAssignment)) {
+ // Check if the servers of the first assignment for the group have any
space left to move the segments of the
+ // group as a whole. If so, let's mark the group as to be moved,
otherwise we mark the group as a whole as not
+ // moveable. It is enough to check the first segment (or partition)
since all segments and partitions in the
+ // group have the same current and target instances.
+ Map<String, String> firstCurrentInstanceStateMap;
+ Map<String, String> firstTargetInstanceStateMap;
+ if (!group._segments.isEmpty()) {
+ Map.Entry<String, Map<String, String>> firstEntry =
group._segments.entrySet().iterator().next();
+ firstCurrentInstanceStateMap = firstEntry.getValue();
+ firstTargetInstanceStateMap =
targetAssignment.get(firstEntry.getKey());
+ } else {
+ int firstPartitionId = group._partitionIds.get(0);
+ firstCurrentInstanceStateMap =
getPartitionInstanceStateMap(currentPartitionAssignment.get(firstPartitionId));
+ firstTargetInstanceStateMap =
getPartitionInstanceStateMap(targetPartitionAssignment.get(firstPartitionId));
+ }
+ SingleSegmentAssignment firstAssignment =
+ getNextSingleSegmentAssignment(firstCurrentInstanceStateMap,
firstTargetInstanceStateMap,
+ minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap,
assignmentMap);
+ Set<String> serversAdded =
+
getServersAddedInSingleSegmentAssignment(firstCurrentInstanceStateMap,
firstAssignment._instanceStateMap);
+ int numSegments = group._segments.size();
+ boolean anyServerExhaustedBatchSize = false;
+ for (String server : serversAdded) {
+ int segmentsAddedToServerSoFar =
serverToNumSegmentsAddedSoFar.getOrDefault(server, 0);
+ // Case I: We already exceeded the batchSizePerServer for this server,
cannot add any more segments
+ // Case II: We have not yet exceeded the batchSizePerServer for this
server, but we don't have sufficient
+ // space to host the segments for this assignment on the server, and
we have allocated some partitions so
+ // far. If the batchSizePerServer is less than the number of segments
in a given partitionId, we must host
+ // at least 1 partition and exceed the batchSizePerServer to ensure
progress is made. Thus, performing this
+ // check only if segmentsAddedToServerSoFar > 0 is necessary.
+ if ((segmentsAddedToServerSoFar >= batchSizePerServer) ||
(segmentsAddedToServerSoFar > 0
+ && (segmentsAddedToServerSoFar + numSegments) >
batchSizePerServer)) {
+ anyServerExhaustedBatchSize = true;
+ break;
+ }
+ }
+ moveStrictReplicaGroup(group, anyServerExhaustedBatchSize,
targetAssignment, currentPartitionAssignment,
+ targetPartitionAssignment, nextAssignment, nextPartitionAssignment,
minAvailableReplicas, lowDiskMode,
+ numSegmentsToOffloadMap, assignmentMap, availableInstancesMap,
serverToNumSegmentsAddedSoFar,
+ dataLossRiskAssessor);
}
checkIfAnyServersAssignedMoreSegmentsThanBatchSize(batchSizePerServer,
serverToNumSegmentsAddedSoFar,
tableRebalanceLogger);
- return nextAssignment;
+ return new TableAssignment(nextAssignment, nextPartitionAssignment);
}
- private static void updateNextAssignmentForPartitionIdStrictReplicaGroup(
- Map<String, Map<String, String>> currentAssignment, Map<String,
Map<String, String>> targetAssignment,
- Map<String, Map<String, String>> nextAssignment, boolean
anyServerExhaustedBatchSize, int minAvailableReplicas,
- boolean lowDiskMode, Map<String, Integer> numSegmentsToOffloadMap,
- Map<Pair<Set<String>, Set<String>>, Set<String>> assignmentMap,
+ /// Moves the segments and the partitions of the given group to their next
assignment, or keeps them as they are
+ /// when the batch size is exhausted. The segments and partitions of a group
have the same current and target
+ /// instances, so they take the same step; the minimum available replicas
check is still performed per segment and
+ /// partition because it depends on the other segments assigned to the same
instances.
+ private static void moveStrictReplicaGroup(StrictReplicaGroup group, boolean
anyServerExhaustedBatchSize,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment, Map<String,
Map<String, String>> nextAssignment,
+ Map<Integer, List<String>> nextPartitionAssignment, int
minAvailableReplicas, boolean lowDiskMode,
+ Map<String, Integer> numSegmentsToOffloadMap, Map<Pair<Set<String>,
Set<String>>, Set<String>> assignmentMap,
Map<Set<String>, Set<String>> availableInstancesMap, Map<String,
Integer> serverToNumSegmentsAddedSoFar,
DataLossRiskAssessor dataLossRiskAssessor) {
if (anyServerExhaustedBatchSize) {
- // Exhausted the batch size for at least 1 server, just copy over the
remaining segments as is
- nextAssignment.putAll(currentAssignment);
- } else {
- // Process all the partitionIds even if segmentsAddedSoFar becomes
larger than batchSizePerServer
- // Can only do bestEfforts w.r.t. StrictReplicaGroup since a whole
partition must be moved together for
- // maintaining consistency
- for (Map.Entry<String, Map<String, String>> entry :
currentAssignment.entrySet()) {
- String segmentName = entry.getKey();
- Map<String, String> currentInstanceStateMap = entry.getValue();
- Map<String, String> targetInstanceStateMap =
targetAssignment.get(segmentName);
- SingleSegmentAssignment assignment =
- getNextSingleSegmentAssignment(currentInstanceStateMap,
targetInstanceStateMap, minAvailableReplicas,
- lowDiskMode, numSegmentsToOffloadMap, assignmentMap);
- Set<String> assignedInstances = assignment._instanceStateMap.keySet();
- Set<String> availableInstances = assignment._availableInstances;
- availableInstancesMap.compute(assignedInstances, (k,
currentAvailableInstances) -> {
- if (currentAvailableInstances == null) {
- // First segment assigned to these instances, use the new
assignment and update the available instances
- nextAssignment.put(segmentName, assignment._instanceStateMap);
- updateNumSegmentsToOffloadMap(numSegmentsToOffloadMap,
currentInstanceStateMap.keySet(), k);
- return availableInstances;
- } else {
- // There are other segments assigned to the same instances, check
the available instances to see if
- // adding the new assignment can still hold the minimum available
replicas requirement
- availableInstances.retainAll(currentAvailableInstances);
- if (availableInstances.size() >= minAvailableReplicas) {
- // New assignment can be added
- nextAssignment.put(segmentName, assignment._instanceStateMap);
- updateNumSegmentsToOffloadMap(numSegmentsToOffloadMap,
currentInstanceStateMap.keySet(), k);
- return availableInstances;
- } else {
- // New assignment cannot be added, use the current instance
state map
- nextAssignment.put(segmentName, currentInstanceStateMap);
- return currentAvailableInstances;
- }
- }
- });
+ // Exhausted the batch size for at least 1 server, just copy over the
remaining segments and partitions as is
+ nextAssignment.putAll(group._segments);
+ for (int partitionId : group._partitionIds) {
+ nextPartitionAssignment.put(partitionId,
currentPartitionAssignment.get(partitionId));
+ }
+ return;
+ }
+ // Process all the segments even if segmentsAddedSoFar becomes larger than
batchSizePerServer
+ // Can only do bestEfforts w.r.t. StrictReplicaGroup since a whole
partition must be moved together for
+ // maintaining consistency
+ for (Map.Entry<String, Map<String, String>> entry :
group._segments.entrySet()) {
+ String segmentName = entry.getKey();
+ Map<String, String> currentInstanceStateMap = entry.getValue();
+ Map<String, String> targetInstanceStateMap =
targetAssignment.get(segmentName);
+ SingleSegmentAssignment assignment =
+ getNextSingleSegmentAssignment(currentInstanceStateMap,
targetInstanceStateMap, minAvailableReplicas,
+ lowDiskMode, numSegmentsToOffloadMap, assignmentMap);
+ Map<String, String> nextInstanceStateMap;
+ if (canAssignWithStrictReplicaGroup(assignment, minAvailableReplicas,
availableInstancesMap)) {
+ nextInstanceStateMap = assignment._instanceStateMap;
+ updateNumSegmentsToOffloadMap(numSegmentsToOffloadMap,
currentInstanceStateMap.keySet(),
+ nextInstanceStateMap.keySet());
+ } else {
+ nextInstanceStateMap = currentInstanceStateMap;
+ }
+ nextAssignment.put(segmentName, nextInstanceStateMap);
- if (!nextAssignment.get(segmentName).equals(currentInstanceStateMap)) {
- Set<String> serversAddedForSegment =
-
getServersAddedInSingleSegmentAssignment(currentInstanceStateMap,
nextAssignment.get(segmentName));
- serversAddedForSegment.forEach(server ->
serverToNumSegmentsAddedSoFar.merge(server, 1, Integer::sum));
+ if (!nextInstanceStateMap.equals(currentInstanceStateMap)) {
+ Set<String> serversAddedForSegment =
+ getServersAddedInSingleSegmentAssignment(currentInstanceStateMap,
nextInstanceStateMap);
+ serversAddedForSegment.forEach(server ->
serverToNumSegmentsAddedSoFar.merge(server, 1, Integer::sum));
- // Since next assignment doesn't match current assignment, it means
the segment will be moved. Check if there
- // is a data loss risk
- Pair<Boolean, String> dataLossRiskResult =
dataLossRiskAssessor.assessDataLossRisk(segmentName);
- if (dataLossRiskResult.getLeft()) {
- throw new IllegalStateException(dataLossRiskResult.getRight());
- }
+ // Since next assignment doesn't match current assignment, it means
the segment will be moved. Check if there
+ // is a data loss risk
+ Pair<Boolean, String> dataLossRiskResult =
dataLossRiskAssessor.assessDataLossRisk(segmentName);
+ if (dataLossRiskResult.getLeft()) {
+ throw new IllegalStateException(dataLossRiskResult.getRight());
}
}
}
+ for (int partitionId : group._partitionIds) {
+ List<String> currentInstances =
currentPartitionAssignment.get(partitionId);
+ SingleSegmentAssignment assignment =
+
getNextSingleSegmentAssignment(getPartitionInstanceStateMap(currentInstances),
+
getPartitionInstanceStateMap(targetPartitionAssignment.get(partitionId)),
minAvailableReplicas,
+ lowDiskMode, numSegmentsToOffloadMap, assignmentMap);
+ nextPartitionAssignment.put(partitionId,
+ canAssignWithStrictReplicaGroup(assignment, minAvailableReplicas,
availableInstancesMap)
Review Comment:
**Data-less partition entries take part in the min-available-replicas
accounting like real segments.** `canAssignWithStrictReplicaGroup` registers an
entry's available instances in the shared `availableInstancesMap`, so with
batching on (groups iterated in HashMap order) an entry-only group processed
first can hold back real segments heading to the same next instance set.
Example: minAvailableReplicas=1, partition 1's entry [p,q]->[p,r] with no
segments in that pair, partition 3's segments [r,s]->[p,r]; both step to [p,r]
with available {p} vs {r}; if partition 1's group comes first, partition 3's
segments fail `retainAll` and stay put for an extra round although no real
segment constrains them. Deriving an entry's next step from its segments' next
assignment would avoid this.
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineSegmentAssignment.java:
##########
@@ -21,47 +21,97 @@
import com.google.common.base.Preconditions;
import java.util.List;
import java.util.Map;
+import java.util.TreeMap;
import javax.annotation.Nullable;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.tuple.Pair;
+import org.apache.helix.model.IdealState;
import org.apache.pinot.common.assignment.InstancePartitions;
+import org.apache.pinot.common.assignment.InstancePartitionsUtils;
+import org.apache.pinot.common.assignment.PartitionAssignmentUtils;
import org.apache.pinot.common.restlet.resources.RebalanceConfig;
import org.apache.pinot.common.tier.Tier;
+import org.apache.pinot.common.utils.SegmentUtils;
import
org.apache.pinot.controller.helix.core.assignment.segment.strategy.AllServersSegmentAssignmentStrategy;
import
org.apache.pinot.controller.helix.core.assignment.segment.strategy.SegmentAssignmentStrategy;
import
org.apache.pinot.controller.helix.core.assignment.segment.strategy.SegmentAssignmentStrategyFactory;
+import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
+import
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel;
/// Segment assignment for offline table.
+///
+/// For a partition-pinned table (strict replica-group routing), all segments
of a value partition are assigned to
+/// the same instances, recorded in the partition assignment of the ideal
state (see
+/// [org.apache.pinot.common.assignment.PartitionAssignmentUtils]), even when
the instance partitions have changed.
+/// The partition id of a new segment is read from its ZK metadata. This
requires a partition column, a table that is
+/// not a dimension table (whose segments are assigned to all servers of the
tenant), and OFFLINE instance partitions
+/// holding exactly one instance per partition within each replica-group;
otherwise a warning is logged and the
+/// segments are assigned with the segment assignment strategy, as are
segments without partition id and segments
+/// eligible for a tier.
public class OfflineSegmentAssignment extends BaseSegmentAssignment {
@Override
- public List<String> assignSegment(String segmentName, Map<String,
Map<String, String>> currentAssignment,
+ protected boolean isPartitionPinned(TableConfig tableConfig) {
+ if (!super.isPartitionPinned(tableConfig)) {
+ return false;
+ }
+ if (_partitionColumn == null) {
+ _logger.warn("Table: {} with strict replica-group routing has no
partition column, cannot pin its segments to "
+ + "their partition, assigning them with the segment assignment
strategy", _tableNameWithType);
+ return false;
+ }
+ if (tableConfig.isDimTable()) {
+ _logger.warn("Dimension table: {} with strict replica-group routing
cannot pin its segments to their partition "
+ + "because they are assigned to all servers, assigning them with the
segment assignment strategy",
+ _tableNameWithType);
+ return false;
+ }
+ return true;
+ }
+
+ @Override
+ public List<String> assignSegment(String segmentName, IdealState idealState,
Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
InstancePartitions instancePartitions =
instancePartitionsMap.get(InstancePartitionsType.OFFLINE);
Preconditions.checkState(instancePartitions != null, "Failed to find
OFFLINE instance partitions for table: %s",
_tableNameWithType);
- // Gets Segment assignment strategy for instance partitions
- SegmentAssignmentStrategy segmentAssignmentStrategy =
-
SegmentAssignmentStrategyFactory.getSegmentAssignmentStrategy(_helixManager,
_tableConfig,
- InstancePartitionsType.OFFLINE.toString(), instancePartitions);
_logger.info("Assigning segment: {} with instance partitions: {} for
table: {}", segmentName, instancePartitions,
_tableNameWithType);
- List<String> instancesAssigned =
- segmentAssignmentStrategy.assignSegment(segmentName,
currentAssignment, instancePartitions);
+ List<String> instancesAssigned = null;
+ if (isPinnedWith(instancePartitions)) {
+ Integer partitionId =
+ SegmentUtils.getSegmentPartitionId(segmentName, _tableNameWithType,
_helixManager, _partitionColumn);
+ if (partitionId != null) {
+ instancesAssigned = getOrCreatePartitionInstances(idealState,
partitionId,
Review Comment:
**OfflineSegmentAssignment does not override `getExistingAssignment`, so for
offline tables created before this PR a partition's first entry is seeded from
the current instance partitions rather than from where its segments are.** That
contradicts the class doc ('even when the instance partitions have changed'),
and the realtime side has this handling. Example: pre-existing pinned OFFLINE
(e.g. upsert) table without entries: replaceInstance X->Y without a rebalance,
then the first upload of partition P stores entry [Y]; the operator reverts the
instance partitions to X, but every later upload of P still goes to Y (base:
back to X), so P stays split with a growing share on Y until a rebalance.
##########
pinot-common/src/main/java/org/apache/pinot/common/assignment/PartitionAssignmentUtils.java:
##########
@@ -0,0 +1,102 @@
+/**
+ * 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.common.assignment;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import javax.annotation.Nullable;
+import org.apache.helix.model.IdealState;
+
+
+/// Partition assignment of a table: the instances hosting all segments of
each segment partition, stored in the list
+/// fields of the table IdealState. List fields whose key is not a partition
id are not part of it and are left alone.
+///
+/// Each list field maps a segment partition id (the stream partition for a
real-time table, the value partition for an
+/// offline table) to the sorted list of instances hosting all segments of
that partition. The partition assignment is
+/// maintained for partition-pinned tables only, i.e. tables whose segments
must be collocated per partition: tables
+/// with strict replica-group routing, which upsert and dedup tables use. For
these tables:
+/// - A new segment is assigned to the instances of its partition when the
partition has an entry, even when the
+/// instance partitions have changed since the entry was created. The entry
is created from the instance partitions
+/// otherwise.
+/// - Table rebalance moves the entries together with the segments, so that a
segment added during the rebalance is
+/// assigned to the same instances as the other segments of its partition at
that point of the rebalance.
+///
+/// The entries are read and written together with the map fields of the
IdealState, so they are never observed out of
+/// sync with the segment assignment. Helix ignores the list fields of an
IdealState in CUSTOMIZED rebalance mode.
+public class PartitionAssignmentUtils {
+ private PartitionAssignmentUtils() {
+ }
+
+ /// Returns the instances of the given partition, or `null` if the partition
has no entry.
+ @Nullable
+ public static List<String> getInstances(IdealState idealState, int
partitionId) {
+ return idealState.getRecord().getListField(Integer.toString(partitionId));
+ }
+
+ /// Sets the instances of the given partition.
+ public static void setInstances(IdealState idealState, int partitionId,
Collection<String> instances) {
+ idealState.getRecord().setListField(Integer.toString(partitionId),
toInstanceList(instances));
+ }
+
+ /// Returns the partition assignment as a sorted map from partition id to
the instances of the partition. The lists
+ /// are the ones stored in the IdealState record and must not be modified.
+ public static Map<Integer, List<String>> getPartitionAssignment(IdealState
idealState) {
+ Map<Integer, List<String>> partitionAssignment = new TreeMap<>();
+ for (Map.Entry<String, List<String>> entry :
idealState.getRecord().getListFields().entrySet()) {
+ try {
+ partitionAssignment.put(Integer.parseInt(entry.getKey()),
entry.getValue());
Review Comment:
**Any list field whose key passes lenient `Integer.parseInt` is treated as a
partition entry, including negative keys.** A hand-edited `-1` key makes every
`rebalanceTable` throw while seeding the target:
`OfflineSegmentAssignment.getPartitionInstances(-1)` computes `-1 %
numPartitions = -1`, so `getInstances(-1, 0)` is null and NPEs;
`RealtimeSegmentAssignment.assignConsumingSegment(-1, ...)` throws
IndexOutOfBounds. The rebalance returns FAILED and nothing removes the key. A
JSON-null list value similarly NPEs at `new HashSet<>(entry.getValue())`. Only
canonical non-negative keys (`Integer.toString(id)`, as `getInstances` writes)
should be accepted.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
