xiangfu0 commented on code in PR #19916:
URL: https://github.com/apache/pinot/pull/19916#discussion_r4332395498
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -82,21 +101,95 @@ public abstract class BaseSegmentAssignment implements
SegmentAssignment {
@Nullable
private Map<String, String> _segmentToTierName;
+ // Cache of segment name to partition id (null value means the partition id
cannot be determined) for the segments
+ // whose partition id is not in the segment name, loaded with a single bulk
ZK read on first use. Same scope and
+ // limitations as _segmentToTierName; a segment absent from the cache is
resolved on demand and merged in.
+ @Nullable
+ private Map<String, Integer> _segmentToPartitionId;
+
@Override
public void init(HelixManager helixManager, TableConfig tableConfig,
@Nullable ControllerMetrics controllerMetrics) {
_helixManager = helixManager;
_tableNameWithType = tableConfig.getTableName();
_tableConfig = tableConfig;
_replication = tableConfig.getReplication();
_partitionColumn = TableConfigUtils.getPartitionColumn(_tableConfig);
- _controllerMetrics = controllerMetrics;
+ _partitionPinned = isPartitionPinned(tableConfig);
if (_partitionColumn == null) {
- _logger.info("Initialized with replication: {} without partition column
for table: {} ", _replication,
- _tableNameWithType);
+ _logger.info("Initialized with replication: {} without partition column,
partitionPinned: {} for table: {}",
+ _replication, _partitionPinned, _tableNameWithType);
} else {
- _logger.info("Initialized with replication: {} and partition column: {}
for table: {}", _replication,
- _partitionColumn, _tableNameWithType);
+ _logger.info("Initialized with replication: {} and partition column: {},
partitionPinned: {} for table: {}",
+ _replication, _partitionColumn, _partitionPinned,
_tableNameWithType);
+ }
+ }
+
+ /// Returns whether the segments of the table are pinned to their partition
(see the class Javadoc): the table uses
+ /// strict replica-group routing. Subclasses add the conditions they need to
pin the segments, and log a warning
+ /// when such a table cannot be pinned, in which case its segments are
assigned with the segment assignment strategy.
+ protected boolean isPartitionPinned(TableConfig tableConfig) {
+ return SegmentAssignmentUtils.isStrictReplicaGroupRouting(tableConfig);
+ }
+
+ /// Returns the instances of the given partition of a partition-pinned
table, creating its entry in the partition
+ /// assignment of the given ideal state when the partition has no entry yet,
or an empty one (a rebalance in low disk
+ /// mode without minimum available replicas empties the entry while it drops
the partition from its current
+ /// instances). The entry is created from the existing assignment of the
partition (see [#getExistingAssignment])
+ /// when there is one, else from the given candidate instances computed from
the instance partitions. An existing
+ /// entry wins over the candidate instances, so that the segments of the
partition stay collocated when the instance
+ /// partitions change: the partition moves on the next rebalance.
+ protected List<String> getOrCreatePartitionInstances(IdealState idealState,
int partitionId,
+ List<String> candidateInstances) {
+ List<String> instances = PartitionAssignmentUtils.getInstances(idealState,
partitionId);
+ if (CollectionUtils.isEmpty(instances)) {
Review Comment:
**A non-empty entry is trusted unconditionally, but only the new
TableRebalancer keeps entries in sync with segment moves.** Entries are written
on the first commit per partition (not only on rebalance), and any move made
without the new rebalancer leaves them stale; they then redirect every new
segment of the partition. Example: a new-version controller writes entry
`0`=[S1,S2]; controllers are rolled back (or an old-version controller runs a
rebalance during a rolling upgrade) and the old TableRebalancer, which only
calls `setMapFields`, moves partition 0 to [S2,S3]; after (re-)upgrade every
new CONSUMING segment of partition 0 goes to [S1,S2] (S1 possibly
decommissioned) while its history is on [S2,S3], splitting the upsert/dedup
partition until a new-version rebalance. The removed mismatch meter was the
only signal. CLAUDE.md asks to preserve mixed-version compatibility; consider
validating the entry against the partition's existing segments (e.g. its
CONSUMING segment) before
trusting it.
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java:
##########
@@ -82,21 +101,95 @@ public abstract class BaseSegmentAssignment implements
SegmentAssignment {
@Nullable
private Map<String, String> _segmentToTierName;
+ // Cache of segment name to partition id (null value means the partition id
cannot be determined) for the segments
+ // whose partition id is not in the segment name, loaded with a single bulk
ZK read on first use. Same scope and
+ // limitations as _segmentToTierName; a segment absent from the cache is
resolved on demand and merged in.
+ @Nullable
+ private Map<String, Integer> _segmentToPartitionId;
+
@Override
public void init(HelixManager helixManager, TableConfig tableConfig,
@Nullable ControllerMetrics controllerMetrics) {
_helixManager = helixManager;
_tableNameWithType = tableConfig.getTableName();
_tableConfig = tableConfig;
_replication = tableConfig.getReplication();
_partitionColumn = TableConfigUtils.getPartitionColumn(_tableConfig);
- _controllerMetrics = controllerMetrics;
+ _partitionPinned = isPartitionPinned(tableConfig);
if (_partitionColumn == null) {
- _logger.info("Initialized with replication: {} without partition column
for table: {} ", _replication,
- _tableNameWithType);
+ _logger.info("Initialized with replication: {} without partition column,
partitionPinned: {} for table: {}",
+ _replication, _partitionPinned, _tableNameWithType);
} else {
- _logger.info("Initialized with replication: {} and partition column: {}
for table: {}", _replication,
- _partitionColumn, _tableNameWithType);
+ _logger.info("Initialized with replication: {} and partition column: {},
partitionPinned: {} for table: {}",
+ _replication, _partitionColumn, _partitionPinned,
_tableNameWithType);
+ }
+ }
+
+ /// Returns whether the segments of the table are pinned to their partition
(see the class Javadoc): the table uses
+ /// strict replica-group routing. Subclasses add the conditions they need to
pin the segments, and log a warning
+ /// when such a table cannot be pinned, in which case its segments are
assigned with the segment assignment strategy.
+ protected boolean isPartitionPinned(TableConfig tableConfig) {
+ return SegmentAssignmentUtils.isStrictReplicaGroupRouting(tableConfig);
Review Comment:
**Upsert/dedup tables on `multiStageReplicaGroup` routing silently lose
partition pinning.** `TableConfigUtils.UPSERT_DEDUP_ALLOWED_ROUTING_STRATEGIES`
still accepts that routing, and the deleted factory pinned every upsert/dedup
table regardless of routing. Example: such a table's CONSUMING instance
partitions change without moving segments (assignInstances, PUT
instancePartitions, replaceInstance, or a reassignInstances rebalance that
persists new instance partitions and then fails/stops); the next commit of
partition P runs `assignConsumingSegment` from the instance partitions and puts
the new CONSUMING segment on the new server while P's completed segments stay
on the old one, so upsert serves stale and new rows for the same key and dedup
admits duplicates. The PR calls this deliberate, but nothing rejects, warns
about or migrates the still-valid config.
##########
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) {
Review Comment:
**With includeConsuming=false on a pinned non-upsert table, CONSUMING
segments created mid-rebalance are stranded on intermediate instances.** The
partition entry still moves step by step (its target comes from the instance
partitions), new CONSUMING segments are assigned to the entry at that moment,
and both this `continue` and the full recompute
(`newAssignment.putAll(consumingSegmentAssignment)`) leave them there. Example:
non-upsert strictReplicaGroup table rebalanced [A,B]->[C,D] with
includeConsuming=false (RebalanceConfig and SegmentRelocator default) and
minAvailableReplicas=1: step 1 writes entry [A,C] ([A] in lowDiskMode); p0__6
is created CONSUMING on {A,C}; the rebalance reports DONE with p0__6 on A
(being decommissioned) or with a single replica, and after it commits it stays
split from the rest of p0 on [C,D]. Before the PR the new segment went straight
to the new instance partitions {C,D}.
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -1803,71 +1943,63 @@ private static void
checkIfAnyServersAssignedMoreSegmentsThanBatchSize(int batch
}
}
- /// Create a mapping of Pair(currentInstances, targetInstances) to
partitionId to the current assignment of segments.
- /// This is to be used for batching purposes for StrictReplicaGroup routing,
for all segment assignment types:
- /// RealtimeSegmentAssignment, StrictRealtimeSegmentAssignment and
OfflineSegmentAssignment
- /// @param currentAssignment the current assignment
- /// @param targetAssignment the target assignment
- /// @param segmentPartitionIdMap cache to store the partition ids to avoid
fetching ZK segment metadata
- /// @param partitionIdFetcher function to fetch the partition id
- /// @return a mapping from Pair(currentInstances, targetInstances) to the
partitionId to the segment assignment map of
- /// all segments that fall in that category
- private static Map<Pair<Set<String>, Set<String>>, Map<Integer, Map<String,
Map<String, String>>>>
- getCurrentAndTargetInstancesToPartitionIdToCurrentAssignmentMap(
- Map<String, Map<String, String>> currentAssignment, Map<String,
Map<String, String>> targetAssignment,
- Object2IntOpenHashMap<String> segmentPartitionIdMap,
PartitionIdFetcher partitionIdFetcher) {
- Map<Pair<Set<String>, Set<String>>, Map<Integer, Map<String, Map<String,
String>>>>
- currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap = new
HashMap<>();
-
- for (Map.Entry<String, Map<String, String>> assignment :
currentAssignment.entrySet()) {
- String segmentName = assignment.getKey();
- Map<String, String> currentInstanceStateMap = assignment.getValue();
- Map<String, String> targetInstanceStateMap =
targetAssignment.get(segmentName);
-
- int partitionId = segmentPartitionIdMap.computeIfAbsent(segmentName, v
-> partitionIdFetcher.fetch(segmentName));
- Pair<Set<String>, Set<String>> currentAndTargetInstances =
- Pair.of(currentInstanceStateMap.keySet(),
targetInstanceStateMap.keySet());
-
currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap.computeIfAbsent(currentAndTargetInstances,
- k -> new TreeMap<>())
- .computeIfAbsent(partitionId, k -> new TreeMap<>())
- .put(segmentName, currentInstanceStateMap);
- }
-
- return currentAndTargetInstancesToPartitionIdToCurrentAssignmentMap;
- }
-
- @VisibleForTesting
- @FunctionalInterface
- interface PartitionIdFetcher {
- int fetch(String segmentName);
+ /// Segments and partitions with the same current and target instances,
which are moved together for strict
+ /// replica-group routing.
+ private static class StrictReplicaGroup {
+ final Map<String, Map<String, String>> _segments = new TreeMap<>();
+ final List<Integer> _partitionIds = new ArrayList<>();
}
- private static class PartitionIdFetcherImpl implements PartitionIdFetcher {
- private final String _tableNameWithType;
- private final String _partitionColumn;
- private final HelixManager _helixManager;
- private final boolean _isStrictRealtimeSegmentAssignment;
-
- private PartitionIdFetcherImpl(String tableNameWithType, @Nullable String
partitionColumn,
- HelixManager helixManager, boolean isStrictRealtimeSegmentAssignment) {
- _tableNameWithType = tableNameWithType;
- _partitionColumn = partitionColumn;
- _helixManager = helixManager;
- _isStrictRealtimeSegmentAssignment = isStrictRealtimeSegmentAssignment;
- }
-
- @Override
- public int fetch(String segmentName) {
- Integer partitionId =
- SegmentUtils.getSegmentPartitionId(segmentName, _tableNameWithType,
_helixManager, _partitionColumn);
- if (partitionId != null) {
- return partitionId;
+ /// Groups the segments and the partitions by their current and target
instances for batching with strict
+ /// replica-group routing, for all table types. Within a pair of current and
target instances, the segments are
+ /// further grouped by the partition id in their name, so that partitions
sharing the same instances can be moved in
+ /// different batches. When a segment has no partition id in its name (a
segment of an offline table, or an uploaded
+ /// segment with a custom name), the whole pair is moved together so that
the segment stays with the other segments
+ /// of its partition. The partitions of a partition-pinned table join the
group of their current and target
+ /// instances. Partitions absent from the target partition assignment are
dropped.
+ private static List<StrictReplicaGroup> getStrictReplicaGroups(Map<String,
Map<String, String>> currentAssignment,
+ Map<String, Map<String, String>> targetAssignment, Map<Integer,
List<String>> currentPartitionAssignment,
+ Map<Integer, List<String>> targetPartitionAssignment) {
+ Map<Pair<Set<String>, Set<String>>, Map<Integer, StrictReplicaGroup>>
groupsMap = new HashMap<>();
+ for (Map.Entry<String, Map<String, String>> entry :
currentAssignment.entrySet()) {
+ String segmentName = entry.getKey();
+ Map<String, String> currentInstanceStateMap = entry.getValue();
+ Pair<Set<String>, Set<String>> currentAndTargetInstances =
+ Pair.of(currentInstanceStateMap.keySet(),
targetAssignment.get(segmentName).keySet());
+ Integer partitionId =
SegmentUtils.getPartitionIdFromSegmentName(segmentName);
Review Comment:
**batchSizePerServer no longer limits OFFLINE strictReplicaGroup
rebalances.** The partition id comes from the segment name only; offline names
never carry one, so every segment maps to UNKNOWN_PARTITION_ID and the merge
below moves the whole (current,target) pair as one group, which is always
admitted because the first group reaching a server has
`segmentsAddedToServerSoFar == 0`. Example: offline strictReplicaGroup table
with replica-group instance partitions, batchSizePerServer=100, replacing
server X with X': all ~2,500 segments of the pair ({X,Y},{X',Y}) are added to
X' in one IdealState step (only a WARN is logged), where the removed
PartitionIdFetcher grouped by ZK partition id or a per-segment hash and moved
~100 per step. One custom-named segment does the same to a realtime pinned
pair. The partition ids that rebalanceTable already resolved
(`_segmentToPartitionId`) could be returned and used here, with a per-segment
default for unpinned tables.
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java:
##########
@@ -396,6 +402,30 @@ private RebalanceResult doRebalance(TableConfig
tableConfig, RebalanceConfig reb
if (segmentAssignmentUnchanged) {
tableRebalanceLogger.info("Table is already balanced");
+ if (!partitionAssignmentUnchanged && !dryRun) {
+ // Only the partition assignment changed (e.g. a table created before
the partition assignment was introduced,
+ // or partitions without segment): update it without moving any segment
+ tableRebalanceLogger.info("Updating the partition assignment from: {}
to: {}", currentPartitionAssignment,
+ targetPartitionAssignment);
+ PartitionAssignmentUtils.setPartitionAssignment(currentIdealState,
targetPartitionAssignment);
+ ZNRecord idealStateRecord = currentIdealState.getRecord();
+ try {
+ Preconditions.checkState(_helixDataAccessor.getBaseDataAccessor()
+ .set(idealStatePropertyKey.getPath(), idealStateRecord,
idealStateRecord.getVersion(),
Review Comment:
**The partition-only update writes with the version read at the start of
doRebalance and turns a conflict into FAILED with no retry; the dry-run
meanwhile reports NO_OP.** Between the read (~287) and this write run
instance-partition computation, rebalanceTable (bulk ZK reads),
fetchTableSizeDetails (up to 30s) and the summary, so on a busy table any
commit makes `set(...)` throw ZkBadVersionException and the catch returns
FAILED, where the base returned NO_OP. Example: pre-PR upsert table, segments
balanced, an idle or paused partition without entry: a tenant rebalance or REST
downtime=true rebalance fails and every rerun races again. The default REST
flow runs a dry-run first, gets NO_OP and stops, so the PR's 'first rebalance
populates them' never happens through it. Consider routing partition-only
changes through the retrying main loop.
--
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]
