diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java index d264fb320315..ef4dd7361cfb 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java @@ -142,6 +142,10 @@ public class TableRebalancer { // TODO: Consider making the timeoutMs below table rebalancer configurable private static final int TABLE_SIZE_READER_TIMEOUT_MS = 30_000; private static final int STREAM_PARTITION_OFFSET_READ_TIMEOUT_MS = 10_000; + // Max number of times to rebase a single IdealState batch update onto a newer version and retry the compare-and-set + // in place (without waiting for the ExternalView to converge again) when the version changed only due to concurrent + // writes that do not touch the segments being moved. Beyond this, fall back to the regular convergence loop. + private static final int MAX_IDEAL_STATE_UPDATE_REBASE_ATTEMPTS = 5; private static final AtomicInteger REBALANCE_JOB_COUNTER = new AtomicInteger(0); private final HelixManager _helixManager; private final HelixDataAccessor _helixDataAccessor; @@ -801,27 +805,105 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb idealState.setNumPartitions(nextAssignment.size()); idealState.setReplicas(Integer.toString(nextAssignment.values().iterator().next().size())); - // Check version and update IdealState - try { - Preconditions.checkState(_helixDataAccessor.getBaseDataAccessor() - .set(idealStatePropertyKey.getPath(), idealStateRecord, expectedVersion, AccessOption.PERSISTENT), - "Failed to update IdealState"); + // Segments this batch changes relative to the current assignment it was computed against. Captured before the + // compare-and-set so that, on a version conflict, we can tell whether a concurrent write touched any segment + // this batch moves. + List batchMovedSegments = SegmentAssignmentUtils.getSegmentsToMove(currentAssignment, nextAssignment); + + // Check version and update the IdealState. If the compare-and-set fails only because a concurrent write bumped + // the version without touching the segments this batch moves (e.g. consuming segment commits on a continuously + // ingesting table), rebase this batch onto the latest IdealState and retry the compare-and-set in place, without + // waiting for the ExternalView to converge again (this batch never landed, so there is nothing new to wait for) + // or recomputing the full target assignment. This keeps the rebalance from live-locking against a steady stream + // of version bumps. Only attempted when this rebalance moves only tier segments: the base placements are then + // unchanged, so a segment added concurrently keeps its correct placement and can be carried over as-is. + boolean rebasable = isMovingOnlyTierSegments(segmentsToMove, providedTierToSegmentsMap); + boolean updated = false; + int rebaseAttempts = 0; + while (true) { + try { + Preconditions.checkState(_helixDataAccessor.getBaseDataAccessor() + .set(idealStatePropertyKey.getPath(), idealStateRecord, expectedVersion, AccessOption.PERSISTENT), + "Failed to update IdealState"); + updated = true; + break; + } catch (ZkBadVersionException e) { + if (!rebasable || rebaseAttempts >= MAX_IDEAL_STATE_UPDATE_REBASE_ATTEMPTS) { + break; + } + IdealState latestIdealState; + try { + latestIdealState = _helixDataAccessor.getProperty(idealStatePropertyKey); + } catch (Exception re) { + tableRebalanceLogger.warn("Failed to re-read IdealState for rebasing after a version conflict", re); + break; + } + if (latestIdealState == null) { + break; + } + Map> latestAssignment = latestIdealState.getRecord().getMapFields(); + boolean concurrentChangeTouchesBatch = false; + for (String segment : batchMovedSegments) { + if (!Objects.equals(currentAssignment.get(segment), latestAssignment.get(segment))) { + concurrentChangeTouchesBatch = true; + break; + } + } + if (concurrentChangeTouchesBatch) { + break; + } + // Rebase the batch onto the latest IdealState, preserving concurrent changes to the other segments. + Map> rebasedAssignment = new TreeMap<>(latestAssignment); + for (String segment : batchMovedSegments) { + rebasedAssignment.put(segment, nextAssignment.get(segment)); + } + nextAssignment = rebasedAssignment; + idealState = latestIdealState; + idealStateRecord = latestIdealState.getRecord(); + idealStateRecord.setMapFields(nextAssignment); + idealState.setNumPartitions(nextAssignment.size()); + idealState.setReplicas(Integer.toString(nextAssignment.values().iterator().next().size())); + expectedVersion = idealStateRecord.getVersion(); + rebaseAttempts++; + tableRebalanceLogger.info("Rebasing IdealState update onto version {} after a concurrent change that does " + + "not affect the segments being moved (attempt {})", expectedVersion, rebaseAttempts); + } catch (Exception e) { + onReturnFailure("Caught exception while updating IdealState, aborting the rebalance", e, + tableRebalanceLogger); + return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, + "Caught exception while updating IdealState: " + e, instancePartitionsMap, tierToInstancePartitionsMap, + targetAssignment, preChecksResult, summaryResult); + } + } + + if (updated) { expectedVersion++; currentAssignment = nextAssignment; + if (rebaseAttempts > 0) { + // A rebase adopts the latest IdealState, which may have segments added (e.g. new consuming segments) or + // removed (e.g. by retention) concurrently that the target assignment, computed earlier, does not match. + // Since a rebase is only done when the rebalance moves only tier segments, added segments keep their current + // placement, so rebuild the target from the adopted current assignment while preserving the target only for + // the segments still moving that still exist. This keeps the convergence check + // (currentAssignment.equals(targetAssignment)) well-defined and its key set aligned with the current + // assignment. + Map> refreshedTarget = new TreeMap<>(currentAssignment); + for (String segment : SegmentAssignmentUtils.getSegmentsToMove(currentAssignment, targetAssignment)) { + if (currentAssignment.containsKey(segment)) { + refreshedTarget.put(segment, targetAssignment.get(segment)); + } + } + targetAssignment = refreshedTarget; + } segmentsToMove = SegmentAssignmentUtils.getSegmentsToMove(currentAssignment, targetAssignment); // IdealState update is successful. Update the segment list as the IDEAL_STATE_CHANGE_TRIGGER should have // captured the newly added / deleted segments allSegmentsFromIdealState = currentAssignment.keySet(); tableRebalanceLogger.info("Successfully updated the IdealState"); - } catch (ZkBadVersionException e) { + } else { tableRebalanceLogger.info("Version changed while updating IdealState"); // Since IdealState wasn't updated, rollback the stats changes made _tableRebalanceObserver.onRollback(); - } catch (Exception e) { - onReturnFailure("Caught exception while updating IdealState, aborting the rebalance", e, tableRebalanceLogger); - return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, - "Caught exception while updating IdealState: " + e, instancePartitionsMap, tierToInstancePartitionsMap, - targetAssignment, preChecksResult, summaryResult); } } }