diff --git a/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/RebalanceStatus.java b/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/RebalanceStatus.java index 5edbafd0964..0434940ea85 100644 --- a/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/RebalanceStatus.java +++ b/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/RebalanceStatus.java @@ -38,7 +38,7 @@ public enum RebalanceStatus { TIMEOUT(5); public static final Set FINAL_STATUSES = - new HashSet<>(Arrays.asList(COMPLETED, CANCELED, FAILED, TIMEOUT)); + new HashSet<>(Arrays.asList(COMPLETED, CANCELED, FAILED)); private final int code; diff --git a/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java b/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java index 506ae476c2c..f9b769aee6b 100644 --- a/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java +++ b/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java @@ -241,6 +241,39 @@ public class ConfigOptions { + "TableLifecycleThrottler scans in-flight drops for " + "timeouts."); + public static final ConfigOption COORDINATOR_REBALANCE_TARGET_UNAVAILABLE_TIMEOUT = + key("coordinator.rebalance.target-unavailable-timeout") + .durationType() + .defaultValue(Duration.ofMinutes(30)) + .withDescription( + "The time after which a timed-out bucket migration is marked FAILED " + + "if its target tablet servers remain unavailable and no " + + "bucket-state progress is observed. Observed progress or all " + + "target servers becoming live resets this timer. Must be at " + + "least 1 ms and fit in a signed 64-bit millisecond value."); + + public static final ConfigOption COORDINATOR_REBALANCE_NO_PROGRESS_TIMEOUT = + key("coordinator.rebalance.no-progress-timeout") + .durationType() + .defaultValue(Duration.ofHours(24)) + .withDescription( + "The time after which a timed-out bucket migration is marked FAILED " + + "if no bucket-state progress is observed, even when its target " + + "tablet servers are live. Observed progress resets this timer; " + + "this does not limit the total migration duration. Must be at " + + "least 1 ms and fit in a signed 64-bit millisecond value."); + + public static final ConfigOption COORDINATOR_REBALANCE_MAX_TRACKED_TIMED_OUT_TASKS = + key("coordinator.rebalance.max-tracked-timed-out-tasks") + .intType() + .defaultValue(8) + .withDescription( + "The maximum number of non-final timed-out bucket migrations tracked " + + "before the coordinator stops admitting new bucket migrations. " + + "Admission resumes when a tracked migration reaches a final " + + "status. This limits concurrent migrations and reconciliation " + + "work on the coordinator and ZooKeeper. Must be at least 1."); + public static final ConfigOption COORDINATOR_OFFLINE_LEADER_RETRY_DELAY = key("coordinator.offline-leader.retry-delay") .durationType() diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java index 2df749781b0..027237fb6a2 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java @@ -82,14 +82,17 @@ import org.apache.fluss.server.coordinator.event.DropTableEvent; import org.apache.fluss.server.coordinator.event.EventProcessor; import org.apache.fluss.server.coordinator.event.FencedCoordinatorEvent; +import org.apache.fluss.server.coordinator.event.FinalizeRebalanceEvent; import org.apache.fluss.server.coordinator.event.ListRebalanceProgressEvent; import org.apache.fluss.server.coordinator.event.NewCoordinatorEvent; import org.apache.fluss.server.coordinator.event.NewTabletServerEvent; import org.apache.fluss.server.coordinator.event.NotifyKvSnapshotOffsetEvent; import org.apache.fluss.server.coordinator.event.NotifyLakeTableOffsetEvent; +import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrRequestContext; import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrResponseReceivedEvent; import org.apache.fluss.server.coordinator.event.RebalanceEvent; import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent; +import org.apache.fluss.server.coordinator.event.ReconcileRebalanceTaskEvent; import org.apache.fluss.server.coordinator.event.RecoverRebalanceEvent; import org.apache.fluss.server.coordinator.event.RemoveServerTagByRackEvent; import org.apache.fluss.server.coordinator.event.RemoveServerTagEvent; @@ -101,6 +104,7 @@ import org.apache.fluss.server.coordinator.event.watcher.TableChangeWatcher; import org.apache.fluss.server.coordinator.event.watcher.TabletServerChangeWatcher; import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager; +import org.apache.fluss.server.coordinator.rebalance.RebalanceExecutor; import org.apache.fluss.server.coordinator.rebalance.RebalanceManager; import org.apache.fluss.server.coordinator.statemachine.ReplicaLeaderElection.ControlledShutdownLeaderElection; import org.apache.fluss.server.coordinator.statemachine.ReplicaLeaderElection.ReassignmentLeaderElection; @@ -132,7 +136,6 @@ import org.apache.fluss.server.zk.data.lake.LakeTableSnapshot; import org.apache.fluss.utils.AutoPartitionStrategy; import org.apache.fluss.utils.clock.Clock; -import org.apache.fluss.utils.clock.SystemClock; import org.apache.fluss.utils.concurrent.Scheduler; import org.apache.fluss.utils.types.Tuple2; @@ -173,7 +176,7 @@ /** An implementation for {@link EventProcessor}. */ @NotThreadSafe -public class CoordinatorEventProcessor implements EventProcessor { +public class CoordinatorEventProcessor implements EventProcessor, RebalanceExecutor { private static final Logger LOG = LoggerFactory.getLogger(CoordinatorEventProcessor.class); @@ -225,22 +228,17 @@ public CoordinatorEventProcessor( this.coordinatorContext = coordinatorContext; this.replicaCapacityController = replicaCapacityController; this.coordinatorEventManager = new CoordinatorEventManager(this, coordinatorMetricGroup); + CoordinatorRequestBatch replicaRequestBatch = + new CoordinatorRequestBatch( + coordinatorChannelManager, coordinatorEventManager, coordinatorContext); this.replicaStateMachine = - new ReplicaStateMachine( - coordinatorContext, - new CoordinatorRequestBatch( - coordinatorChannelManager, - coordinatorEventManager, - coordinatorContext), - zooKeeperClient); + new ReplicaStateMachine(coordinatorContext, replicaRequestBatch, zooKeeperClient); + CoordinatorRequestBatch bucketRequestBatch = + new CoordinatorRequestBatch( + coordinatorChannelManager, coordinatorEventManager, coordinatorContext); this.tableBucketStateMachine = new TableBucketStateMachine( - coordinatorContext, - new CoordinatorRequestBatch( - coordinatorChannelManager, - coordinatorEventManager, - coordinatorContext), - zooKeeperClient); + coordinatorContext, bucketRequestBatch, zooKeeperClient); this.metadataManager = metadataManager; this.lifecycleThrottler = new TableLifecycleThrottler(coordinatorEventManager, clock, conf); @@ -277,8 +275,10 @@ public CoordinatorEventProcessor( this.coordinatorMetricGroup = coordinatorMetricGroup; this.internalListenerName = conf.getString(ConfigOptions.INTERNAL_LISTENER_NAME); this.rebalanceManager = - new RebalanceManager( - this, zooKeeperClient, coordinatorEventManager, SystemClock.getInstance()); + new RebalanceManager(this, zooKeeperClient, coordinatorEventManager, clock, conf); + replicaRequestBatch.setRebalanceExecutionKeyResolver(rebalanceManager::getExecutionKey); + bucketRequestBatch.setRebalanceExecutionKeyResolver(rebalanceManager::getExecutionKey); + coordinatorRequestBatch.setRebalanceExecutionKeyResolver(rebalanceManager::getExecutionKey); this.offlineLeaderRetryDelayMs = conf.get(ConfigOptions.COORDINATOR_OFFLINE_LEADER_RETRY_DELAY).toMillis(); if (offlineLeaderRetryDelayMs <= 0) { @@ -340,7 +340,6 @@ public void startup() { // start rebalance manager. rebalanceManager.startup(); - rebalanceManager.start(); } public void shutdown() { @@ -767,12 +766,7 @@ public void process(CoordinatorEvent event) { completeFromCallable( rebalanceEvent.getRespCallback(), () -> processRebalance(rebalanceEvent)); } else if (event instanceof RecoverRebalanceEvent) { - RecoverRebalanceEvent recoverRebalanceEvent = (RecoverRebalanceEvent) event; - RebalanceTask rebalanceTask = recoverRebalanceEvent.getRebalanceTask(); - rebalanceManager.registerRebalance( - rebalanceTask.getRebalanceId(), - rebalanceTask.getExecutePlan(), - rebalanceTask.getRebalanceStatus()); + rebalanceManager.recoverRebalance(((RecoverRebalanceEvent) event).getRebalanceTask()); } else if (event instanceof CancelRebalanceEvent) { CancelRebalanceEvent cancelRebalanceEvent = (CancelRebalanceEvent) event; completeFromCallable( @@ -780,11 +774,21 @@ public void process(CoordinatorEvent event) { () -> processCancelRebalance(cancelRebalanceEvent)); } else if (event instanceof RebalanceTaskTimeoutEvent) { RebalanceTaskTimeoutEvent timeoutEvent = (RebalanceTaskTimeoutEvent) event; - LOG.warn( - "Rebalance task for {} timed out. Treating as timeout.", - timeoutEvent.getTableBucket()); - rebalanceManager.finishRebalanceTask( - timeoutEvent.getTableBucket(), RebalanceStatus.TIMEOUT); + if (rebalanceManager.timeoutRebalanceTask(timeoutEvent.getExecutionKey())) { + LOG.warn( + "Rebalance task {} timed out and remains under reconciliation.", + timeoutEvent.getExecutionKey()); + } + } else if (event instanceof ReconcileRebalanceTaskEvent) { + ReconcileRebalanceTaskEvent reconcileEvent = (ReconcileRebalanceTaskEvent) event; + RebalancePlanForBucket plan = + rebalanceManager.getPlanForReconciliation(reconcileEvent.getExecutionKey()); + if (plan != null) { + tryToExecuteRebalanceTask(plan); + } + } else if (event instanceof FinalizeRebalanceEvent) { + rebalanceManager.retryFinalizeRebalance( + ((FinalizeRebalanceEvent) event).getRebalanceId()); } else if (event instanceof ResumeDropEvent) { // Resume-mode reconciliation queued by TableLifecycleThrottler: dispatch on the event // thread so the @NotThreadSafe CoordinatorContext / state machines are only mutated @@ -1241,9 +1245,12 @@ private void processNotifyLeaderAndIsrResponseReceivedEvent( notifyLeaderAndIsrResponseReceivedEvent.getNotifyLeaderAndIsrResultForBuckets(); for (NotifyLeaderAndIsrResultForBucket notifyLeaderAndIsrResultForBucket : notifyLeaderAndIsrResultForBuckets) { + TableBucket tableBucket = notifyLeaderAndIsrResultForBucket.getTableBucket(); + // Replica liveness bookkeeping must not depend on how fresh the request was: a replica + // that failed to apply the state is offline even when the leader has bumped the bucket + // epoch in the meantime. // if the error code is not none, we will consider it as offline if (notifyLeaderAndIsrResultForBucket.failed()) { - TableBucket tableBucket = notifyLeaderAndIsrResultForBucket.getTableBucket(); if (coordinatorContext.isToBeDeleted(tableBucket) || !coordinatorContext.getAssignment(tableBucket).contains(serverId)) { // A late activation failure must not reset a deletion state to OfflineReplica @@ -1252,7 +1259,7 @@ private void processNotifyLeaderAndIsrResponseReceivedEvent( } offlineReplicas.add(new TableBucketReplica(tableBucket, serverId)); } else { - succeededBuckets.add(notifyLeaderAndIsrResultForBucket.getTableBucket()); + succeededBuckets.add(tableBucket); } } for (TableBucket tb : succeededBuckets) { @@ -1269,9 +1276,15 @@ private void processNotifyLeaderAndIsrResponseReceivedEvent( // Try to complete rebalance tasks for the buckets in the response. // This is essential for leader-only migrations to ensure they wait for the tablet // server to acknowledge the leader change before proceeding to the next migration. + // Only this part is gated on the response belonging to the current request, because only + // completing a rebalance task requires an acknowledgement of the state we last sent. for (NotifyLeaderAndIsrResultForBucket notifyLeaderAndIsrResultForBucket : notifyLeaderAndIsrResultForBuckets) { - tryToCompleteRebalanceTask(notifyLeaderAndIsrResultForBucket, serverId); + tryToCompleteRebalanceTask( + notifyLeaderAndIsrResultForBucket, + serverId, + notifyLeaderAndIsrResponseReceivedEvent.getRequestContext( + notifyLeaderAndIsrResultForBucket.getTableBucket())); } } @@ -1623,8 +1636,7 @@ private RebalanceResponse processRebalance(RebalanceEvent rebalanceEvent) { // 2. execute rebalance plan. Map executePlan = rebalanceTask.getExecutePlan(); zooKeeperClient.registerRebalanceTask(rebalanceTask); - rebalanceManager.registerRebalance( - rebalanceTask.getRebalanceId(), executePlan, RebalanceStatus.NOT_STARTED); + rebalanceManager.registerRebalance(rebalanceTask.getRebalanceId(), executePlan); } catch (Exception e) { throw new RebalanceFailureException( String.format( @@ -1664,70 +1676,82 @@ private ListRebalanceProgressResponse processListRebalanceProgress( * */ public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) { - Set allBuckets = coordinatorContext.getAllBuckets(); TableBucket tableBucket = planForBucket.getTableBucket(); - if (!allBuckets.contains(tableBucket)) { - LOG.warn( - "Skipping rebalance task of tableBucket {} since it doesn't exist.", - tableBucket); - rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED); - return; - } - - if (coordinatorContext.isTableQueuedForDeletion(tableBucket.getTableId())) { - LOG.warn( - "Skipping rebalance task of tableBucket {} since the respective " - + "tables are being deleted.", - tableBucket); - rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED); + if (completeRebalanceTaskIfDeleted(tableBucket)) { return; } List newReplicas = planForBucket.getNewReplicas(); ReplicaReassignment reassignment = - ReplicaReassignment.build( - coordinatorContext.getAssignment(tableBucket), newReplicas); - - if (planForBucket.isLeaderChanged() && !reassignment.isBeingReassigned()) { - // buckets only need to change leader like leader replica rebalance. - // Don't finish the task immediately; wait for the NotifyLeaderAndIsr response - // from the tablet server to confirm the leader change has been applied. - // This ensures leader migrations are executed sequentially, avoiding excessive - // pressure on tablet servers (especially for KV tables). - LOG.info("trigger leader election for tableBucket {}.", tableBucket); - tableBucketStateMachine.handleStateChange( - Collections.singleton(tableBucket), - OnlineBucket, - new ReassignmentLeaderElection(newReplicas)); - } else { - try { - LOG.info( - "Try to processing bucket reassignment for tableBucket {} with assignment: {}.", - tableBucket, - reassignment); - onBucketReassignment(tableBucket, reassignment, false); - } catch (Exception e) { - LOG.error("Error when processing bucket reassignment.", e); - rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED); + ReplicaReassignment.build(planForBucket.getOriginReplicas(), newReplicas); + + try { + if (reassignment.isBeingReassigned()) { + if (advanceReassignment(tableBucket, reassignment, true)) { + LOG.info( + "Re-send final state for rebalance task of {} and wait for leader ack.", + tableBucket); + sendCurrentLeaderAndIsrRequest(tableBucket, reassignment.getTargetReplicas()); + } + } else if (planForBucket.isLeaderChanged()) { + LeaderAndIsr leaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket).get(); + if (leaderAndIsr.leader() == planForBucket.getNewLeader()) { + // Recovery and retries must confirm the applied state without re-electing + // the same leader. A leader-only move does not reorder the assignment. + sendCurrentLeaderAndIsrRequest( + tableBucket, coordinatorContext.getAssignment(tableBucket)); + } else { + tableBucketStateMachine.handleStateChange( + Collections.singleton(tableBucket), + OnlineBucket, + new ReassignmentLeaderElection(newReplicas)); + } + } else { + rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.COMPLETED); } + } catch (Exception e) { + LOG.warn( + "Failed to process rebalance task for {}. It remains active for timeout " + + "reconciliation.", + tableBucket, + e); } } /** try to finish rebalance tasks after receive notify leader and isr response. */ private void tryToCompleteRebalanceTask( NotifyLeaderAndIsrResultForBucket notifyLeaderAndIsrResultForBucket, - int responseServerId) { + int responseServerId, + @Nullable NotifyLeaderAndIsrRequestContext requestContext) { TableBucket tableBucket = notifyLeaderAndIsrResultForBucket.getTableBucket(); + if (requestContext != null && !isCurrentRequestContext(tableBucket, requestContext)) { + LOG.debug( + "Ignore stale NotifyLeaderAndIsr response for {} with context {}.", + tableBucket, + requestContext); + return; + } RebalancePlanForBucket planForBucket = rebalanceManager.getRebalancePlanForBucket(tableBucket); if (planForBucket != null) { + if (requestContext != null + && !Objects.equals( + rebalanceManager.getExecutionKey(tableBucket), + requestContext.getRebalanceExecutionKey())) { + LOG.debug( + "Ignore NotifyLeaderAndIsr response from another rebalance attempt for {} " + + "with context {}.", + tableBucket, + requestContext); + return; + } ReplicaReassignment reassignment = ReplicaReassignment.build( planForBucket.getOriginReplicas(), planForBucket.getNewReplicas()); try { if (planForBucket.isLeaderChanged() && !reassignment.isBeingReassigned()) { - if (!isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader( - notifyLeaderAndIsrResultForBucket, responseServerId, planForBucket)) { + if (notifyLeaderAndIsrResultForBucket.failed() + || responseServerId != planForBucket.getNewLeader()) { return; } LeaderAndIsr leaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket).get(); @@ -1737,17 +1761,83 @@ private void tryToCompleteRebalanceTask( rebalanceManager.finishRebalanceTask( tableBucket, RebalanceStatus.COMPLETED); } - } else if (notifyLeaderAndIsrResultForBucket.succeeded()) { - tryToCompleteReassignmentTask(tableBucket, reassignment); + } else if (notifyLeaderAndIsrResultForBucket.succeeded() + && advanceReassignment(tableBucket, reassignment, false) + && isSuccessfulReassignmentResponseFromLeader( + responseServerId, requestContext)) { + rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.COMPLETED); } } catch (Exception e) { - LOG.error( - "Failed to complete the reassignment for table bucket {}", tableBucket, e); - rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED); + LOG.warn( + "Failed to inspect or complete the reassignment for table bucket {}. It " + + "remains active for timeout reconciliation.", + tableBucket, + e); } } } + private boolean isCurrentRequestContext( + TableBucket tableBucket, NotifyLeaderAndIsrRequestContext requestContext) { + Optional leaderAndIsrOpt = + coordinatorContext.getBucketLeaderAndIsr(tableBucket); + if (!leaderAndIsrOpt.isPresent()) { + return false; + } + LeaderAndIsr leaderAndIsr = leaderAndIsrOpt.get(); + return coordinatorContext.getCoordinatorEpoch() == requestContext.getCoordinatorEpoch() + && leaderAndIsr.leader() == requestContext.getLeader() + && leaderAndIsr.leaderEpoch() == requestContext.getLeaderEpoch() + && leaderAndIsr.bucketEpoch() == requestContext.getBucketEpoch(); + } + + /** Returns whether a persisted plan is still at its clean origin state. */ + public boolean isRebalanceTaskAtOrigin(RebalancePlanForBucket planForBucket) { + TableBucket tableBucket = planForBucket.getTableBucket(); + Optional leaderAndIsrOpt = + coordinatorContext.getBucketLeaderAndIsr(tableBucket); + if (!leaderAndIsrOpt.isPresent() + || !coordinatorContext + .getAssignment(tableBucket) + .equals(planForBucket.getOriginReplicas())) { + return false; + } + LeaderAndIsr leaderAndIsr = leaderAndIsrOpt.get(); + // Losing an original follower from ISR does not mean the migration has started. Only an + // ISR member outside the origin indicates leftover work that cancellation must drain. + if (!new HashSet<>(planForBucket.getOriginReplicas()).containsAll(leaderAndIsr.isr())) { + return false; + } + ReplicaReassignment reassignment = + ReplicaReassignment.build( + planForBucket.getOriginReplicas(), planForBucket.getNewReplicas()); + // An unrelated leader election is not evidence that a replica migration started. A + // leader-only plan has no intermediate state to drain either: only once its target leader + // is elected does it have an applied state that still needs the leader acknowledgement. + return reassignment.isBeingReassigned() + || leaderAndIsr.leader() != planForBucket.getNewLeader(); + } + + @Override + public boolean isRebalanceTaskDeleted(TableBucket tableBucket) { + return coordinatorContext.getAssignment(tableBucket).isEmpty() + || coordinatorContext.isTableQueuedForDeletion(tableBucket.getTableId()) + || coordinatorContext.isToBeDeleted(tableBucket); + } + + private boolean completeRebalanceTaskIfDeleted(TableBucket tableBucket) { + if (!isRebalanceTaskDeleted(tableBucket)) { + return false; + } + + LOG.info( + "Complete rebalance task {} because its table bucket was deleted or is being " + + "deleted.", + tableBucket); + rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.COMPLETED); + return true; + } + private void tryToCompleteRebalanceTaskOnLeaderAndIsrChange(TableBucket tableBucket) { RebalancePlanForBucket planForBucket = rebalanceManager.getRebalancePlanForBucket(tableBucket); @@ -1759,34 +1849,73 @@ private void tryToCompleteRebalanceTaskOnLeaderAndIsrChange(TableBucket tableBuc return; } try { - tryToCompleteReassignmentTask(tableBucket, reassignment); + if (advanceReassignment(tableBucket, reassignment, false)) { + // ZK and ISR already reached the target, but completion still requires a + // successful response to a current final-state request. Re-send without + // changing bucket epoch. + sendCurrentLeaderAndIsrRequest(tableBucket, reassignment.getTargetReplicas()); + } } catch (Exception e) { - LOG.error( - "Failed to complete the reassignment for table bucket {}", tableBucket, e); - rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED); + LOG.warn( + "Failed to complete the reassignment for table bucket {}. It remains " + + "active for timeout reconciliation.", + tableBucket, + e); } } } - private void tryToCompleteReassignmentTask( - TableBucket tableBucket, ReplicaReassignment reassignment) throws Exception { - boolean isReassignmentComplete = isReassignmentComplete(tableBucket, reassignment); - if (isReassignmentComplete) { + /** + * Moves a replica reassignment forward from its persisted state without replaying finished + * steps. Returns true, without sending anything, if the final state has already been reached; + * the caller decides whether to complete the task or re-send that state. + * + *

Only execution and timeout reconciliation ({@code reconcile}) (re)start phase A or probe a + * persisted phase B whose targets have not caught up. A successful Notify response alone must + * not start an immediate resend loop while a target is still offline; responses and leader/ISR + * changes only advance caught-up targets. + */ + private boolean advanceReassignment( + TableBucket tableBucket, ReplicaReassignment reassignment, boolean reconcile) + throws Exception { + List targetReplicas = reassignment.getTargetReplicas(); + LeaderAndIsr leaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket).get(); + boolean targetsCaughtUp = isReassignmentComplete(leaderAndIsr, targetReplicas); + if (isReplicaAssignmentAtTarget(tableBucket, targetReplicas)) { + if (targetsCaughtUp && isFinalReassignmentState(leaderAndIsr, targetReplicas)) { + return true; + } + if (targetsCaughtUp) { + // Phase B persisted the target assignment before an acknowledgement was observed. + // Do not replay replica deletion. Repair any remaining ISR/leader transition + // directly and re-send the final state. + resumePersistedPhaseB(tableBucket, targetReplicas, leaderAndIsr); + } else if (reconcile) { + sendCurrentLeaderAndIsrRequest(tableBucket, targetReplicas); + } + } else if (targetsCaughtUp) { LOG.info( "Target replicas {} have all caught up with the leader for reassigning bucket {}", - reassignment.getTargetReplicas(), + targetReplicas, tableBucket); onBucketReassignment(tableBucket, reassignment, true); + } else if (reconcile) { + if (coordinatorContext.getAssignment(tableBucket).equals(reassignment.replicas)) { + LOG.info( + "Retry phase A for rebalance task of {} without advancing epoch.", + tableBucket); + retryBucketReassignmentPhaseA(tableBucket, reassignment); + } else { + LOG.info("Start phase A for rebalance task of {}.", tableBucket); + onBucketReassignment(tableBucket, reassignment, false); + } } + return false; } - @VisibleForTesting - static boolean isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader( - NotifyLeaderAndIsrResultForBucket notifyLeaderAndIsrResultForBucket, - int responseServerId, - RebalancePlanForBucket planForBucket) { - return notifyLeaderAndIsrResultForBucket.succeeded() - && responseServerId == planForBucket.getNewLeader(); + private static boolean isSuccessfulReassignmentResponseFromLeader( + int responseServerId, @Nullable NotifyLeaderAndIsrRequestContext requestContext) { + return requestContext != null && responseServerId == requestContext.getLeader(); } /** @@ -1833,7 +1962,8 @@ static boolean isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader( *

  • B7. Update ZK with RS=TRS, AR=[], RR=[]. *
  • B8. After electing leader, the replicas and isr information changes. So resend the * update metadata request to every tabletServer. - *
  • B8. Mark the ongoing rebalance task to finish. + *
  • B9. Wait for the current leader to acknowledge the final LeaderAndIsr state, then mark + * the ongoing rebalance task as finished. * * *

    In general, there are two goals we want to aim for: @@ -1918,17 +2048,95 @@ private void onBucketReassignment( null, null, Collections.singleton(tableBucket)); - // B8. Mark the ongoing rebalance task to finish. - rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.COMPLETED); + // Completion is intentionally deferred until the current leader acknowledges the + // final NotifyLeaderAndIsr state. This prevents a target tablet-server failure during + // phase B from being recorded as a successful rebalance. } } - private boolean isReassignmentComplete( + private void retryBucketReassignmentPhaseA( TableBucket tableBucket, ReplicaReassignment reassignment) throws Exception { - LeaderAndIsr leaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket).get(); - List isr = leaderAndIsr.isr(); - List targetReplicas = reassignment.getTargetReplicas(); - return targetReplicas.isEmpty() || new HashSet<>(isr).containsAll(targetReplicas); + // Persist the idempotent union assignment again in case the previous attempt updated + // memory but failed before its ZK write. Re-sending the current LeaderAndIsr must not bump + // bucket epoch: otherwise a slow AdjustIsr response can be fenced forever by retries. + coordinatorContext.updateBucketReplicaAssignment(tableBucket, reassignment.replicas); + updateReplicaAssignmentForBucket(tableBucket, reassignment.replicas); + for (Integer replica : reassignment.addingReplicas) { + TableBucketReplica tableBucketReplica = new TableBucketReplica(tableBucket, replica); + if (coordinatorContext.getReplicaState(tableBucketReplica) == null + || coordinatorContext.getReplicaState(tableBucketReplica) + == NonExistentReplica) { + replicaStateMachine.handleStateChanges( + Collections.singleton(tableBucketReplica), NewReplica); + } + } + sendCurrentLeaderAndIsrRequest(tableBucket, reassignment.replicas); + } + + private void resumePersistedPhaseB( + TableBucket tableBucket, List targetReplicas, LeaderAndIsr currentLeaderAndIsr) + throws Exception { + LeaderAndIsr leaderAndIsr = currentLeaderAndIsr; + if (!targetReplicas.contains(leaderAndIsr.leader())) { + tableBucketStateMachine.handleStateChange( + Collections.singleton(tableBucket), + OnlineBucket, + new ReassignmentLeaderElection(targetReplicas)); + leaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket).get(); + } + + Set targetReplicaSet = new HashSet<>(targetReplicas); + if (!new HashSet<>(leaderAndIsr.isr()).equals(targetReplicaSet)) { + List targetIsr = + leaderAndIsr.isr().stream() + .filter(targetReplicaSet::contains) + .collect(Collectors.toList()); + LeaderAndIsr newLeaderAndIsr = leaderAndIsr.newLeaderAndIsr(targetIsr); + zooKeeperClient.updateLeaderAndIsr( + tableBucket, newLeaderAndIsr, coordinatorContext.getCoordinatorZkVersion()); + coordinatorContext.putBucketLeaderAndIsr(tableBucket, newLeaderAndIsr); + } + sendCurrentLeaderAndIsrRequest(tableBucket, targetReplicas); + } + + private boolean isReassignmentComplete( + LeaderAndIsr leaderAndIsr, List targetReplicas) { + return coordinatorContext.liveTabletServerSet().containsAll(targetReplicas) + && (targetReplicas.isEmpty() + || new HashSet<>(leaderAndIsr.isr()).containsAll(targetReplicas)); + } + + private static boolean isFinalReassignmentState( + LeaderAndIsr leaderAndIsr, List targetReplicas) { + return targetReplicas.contains(leaderAndIsr.leader()) + && new HashSet<>(leaderAndIsr.isr()).equals(new HashSet<>(targetReplicas)); + } + + private boolean isReplicaAssignmentAtTarget( + TableBucket tableBucket, List targetReplicas) throws Exception { + return coordinatorContext.getAssignment(tableBucket).equals(targetReplicas) + && isReplicaAssignmentPersisted(tableBucket, targetReplicas); + } + + private boolean isReplicaAssignmentPersisted( + TableBucket tableBucket, List expectedReplicas) throws Exception { + BucketAssignment bucketAssignment; + if (tableBucket.getPartitionId() == null) { + Optional assignment = + zooKeeperClient.getTableAssignment(tableBucket.getTableId()); + bucketAssignment = + assignment.isPresent() + ? assignment.get().getBucketAssignment(tableBucket.getBucket()) + : null; + } else { + Optional assignment = + zooKeeperClient.getPartitionAssignment(tableBucket.getPartitionId()); + bucketAssignment = + assignment.isPresent() + ? assignment.get().getBucketAssignment(tableBucket.getBucket()) + : null; + } + return bucketAssignment != null && bucketAssignment.getReplicas().equals(expectedReplicas); } private void maybeReassignedBucketLeaderIfRequired( @@ -2641,6 +2849,29 @@ private void updateBucketEpochAndSendRequest(TableBucket tableBucket, List replicas) + throws Exception { + Optional leaderAndIsrOpt = zooKeeperClient.getLeaderAndIsr(tableBucket); + if (!leaderAndIsrOpt.isPresent()) { + return; + } + LeaderAndIsr leaderAndIsr = leaderAndIsrOpt.get(); + coordinatorContext.putBucketLeaderAndIsr(tableBucket, leaderAndIsr); + sendLeaderAndIsrRequest(tableBucket, replicas, leaderAndIsr); + } + + private void sendLeaderAndIsrRequest( + TableBucket tableBucket, List replicas, LeaderAndIsr leaderAndIsr) { String partitionName = null; if (tableBucket.getPartitionId() != null) { partitionName = coordinatorContext.getPartitionName(tableBucket.getPartitionId()); @@ -2650,22 +2881,15 @@ private void updateBucketEpochAndSendRequest(TableBucket tableBucket, List(newReplicas), + new HashSet<>(replicas), PhysicalTablePath.of( coordinatorContext.getTablePathById(tableBucket.getTableId()), partitionName), tableBucket, - newReplicas, - newLeaderAndIsr); + replicas, + leaderAndIsr); coordinatorRequestBatch.sendRequestToTabletServers( coordinatorContext.getCoordinatorEpoch()); } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatch.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatch.java index b4bf7e90d41..eb91b5c096e 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatch.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatch.java @@ -38,7 +38,9 @@ import org.apache.fluss.server.coordinator.event.AccessContextEvent; import org.apache.fluss.server.coordinator.event.DeleteReplicaResponseReceivedEvent; import org.apache.fluss.server.coordinator.event.EventManager; +import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrRequestContext; import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrResponseReceivedEvent; +import org.apache.fluss.server.coordinator.rebalance.RebalanceExecutionKey; import org.apache.fluss.server.entity.DeleteReplicaResultForBucket; import org.apache.fluss.server.entity.NotifyLeaderAndIsrData; import org.apache.fluss.server.metadata.BucketMetadata; @@ -60,6 +62,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.function.Function; import java.util.stream.Collectors; import static org.apache.fluss.server.metadata.PartitionMetadata.DELETED_PARTITION_ID; @@ -114,6 +117,8 @@ public class CoordinatorRequestBatch { private final CoordinatorChannelManager coordinatorChannelManager; private final EventManager eventManager; private final CoordinatorContext coordinatorContext; + private Function rebalanceExecutionKeyResolver = + ignored -> null; public CoordinatorRequestBatch( CoordinatorChannelManager coordinatorChannelManager, @@ -124,6 +129,11 @@ public CoordinatorRequestBatch( this.coordinatorContext = coordinatorContext; } + void setRebalanceExecutionKeyResolver( + Function rebalanceExecutionKeyResolver) { + this.rebalanceExecutionKeyResolver = rebalanceExecutionKeyResolver; + } + public void newBatch() { if (!notifyLeaderAndIsrRequestMap.isEmpty()) { throw new IllegalStateException( @@ -456,6 +466,19 @@ private void sendNotifyLeaderAndIsrRequest(int coordinatorEpoch) { NotifyLeaderAndIsrRequest notifyLeaderAndIsrRequest = makeNotifyLeaderAndIsrRequest( coordinatorEpoch, notifyRequestEntry.getValue().values()); + Map requestContexts = new HashMap<>(); + for (Map.Entry entry : + notifyRequestEntry.getValue().entrySet()) { + PbNotifyLeaderAndIsrReqForBucket request = entry.getValue(); + requestContexts.put( + entry.getKey(), + new NotifyLeaderAndIsrRequestContext( + coordinatorEpoch, + request.getLeader(), + request.getLeaderEpoch(), + request.getBucketEpoch(), + rebalanceExecutionKeyResolver.apply(entry.getKey()))); + } // Track exactly which buckets THIS request marked as pending leader activation. Only // those entries (where leader == serverId) need to be cleared if the request fails @@ -504,7 +527,9 @@ private void sendNotifyLeaderAndIsrRequest(int coordinatorEpoch) { // put the response receive event into the event manager eventManager.put( new NotifyLeaderAndIsrResponseReceivedEvent( - getNotifyLeaderAndIsrResponseData(response), serverId)); + getNotifyLeaderAndIsrResponseData(response), + serverId, + requestContexts)); }); } notifyLeaderAndIsrRequestMap.clear(); diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/FinalizeRebalanceEvent.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/FinalizeRebalanceEvent.java new file mode 100644 index 00000000000..1e4ca33c865 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/FinalizeRebalanceEvent.java @@ -0,0 +1,36 @@ +/* + * 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.fluss.server.coordinator.event; + +/** An event that retries persisting the final state of a rebalance. */ +public final class FinalizeRebalanceEvent implements CoordinatorEvent { + private final String rebalanceId; + + public FinalizeRebalanceEvent(String rebalanceId) { + this.rebalanceId = rebalanceId; + } + + public String getRebalanceId() { + return rebalanceId; + } + + @Override + public String toString() { + return "FinalizeRebalanceEvent{rebalanceId='" + rebalanceId + "'}"; + } +} diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/NotifyLeaderAndIsrRequestContext.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/NotifyLeaderAndIsrRequestContext.java new file mode 100644 index 00000000000..90cdd0a4b0a --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/NotifyLeaderAndIsrRequestContext.java @@ -0,0 +1,86 @@ +/* + * 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.fluss.server.coordinator.event; + +import org.apache.fluss.server.coordinator.rebalance.RebalanceExecutionKey; + +import javax.annotation.Nullable; + +/** The coordinator and bucket epochs of a sent NotifyLeaderAndIsr request. */ +public final class NotifyLeaderAndIsrRequestContext { + + private final int coordinatorEpoch; + private final int leader; + private final int leaderEpoch; + private final int bucketEpoch; + private final @Nullable RebalanceExecutionKey rebalanceExecutionKey; + + public NotifyLeaderAndIsrRequestContext( + int coordinatorEpoch, int leader, int leaderEpoch, int bucketEpoch) { + this(coordinatorEpoch, leader, leaderEpoch, bucketEpoch, null); + } + + public NotifyLeaderAndIsrRequestContext( + int coordinatorEpoch, + int leader, + int leaderEpoch, + int bucketEpoch, + @Nullable RebalanceExecutionKey rebalanceExecutionKey) { + this.coordinatorEpoch = coordinatorEpoch; + this.leader = leader; + this.leaderEpoch = leaderEpoch; + this.bucketEpoch = bucketEpoch; + this.rebalanceExecutionKey = rebalanceExecutionKey; + } + + public int getCoordinatorEpoch() { + return coordinatorEpoch; + } + + public int getLeader() { + return leader; + } + + public int getLeaderEpoch() { + return leaderEpoch; + } + + public int getBucketEpoch() { + return bucketEpoch; + } + + public @Nullable RebalanceExecutionKey getRebalanceExecutionKey() { + return rebalanceExecutionKey; + } + + @Override + public String toString() { + return "NotifyLeaderAndIsrRequestContext{" + + "coordinatorEpoch=" + + coordinatorEpoch + + ", leader=" + + leader + + ", leaderEpoch=" + + leaderEpoch + + ", bucketEpoch=" + + bucketEpoch + + ", rebalanceExecutionKey=" + + rebalanceExecutionKey + + '}'; + } +} diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/NotifyLeaderAndIsrResponseReceivedEvent.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/NotifyLeaderAndIsrResponseReceivedEvent.java index fcf656225ba..a192c60f92a 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/NotifyLeaderAndIsrResponseReceivedEvent.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/NotifyLeaderAndIsrResponseReceivedEvent.java @@ -17,10 +17,14 @@ package org.apache.fluss.server.coordinator.event; +import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrRequest; import org.apache.fluss.server.entity.NotifyLeaderAndIsrResultForBucket; +import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; /** An event for receive the response of {@link NotifyLeaderAndIsrRequest} from tablet server. */ public class NotifyLeaderAndIsrResponseReceivedEvent implements CoordinatorEvent { @@ -30,11 +34,21 @@ public class NotifyLeaderAndIsrResponseReceivedEvent implements CoordinatorEvent // the server id that return the response private final int responseServerId; + private final Map requestContexts; + public NotifyLeaderAndIsrResponseReceivedEvent( List notifyLeaderAndIsrResultForBuckets, int responseServerId) { + this(notifyLeaderAndIsrResultForBuckets, responseServerId, Collections.emptyMap()); + } + + public NotifyLeaderAndIsrResponseReceivedEvent( + List notifyLeaderAndIsrResultForBuckets, + int responseServerId, + Map requestContexts) { this.notifyLeaderAndIsrResultForBuckets = notifyLeaderAndIsrResultForBuckets; this.responseServerId = responseServerId; + this.requestContexts = Collections.unmodifiableMap(new HashMap<>(requestContexts)); } public int getResponseServerId() { @@ -44,4 +58,8 @@ public int getResponseServerId() { public List getNotifyLeaderAndIsrResultForBuckets() { return notifyLeaderAndIsrResultForBuckets; } + + public NotifyLeaderAndIsrRequestContext getRequestContext(TableBucket tableBucket) { + return requestContexts.get(tableBucket); + } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RebalanceTaskTimeoutEvent.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RebalanceTaskTimeoutEvent.java index c5d962d02cd..a84de0696ca 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RebalanceTaskTimeoutEvent.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RebalanceTaskTimeoutEvent.java @@ -17,23 +17,23 @@ package org.apache.fluss.server.coordinator.event; -import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.server.coordinator.rebalance.RebalanceExecutionKey; /** An event fired when a rebalance task exceeds the timeout without completing. */ public class RebalanceTaskTimeoutEvent implements CoordinatorEvent { - private final TableBucket tableBucket; + private final RebalanceExecutionKey executionKey; - public RebalanceTaskTimeoutEvent(TableBucket tableBucket) { - this.tableBucket = tableBucket; + public RebalanceTaskTimeoutEvent(RebalanceExecutionKey executionKey) { + this.executionKey = executionKey; } - public TableBucket getTableBucket() { - return tableBucket; + public RebalanceExecutionKey getExecutionKey() { + return executionKey; } @Override public String toString() { - return "RebalanceTaskTimeoutEvent{tableBucket=" + tableBucket + "}"; + return "RebalanceTaskTimeoutEvent{executionKey=" + executionKey + "}"; } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/ReconcileRebalanceTaskEvent.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/ReconcileRebalanceTaskEvent.java new file mode 100644 index 00000000000..e37ed69e69c --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/ReconcileRebalanceTaskEvent.java @@ -0,0 +1,38 @@ +/* + * 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.fluss.server.coordinator.event; + +import org.apache.fluss.server.coordinator.rebalance.RebalanceExecutionKey; + +/** An event that reconciles a timed-out rebalance task against coordinator state. */ +public final class ReconcileRebalanceTaskEvent implements CoordinatorEvent { + private final RebalanceExecutionKey executionKey; + + public ReconcileRebalanceTaskEvent(RebalanceExecutionKey executionKey) { + this.executionKey = executionKey; + } + + public RebalanceExecutionKey getExecutionKey() { + return executionKey; + } + + @Override + public String toString() { + return "ReconcileRebalanceTaskEvent{executionKey=" + executionKey + "}"; + } +} diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RecoverRebalanceEvent.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RecoverRebalanceEvent.java index b891423513a..2d329ef6b4e 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RecoverRebalanceEvent.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RecoverRebalanceEvent.java @@ -19,9 +19,8 @@ import org.apache.fluss.server.zk.data.RebalanceTask; -/** Fired during startup to recover a rebalance task that was persisted in ZooKeeper. */ -public class RecoverRebalanceEvent implements CoordinatorEvent { - +/** An event that recovers and reconciles a persisted rebalance task. */ +public final class RecoverRebalanceEvent implements CoordinatorEvent { private final RebalanceTask rebalanceTask; public RecoverRebalanceEvent(RebalanceTask rebalanceTask) { diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceExecutionKey.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceExecutionKey.java new file mode 100644 index 00000000000..624d954f1a3 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceExecutionKey.java @@ -0,0 +1,70 @@ +/* + * 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.fluss.server.coordinator.rebalance; + +import org.apache.fluss.metadata.TableBucket; + +import java.util.Objects; + +/** Identifies the execution of a bucket-level task within a rebalance. */ +public final class RebalanceExecutionKey { + private final String rebalanceId; + private final TableBucket tableBucket; + + public RebalanceExecutionKey(String rebalanceId, TableBucket tableBucket) { + this.rebalanceId = rebalanceId; + this.tableBucket = tableBucket; + } + + public String getRebalanceId() { + return rebalanceId; + } + + public TableBucket getTableBucket() { + return tableBucket; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + RebalanceExecutionKey that = (RebalanceExecutionKey) o; + return Objects.equals(rebalanceId, that.rebalanceId) + && Objects.equals(tableBucket, that.tableBucket); + } + + @Override + public int hashCode() { + return Objects.hash(rebalanceId, tableBucket); + } + + @Override + public String toString() { + return "RebalanceExecutionKey{" + + "rebalanceId='" + + rebalanceId + + '\'' + + ", tableBucket=" + + tableBucket + + '}'; + } +} diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceExecutor.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceExecutor.java new file mode 100644 index 00000000000..e83732a3300 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceExecutor.java @@ -0,0 +1,38 @@ +/* + * 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.fluss.server.coordinator.rebalance; + +import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.server.coordinator.CoordinatorContext; + +/** Coordinator operations needed to execute and reconcile rebalance bucket plans. */ +public interface RebalanceExecutor { + + /** Returns the coordinator state used to build a cluster model. */ + CoordinatorContext getCoordinatorContext(); + + /** Starts or resumes one bucket plan. */ + void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket); + + /** Returns whether the plan remains at its clean origin state. */ + boolean isRebalanceTaskAtOrigin(RebalancePlanForBucket planForBucket); + + /** Returns whether the bucket was deleted or is being deleted. */ + boolean isRebalanceTaskDeleted(TableBucket tableBucket); +} diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java index ce95dd862a0..56164d860a8 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java @@ -23,12 +23,18 @@ import org.apache.fluss.cluster.rebalance.RebalanceResultForBucket; import org.apache.fluss.cluster.rebalance.RebalanceStatus; import org.apache.fluss.cluster.rebalance.ServerTag; +import org.apache.fluss.config.ConfigOption; +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.config.Configuration; import org.apache.fluss.exception.NoRebalanceInProgressException; +import org.apache.fluss.exception.RebalanceFailureException; import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableInfo; import org.apache.fluss.server.coordinator.CoordinatorContext; -import org.apache.fluss.server.coordinator.CoordinatorEventProcessor; import org.apache.fluss.server.coordinator.event.EventManager; +import org.apache.fluss.server.coordinator.event.FinalizeRebalanceEvent; import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent; +import org.apache.fluss.server.coordinator.event.ReconcileRebalanceTaskEvent; import org.apache.fluss.server.coordinator.event.RecoverRebalanceEvent; import org.apache.fluss.server.coordinator.rebalance.goal.Goal; import org.apache.fluss.server.coordinator.rebalance.goal.GoalOptimizer; @@ -47,9 +53,15 @@ import org.slf4j.LoggerFactory; import javax.annotation.Nullable; +import javax.annotation.concurrent.ThreadSafe; +import java.time.Duration; import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Comparator; import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Optional; @@ -58,24 +70,57 @@ import java.util.SortedSet; import java.util.TreeSet; import java.util.UUID; -import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import static org.apache.fluss.cluster.rebalance.RebalanceStatus.CANCELED; import static org.apache.fluss.cluster.rebalance.RebalanceStatus.COMPLETED; +import static org.apache.fluss.cluster.rebalance.RebalanceStatus.FAILED; import static org.apache.fluss.cluster.rebalance.RebalanceStatus.FINAL_STATUSES; import static org.apache.fluss.cluster.rebalance.RebalanceStatus.NOT_STARTED; import static org.apache.fluss.cluster.rebalance.RebalanceStatus.REBALANCING; +import static org.apache.fluss.cluster.rebalance.RebalanceStatus.TIMEOUT; +import static org.apache.fluss.server.coordinator.rebalance.goal.GoalOptimizerUtils.getDiff; import static org.apache.fluss.utils.Preconditions.checkArgument; -import static org.apache.fluss.utils.Preconditions.checkNotNull; +import static org.apache.fluss.utils.Preconditions.checkState; /** - * A rebalance manager to generate rebalance plan, and execution rebalance plan. + * Generates rebalance plans and tracks their bucket migrations through completion. * - *

    This manager can only be used in {@link CoordinatorEventProcessor} as a single threaded model. + *

    Bucket task state transitions: + * + * 

    + * NOT_STARTED --admit------------------------> REBALANCING
+ * REBALANCING --current leader acknowledges--> COMPLETED
+ * REBALANCING --execution timeout------------> TIMEOUT
+ * TIMEOUT     --reconcile with backoff-------> TIMEOUT
+ * TIMEOUT     --current leader acknowledges--> COMPLETED
+ * TIMEOUT     --unavailable/no-progress limit-> FAILED
+ * NOT_STARTED / REBALANCING / TIMEOUT
+ *             --cancel before migration-----> CANCELED
+ * Any non-final task --bucket deleted--------> COMPLETED
+ *
    + * + *

    Only one bucket occupies the normal execution slot. TIMEOUT releases that slot but remains + * non-final and tracked; reaching the configured tracking limit pauses admission of new buckets. + * Cancellation is persisted before local state changes, including which buckets were canceled. New + * pending tasks are canceled directly; recovered or admitted tasks must still be at their origin to + * be canceled. Recovery admits intermediate tasks into reconciliation up to the configured tracking + * limit, and still requires a current leader acknowledgement even when the persisted state already + * matches the target. + * + *

    The overall rebalance remains REBALANCING while buckets are pending, running, or timed out, + * and while its final state cannot be persisted. Once all buckets are final, any failed bucket + * yields FAILED; otherwise a cancellation request yields CANCELED, and all completed buckets yield + * COMPLETED. Final-state persistence is retried before another rebalance can be admitted. A failed + * bucket can retain an intermediate assignment for a subsequent rebalance to clean up. + * + *

    Rebalance progress is driven from the coordinator event loop, but the periodic {@link + * #checkTimeout()} runs on its own scheduler thread. All mutable state is therefore guarded by this + * manager's monitor. */ +@ThreadSafe public class RebalanceManager { private static final Logger LOG = LoggerFactory.getLogger(RebalanceManager.class); @@ -85,53 +130,68 @@ public class RebalanceManager { /** Hardcoded interval for the periodic timeout check: 30 seconds. */ private static final long TIMEOUT_CHECK_INTERVAL_MS = 30 * 1000L; + /** Hardcoded upper bound for the exponential reconciliation backoff: 5 minutes. */ + private static final long MAX_RECONCILE_BACKOFF_MS = 5 * 60 * 1000L; + private final ZooKeeperClient zkClient; - private final CoordinatorEventProcessor eventProcessor; + private final RebalanceExecutor rebalanceExecutor; private final EventManager eventManager; private final Clock clock; private final ScheduledExecutorService timeoutChecker; + private final long targetUnavailableTimeoutMs; + private final long noProgressTimeoutMs; + private final int maxTrackedTimedOutTasks; + + /** Tasks awaiting admission. */ + private final Queue pendingRebalanceTasks = new ArrayDeque<>(); - /** A queue of in progress table bucket to rebalance. */ - private final Queue inProgressRebalanceTasksQueue = new ArrayDeque<>(); + /** Recovered intermediate tasks awaiting bounded reconciliation admission. */ + private final Queue pendingRecoveredTasks = new ArrayDeque<>(); - /** A mapping from table bucket to rebalance status of pending and running tasks. */ + /** A mapping from table bucket to rebalance status of all non-final tasks. */ private final Map inProgressRebalanceTasks = - new ConcurrentHashMap<>(); + new HashMap<>(); + + /** Normally running tasks. This map contains at most one entry until concurrency is added. */ + private final Map runningRebalanceTasks = new HashMap<>(); + + /** Soft-timed-out tasks that no longer occupy the normal execution slot. */ + private final Map timedOutRebalanceTasks = new HashMap<>(); + + private final Set queuedTimeoutEvents = new HashSet<>(); + private final Set queuedReconcileEvents = new HashSet<>(); /** A mapping from table bucket to rebalance status of failed or completed tasks. */ private final Map finishedRebalanceTasks = - new ConcurrentHashMap<>(); + new HashMap<>(); private final GoalOptimizer goalOptimizer; - private volatile long registerTime; - private volatile @Nullable RebalanceStatus rebalanceStatus; - private volatile @Nullable String currentRebalanceId; - private volatile boolean isClosed = false; + private long registerTime; + private @Nullable RebalanceStatus rebalanceStatus; + private @Nullable String currentRebalanceId; + private boolean recoveryPending; - /** - * Timestamp when the current in-flight task was started, or -1 if idle. - * - *

    Write ordering contract (volatile publication idiom): always write {@code - * inflightTaskStartMs} BEFORE {@code inflightTaskBucket} when setting, and clear {@code - * inflightTaskBucket} BEFORE {@code inflightTaskStartMs} when resetting. The timeout checker - * reads in reverse order (bucket first, then startMs), ensuring it never observes a stale - * startMs paired with a new bucket. - */ - private volatile long inflightTaskStartMs = -1; + /** Pending tasks loaded from ZooKeeper may have started before the coordinator restarted. */ + private boolean recoveredRebalance; - /** The bucket of the current in-flight task, or null if idle. Acts as the "gate" variable. */ - private volatile @Nullable TableBucket inflightTaskBucket; + private boolean cancelRequested; + private boolean finalizationPending; + private boolean finalizationEventQueued; + private boolean isClosed = false; + /** Creates a manager with validated coordinator rebalance settings. */ public RebalanceManager( - CoordinatorEventProcessor eventProcessor, + RebalanceExecutor rebalanceExecutor, ZooKeeperClient zkClient, EventManager eventManager, - Clock clock) { + Clock clock, + Configuration conf) { this( - eventProcessor, + rebalanceExecutor, zkClient, eventManager, clock, + conf, // TODO: Reuse the CoordinatorServer shared scheduler for this lightweight // coordinator timeout checker instead of creating a component-owned scheduler. Executors.newScheduledThreadPool( @@ -140,12 +200,25 @@ public RebalanceManager( @VisibleForTesting RebalanceManager( - CoordinatorEventProcessor eventProcessor, + RebalanceExecutor rebalanceExecutor, ZooKeeperClient zkClient, EventManager eventManager, Clock clock, + Configuration conf, ScheduledExecutorService timeoutChecker) { - this.eventProcessor = eventProcessor; + this.targetUnavailableTimeoutMs = + getTimeoutMillis( + conf, ConfigOptions.COORDINATOR_REBALANCE_TARGET_UNAVAILABLE_TIMEOUT); + this.noProgressTimeoutMs = + getTimeoutMillis(conf, ConfigOptions.COORDINATOR_REBALANCE_NO_PROGRESS_TIMEOUT); + this.maxTrackedTimedOutTasks = + conf.get(ConfigOptions.COORDINATOR_REBALANCE_MAX_TRACKED_TIMED_OUT_TASKS); + checkArgument( + maxTrackedTimedOutTasks > 0, + "%s must be at least 1, but was %s.", + ConfigOptions.COORDINATOR_REBALANCE_MAX_TRACKED_TIMED_OUT_TASKS.key(), + maxTrackedTimedOutTasks); + this.rebalanceExecutor = rebalanceExecutor; this.zkClient = zkClient; this.eventManager = eventManager; this.clock = clock == null ? SystemClock.getInstance() : clock; @@ -153,13 +226,11 @@ public RebalanceManager( this.goalOptimizer = new GoalOptimizer(); } - public void startup() { + public synchronized void startup() { LOG.info("Start up rebalance manager."); - initialize(); - } - - /** Starts the periodic timeout checker. Call after {@link #startup()}. */ - public void start() { + // Keep initialization in one critical section. Otherwise the queued recovery event could + // acquire this manager's monitor before a separate checker-start call and delay coordinator + // startup while it scans a large persisted task. timeoutChecker.scheduleWithFixedDelay( this::checkTimeoutSafely, TIMEOUT_CHECK_INTERVAL_MS, @@ -169,9 +240,10 @@ public void start() { "RebalanceManager timeout checker started: timeoutMs={}, checkIntervalMs={}", REBALANCE_TASK_TIMEOUT_MS, TIMEOUT_CHECK_INTERVAL_MS); + initialize(); } - public @Nullable String getRebalanceId() { + public synchronized @Nullable String getRebalanceId() { return currentRebalanceId; } @@ -179,87 +251,126 @@ private void initialize() { try { zkClient.getRebalanceTask() .ifPresent( - rebalanceTask -> - eventManager.put(new RecoverRebalanceEvent(rebalanceTask))); + rebalanceTask -> { + recoveryPending = true; + eventManager.put(new RecoverRebalanceEvent(rebalanceTask)); + }); } catch (Exception e) { + // TODO: Retry loading the persisted rebalance task instead of skipping recovery. + recoveryPending = false; LOG.error( - "Failed to get rebalance plan from zookeeper, it will be treated as no" - + "rebalance tasks.", + "Failed to get rebalance plan from zookeeper. Skipping recovery and allowing " + + "new rebalance requests; any unfinished task will not be resumed.", e); } } - public void registerRebalance( - String rebalanceId, - Map rebalancePlan, - RebalanceStatus newStatus) { + /** Registers a new rebalance; persisted tasks are handled by {@link #recoverRebalance}. */ + public synchronized void registerRebalance( + String rebalanceId, Map rebalancePlan) { checkNotClosed(); - registerTime = System.currentTimeMillis(); - // first clear all exists tasks. - inProgressRebalanceTasks.clear(); - inProgressRebalanceTasksQueue.clear(); - finishedRebalanceTasks.clear(); - // Clear gate (bucket) first, then data (startMs). - inflightTaskBucket = null; - inflightTaskStartMs = -1; - - currentRebalanceId = rebalanceId; + resetRebalance(rebalanceId, false); if (rebalancePlan.isEmpty()) { - completeRebalance(); + finalizeRebalance(); return; } - rebalancePlan.forEach( - ((tableBucket, planForBucket) -> { - if (FINAL_STATUSES.contains(newStatus)) { - finishedRebalanceTasks.put( - tableBucket, RebalanceResultForBucket.of(planForBucket, newStatus)); - } else { - inProgressRebalanceTasksQueue.add(tableBucket); - inProgressRebalanceTasks.put( - tableBucket, - RebalanceResultForBucket.of(planForBucket, NOT_STARTED)); - } - })); - - if (!inProgressRebalanceTasksQueue.isEmpty()) { - // Trigger one rebalance task to execute. - rebalanceStatus = REBALANCING; - processNewRebalanceTask(); - } else { - rebalanceStatus = newStatus; + for (Map.Entry entry : rebalancePlan.entrySet()) { + addPendingTask(entry.getKey(), entry.getValue()); } + + rebalanceStatus = REBALANCING; + processNewRebalanceTask(); } - public void finishRebalanceTask(TableBucket tableBucket, RebalanceStatus statusForBucket) { + /** Recovers a persisted task by comparing every bucket plan with current coordinator state. */ + public synchronized void recoverRebalance(RebalanceTask rebalanceTask) { checkNotClosed(); - if (inProgressRebalanceTasksQueue.contains(tableBucket)) { - inProgressRebalanceTasksQueue.remove(tableBucket); - RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.remove(tableBucket); - checkNotNull(resultForBucket, "RebalanceResultForBucket is null."); - finishedRebalanceTasks.put( - tableBucket, - RebalanceResultForBucket.of(resultForBucket.plan(), statusForBucket)); - // Clear gate (bucket) first, then data (startMs). - inflightTaskBucket = null; - inflightTaskStartMs = -1; - LOG.info( - "Rebalance task {} in progress: {} tasks pending, {} completed.", - currentRebalanceId, - inProgressRebalanceTasksQueue.size(), - finishedRebalanceTasks.size()); + resetRebalance(rebalanceTask.getRebalanceId(), rebalanceTask.isCancelRequested()); + recoveredRebalance = true; + if (FINAL_STATUSES.contains(rebalanceTask.getRebalanceStatus())) { + for (Map.Entry entry : + rebalanceTask.getExecutePlan().entrySet()) { + finishedRebalanceTasks.put( + entry.getKey(), + RebalanceResultForBucket.of( + entry.getValue(), rebalanceTask.getRebalanceStatus())); + } + rebalanceStatus = rebalanceTask.getRebalanceStatus(); + return; + } - if (inProgressRebalanceTasksQueue.isEmpty()) { - // All rebalance tasks are completed. - completeRebalance(); + for (Map.Entry entry : + rebalanceTask.getExecutePlan().entrySet()) { + if (rebalanceTask.getCanceledBuckets().contains(entry.getKey())) { + finishedRebalanceTasks.put( + entry.getKey(), RebalanceResultForBucket.of(entry.getValue(), CANCELED)); + } else if (rebalanceExecutor.isRebalanceTaskAtOrigin(entry.getValue())) { + addPendingTask(entry.getKey(), entry.getValue()); } else { - // Trigger one rebalance task to execute. - processNewRebalanceTask(); + addRecoveredIntermediateTask(entry.getKey(), entry.getValue()); } } + if (cancelRequested) { + cancelPendingTasks(findTasksToCancel()); + } + + if (inProgressRebalanceTasks.isEmpty()) { + finalizeRebalance(); + } else { + rebalanceStatus = REBALANCING; + processNewRebalanceTask(); + } } - public @Nullable RebalanceProgress listRebalanceProgress(@Nullable String rebalanceId) { + public synchronized void finishRebalanceTask( + TableBucket tableBucket, RebalanceStatus statusForBucket) { + RebalanceExecutionKey executionKey = getExecutionKey(tableBucket); + if (executionKey != null) { + finishRebalanceTask(executionKey, statusForBucket); + } + } + + public synchronized boolean finishRebalanceTask( + RebalanceExecutionKey executionKey, RebalanceStatus statusForBucket) { + checkNotClosed(); + checkArgument(statusForBucket != TIMEOUT, "Use timeoutRebalanceTask for soft timeouts."); + RebalanceTaskAttempt attempt = findActiveAttempt(executionKey); + if (attempt == null) { + LOG.debug("Ignore stale completion for {}.", executionKey); + return false; + } + + TableBucket tableBucket = executionKey.getTableBucket(); + runningRebalanceTasks.remove(tableBucket); + timedOutRebalanceTasks.remove(tableBucket); + queuedTimeoutEvents.remove(executionKey); + queuedReconcileEvents.remove(executionKey); + RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.remove(tableBucket); + if (resultForBucket == null) { + return false; + } + finishedRebalanceTasks.put( + tableBucket, RebalanceResultForBucket.of(resultForBucket.plan(), statusForBucket)); + LOG.info( + "Rebalance {} progress: {} pending, {} running, {} timed out and tracking, " + + "{} finished.", + currentRebalanceId, + pendingRebalanceTasks.size(), + runningRebalanceTasks.size(), + timedOutRebalanceTasks.size(), + finishedRebalanceTasks.size()); + + if (inProgressRebalanceTasks.isEmpty()) { + finalizeRebalance(); + } else { + processNewRebalanceTask(); + } + return true; + } + + public synchronized @Nullable RebalanceProgress listRebalanceProgress( + @Nullable String rebalanceId) { checkNotClosed(); if (rebalanceId != null && currentRebalanceId != null @@ -285,9 +396,13 @@ public void finishRebalanceTask(TableBucket tableBucket, RebalanceStatus statusF currentRebalanceId, rebalanceStatus, 0.0, progressForBucketMap); } - public void cancelRebalance(@Nullable String rebalanceId) { + public synchronized void cancelRebalance(@Nullable String rebalanceId) { checkNotClosed(); + if (currentRebalanceId == null) { + return; + } + if (rebalanceId != null && currentRebalanceId != null && !rebalanceId.equals(currentRebalanceId)) { @@ -306,45 +421,96 @@ public void cancelRebalance(@Nullable String rebalanceId) { return; } + Set tasksToCancel = findTasksToCancel(); + Set canceledBuckets = canceledBuckets(); + canceledBuckets.addAll(tasksToCancel); try { - Optional rebalanceTaskOpt = zkClient.getRebalanceTask(); - if (rebalanceTaskOpt.isPresent()) { - RebalanceTask rebalanceTask = rebalanceTaskOpt.get(); - zkClient.registerRebalanceTask( - new RebalanceTask( - rebalanceTask.getRebalanceId(), - CANCELED, - rebalanceTask.getExecutePlan())); - } + zkClient.registerRebalanceTask( + new RebalanceTask( + currentRebalanceId, + REBALANCING, + allRebalancePlans(), + true, + canceledBuckets)); } catch (Exception e) { - LOG.error("Error when delete rebalance plan from zookeeper.", e); + throw new RebalanceFailureException( + "Failed to persist rebalance cancellation request.", e); } - rebalanceStatus = CANCELED; - inProgressRebalanceTasksQueue.clear(); - inProgressRebalanceTasks.clear(); - // Clear gate (bucket) first, then data (startMs). - inflightTaskBucket = null; - inflightTaskStartMs = -1; - // Here, it will not clear finishedRebalanceTasks, because it will be used by - // listRebalanceProgress. It will be cleared when next register. + cancelRequested = true; + cancelPendingTasks(tasksToCancel); + + // Admitted tasks that have not changed anything yet can be given up on right away: there + // is no half-applied assignment to drain, so cancellation does not have to wait for them. + for (RebalanceTaskAttempt attempt : activeAttempts()) { + if (tasksToCancel.contains(attempt.executionKey.getTableBucket())) { + finishRebalanceTask(attempt.executionKey, CANCELED); + } + } - LOG.info("Cancel rebalance task success."); + if (inProgressRebalanceTasks.isEmpty() && !FINAL_STATUSES.contains(rebalanceStatus)) { + finalizeRebalance(); + } + LOG.info( + "Accepted cancellation for rebalance {}. Running and timed-out tasks will be " + + "drained before the rebalance becomes canceled.", + currentRebalanceId); } - public boolean hasInProgressRebalance() { + private Set findTasksToCancel() { + Set tasksToCancel = new HashSet<>(); + for (TableBucket bucket : pendingRebalanceTasks) { + RebalanceResultForBucket result = inProgressRebalanceTasks.get(bucket); + if (result != null + && (!recoveredRebalance + || rebalanceExecutor.isRebalanceTaskAtOrigin(result.plan()))) { + tasksToCancel.add(bucket); + } + } + for (RebalanceTaskAttempt attempt : activeAttempts()) { + TableBucket bucket = attempt.executionKey.getTableBucket(); + RebalanceResultForBucket result = inProgressRebalanceTasks.get(bucket); + if (result != null && rebalanceExecutor.isRebalanceTaskAtOrigin(result.plan())) { + tasksToCancel.add(bucket); + } + } + return tasksToCancel; + } + + private void cancelPendingTasks(Set tasksToCancel) { + Iterator pending = pendingRebalanceTasks.iterator(); + while (pending.hasNext()) { + TableBucket bucket = pending.next(); + if (tasksToCancel.contains(bucket)) { + RebalanceResultForBucket result = inProgressRebalanceTasks.get(bucket); + pending.remove(); + inProgressRebalanceTasks.remove(bucket); + finishedRebalanceTasks.put( + bucket, RebalanceResultForBucket.of(result.plan(), CANCELED)); + } + } + } + + public synchronized boolean hasInProgressRebalance() { checkNotClosed(); - return !inProgressRebalanceTasks.isEmpty() || !inProgressRebalanceTasksQueue.isEmpty(); + return recoveryPending || finalizationPending || !inProgressRebalanceTasks.isEmpty(); } - public RebalanceTask generateRebalanceTask(List goalsByPriority) { + public synchronized RebalanceTask generateRebalanceTask(List goalsByPriority) { checkNotClosed(); List rebalancePlanForBuckets; String rebalanceId = UUID.randomUUID().toString(); try { // Generate the latest cluster model. long startTime = System.currentTimeMillis(); - ClusterModel clusterModel = buildClusterModel(eventProcessor.getCoordinatorContext()); + CoordinatorContext context = rebalanceExecutor.getCoordinatorContext(); + ClusterModel clusterModel = buildClusterModel(context); + Map> actualAssignments = new HashMap<>(); + Map actualLeaders = new HashMap<>(); + for (TableBucket bucket : clusterModel.getReplicaDistribution().keySet()) { + actualAssignments.put(bucket, new ArrayList<>(context.getAssignment(bucket))); + actualLeaders.put(bucket, context.getBucketLeaderAndIsr(bucket).get().leader()); + } LOG.info( "Build cluster model for rebalance id {} with {} ms.", rebalanceId, @@ -352,7 +518,11 @@ public RebalanceTask generateRebalanceTask(List goalsByPriority) { // do optimize. startTime = System.currentTimeMillis(); - rebalancePlanForBuckets = goalOptimizer.doOptimizeOnce(clusterModel, goalsByPriority); + goalOptimizer.doOptimizeOnce(clusterModel, goalsByPriority); + // The optimization model excludes excess replicas left by a failed migration. Diff + // against the real assignment so that cleanup is still executed and persisted, even + // when the requested goals make no further changes to the model. + rebalancePlanForBuckets = getDiff(actualAssignments, actualLeaders, clusterModel); LOG.info( "Do optimize for rebalance id {} with {} ms.", rebalanceId, @@ -366,8 +536,13 @@ public RebalanceTask generateRebalanceTask(List goalsByPriority) { return buildRebalanceTask(rebalanceId, rebalancePlanForBuckets); } - public @Nullable RebalancePlanForBucket getRebalancePlanForBucket(TableBucket tableBucket) { + public synchronized @Nullable RebalancePlanForBucket getRebalancePlanForBucket( + TableBucket tableBucket) { checkNotClosed(); + if (!runningRebalanceTasks.containsKey(tableBucket) + && !timedOutRebalanceTasks.containsKey(tableBucket)) { + return null; + } RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.get(tableBucket); if (resultForBucket != null) { return resultForBucket.plan(); @@ -375,46 +550,316 @@ public RebalanceTask generateRebalanceTask(List goalsByPriority) { return null; } + public synchronized @Nullable RebalanceExecutionKey getExecutionKey(TableBucket tableBucket) { + RebalanceTaskAttempt attempt = runningRebalanceTasks.get(tableBucket); + if (attempt == null) { + attempt = timedOutRebalanceTasks.get(tableBucket); + } + return attempt == null ? null : attempt.executionKey; + } + + public synchronized boolean timeoutRebalanceTask(RebalanceExecutionKey executionKey) { + checkNotClosed(); + queuedTimeoutEvents.remove(executionKey); + RebalanceTaskAttempt attempt = runningRebalanceTasks.get(executionKey.getTableBucket()); + if (attempt == null || !attempt.executionKey.equals(executionKey)) { + LOG.debug("Ignore stale timeout for {}.", executionKey); + return false; + } + + TableBucket tableBucket = executionKey.getTableBucket(); + runningRebalanceTasks.remove(tableBucket); + timedOutRebalanceTasks.put(tableBucket, attempt); + RebalanceResultForBucket result = inProgressRebalanceTasks.get(tableBucket); + if (result == null) { + timedOutRebalanceTasks.remove(tableBucket); + return false; + } + inProgressRebalanceTasks.put( + tableBucket, RebalanceResultForBucket.of(result.plan(), TIMEOUT)); + attempt.onTimedOut(clock.milliseconds(), observeBucketState(tableBucket)); + enqueueReconciliation(attempt); + processNewRebalanceTask(); + return true; + } + + /** + * Returns the plan to reconcile for the given attempt, or null if the attempt is stale or has + * just been given up on. + * + *

    Reconciliation has to terminate. Otherwise a bucket that can never converge, for example + * because a target server is gone for good, keeps the overall rebalance in a non-final status + * and every later rebalance request is rejected forever. + */ + public synchronized @Nullable RebalancePlanForBucket getPlanForReconciliation( + RebalanceExecutionKey executionKey) { + queuedReconcileEvents.remove(executionKey); + TableBucket tableBucket = executionKey.getTableBucket(); + RebalanceTaskAttempt attempt = timedOutRebalanceTasks.get(tableBucket); + if (attempt == null || !attempt.executionKey.equals(executionKey)) { + return null; + } + RebalanceResultForBucket result = inProgressRebalanceTasks.get(tableBucket); + if (result == null) { + return null; + } + if (rebalanceExecutor.isRebalanceTaskDeleted(tableBucket)) { + finishRebalanceTask(executionKey, COMPLETED); + return null; + } + + // Called on the coordinator event loop, so reading the coordinator state is safe here. + long now = clock.milliseconds(); + String observedState = observeBucketState(tableBucket); + boolean targetsLive = + rebalanceExecutor + .getCoordinatorContext() + .liveTabletServerSet() + .containsAll(result.plan().getNewReplicas()); + if (!observedState.equals(attempt.observedState)) { + attempt.onProgress(now, observedState); + } else if (targetsLive) { + attempt.onTargetsAvailable(); + } else { + attempt.onTargetsUnavailable(now); + } + + if (attempt.blockedForMs(now) > targetUnavailableTimeoutMs + || now - attempt.lastProgressMs > noProgressTimeoutMs) { + LOG.error( + "Giving up on rebalance task {} after {} ms without progress, target replicas " + + "live: {}. The bucket may be left with the intermediate assignment " + + "and can be moved again by a new rebalance.", + executionKey, + now - attempt.lastProgressMs, + targetsLive); + finishRebalanceTask(executionKey, FAILED); + return null; + } + + attempt.onReconcileDispatched(now); + return result.plan(); + } + + public synchronized void retryFinalizeRebalance(String rebalanceId) { + finalizationEventQueued = false; + if (finalizationPending && rebalanceId.equals(currentRebalanceId)) { + persistFinalStatus(); + } + } + private void processNewRebalanceTask() { - TableBucket tableBucket = inProgressRebalanceTasksQueue.peek(); - if (tableBucket != null && inProgressRebalanceTasks.containsKey(tableBucket)) { - // Write data (startMs) first, then publish gate (bucket). - inflightTaskStartMs = clock.milliseconds(); - inflightTaskBucket = tableBucket; + TableBucket tableBucket; + while (timedOutRebalanceTasks.size() < maxTrackedTimedOutTasks + && (tableBucket = pendingRecoveredTasks.poll()) != null) { RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.get(tableBucket); - RebalanceResultForBucket rebalanceResultForBucket = - RebalanceResultForBucket.of(resultForBucket.plan(), REBALANCING); - eventProcessor.tryToExecuteRebalanceTask(rebalanceResultForBucket.plan()); + if (resultForBucket == null || resultForBucket.status() != TIMEOUT) { + continue; + } + long now = clock.milliseconds(); + RebalanceTaskAttempt attempt = + new RebalanceTaskAttempt( + new RebalanceExecutionKey(currentRebalanceId, tableBucket), now); + attempt.onTimedOut(now, observeBucketState(tableBucket)); + timedOutRebalanceTasks.put(tableBucket, attempt); + enqueueReconciliation(attempt); + } + if (!pendingRecoveredTasks.isEmpty()) { + return; + } + if (!runningRebalanceTasks.isEmpty()) { + return; + } + if (timedOutRebalanceTasks.size() >= maxTrackedTimedOutTasks) { + // Stop admitting work until some of the timed-out tasks reach a final status, so that + // a long cluster operation cannot grow the tracked set, and with it the reconciliation + // work and the number of concurrent replica migrations, without bound. + LOG.info( + "Hold back new tasks of rebalance {} because {} timed-out tasks are still " + + "being reconciled.", + currentRebalanceId, + timedOutRebalanceTasks.size()); + return; } + while ((tableBucket = pendingRebalanceTasks.poll()) != null) { + RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.get(tableBucket); + if (resultForBucket == null || resultForBucket.status() != NOT_STARTED) { + continue; + } + RebalanceExecutionKey executionKey = + new RebalanceExecutionKey(currentRebalanceId, tableBucket); + runningRebalanceTasks.put( + tableBucket, new RebalanceTaskAttempt(executionKey, clock.milliseconds())); + inProgressRebalanceTasks.put( + tableBucket, RebalanceResultForBucket.of(resultForBucket.plan(), REBALANCING)); + rebalanceExecutor.tryToExecuteRebalanceTask(resultForBucket.plan()); + return; + } + } + + private void finalizeRebalance() { + finalizationPending = true; + persistFinalStatus(); } - private void completeRebalance() { + private void persistFinalStatus() { checkNotClosed(); + RebalanceStatus finalStatus = aggregateFinalStatus(); try { - Optional rebalanceTaskOpt = zkClient.getRebalanceTask(); - Map bucketPlan; - if (rebalanceTaskOpt.isPresent()) { - bucketPlan = rebalanceTaskOpt.get().getExecutePlan(); - } else { - LOG.warn( - "Rebalance task is empty in zk when complete rebalance. " - + "It will be treated as no rebalance tasks."); - bucketPlan = new HashMap<>(); - } zkClient.registerRebalanceTask( - new RebalanceTask(currentRebalanceId, COMPLETED, bucketPlan)); + new RebalanceTask( + currentRebalanceId, + finalStatus, + allRebalancePlans(), + cancelRequested, + canceledBuckets())); } catch (Exception e) { - LOG.error("Error when update rebalance plan from zookeeper.", e); + rebalanceStatus = REBALANCING; + finalizationPending = true; + LOG.error( + "Failed to persist final state for rebalance {}. It will be retried.", + currentRebalanceId, + e); + return; } - rebalanceStatus = COMPLETED; + rebalanceStatus = finalStatus; + finalizationPending = false; + finalizationEventQueued = false; inProgressRebalanceTasks.clear(); - inProgressRebalanceTasksQueue.clear(); + pendingRebalanceTasks.clear(); + pendingRecoveredTasks.clear(); + runningRebalanceTasks.clear(); + timedOutRebalanceTasks.clear(); + queuedTimeoutEvents.clear(); + queuedReconcileEvents.clear(); + + LOG.info( + "Rebalance {} reached final status {} in {} ms.", + currentRebalanceId, + finalStatus, + System.currentTimeMillis() - registerTime); + } + + private void resetRebalance(String rebalanceId, boolean cancelRequested) { + registerTime = System.currentTimeMillis(); + currentRebalanceId = rebalanceId; + recoveryPending = false; + recoveredRebalance = false; + this.cancelRequested = cancelRequested; + finalizationPending = false; + finalizationEventQueued = false; + inProgressRebalanceTasks.clear(); + pendingRebalanceTasks.clear(); + pendingRecoveredTasks.clear(); + runningRebalanceTasks.clear(); + timedOutRebalanceTasks.clear(); + finishedRebalanceTasks.clear(); + queuedTimeoutEvents.clear(); + queuedReconcileEvents.clear(); + } + + private void addPendingTask(TableBucket tableBucket, RebalancePlanForBucket plan) { + pendingRebalanceTasks.add(tableBucket); + inProgressRebalanceTasks.put(tableBucket, RebalanceResultForBucket.of(plan, NOT_STARTED)); + } + + private void addRecoveredIntermediateTask( + TableBucket tableBucket, RebalancePlanForBucket plan) { + pendingRecoveredTasks.add(tableBucket); + inProgressRebalanceTasks.put(tableBucket, RebalanceResultForBucket.of(plan, TIMEOUT)); + } + + private @Nullable RebalanceTaskAttempt findActiveAttempt(RebalanceExecutionKey executionKey) { + RebalanceTaskAttempt attempt = runningRebalanceTasks.get(executionKey.getTableBucket()); + if (attempt == null) { + attempt = timedOutRebalanceTasks.get(executionKey.getTableBucket()); + } + return attempt != null && attempt.executionKey.equals(executionKey) ? attempt : null; + } + + private List activeAttempts() { + List attempts = new ArrayList<>(runningRebalanceTasks.values()); + attempts.addAll(timedOutRebalanceTasks.values()); + return attempts; + } + + /** + * Returns the observable state of a bucket, used to detect whether a timed-out task is still + * making progress. + * + *

    The leader and bucket epochs are deliberately left out: a reconciliation re-sends the + * current state and can bump them without the migration moving forward at all. + */ + private String observeBucketState(TableBucket tableBucket) { + CoordinatorContext coordinatorContext = rebalanceExecutor.getCoordinatorContext(); + StringBuilder observed = + new StringBuilder(coordinatorContext.getAssignment(tableBucket).toString()); + coordinatorContext + .getBucketLeaderAndIsr(tableBucket) + .ifPresent( + leaderAndIsr -> + observed.append("|leader=") + .append(leaderAndIsr.leader()) + .append("|isr=") + .append(new TreeSet<>(leaderAndIsr.isr()))); + return observed.toString(); + } + + private static long getTimeoutMillis(Configuration conf, ConfigOption option) { + Duration timeout = conf.get(option); + checkArgument( + timeout.compareTo(Duration.ofMillis(1)) >= 0 + && timeout.compareTo(Duration.ofMillis(Long.MAX_VALUE)) <= 0, + "%s must be between 1 ms and %s ms, but was %s.", + option.key(), + Long.MAX_VALUE, + timeout); + return timeout.toMillis(); + } + + private static long reconcileBackoffMs(int dispatchedAttempts) { + long backoff = TIMEOUT_CHECK_INTERVAL_MS << Math.min(dispatchedAttempts, 8); + return Math.min(backoff, MAX_RECONCILE_BACKOFF_MS); + } + + private void enqueueReconciliation(RebalanceTaskAttempt attempt) { + if (queuedReconcileEvents.add(attempt.executionKey)) { + eventManager.put(new ReconcileRebalanceTaskEvent(attempt.executionKey)); + } + } - // Here, it will not clear finishedRebalanceTasks, because it will be used by - // listRebalanceProgress. It will be cleared when next register. + private Map allRebalancePlans() { + Map plans = new HashMap<>(); + for (Map.Entry entry : + inProgressRebalanceTasks.entrySet()) { + plans.put(entry.getKey(), entry.getValue().plan()); + } + for (Map.Entry entry : + finishedRebalanceTasks.entrySet()) { + plans.put(entry.getKey(), entry.getValue().plan()); + } + return plans; + } + + private Set canceledBuckets() { + Set canceledBuckets = new HashSet<>(); + for (Map.Entry entry : + finishedRebalanceTasks.entrySet()) { + if (entry.getValue().status() == CANCELED) { + canceledBuckets.add(entry.getKey()); + } + } + return canceledBuckets; + } - LOG.info("Rebalance complete with {} ms.", System.currentTimeMillis() - registerTime); + private RebalanceStatus aggregateFinalStatus() { + for (RebalanceResultForBucket result : finishedRebalanceTasks.values()) { + if (result.status() == FAILED) { + return FAILED; + } + } + return cancelRequested ? CANCELED : COMPLETED; } private ClusterModel buildClusterModel(CoordinatorContext coordinatorContext) { @@ -433,10 +878,18 @@ private ClusterModel buildClusterModel(CoordinatorContext coordinatorContext) { } } + // Failed migrations can leave an unavailable target in the persisted assignment. Keep + // such servers in the model as ineligible sources so the goals can evacuate their replicas. + Set allBuckets = coordinatorContext.getAllBuckets(); + for (TableBucket tableBucket : allBuckets) { + for (Integer replica : coordinatorContext.getAssignment(tableBucket)) { + serverModelMap.computeIfAbsent( + replica, id -> new ServerModel(id, RackModel.DEFAULT_RACK, true)); + } + } ClusterModel clusterModel = initialClusterModel(serverModelMap); // Try to update the cluster model with the latest bucket states. - Set allBuckets = coordinatorContext.getAllBuckets(); for (TableBucket tableBucket : allBuckets) { List assignment = coordinatorContext.getAssignment(tableBucket); Optional bucketLeaderAndIsrOpt = @@ -453,6 +906,8 @@ private ClusterModel buildClusterModel(CoordinatorContext coordinatorContext) { if (leader == -1 || !assignment.contains(leader)) { continue; } + assignment = + assignmentForOptimization(coordinatorContext, tableBucket, assignment, isr); for (int i = 0; i < assignment.size(); i++) { int replica = assignment.get(i); clusterModel.createReplica(replica, tableBucket, i, leader == replica); @@ -461,6 +916,39 @@ private ClusterModel buildClusterModel(CoordinatorContext coordinatorContext) { return clusterModel; } + private List assignmentForOptimization( + CoordinatorContext context, + TableBucket bucket, + List assignment, + LeaderAndIsr leaderAndIsr) { + TableInfo tableInfo = context.getTableInfoById(bucket.getTableId()); + int replicationFactor = + tableInfo == null + ? assignment.size() + : tableInfo + .getProperties() + .getOptional(ConfigOptions.TABLE_REPLICATION_FACTOR) + .orElse(assignment.size()); + if (assignment.size() <= replicationFactor) { + return assignment; + } + + // A Phase A union is not a new replication factor. Prefer the current leader and live, + // caught-up replicas when selecting the model's starting assignment. The requested goals + // can then relocate these replicas; execution retains the full real origin for cleanup. + List replicas = new ArrayList<>(assignment); + replicas.sort( + Comparator.comparing((Integer replica) -> replica != leaderAndIsr.leader()) + .thenComparing(replica -> !context.liveTabletServerSet().contains(replica)) + .thenComparing(replica -> !leaderAndIsr.isr().contains(replica)) + .thenComparing( + replica -> + context.getServerTag(replica) + .map(this::isOfflineTagged) + .orElse(false))); + return new ArrayList<>(replicas.subList(0, replicationFactor)); + } + private RebalanceTask buildRebalanceTask( String rebalanceId, List rebalancePlanForBuckets) { Map bucketPlan = new HashMap<>(); @@ -488,47 +976,117 @@ private void checkTimeoutSafely() { } @VisibleForTesting - void checkTimeout() { - // Read gate (bucket) first, then data (startMs). - // If bucket is non-null, happens-before guarantees startMs is at least as - // fresh as the value written before bucket was published. - TableBucket bucket = inflightTaskBucket; - long startMs = inflightTaskStartMs; - if (bucket == null || startMs < 0) { - return; + synchronized void checkTimeout() { + long now = clock.milliseconds(); + for (RebalanceTaskAttempt attempt : runningRebalanceTasks.values()) { + long elapsed = now - attempt.startMs; + if (elapsed > REBALANCE_TASK_TIMEOUT_MS + && queuedTimeoutEvents.add(attempt.executionKey)) { + LOG.warn( + "In-flight rebalance task {} timed out after {}ms. It will continue to be " + + "tracked while the next pending task is admitted.", + attempt.executionKey, + elapsed); + eventManager.put(new RebalanceTaskTimeoutEvent(attempt.executionKey)); + } } - long elapsed = clock.milliseconds() - startMs; - if (elapsed > REBALANCE_TASK_TIMEOUT_MS) { - LOG.warn( - "In-flight rebalance task for {} timed out after {}ms. " - + "Treating it as timed out and advancing to the next task.", - bucket, - elapsed); - // Clear gate (bucket) first, then data (startMs), matching the - // publication idiom so the next checkTimeout sees bucket==null. - inflightTaskBucket = null; - inflightTaskStartMs = -1; - eventManager.put(new RebalanceTaskTimeoutEvent(bucket)); + + // Reconcile timed-out tasks on a growing backoff, so that a long cluster operation such as + // a rolling upgrade does not turn into a constant retry storm on the event loop. + for (RebalanceTaskAttempt attempt : timedOutRebalanceTasks.values()) { + if (now >= attempt.nextReconcileMs) { + enqueueReconciliation(attempt); + } + } + + String rebalanceId = currentRebalanceId; + if (finalizationPending && rebalanceId != null && !finalizationEventQueued) { + finalizationEventQueued = true; + eventManager.put(new FinalizeRebalanceEvent(rebalanceId)); } } private void checkNotClosed() { - checkArgument(!isClosed, "RebalanceManager is already closed."); + checkState(!isClosed, "RebalanceManager is already closed."); } - public void close() { + public synchronized void close() { isClosed = true; timeoutChecker.shutdownNow(); } @VisibleForTesting public ClusterModel buildClusterModel() { - return buildClusterModel(eventProcessor.getCoordinatorContext()); + return buildClusterModel(rebalanceExecutor.getCoordinatorContext()); } @VisibleForTesting @Nullable - RebalanceStatus getRebalanceStatus() { + synchronized RebalanceStatus getRebalanceStatus() { return rebalanceStatus; } + + @VisibleForTesting + synchronized boolean isCancelRequested() { + return cancelRequested; + } + + private static final class RebalanceTaskAttempt { + private final RebalanceExecutionKey executionKey; + private final long startMs; + + /** The last time this task was observed to change any bucket state. */ + private long lastProgressMs; + + /** The bucket state observed at {@link #lastProgressMs}. */ + private String observedState = ""; + + /** Since when the target replicas are not all live, or -1 if they are. */ + private long blockedSinceMs = -1; + + /** The number of reconciliations already dispatched, used to grow the backoff. */ + private int reconcileAttempts; + + private long nextReconcileMs; + + private RebalanceTaskAttempt(RebalanceExecutionKey executionKey, long startMs) { + this.executionKey = executionKey; + this.startMs = startMs; + } + + private void onTimedOut(long nowMs, String observedState) { + this.lastProgressMs = nowMs; + this.observedState = observedState; + this.blockedSinceMs = -1; + this.reconcileAttempts = 0; + this.nextReconcileMs = nowMs; + } + + private void onProgress(long nowMs, String observedState) { + this.lastProgressMs = nowMs; + this.observedState = observedState; + this.blockedSinceMs = -1; + // A task that moves forward is worth probing at the base interval again. + this.reconcileAttempts = 0; + } + + private void onTargetsAvailable() { + this.blockedSinceMs = -1; + } + + private void onTargetsUnavailable(long nowMs) { + if (blockedSinceMs < 0) { + this.blockedSinceMs = nowMs; + } + } + + private long blockedForMs(long nowMs) { + return blockedSinceMs < 0 ? 0 : nowMs - blockedSinceMs; + } + + private void onReconcileDispatched(long nowMs) { + this.nextReconcileMs = nowMs + reconcileBackoffMs(reconcileAttempts); + this.reconcileAttempts++; + } + } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/RebalanceTask.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/RebalanceTask.java index 2e342123a1a..1ce0b3a1368 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/RebalanceTask.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/RebalanceTask.java @@ -24,10 +24,13 @@ import org.apache.fluss.metadata.TablePartition; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; /** * The generated rebalance task for this cluster. @@ -45,6 +48,12 @@ public class RebalanceTask { /** The rebalance status for the overall rebalance. */ private final RebalanceStatus rebalanceStatus; + /** Whether cancellation has been requested and unfinished tasks should only be drained. */ + private final boolean cancelRequested; + + /** Buckets that were explicitly canceled before the task reached its final status. */ + private final Set canceledBuckets; + /** A mapping from tableBucket to RebalancePlanForBuckets of none-partitioned table. */ private final Map> planForBuckets; @@ -56,8 +65,27 @@ public RebalanceTask( String rebalanceId, RebalanceStatus rebalanceStatus, Map bucketPlan) { + this(rebalanceId, rebalanceStatus, bucketPlan, false); + } + + public RebalanceTask( + String rebalanceId, + RebalanceStatus rebalanceStatus, + Map bucketPlan, + boolean cancelRequested) { + this(rebalanceId, rebalanceStatus, bucketPlan, cancelRequested, Collections.emptySet()); + } + + public RebalanceTask( + String rebalanceId, + RebalanceStatus rebalanceStatus, + Map bucketPlan, + boolean cancelRequested, + Set canceledBuckets) { this.rebalanceId = rebalanceId; this.rebalanceStatus = rebalanceStatus; + this.cancelRequested = cancelRequested; + this.canceledBuckets = new HashSet<>(canceledBuckets); this.planForBuckets = new HashMap<>(); this.planForBucketsOfPartitionedTable = new HashMap<>(); @@ -86,6 +114,14 @@ public RebalanceStatus getRebalanceStatus() { return rebalanceStatus; } + public boolean isCancelRequested() { + return cancelRequested; + } + + public Set getCanceledBuckets() { + return Collections.unmodifiableSet(canceledBuckets); + } + public Map> getPlanForBuckets() { return planForBuckets; } @@ -121,6 +157,10 @@ public String toString() { + rebalanceId + ", rebalanceStatus=" + rebalanceStatus + + ", cancelRequested=" + + cancelRequested + + ", canceledBuckets=" + + canceledBuckets + ", planForBuckets=" + planForBuckets + ", planForBucketsOfPartitionedTable=" @@ -138,8 +178,10 @@ public boolean equals(Object o) { } RebalanceTask that = (RebalanceTask) o; - return rebalanceStatus == that.rebalanceStatus + return cancelRequested == that.cancelRequested + && rebalanceStatus == that.rebalanceStatus && Objects.equals(rebalanceId, that.rebalanceId) + && Objects.equals(canceledBuckets, that.canceledBuckets) && Objects.equals(planForBuckets, that.planForBuckets) && Objects.equals( planForBucketsOfPartitionedTable, that.planForBucketsOfPartitionedTable); @@ -148,6 +190,11 @@ public boolean equals(Object o) { @Override public int hashCode() { return Objects.hash( - rebalanceId, rebalanceStatus, planForBuckets, planForBucketsOfPartitionedTable); + rebalanceId, + rebalanceStatus, + cancelRequested, + canceledBuckets, + planForBuckets, + planForBucketsOfPartitionedTable); } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/RebalanceTaskJsonSerde.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/RebalanceTaskJsonSerde.java index dfb920125a1..f7a34328254 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/RebalanceTaskJsonSerde.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/RebalanceTaskJsonSerde.java @@ -29,9 +29,11 @@ import java.io.IOException; import java.util.ArrayList; import java.util.HashMap; +import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.Set; /** Json serializer and deserializer for {@link RebalanceTask}. */ public class RebalanceTaskJsonSerde @@ -42,6 +44,7 @@ public class RebalanceTaskJsonSerde private static final String VERSION_KEY = "version"; private static final String REBALANCE_ID = "rebalance_id"; private static final String REBALANCE_STATUS = "rebalance_status"; + private static final String CANCEL_REQUESTED = "cancel_requested"; private static final String REBALANCE_PLAN = "rebalance_plan"; private static final String TABLE_ID = "table_id"; @@ -49,12 +52,13 @@ public class RebalanceTaskJsonSerde private static final String BUCKETS = "buckets"; private static final String BUCKET_ID = "bucket_id"; + private static final String CANCELED = "canceled"; private static final String ORIGINAL_LEADER = "original_leader"; private static final String NEW_LEADER = "new_leader"; private static final String ORIGIN_REPLICAS = "origin_replicas"; private static final String NEW_REPLICAS = "new_replicas"; - private static final int VERSION = 1; + private static final int VERSION = 2; @Override public void serialize(RebalanceTask rebalanceTask, JsonGenerator generator) throws IOException { @@ -62,6 +66,7 @@ public void serialize(RebalanceTask rebalanceTask, JsonGenerator generator) thro generator.writeNumberField(VERSION_KEY, VERSION); generator.writeStringField(REBALANCE_ID, rebalanceTask.getRebalanceId()); generator.writeNumberField(REBALANCE_STATUS, rebalanceTask.getRebalanceStatus().getCode()); + generator.writeBooleanField(CANCEL_REQUESTED, rebalanceTask.isCancelRequested()); generator.writeArrayFieldStart(REBALANCE_PLAN); // first to write none-partitioned tables. @@ -71,7 +76,10 @@ public void serialize(RebalanceTask rebalanceTask, JsonGenerator generator) thro generator.writeNumberField(TABLE_ID, entry.getKey()); generator.writeArrayFieldStart(BUCKETS); for (RebalancePlanForBucket bucketPlan : entry.getValue()) { - serializeRebalancePlanForBucket(generator, bucketPlan); + serializeRebalancePlanForBucket( + generator, + bucketPlan, + rebalanceTask.getCanceledBuckets().contains(bucketPlan.getTableBucket())); } generator.writeEndArray(); generator.writeEndObject(); @@ -85,7 +93,10 @@ public void serialize(RebalanceTask rebalanceTask, JsonGenerator generator) thro generator.writeNumberField(PARTITION_ID, entry.getKey().getPartitionId()); generator.writeArrayFieldStart(BUCKETS); for (RebalancePlanForBucket bucketPlan : entry.getValue()) { - serializeRebalancePlanForBucket(generator, bucketPlan); + serializeRebalancePlanForBucket( + generator, + bucketPlan, + rebalanceTask.getCanceledBuckets().contains(bucketPlan.getTableBucket())); } generator.writeEndArray(); generator.writeEndObject(); @@ -102,8 +113,11 @@ public RebalanceTask deserialize(JsonNode node) { String rebalanceId = node.get(REBALANCE_ID).asText(); RebalanceStatus rebalanceStatus = RebalanceStatus.of(node.get(REBALANCE_STATUS).asInt()); + boolean cancelRequested = + node.has(CANCEL_REQUESTED) && node.get(CANCEL_REQUESTED).asBoolean(); Map planForBuckets = new HashMap<>(); + Set canceledBuckets = new HashSet<>(); for (JsonNode tablePartitionPlanNode : rebalancePlanNode) { long tableId = tablePartitionPlanNode.get(TABLE_ID).asLong(); @@ -116,6 +130,9 @@ public RebalanceTask deserialize(JsonNode node) { for (JsonNode bucketPlanNode : bucketPlanNodes) { int bucketId = bucketPlanNode.get(BUCKET_ID).asInt(); TableBucket tableBucket = new TableBucket(tableId, partitionId, bucketId); + if (bucketPlanNode.has(CANCELED) && bucketPlanNode.get(CANCELED).asBoolean()) { + canceledBuckets.add(tableBucket); + } int originLeader = bucketPlanNode.get(ORIGINAL_LEADER).asInt(); @@ -140,13 +157,18 @@ public RebalanceTask deserialize(JsonNode node) { } } - return new RebalanceTask(rebalanceId, rebalanceStatus, planForBuckets); + return new RebalanceTask( + rebalanceId, rebalanceStatus, planForBuckets, cancelRequested, canceledBuckets); } private void serializeRebalancePlanForBucket( - JsonGenerator generator, RebalancePlanForBucket bucketPlan) throws IOException { + JsonGenerator generator, RebalancePlanForBucket bucketPlan, boolean canceled) + throws IOException { generator.writeStartObject(); generator.writeNumberField(BUCKET_ID, bucketPlan.getBucketId()); + if (canceled) { + generator.writeBooleanField(CANCELED, true); + } generator.writeNumberField(ORIGINAL_LEADER, bucketPlan.getOriginalLeader()); generator.writeNumberField(NEW_LEADER, bucketPlan.getNewLeader()); generator.writeArrayFieldStart(ORIGIN_REPLICAS); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorAdjustIsrTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorAdjustIsrTest.java index 6fa13ff7d23..f59932a0cea 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorAdjustIsrTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorAdjustIsrTest.java @@ -26,7 +26,6 @@ import org.apache.fluss.metadata.TablePath; import org.apache.fluss.rpc.messages.AdjustIsrResponse; import org.apache.fluss.rpc.protocol.Errors; -import org.apache.fluss.server.coordinator.event.AccessContextEvent; import org.apache.fluss.server.coordinator.event.AdjustIsrReceivedEvent; import org.apache.fluss.server.entity.AdjustIsrResultForBucket; import org.apache.fluss.server.metrics.group.TestingMetricGroups; @@ -54,10 +53,8 @@ import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; -import static org.apache.fluss.server.coordinator.CoordinatorTestUtils.makeSendLeaderAndStopRequestAlwaysSuccess; import static org.apache.fluss.server.utils.ServerRpcMessageUtils.getAdjustIsrResponseData; import static org.apache.fluss.server.utils.TableAssignmentUtils.generateAssignment; import static org.apache.fluss.server.zk.ZooKeeperUtils.startZookeeperClient; @@ -355,17 +352,6 @@ private AdjustIsrResponse adjustIsr(Map leaderAndIsrs return response.get(); } - private AdjustIsrResultForBucket submitAdjustIsr( - TableBucket tableBucket, LeaderAndIsr newLeaderAndIsr) throws Exception { - CompletableFuture response = new CompletableFuture<>(); - eventProcessor - .getCoordinatorEventManager() - .put( - new AdjustIsrReceivedEvent( - Collections.singletonMap(tableBucket, newLeaderAndIsr), response)); - return getAdjustIsrResponseData(response.get()).get(tableBucket); - } - private void assertAdjustIsrResponseCount(AdjustIsrResponse response, int expectedCount) { // Inspect the raw protobuf response: converting it to a map would hide duplicate results // for the same bucket, such as an error followed by an unconditional success. @@ -395,22 +381,6 @@ private void assertAdjustedIsrInstalled( } } - private void initCoordinatorChannel() throws Exception { - makeSendLeaderAndStopRequestAlwaysSuccess( - testCoordinatorChannelManager, - Arrays.stream(zookeeperClient.getSortedTabletServerList()) - .boxed() - .collect(Collectors.toSet()), - Collections.emptySet()); - } - - private T fromCtx(java.util.function.Function retrieveFunction) - throws Exception { - AccessContextEvent event = new AccessContextEvent<>(retrieveFunction); - eventProcessor.getCoordinatorEventManager().put(event); - return event.getResultFuture().get(30, TimeUnit.SECONDS); - } - private CoordinatorEventProcessor buildCoordinatorEventProcessor( ZooKeeperClient coordinatorZooKeeperClient) { Configuration conf = new Configuration(); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java index b26d8cabca2..dc0ca869c3e 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java @@ -40,25 +40,23 @@ import org.apache.fluss.metrics.Gauge; import org.apache.fluss.metrics.MetricNames; import org.apache.fluss.metrics.groups.AbstractMetricGroup; -import org.apache.fluss.rpc.gateway.TabletServerGateway; import org.apache.fluss.rpc.messages.AdjustIsrResponse; import org.apache.fluss.rpc.messages.ApiMessage; import org.apache.fluss.rpc.messages.CommitKvSnapshotResponse; import org.apache.fluss.rpc.messages.CommitRemoteLogManifestResponse; import org.apache.fluss.rpc.messages.NotifyKvSnapshotOffsetRequest; -import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrRequest; -import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrResponse; import org.apache.fluss.rpc.messages.NotifyRemoteLogOffsetsRequest; import org.apache.fluss.rpc.messages.UpdateMetadataRequest; import org.apache.fluss.rpc.messages.UpdateMetadataResponse; import org.apache.fluss.rpc.protocol.ApiError; import org.apache.fluss.rpc.protocol.ApiKeys; import org.apache.fluss.rpc.protocol.Errors; -import org.apache.fluss.server.coordinator.event.AccessContextEvent; +import org.apache.fluss.server.coordinator.TestingControlledNotifyGateway.PendingNotify; import org.apache.fluss.server.coordinator.event.AdjustIsrReceivedEvent; import org.apache.fluss.server.coordinator.event.CommitKvSnapshotEvent; import org.apache.fluss.server.coordinator.event.CommitRemoteLogManifestEvent; import org.apache.fluss.server.coordinator.event.CoordinatorEventManager; +import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrRequestContext; import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrResponseReceivedEvent; import org.apache.fluss.server.coordinator.event.RetryOfflineLeaderEvent; import org.apache.fluss.server.coordinator.remote.RemoteDirDynamicLoader; @@ -89,7 +87,6 @@ import org.apache.fluss.server.zk.data.ZkVersion; import org.apache.fluss.testutils.common.ManuallyTriggeredScheduledExecutorService; import org.apache.fluss.types.DataTypes; -import org.apache.fluss.utils.ExceptionUtils; import org.apache.fluss.utils.clock.SystemClock; import org.apache.fluss.utils.types.Tuple2; @@ -110,10 +107,7 @@ import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; -import java.util.function.Consumer; -import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -1389,10 +1383,53 @@ void testDiskWriteLockedNotifyLeaderResponseMarksReplicaOffline() throws Excepti verifyTableCreated(tableId, tableAssignment, nBuckets, replicationFactor); TableBucket tableBucket = new TableBucket(tableId, 0); - int leader = - tableAssignment.getBucketAssignment(tableBucket.getBucket()).getReplicas().get(0); + List replicas = + tableAssignment.getBucketAssignment(tableBucket.getBucket()).getReplicas(); + int leader = replicas.get(0); TableBucketReplica tableBucketReplica = new TableBucketReplica(tableBucket, leader); + putFailedNotifyResponse(tableBucket, leader, Collections.emptyMap()); + + fromCtx( + ctx -> { + assertThat(ctx.getReplicaState(tableBucketReplica)).isEqualTo(OfflineReplica); + assertThat(ctx.isReplicaOnline(leader, tableBucket)).isFalse(); + return null; + }); + + // The same holds for a response that no longer matches the state we last sent, for example + // because the leader shrank the ISR in the meantime. + int follower = replicas.get(2); + LeaderAndIsr current = fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get()); + putFailedNotifyResponse( + tableBucket, + follower, + Collections.singletonMap( + tableBucket, + new NotifyLeaderAndIsrRequestContext( + eventProcessor.getCoordinatorEpoch(), + current.leader(), + current.leaderEpoch(), + current.bucketEpoch() - 1))); + + retry( + Duration.ofMinutes(1), + () -> + fromCtx( + ctx -> { + assertThat( + ctx.getReplicaState( + new TableBucketReplica( + tableBucket, follower))) + .isEqualTo(OfflineReplica); + return null; + })); + } + + private void putFailedNotifyResponse( + TableBucket tableBucket, + int serverId, + Map requestContexts) { eventProcessor .getCoordinatorEventManager() .put( @@ -1403,14 +1440,8 @@ void testDiskWriteLockedNotifyLeaderResponseMarksReplicaOffline() throws Excepti new ApiError( Errors.DISK_WRITE_LOCKED, "disk write locked"))), - leader)); - - fromCtx( - ctx -> { - assertThat(ctx.getReplicaState(tableBucketReplica)).isEqualTo(OfflineReplica); - assertThat(ctx.isReplicaOnline(leader, tableBucket)).isFalse(); - return null; - }); + serverId, + requestContexts)); } @Test @@ -1516,7 +1547,7 @@ void testRetryOfflineLeaderEventKeepsReplicaOfflineWhenRetryStillFails() throws }); assertThat(eventProcessor.hasOfflineLeaderRetryTaskScheduled()).isTrue(); - CountingFailingNotifyGateway failingGateway = new CountingFailingNotifyGateway(); + TestingFailingNotifyGateway failingGateway = new TestingFailingNotifyGateway(); testCoordinatorChannelManager.setGateways(Collections.singletonMap(leader, failingGateway)); eventProcessor.getCoordinatorEventManager().put(new RetryOfflineLeaderEvent()); @@ -2029,81 +2060,10 @@ void testAlterStandbyReplicaEnabledForLogTable() throws Exception { .hasMessageContaining("can only be altered on primary key tables"); } - @Test - void testDoBucketReassignment() throws Exception { - zookeeperClient.registerTabletServer( - 3, - new TabletServerRegistration( - "rack3", - Collections.singletonList( - new Endpoint("host3", 1001, DEFAULT_LISTENER_NAME)), - System.currentTimeMillis())); - - initCoordinatorChannel(); - TablePath t1 = TablePath.of(defaultDatabase, "test_bucket_reassignment_table"); - // Mock un-balanced table assignment. - Map bucketAssignments = new HashMap<>(); - bucketAssignments.put(0, BucketAssignment.of(0, 1, 3)); - TableAssignment tableAssignment = new TableAssignment(bucketAssignments); - long t1Id = - metadataManager.createTable( - t1, - remoteDataDir, - CoordinatorEventProcessorTest.TEST_TABLE, - tableAssignment, - false); - TableBucket tb0 = new TableBucket(t1Id, 0); - verifyIsr(tb0, 0, Arrays.asList(0, 1, 3)); - - // trigger bucket reassignment for tb0: - // bucket0 -> (0, 1, 2) - Map rebalancePlan = new HashMap<>(); - RebalancePlanForBucket planForBucket0 = - new RebalancePlanForBucket( - tb0, 0, 0, Arrays.asList(0, 1, 3), Arrays.asList(0, 1, 2)); - - rebalancePlan.put(tb0, planForBucket0); - // try to execute. - eventProcessor - .getRebalanceManager() - .registerRebalance( - "rebalance-task-jdsds1", rebalancePlan, RebalanceStatus.NOT_STARTED); - - // Mock to finish rebalance tasks, in production case, this need to be trigged by receiving - // AdjustIsrRequest. - Map leaderAndIsrMap = new HashMap<>(); - CompletableFuture respCallback = new CompletableFuture<>(); - - // This isr list equals originReplicas + addingReplicas. the bucket epoch is 1. - leaderAndIsrMap.put( - tb0, - new LeaderAndIsr(0, 0, Arrays.asList(0, 1, 2, 3), Collections.emptyList(), 0, 1)); - eventProcessor - .getCoordinatorEventManager() - .put(new AdjustIsrReceivedEvent(leaderAndIsrMap, respCallback)); - respCallback.get(); - verifyIsr(tb0, 0, Arrays.asList(0, 1, 2)); - - // clean up the tablet server 3 - ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(ZkData.ServerIdZNode.path(3)); - } - @Test void testLeaderOnlyRebalanceExecutesSequentially() throws Exception { - // Set up controlled gateways that capture NotifyLeaderAndIsr calls. - // Gateways start in pass-through mode for table creation, then switch - // to controlled mode to verify sequential leader migration. - ConcurrentLinkedDeque pendingTriggers = - new ConcurrentLinkedDeque<>(); - int[] servers = zookeeperClient.getSortedTabletServerList(); - Map gateways = new HashMap<>(); - ControlledNotifyGateway[] controlledGateways = new ControlledNotifyGateway[servers.length]; - for (int i = 0; i < servers.length; i++) { - ControlledNotifyGateway gw = new ControlledNotifyGateway(servers[i], pendingTriggers); - gateways.put(servers[i], gw); - controlledGateways[i] = gw; - } - testCoordinatorChannelManager.setGateways(gateways); + initCoordinatorChannel(); + ConcurrentLinkedDeque pendingTriggers = new ConcurrentLinkedDeque<>(); // Create a table with 3 buckets, each assigned to replicas [0, 1, 2] with leader 0. TablePath t1 = TablePath.of(defaultDatabase, "test_leader_rebalance_sequential"); @@ -2124,12 +2084,7 @@ void testLeaderOnlyRebalanceExecutesSequentially() throws Exception { verifyIsr(tb1, 0, Arrays.asList(0, 1, 2)); verifyIsr(tb2, 0, Arrays.asList(0, 1, 2)); - // Switch to controlled mode: from now on, NotifyLeaderAndIsr responses - // are held until the test explicitly releases them. - for (ControlledNotifyGateway gw : controlledGateways) { - gw.enableControlMode(); - } - pendingTriggers.clear(); + installBlockingNotifyGateways(pendingTriggers); // Create leader-only rebalance plan (replicas stay the same, only leaders change): // tb0: leader 0 -> 1 (newReplicas=[1,0,2] puts target leader first) @@ -2153,8 +2108,7 @@ void testLeaderOnlyRebalanceExecutesSequentially() throws Exception { // because subsequent tasks must wait for the NotifyLeaderAndIsr response. eventProcessor .getRebalanceManager() - .registerRebalance( - "rebalance-leader-sequential", rebalancePlan, RebalanceStatus.NOT_STARTED); + .registerRebalance("rebalance-leader-sequential", rebalancePlan); // === Step 1: Verify only the first task started === // registerRebalance() is synchronous, so after it returns, the first task's @@ -2164,6 +2118,14 @@ void testLeaderOnlyRebalanceExecutesSequentially() throws Exception { // All 3 tasks are still in progress (first executing, two waiting). assertThat(countInProgressRebalanceTasks(tb0, tb1, tb2)).isEqualTo(3); + // An acknowledgement from the old leader must not release the execution slot. + retry( + Duration.ofMinutes(1), + () -> assertThat(hasPendingNotifyTrigger(pendingTriggers, 0)).isTrue()); + completePendingNotifyTrigger(pendingTriggers, 0); + fromCtx(ctx -> null); + assertThat(countInProgressRebalanceTasks(tb0, tb1, tb2)).isEqualTo(3); + // Release the first batch - this allows the event processor to complete // the first task and start the second. drainPendingNotifyTriggers(pendingTriggers); @@ -2195,117 +2157,6 @@ void testLeaderOnlyRebalanceExecutesSequentially() throws Exception { verifyIsr(tb2, 1, Arrays.asList(0, 1, 2)); } - @Test - void testLeaderOnlyRebalanceCompletionCheckRequiresSuccessfulResponseFromNewLeader() { - TableBucket tableBucket = new TableBucket(1L, 0); - RebalancePlanForBucket planForBucket = - new RebalancePlanForBucket( - tableBucket, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 0, 2)); - NotifyLeaderAndIsrResultForBucket successResult = - new NotifyLeaderAndIsrResultForBucket(tableBucket); - NotifyLeaderAndIsrResultForBucket failedResult = - new NotifyLeaderAndIsrResultForBucket( - tableBucket, new ApiError(Errors.UNKNOWN_SERVER_ERROR, "failed")); - - assertThat( - CoordinatorEventProcessor - .isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader( - successResult, 1, planForBucket)) - .isTrue(); - assertThat( - CoordinatorEventProcessor - .isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader( - successResult, 0, planForBucket)) - .isFalse(); - assertThat( - CoordinatorEventProcessor - .isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader( - failedResult, 1, planForBucket)) - .isFalse(); - } - - @Test - void testLeaderOnlyRebalanceIgnoresSuccessResponseFromOldLeader() throws Exception { - ConcurrentLinkedDeque pendingTriggers = - new ConcurrentLinkedDeque<>(); - int[] servers = zookeeperClient.getSortedTabletServerList(); - Map gateways = new HashMap<>(); - ControlledNotifyGateway[] controlledGateways = new ControlledNotifyGateway[servers.length]; - for (int i = 0; i < servers.length; i++) { - ControlledNotifyGateway gw = new ControlledNotifyGateway(servers[i], pendingTriggers); - gateways.put(servers[i], gw); - controlledGateways[i] = gw; - } - testCoordinatorChannelManager.setGateways(gateways); - - TablePath t1 = TablePath.of(defaultDatabase, "test_leader_rebalance_wait_new_leader"); - Map bucketAssignments = new HashMap<>(); - bucketAssignments.put(0, BucketAssignment.of(0, 1, 2)); - TableAssignment tableAssignment = new TableAssignment(bucketAssignments); - long t1Id = - metadataManager.createTable(t1, remoteDataDir, TEST_TABLE, tableAssignment, false); - - TableBucket tb0 = new TableBucket(t1Id, 0); - - verifyIsr(tb0, 0, Arrays.asList(0, 1, 2)); - - for (ControlledNotifyGateway gw : controlledGateways) { - gw.enableControlMode(); - } - pendingTriggers.clear(); - - Map rebalancePlan = new HashMap<>(); - rebalancePlan.put( - tb0, - new RebalancePlanForBucket( - tb0, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 0, 2))); - - eventProcessor - .getRebalanceManager() - .registerRebalance( - "rebalance-wait-new-leader-response", - rebalancePlan, - RebalanceStatus.NOT_STARTED); - - retry( - Duration.ofMinutes(1), - () -> assertThat(hasPendingNotifyTrigger(pendingTriggers, 0)).isTrue()); - retry( - Duration.ofMinutes(1), - () -> assertThat(hasPendingNotifyTrigger(pendingTriggers, 1)).isTrue()); - assertThat(countInProgressRebalanceTasks(tb0)).isEqualTo(1); - - completePendingNotifyTrigger(pendingTriggers, 0); - fromCtx(ctx -> null); - - assertThat(countInProgressRebalanceTasks(tb0)).isEqualTo(1); - assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()).isTrue(); - - completePendingNotifyTrigger(pendingTriggers, 1); - retry( - Duration.ofMinutes(1), - () -> - assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) - .isFalse()); - verifyIsr(tb0, 1, Arrays.asList(0, 1, 2)); - } - - private void verifyIsr(TableBucket tb, int expectedLeader, List expectedIsr) - throws Exception { - LeaderAndIsr leaderAndIsr = - waitValue( - () -> fromCtx((ctx) -> ctx.getBucketLeaderAndIsr(tb)), - Duration.ofMinutes(1), - "leader not elected"); - LeaderAndIsr newLeaderAndIsrOfZk = zookeeperClient.getLeaderAndIsr(tb).get(); - assertThat(leaderAndIsr.leader()) - .isEqualTo(newLeaderAndIsrOfZk.leader()) - .isEqualTo(expectedLeader); - assertThat(leaderAndIsr.isr()) - .isEqualTo(newLeaderAndIsrOfZk.isr()) - .hasSameElementsAs(expectedIsr); - } - private static class FailingUpdateMetadataChannelManager extends TestCoordinatorChannelManager { private final CountDownLatch failureObserved = new CountDownLatch(1); @@ -2384,15 +2235,6 @@ private int getLiveTabletServerCountAtInit() { } } - private void initCoordinatorChannel() throws Exception { - makeSendLeaderAndStopRequestAlwaysSuccess( - testCoordinatorChannelManager, - Arrays.stream(zookeeperClient.getSortedTabletServerList()) - .boxed() - .collect(Collectors.toSet()), - Collections.emptySet()); - } - private void initCoordinatorChannel(Set ignoreApiKeys) throws Exception { makeSendLeaderAndStopRequestAlwaysSuccess( testCoordinatorChannelManager, @@ -2679,31 +2521,6 @@ private void verifyMetadataUpdateRequest(int serverCount, TableMetadata tableMet } } - private void retryVerifyContext(Consumer verifyFunction) { - retry( - Duration.ofMinutes(1), - () -> { - AccessContextEvent event = - new AccessContextEvent<>( - ctx -> { - verifyFunction.accept(ctx); - return null; - }); - eventProcessor.getCoordinatorEventManager().put(event); - try { - event.getResultFuture().get(30, TimeUnit.SECONDS); - } catch (Throwable t) { - throw ExceptionUtils.stripExecutionException(t); - } - }); - } - - private T fromCtx(Function retrieveFunction) throws Exception { - AccessContextEvent event = new AccessContextEvent<>(retrieveFunction); - eventProcessor.getCoordinatorEventManager().put(event); - return event.getResultFuture().get(30, TimeUnit.SECONDS); - } - private Map.Entry createTableAndGetLeaderAndIsr(String tableName) throws Exception { initCoordinatorChannel(); @@ -2734,17 +2551,6 @@ private Map.Entry createTableAndGetLeaderAndIsr(Strin return bucketLeaderAndIsrMap.entrySet().iterator().next(); } - private AdjustIsrResultForBucket submitAdjustIsr( - TableBucket tableBucket, LeaderAndIsr newLeaderAndIsr) throws Exception { - CompletableFuture response = new CompletableFuture<>(); - eventProcessor - .getCoordinatorEventManager() - .put( - new AdjustIsrReceivedEvent( - Collections.singletonMap(tableBucket, newLeaderAndIsr), response)); - return getAdjustIsrResponseData(response.get()).get(tableBucket); - } - private long createTable(TablePath tablePath, TabletServerInfo[] servers) { TableAssignment tableAssignment = generateAssignment(N_BUCKETS, REPLICATION_FACTOR, servers); @@ -2791,132 +2597,29 @@ private static List allTableBuckets( .collect(Collectors.toList()); } - private static void drainPendingNotifyTriggers( - ConcurrentLinkedDeque pendingTriggers) { - ControlledNotifyTrigger trigger; - while ((trigger = pendingTriggers.poll()) != null) { - trigger.complete(null); - } - } - - private static boolean hasPendingNotifyTrigger( - ConcurrentLinkedDeque pendingTriggers, int responseServerId) { - for (ControlledNotifyTrigger trigger : pendingTriggers) { - if (trigger.getResponseServerId() == responseServerId) { - return true; - } - } - return false; - } + private static class PartitionIdName { + private final long partitionId; + private final String partitionName; - private static void completePendingNotifyTrigger( - ConcurrentLinkedDeque pendingTriggers, int responseServerId) { - for (ControlledNotifyTrigger trigger : pendingTriggers) { - if (trigger.getResponseServerId() == responseServerId) { - assertThat(pendingTriggers.remove(trigger)).isTrue(); - trigger.complete(null); - return; - } + private PartitionIdName(long partitionId, String partitionName) { + this.partitionId = partitionId; + this.partitionName = partitionName; } - throw new AssertionError( - "No pending NotifyLeaderAndIsr response for server " + responseServerId); } private int countInProgressRebalanceTasks(TableBucket... buckets) { int count = 0; for (TableBucket tb : buckets) { - if (eventProcessor.getRebalanceManager().getRebalancePlanForBucket(tb) != null) { + if (!RebalanceStatus.FINAL_STATUSES.contains( + eventProcessor + .getRebalanceManager() + .listRebalanceProgress(null) + .progressForBucketMap() + .get(tb) + .status())) { count++; } } return count; } - - private static class CountingFailingNotifyGateway extends TestTabletServerGateway { - private final AtomicInteger notifyLeaderAndIsrCount = new AtomicInteger(); - - CountingFailingNotifyGateway() { - super(true, Collections.emptySet()); - } - - int getNotifyLeaderAndIsrCount() { - return notifyLeaderAndIsrCount.get(); - } - - @Override - public CompletableFuture notifyLeaderAndIsr( - NotifyLeaderAndIsrRequest request) { - notifyLeaderAndIsrCount.incrementAndGet(); - return super.notifyLeaderAndIsr(request); - } - } - - /** - * A gateway that intercepts NotifyLeaderAndIsr calls for verifying sequential execution of - * leader migrations. In pass-through mode, it delegates to the parent. In controlled mode, it - * captures the response in a CompletableFuture trigger that the test must explicitly complete - * before the response is delivered. - */ - private static class ControlledNotifyGateway extends TestTabletServerGateway { - private volatile boolean controlMode = false; - private final int responseServerId; - private final ConcurrentLinkedDeque pendingTriggers; - - ControlledNotifyGateway( - int responseServerId, - ConcurrentLinkedDeque pendingTriggers) { - super(false, Collections.emptySet()); - this.responseServerId = responseServerId; - this.pendingTriggers = pendingTriggers; - } - - void enableControlMode() { - controlMode = true; - } - - @Override - public CompletableFuture notifyLeaderAndIsr( - NotifyLeaderAndIsrRequest request) { - if (!controlMode) { - return super.notifyLeaderAndIsr(request); - } - // Build the proper success response using parent's logic. - NotifyLeaderAndIsrResponse response = super.notifyLeaderAndIsr(request).join(); - // Return a future that completes only when the test releases the trigger. - ControlledNotifyTrigger trigger = new ControlledNotifyTrigger(responseServerId); - pendingTriggers.add(trigger); - return trigger.getFuture().thenApply(v -> response); - } - } - - private static class ControlledNotifyTrigger { - private final int responseServerId; - private final CompletableFuture future = new CompletableFuture<>(); - - ControlledNotifyTrigger(int responseServerId) { - this.responseServerId = responseServerId; - } - - int getResponseServerId() { - return responseServerId; - } - - CompletableFuture getFuture() { - return future; - } - - void complete(Void value) { - future.complete(value); - } - } - - private static class PartitionIdName { - private final long partitionId; - private final String partitionName; - - private PartitionIdName(long partitionId, String partitionName) { - this.partitionId = partitionId; - this.partitionName = partitionName; - } - } } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTestBase.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTestBase.java index 372abf19f8e..5da485ba377 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTestBase.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTestBase.java @@ -21,8 +21,15 @@ import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; import org.apache.fluss.metadata.DatabaseDescriptor; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.rpc.gateway.TabletServerGateway; +import org.apache.fluss.rpc.messages.AdjustIsrResponse; +import org.apache.fluss.server.coordinator.TestingControlledNotifyGateway.PendingNotify; +import org.apache.fluss.server.coordinator.event.AccessContextEvent; +import org.apache.fluss.server.coordinator.event.AdjustIsrReceivedEvent; import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager; import org.apache.fluss.server.coordinator.remote.RemoteDirDynamicLoader; +import org.apache.fluss.server.entity.AdjustIsrResultForBucket; import org.apache.fluss.server.metadata.CoordinatorMetadataCache; import org.apache.fluss.server.metrics.group.TestingMetricGroups; import org.apache.fluss.server.zk.NOPErrorHandler; @@ -30,10 +37,13 @@ import org.apache.fluss.server.zk.ZooKeeperClient; import org.apache.fluss.server.zk.ZooKeeperExtension; import org.apache.fluss.server.zk.data.CoordinatorAddress; +import org.apache.fluss.server.zk.data.LeaderAndIsr; import org.apache.fluss.server.zk.data.TabletServerRegistration; import org.apache.fluss.server.zk.data.ZkData.PartitionIdsZNode; import org.apache.fluss.server.zk.data.ZkData.TableIdsZNode; import org.apache.fluss.testutils.common.AllCallbackWrapper; +import org.apache.fluss.utils.ExceptionUtils; +import org.apache.fluss.utils.clock.Clock; import org.apache.fluss.utils.clock.SystemClock; import org.apache.fluss.utils.concurrent.ExecutorThreadFactory; import org.apache.fluss.utils.concurrent.FlussScheduler; @@ -45,10 +55,25 @@ import org.junit.jupiter.api.extension.RegisterExtension; import java.time.Duration; +import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.stream.Collectors; import static org.apache.fluss.config.ConfigOptions.DEFAULT_LISTENER_NAME; +import static org.apache.fluss.server.coordinator.CoordinatorTestUtils.makeSendLeaderAndStopRequestAlwaysSuccess; +import static org.apache.fluss.server.utils.ServerRpcMessageUtils.getAdjustIsrResponseData; +import static org.apache.fluss.testutils.common.CommonTestUtils.retry; +import static org.apache.fluss.testutils.common.CommonTestUtils.waitValue; +import static org.assertj.core.api.Assertions.assertThat; /** * The shared lifecycle harness for {@link CoordinatorEventProcessor} unit tests: a ZooKeeper test @@ -161,6 +186,10 @@ void afterEach() throws Exception { } protected CoordinatorEventProcessor buildCoordinatorEventProcessor() { + return buildCoordinatorEventProcessor(SystemClock.getInstance()); + } + + protected CoordinatorEventProcessor buildCoordinatorEventProcessor(Clock clock) { Configuration conf = new Configuration(); conf.set(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); conf.set(ConfigOptions.COORDINATOR_OFFLINE_LEADER_RETRY_DELAY, Duration.ofDays(1)); @@ -178,6 +207,129 @@ protected CoordinatorEventProcessor buildCoordinatorEventProcessor() { metadataManager, kvSnapshotLeaseManager, scheduler, - SystemClock.getInstance()); + clock); + } + + /** Verifies the elected leader and ISR in memory and ZooKeeper. */ + protected void verifyIsr(TableBucket tb, int expectedLeader, List expectedIsr) + throws Exception { + LeaderAndIsr leaderAndIsr = + waitValue( + () -> fromCtx((ctx) -> ctx.getBucketLeaderAndIsr(tb)), + Duration.ofMinutes(1), + "leader not elected"); + LeaderAndIsr newLeaderAndIsrOfZk = zookeeperClient.getLeaderAndIsr(tb).get(); + assertThat(leaderAndIsr.leader()) + .isEqualTo(newLeaderAndIsrOfZk.leader()) + .isEqualTo(expectedLeader); + assertThat(leaderAndIsr.isr()) + .isEqualTo(newLeaderAndIsrOfZk.isr()) + .hasSameElementsAs(expectedIsr); + } + + /** Installs successful tablet server gateways for registered servers. */ + protected void initCoordinatorChannel() throws Exception { + makeSendLeaderAndStopRequestAlwaysSuccess( + testCoordinatorChannelManager, + Arrays.stream(zookeeperClient.getSortedTabletServerList()) + .boxed() + .collect(Collectors.toSet()), + Collections.emptySet()); + } + + /** Registers an additional tablet server in ZooKeeper. */ + protected void registerTabletServer(int serverId) throws Exception { + zookeeperClient.registerTabletServer( + serverId, + new TabletServerRegistration( + "rack" + serverId, + Collections.singletonList( + new Endpoint("host" + serverId, 1001, DEFAULT_LISTENER_NAME)), + System.currentTimeMillis())); + } + + /** Retries assertions against state read on the coordinator event thread. */ + protected void retryVerifyContext(Consumer verifyFunction) { + retry( + Duration.ofMinutes(1), + () -> { + AccessContextEvent event = + new AccessContextEvent<>( + ctx -> { + verifyFunction.accept(ctx); + return null; + }); + eventProcessor.getCoordinatorEventManager().put(event); + try { + event.getResultFuture().get(30, TimeUnit.SECONDS); + } catch (Throwable t) { + throw ExceptionUtils.stripExecutionException(t); + } + }); + } + + /** Reads or updates state on the coordinator event thread. */ + protected T fromCtx(Function retrieveFunction) throws Exception { + AccessContextEvent event = new AccessContextEvent<>(retrieveFunction); + eventProcessor.getCoordinatorEventManager().put(event); + return event.getResultFuture().get(30, TimeUnit.SECONDS); + } + + /** Submits an ISR update through the coordinator event loop. */ + protected AdjustIsrResultForBucket submitAdjustIsr( + TableBucket tableBucket, LeaderAndIsr newLeaderAndIsr) throws Exception { + CompletableFuture response = new CompletableFuture<>(); + eventProcessor + .getCoordinatorEventManager() + .put( + new AdjustIsrReceivedEvent( + Collections.singletonMap(tableBucket, newLeaderAndIsr), response)); + return getAdjustIsrResponseData(response.get()).get(tableBucket); + } + + /** Installs gateways that hold leader notifications until explicitly released. */ + protected void installBlockingNotifyGateways( + ConcurrentLinkedDeque pendingTriggers) throws Exception { + Map gateways = new HashMap<>(); + for (int server : zookeeperClient.getSortedTabletServerList()) { + TestingControlledNotifyGateway gateway = + new TestingControlledNotifyGateway(server, pendingTriggers); + gateways.put(server, gateway); + } + testCoordinatorChannelManager.setGateways(gateways); + } + + /** Releases queued notification responses. */ + protected static void drainPendingNotifyTriggers( + ConcurrentLinkedDeque pendingTriggers) { + PendingNotify trigger; + while ((trigger = pendingTriggers.poll()) != null) { + trigger.complete(); + } + } + + /** Returns whether a server has a held notification response. */ + protected static boolean hasPendingNotifyTrigger( + ConcurrentLinkedDeque pendingTriggers, int responseServerId) { + for (PendingNotify trigger : pendingTriggers) { + if (trigger.getResponseServerId() == responseServerId) { + return true; + } + } + return false; + } + + /** Releases one held notification response from a server. */ + protected static void completePendingNotifyTrigger( + ConcurrentLinkedDeque pendingTriggers, int responseServerId) { + for (PendingNotify trigger : pendingTriggers) { + if (trigger.getResponseServerId() == responseServerId) { + assertThat(pendingTriggers.remove(trigger)).isTrue(); + trigger.complete(); + return; + } + } + throw new AssertionError( + "No pending NotifyLeaderAndIsr response for server " + responseServerId); } } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorRebalanceTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorRebalanceTest.java new file mode 100644 index 00000000000..b0e1eadb493 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorRebalanceTest.java @@ -0,0 +1,818 @@ +/* + * 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.fluss.server.coordinator; + +import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket; +import org.apache.fluss.cluster.rebalance.RebalanceStatus; +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TablePartition; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrRequest; +import org.apache.fluss.rpc.messages.PbNotifyLeaderAndIsrReqForBucket; +import org.apache.fluss.rpc.messages.RebalanceResponse; +import org.apache.fluss.server.coordinator.TestingControlledNotifyGateway.PendingNotify; +import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrRequestContext; +import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrResponseReceivedEvent; +import org.apache.fluss.server.coordinator.event.RebalanceEvent; +import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent; +import org.apache.fluss.server.coordinator.event.ReconcileRebalanceTaskEvent; +import org.apache.fluss.server.coordinator.rebalance.RebalanceExecutionKey; +import org.apache.fluss.server.coordinator.rebalance.goal.ReplicaDistributionGoal; +import org.apache.fluss.server.entity.NotifyLeaderAndIsrResultForBucket; +import org.apache.fluss.server.zk.data.BucketAssignment; +import org.apache.fluss.server.zk.data.LeaderAndIsr; +import org.apache.fluss.server.zk.data.RebalanceTask; +import org.apache.fluss.server.zk.data.TableAssignment; +import org.apache.fluss.server.zk.data.ZkData; +import org.apache.fluss.types.DataTypes; +import org.apache.fluss.utils.clock.Clock; +import org.apache.fluss.utils.clock.ManualClock; +import org.apache.fluss.utils.clock.SystemClock; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import javax.annotation.Nullable; + +import java.time.Duration; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentLinkedDeque; + +import static org.apache.fluss.testutils.common.CommonTestUtils.retry; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for rebalance execution and recovery through the coordinator event loop. */ +class CoordinatorRebalanceTest extends CoordinatorEventProcessorTestBase { + + private static final int REPLICATION_FACTOR = 3; + + private static final TableDescriptor TEST_TABLE = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("a", DataTypes.INT()) + .primaryKey("a") + .build()) + .distributedBy(3, "a") + .property(ConfigOptions.TABLE_KV_STANDBY_REPLICA_ENABLED.key(), "true") + .build() + .withReplicationFactor(REPLICATION_FACTOR); + + private final ConcurrentLinkedDeque pendingResponses = + new ConcurrentLinkedDeque<>(); + private final ConcurrentLinkedDeque recoveredResponses = + new ConcurrentLinkedDeque<>(); + private Clock testClock = SystemClock.getInstance(); + + @AfterEach + void cleanupRebalance() throws Exception { + initCoordinatorChannel(); + drainPendingNotifyTriggers(pendingResponses); + drainPendingNotifyTriggers(recoveredResponses); + zookeeperClient.deleteRebalanceTask(); + if (Arrays.stream(zookeeperClient.getSortedTabletServerList()).noneMatch(id -> id == 2)) { + registerTabletServer(2); + } + ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(ZkData.ServerIdZNode.path(3)); + } + + @Test + void testTimedOutReassignmentRetriesPhaseAIdempotently() throws Exception { + registerTabletServer(3); + + initCoordinatorChannel(); + + long tableId = createTableWithReplicas(Arrays.asList(0, 1, 3), 1); + TableBucket tableBucket = new TableBucket(tableId, 0); + installBlockingNotifyGateways(pendingResponses); + + RebalancePlanForBucket plan = + new RebalancePlanForBucket( + tableBucket, 0, 0, Arrays.asList(0, 1, 3), Arrays.asList(0, 1, 2)); + eventProcessor + .getRebalanceManager() + .registerRebalance( + "timed-out-retry-test", Collections.singletonMap(tableBucket, plan)); + RebalanceExecutionKey executionKey = + eventProcessor.getRebalanceManager().getExecutionKey(tableBucket); + retry(Duration.ofMinutes(1), () -> assertThat(pendingResponses).isNotEmpty()); + int requestsAfterInitialPhaseA = pendingResponses.size(); + int epochAfterInitialPhaseA = + fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get().bucketEpoch()); + + eventProcessor + .getCoordinatorEventManager() + .put(new RebalanceTaskTimeoutEvent(executionKey)); + retry( + Duration.ofMinutes(1), + () -> assertThat(rebalanceStatus(tableBucket)).isEqualTo(RebalanceStatus.TIMEOUT)); + retry( + Duration.ofMinutes(1), + () -> + assertThat(pendingResponses.size()) + .isGreaterThan(requestsAfterInitialPhaseA)); + int requestsAfterFirstRetry = pendingResponses.size(); + int epochAfterFirstRetry = + fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get().bucketEpoch()); + assertThat(epochAfterFirstRetry).isEqualTo(epochAfterInitialPhaseA); + List assignmentAfterFirstRetry = fromCtx(ctx -> ctx.getAssignment(tableBucket)); + assertThat(assignmentAfterFirstRetry).containsExactly(0, 1, 2, 3); + + eventProcessor + .getCoordinatorEventManager() + .put(new ReconcileRebalanceTaskEvent(executionKey)); + retry( + Duration.ofMinutes(1), + () -> assertThat(pendingResponses.size()).isGreaterThan(requestsAfterFirstRetry)); + assertThat(eventProcessor.getRebalanceManager().getExecutionKey(tableBucket)) + .isEqualTo(executionKey); + int epochAfterDuplicateRetry = + fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get().bucketEpoch()); + assertThat(epochAfterDuplicateRetry).isEqualTo(epochAfterInitialPhaseA); + List assignmentAfterDuplicateRetry = + fromCtx(ctx -> ctx.getAssignment(tableBucket)); + assertThat(assignmentAfterDuplicateRetry).containsExactly(0, 1, 2, 3); + + drainPendingNotifyTriggers(pendingResponses); + fromCtx(ctx -> null); + adjustRebalanceIsr(tableBucket, Arrays.asList(0, 1, 2, 3)); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()).isTrue(); + + ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(ZkData.ServerIdZNode.path(2)); + retryVerifyContext(ctx -> assertThat(ctx.liveTabletServerSet()).doesNotContain(2)); + drainPendingNotifyTriggers(pendingResponses); + fromCtx(ctx -> null); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()).isTrue(); + assertThat(pendingResponses).isEmpty(); + List assignmentWhileTargetOffline = fromCtx(ctx -> ctx.getAssignment(tableBucket)); + assertThat(assignmentWhileTargetOffline).containsExactly(0, 1, 2); + + registerTabletServer(2); + retryVerifyContext(ctx -> assertThat(ctx.liveTabletServerSet()).contains(2)); + drainPendingNotifyTriggers(pendingResponses); + fromCtx(ctx -> null); + adjustRebalanceIsr(tableBucket, Arrays.asList(0, 1, 2)); + drainPendingNotifyTriggers(pendingResponses); + retry( + Duration.ofMinutes(1), + () -> + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) + .isFalse()); + fromCtx(ctx -> null); + List finalAssignment = fromCtx(ctx -> ctx.getAssignment(tableBucket)); + assertThat(finalAssignment).containsExactly(0, 1, 2); + verifyIsr(tableBucket, 0, Arrays.asList(0, 1, 2)); + } + + @Test + void testTimedOutRebalanceCompletesWhenTableIsBeingDeleted() throws Exception { + registerTabletServer(3); + initCoordinatorChannel(); + long tableId = createTableWithReplicas(Arrays.asList(0, 1, 2), 1); + TableBucket tableBucket = new TableBucket(tableId, 0); + installBlockingNotifyGateways(pendingResponses); + RebalancePlanForBucket plan = + new RebalancePlanForBucket( + tableBucket, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 3)); + eventProcessor + .getRebalanceManager() + .registerRebalance( + "delete-during-rebalance", Collections.singletonMap(tableBucket, plan)); + RebalanceExecutionKey executionKey = + eventProcessor.getRebalanceManager().getExecutionKey(tableBucket); + + retry(Duration.ofMinutes(1), () -> assertThat(pendingResponses).isNotEmpty()); + fromCtx( + ctx -> { + ctx.queueTableDeletion(Collections.singleton(tableId)); + return null; + }); + eventProcessor + .getCoordinatorEventManager() + .put(new RebalanceTaskTimeoutEvent(executionKey)); + + retry( + Duration.ofMinutes(1), + () -> + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) + .isFalse()); + assertThat(rebalanceStatus(tableBucket)).isEqualTo(RebalanceStatus.COMPLETED); + } + + @Test + void testRebalanceCompletesWhenPartitionIsBeingDeleted() throws Exception { + TableBucket tableBucket = new TableBucket(123L, 456L, 0); + fromCtx( + ctx -> { + ctx.updateBucketReplicaAssignment(tableBucket, Arrays.asList(0, 1, 2)); + ctx.queuePartitionDeletion( + Collections.singleton(new TablePartition(123L, 456L))); + return null; + }); + RebalancePlanForBucket plan = + new RebalancePlanForBucket( + tableBucket, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(1, 2, 0)); + eventProcessor + .getRebalanceManager() + .registerRebalance( + "partition-delete-during-rebalance", + Collections.singletonMap(tableBucket, plan)); + + retry( + Duration.ofMinutes(1), + () -> + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) + .isFalse()); + assertThat(rebalanceStatus(tableBucket)).isEqualTo(RebalanceStatus.COMPLETED); + } + + @ParameterizedTest(name = "recover {0}") + @ValueSource(strings = {"A", "B", "FINAL", "LEADER"}) + void testRestartResumesRebalanceAndWaitsForLeaderAck(String phase) throws Exception { + registerTabletServer(3); + initCoordinatorChannel(); + List originReplicas = Arrays.asList(0, 1, 3); + boolean leaderOnly = "LEADER".equals(phase); + int targetLeader = leaderOnly ? 1 : 0; + List targetReplicas = leaderOnly ? Arrays.asList(1, 0, 3) : Arrays.asList(0, 1, 2); + List finalAssignment = leaderOnly ? originReplicas : targetReplicas; + List unionReplicas = Arrays.asList(0, 1, 2, 3); + long tableId = createTableWithReplicas(originReplicas, 1); + TableBucket tableBucket = new TableBucket(tableId, 0); + installBlockingNotifyGateways(pendingResponses); + + String rebalanceId = "restart-phase-" + phase; + Map plan = + Collections.singletonMap( + tableBucket, + new RebalancePlanForBucket( + tableBucket, 0, targetLeader, originReplicas, targetReplicas)); + RebalanceTask persistedTask = + new RebalanceTask(rebalanceId, RebalanceStatus.NOT_STARTED, plan); + registerRebalance(rebalanceId, plan); + assertThat(zookeeperClient.getTableAssignment(tableId).get().getBucketAssignment(0)) + .isEqualTo(new BucketAssignment(leaderOnly ? originReplicas : unionReplicas)); + + if ("B".equals(phase) || "FINAL".equals(phase)) { + // Reach Phase B through ISR expansion, but hold back the final leader ACK. + adjustRebalanceIsr(tableBucket, unionReplicas); + verifyIsr(tableBucket, targetLeader, targetReplicas); + // A target falls out of ISR before completion. The persisted assignment is at + // the target, but the next coordinator must still resume this unfinished phase. + if ("B".equals(phase)) { + adjustRebalanceIsr(tableBucket, Arrays.asList(0, 1)); + } + } + List intermediateAssignment = + leaderOnly ? originReplicas : "A".equals(phase) ? unionReplicas : targetReplicas; + List intermediateIsr = + "B".equals(phase) + ? Arrays.asList(0, 1) + : "FINAL".equals(phase) ? targetReplicas : originReplicas; + verifyIsr(tableBucket, targetLeader, intermediateIsr); + LeaderAndIsr beforeRestart = fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get()); + assertThat(zookeeperClient.getTableAssignment(tableId).get().getBucketAssignment(0)) + .isEqualTo(new BucketAssignment(intermediateAssignment)); + assertThat(rebalanceStatus(tableBucket)).isEqualTo(RebalanceStatus.REBALANCING); + assertThat(zookeeperClient.getRebalanceTask()).hasValue(persistedTask); + + retry( + Duration.ofMinutes(1), + () -> assertThat(hasPendingNotifyTrigger(pendingResponses, targetLeader)).isTrue()); + PendingNotify oldResponse = + pendingResponses.stream() + .filter(response -> response.getResponseServerId() == targetLeader) + .findFirst() + .get(); + NotifyLeaderAndIsrRequest oldRequest = oldResponse.getRequest(); + assertThat(oldRequest.getNotifyBucketsLeaderReqsList()).hasSize(1); + PbNotifyLeaderAndIsrReqForBucket oldBucketRequest = + oldRequest.getNotifyBucketsLeaderReqsList().get(0); + RebalanceExecutionKey executionKey = + eventProcessor.getRebalanceManager().getExecutionKey(tableBucket); + + // Restart must resume the persisted plan without explicit registration. + restartCoordinator(recoveredResponses); + assertThat(eventProcessor.getCoordinatorEpoch()) + .isGreaterThan(oldRequest.getCoordinatorEpoch()); + fromCtx(ctx -> null); + assertThat(eventProcessor.getRebalanceManager().getExecutionKey(tableBucket)) + .isEqualTo(executionKey); + verifyIsr(tableBucket, targetLeader, intermediateIsr); + LeaderAndIsr afterRecovery = fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get()); + assertThat(afterRecovery.leaderEpoch()).isEqualTo(beforeRestart.leaderEpoch()); + assertThat(afterRecovery.bucketEpoch()).isEqualTo(beforeRestart.bucketEpoch()); + List recoveredAssignment = fromCtx(ctx -> ctx.getAssignment(tableBucket)); + assertThat(recoveredAssignment).containsExactlyElementsOf(intermediateAssignment); + + // The new leader reports catch-up. Keep its final ACK blocked so a stale success + // would visibly and incorrectly complete the recovered attempt. + if ("A".equals(phase) || "B".equals(phase)) { + adjustRebalanceIsr(tableBucket, "B".equals(phase) ? targetReplicas : unionReplicas); + } + verifyIsr(tableBucket, targetLeader, targetReplicas); + LeaderAndIsr current = fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get()); + NotifyLeaderAndIsrRequestContext oldContext = + new NotifyLeaderAndIsrRequestContext( + oldRequest.getCoordinatorEpoch(), + oldBucketRequest.getLeader(), + oldBucketRequest.getLeaderEpoch(), + oldBucketRequest.getBucketEpoch(), + executionKey); + // Also isolate the coordinator-epoch fence: after restart the rebalance id and + // bucket are unchanged, so rejecting a different execution key is insufficient. + NotifyLeaderAndIsrRequestContext oldEpochContext = + new NotifyLeaderAndIsrRequestContext( + oldRequest.getCoordinatorEpoch(), + current.leader(), + current.leaderEpoch(), + current.bucketEpoch(), + executionKey); + oldResponse.complete(); + for (NotifyLeaderAndIsrRequestContext staleContext : + Arrays.asList(oldContext, oldEpochContext)) { + notifyLeaderResponse(tableBucket, current.leader(), staleContext); + assertThat(rebalanceStatus(tableBucket)).isEqualTo(RebalanceStatus.TIMEOUT); + assertThat(zookeeperClient.getRebalanceTask()).hasValue(persistedTask); + } + + retry( + Duration.ofMinutes(1), + () -> { + drainPendingNotifyTriggers(recoveredResponses); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) + .isFalse(); + assertThat(zookeeperClient.getRebalanceTask()) + .hasValue( + new RebalanceTask( + rebalanceId, RebalanceStatus.COMPLETED, plan)); + }); + assertThat(rebalanceStatus(tableBucket)).isEqualTo(RebalanceStatus.COMPLETED); + List assignment = fromCtx(ctx -> ctx.getAssignment(tableBucket)); + assertThat(assignment).containsExactlyElementsOf(finalAssignment); + assertThat(zookeeperClient.getTableAssignment(tableId).get().getBucketAssignment(0)) + .isEqualTo(new BucketAssignment(finalAssignment)); + verifyIsr(tableBucket, targetLeader, targetReplicas); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testNewRebalanceConvergesAfterFailureAndRestart(boolean targetUnavailable) + throws Exception { + ManualClock clock = new ManualClock(0L); + testClock = clock; + eventProcessor.shutdown(); + zookeeperClient.deleteRebalanceTask(); + eventProcessor = buildCoordinatorEventProcessor(clock); + eventProcessor.startup(); + registerTabletServer(3); + initCoordinatorChannel(); + List origin = Arrays.asList(0, 1, 2); + long tableId = createTableWithReplicas(origin, 1); + TableBucket bucket = new TableBucket(tableId, 0); + installBlockingNotifyGateways(pendingResponses); + Map plan = + Collections.singletonMap( + bucket, + new RebalancePlanForBucket(bucket, 0, 0, origin, Arrays.asList(0, 1, 3))); + String failedId = "failed-before-restart"; + registerRebalance(failedId, plan); + if (targetUnavailable) { + ZOO_KEEPER_EXTENSION_WRAPPER + .getCustomExtension() + .cleanupPath(ZkData.ServerIdZNode.path(3)); + retryVerifyContext(ctx -> assertThat(ctx.liveTabletServerSet()).doesNotContain(3)); + } + RebalanceExecutionKey key = eventProcessor.getRebalanceManager().getExecutionKey(bucket); + eventProcessor.getCoordinatorEventManager().put(new RebalanceTaskTimeoutEvent(key)); + retry( + Duration.ofMinutes(1), + () -> assertThat(rebalanceStatus(bucket)).isEqualTo(RebalanceStatus.TIMEOUT)); + // Wait for the immediate reconciliation to establish the unavailable-target timer. + fromCtx(ctx -> null); + clock.advanceTime(targetUnavailable ? Duration.ofMinutes(31) : Duration.ofHours(25)); + eventProcessor.getCoordinatorEventManager().put(new ReconcileRebalanceTaskEvent(key)); + retry( + Duration.ofMinutes(1), + () -> assertThat(rebalanceStatus(bucket)).isEqualTo(RebalanceStatus.FAILED)); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()).isFalse(); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(RebalanceStatus.FAILED); + assertThat( + zookeeperClient + .getTableAssignment(tableId) + .get() + .getBucketAssignment(0) + .getReplicas()) + .containsExactlyInAnyOrder(0, 1, 2, 3); + + restartCoordinator(null); + // Use the normal request path, including model construction, optimization and + // execution. + CompletableFuture response = new CompletableFuture<>(); + eventProcessor + .getCoordinatorEventManager() + .put( + new RebalanceEvent( + Collections.singletonList(new ReplicaDistributionGoal()), + response)); + response.get(); + assertThat(eventProcessor.getRebalanceManager().getRebalanceId()).isNotEqualTo(failedId); + retry( + Duration.ofMinutes(1), + () -> { + assertThat(rebalanceStatus(bucket)).isEqualTo(RebalanceStatus.COMPLETED); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) + .isFalse(); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(RebalanceStatus.COMPLETED); + }); + assertThat(zookeeperClient.getTableAssignment(tableId).get().getBucketAssignment(0)) + .isEqualTo(new BucketAssignment(origin)); + verifyIsr(bucket, 0, origin); + } + + @ParameterizedTest(name = "leaderOnly={0}") + @ValueSource(booleans = {false, true}) + void testRestartAfterCancellationDoesNotStartPendingBucketAfterLeaderChange(boolean leaderOnly) + throws Exception { + registerTabletServer(3); + initCoordinatorChannel(); + List origin = Arrays.asList(0, 1, 2); + List target = Arrays.asList(0, 1, 3); + long tableId = createTableWithReplicas(origin, 2); + TableBucket running = new TableBucket(tableId, 0); + TableBucket pending = new TableBucket(tableId, 1); + // A leader election must not make the pending plan look started. For a leader-only plan, + // the ordinary election deliberately chooses the plan's target before it is admitted. + int pendingLeader = 1; + List pendingIsr = leaderOnly ? origin : Arrays.asList(0, 1); + if (!leaderOnly) { + adjustRebalanceIsr(pending, pendingIsr); + } + changeLeader(pending, pendingLeader); + installBlockingNotifyGateways(pendingResponses); + Map plan = new LinkedHashMap<>(); + plan.put(running, new RebalancePlanForBucket(running, 0, 0, origin, target)); + plan.put( + pending, + leaderOnly + ? new RebalancePlanForBucket(pending, 0, 1, origin, Arrays.asList(1, 0, 2)) + : new RebalancePlanForBucket(pending, 0, 0, origin, target)); + String rebalanceId = "cancel-before-restart"; + registerRebalance(rebalanceId, plan); + fromCtx( + ctx -> { + eventProcessor.getRebalanceManager().cancelRebalance(rebalanceId); + return null; + }); + assertThat(rebalanceStatus(pending)).isEqualTo(RebalanceStatus.CANCELED); + assertThat(zookeeperClient.getRebalanceTask().get().isCancelRequested()).isTrue(); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()).isTrue(); + + restartCoordinator(recoveredResponses); + retry( + Duration.ofMinutes(1), + () -> + assertThat(eventProcessor.getRebalanceManager().getExecutionKey(running)) + .isNotNull()); + assertThat(rebalanceStatus(pending)).isEqualTo(RebalanceStatus.CANCELED); + assertThat(eventProcessor.getRebalanceManager().getExecutionKey(pending)).isNull(); + assertThat(zookeeperClient.getTableAssignment(tableId).get().getBucketAssignment(1)) + .isEqualTo(new BucketAssignment(origin)); + verifyIsr(pending, pendingLeader, pendingIsr); + + // A later failover must not revive the canceled plan and move the leader back to its old + // target. + int expectedFinalPendingLeader = pendingLeader; + if (leaderOnly) { + expectedFinalPendingLeader = 2; + changeLeader(pending, expectedFinalPendingLeader); + } + + adjustRebalanceIsr(running, Arrays.asList(0, 1, 2, 3)); + retry( + Duration.ofMinutes(1), + () -> { + drainPendingNotifyTriggers(recoveredResponses); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) + .isFalse(); + }); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(RebalanceStatus.CANCELED); + assertThat(zookeeperClient.getTableAssignment(tableId).get().getBucketAssignment(1)) + .isEqualTo(new BucketAssignment(origin)); + verifyIsr(pending, expectedFinalPendingLeader, pendingIsr); + verifyIsr(running, 0, target); + } + + @Test + void testCancelAfterRestartDrainsEveryIntermediateBucket() throws Exception { + registerTabletServer(3); + initCoordinatorChannel(); + List origin = Arrays.asList(0, 1, 3); + List target = Arrays.asList(0, 1, 2); + List union = Arrays.asList(0, 1, 2, 3); + long tableId = createTableWithReplicas(origin, 2); + TableBucket first = new TableBucket(tableId, 0); + TableBucket second = new TableBucket(tableId, 1); + installBlockingNotifyGateways(pendingResponses); + Map plans = new LinkedHashMap<>(); + for (TableBucket bucket : Arrays.asList(first, second)) { + plans.put(bucket, new RebalancePlanForBucket(bucket, 0, 0, origin, target)); + } + String rebalanceId = "review-cancel-after-restart"; + registerRebalance(rebalanceId, plans); + RebalanceExecutionKey key = eventProcessor.getRebalanceManager().getExecutionKey(first); + eventProcessor.getCoordinatorEventManager().put(new RebalanceTaskTimeoutEvent(key)); + fromCtx(ctx -> null); + for (TableBucket bucket : plans.keySet()) { + assertThat( + zookeeperClient + .getTableAssignment(tableId) + .get() + .getBucketAssignment(bucket.getBucket())) + .isEqualTo(new BucketAssignment(union)); + } + + restartCoordinator(recoveredResponses); + retry( + Duration.ofMinutes(1), + () -> + assertThat(eventProcessor.getRebalanceManager().getRebalanceId()) + .isEqualTo(rebalanceId)); + TableBucket recoveredIntermediate = + plans.keySet().stream() + .filter(bucket -> rebalanceStatus(bucket) == RebalanceStatus.TIMEOUT) + .findFirst() + .get(); + assertThat( + zookeeperClient + .getTableAssignment(tableId) + .get() + .getBucketAssignment(recoveredIntermediate.getBucket())) + .isEqualTo(new BucketAssignment(union)); + fromCtx( + ctx -> { + assertThat( + eventProcessor.isRebalanceTaskAtOrigin( + plans.get(recoveredIntermediate))) + .isFalse(); + eventProcessor.getRebalanceManager().cancelRebalance(rebalanceId); + return null; + }); + assertThat(rebalanceStatus(recoveredIntermediate)) + .as("A previously started migration still needs draining after recovery") + .isNotEqualTo(RebalanceStatus.CANCELED); + for (TableBucket bucket : plans.keySet()) { + adjustRebalanceIsr(bucket, union); + } + retry( + Duration.ofMinutes(1), + () -> { + drainPendingNotifyTriggers(recoveredResponses); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) + .isFalse(); + }); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(RebalanceStatus.CANCELED); + for (TableBucket bucket : plans.keySet()) { + assertThat(rebalanceStatus(bucket)).isEqualTo(RebalanceStatus.COMPLETED); + assertThat( + zookeeperClient + .getTableAssignment(tableId) + .get() + .getBucketAssignment(bucket.getBucket())) + .isEqualTo(new BucketAssignment(target)); + verifyIsr(bucket, 0, target); + } + } + + private void adjustRebalanceIsr(TableBucket tableBucket, List isr) throws Exception { + LeaderAndIsr current = fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get()); + assertThat( + submitAdjustIsr( + tableBucket, + new LeaderAndIsr( + current.leader(), + current.leaderEpoch(), + isr, + Collections.emptyList(), + current.coordinatorEpoch(), + current.bucketEpoch())) + .succeeded()) + .isTrue(); + } + + @Test + void testRebalanceAtOriginAllowsShrunkIsr() throws Exception { + TableBucket tableBucket = new TableBucket(987L, 0); + RebalancePlanForBucket plan = + new RebalancePlanForBucket( + tableBucket, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 0, 3)); + putBucketState(tableBucket, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2)); + assertThat(eventProcessor.isRebalanceTaskAtOrigin(plan)).isTrue(); + putBucketState(tableBucket, 0, Arrays.asList(0, 1), Arrays.asList(0, 1, 2)); + assertThat(eventProcessor.isRebalanceTaskAtOrigin(plan)).isTrue(); + putBucketState(tableBucket, 0, Arrays.asList(0, 1, 2, 3), Arrays.asList(0, 1, 2)); + assertThat(eventProcessor.isRebalanceTaskAtOrigin(plan)).isFalse(); + putBucketState(tableBucket, 1, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2)); + assertThat(eventProcessor.isRebalanceTaskAtOrigin(plan)).isTrue(); + + RebalancePlanForBucket leaderOnlyPlan = + new RebalancePlanForBucket( + tableBucket, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2)); + // The target leader is elected: wait for its acknowledgement instead of canceling. + assertThat(eventProcessor.isRebalanceTaskAtOrigin(leaderOnlyPlan)).isFalse(); + // An unrelated election does not apply a leader-only plan. + putBucketState(tableBucket, 2, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2)); + assertThat(eventProcessor.isRebalanceTaskAtOrigin(leaderOnlyPlan)).isTrue(); + } + + private void changeLeader(TableBucket tableBucket, int newLeader) throws Exception { + LeaderAndIsr current = zookeeperClient.getLeaderAndIsr(tableBucket).get(); + LeaderAndIsr changed = + new LeaderAndIsr( + newLeader, + current.leaderEpoch() + 1, + current.isr(), + current.standbyReplicas(), + current.coordinatorEpoch(), + current.bucketEpoch() + 1); + int coordinatorZkVersion = fromCtx(CoordinatorContext::getCoordinatorZkVersion); + zookeeperClient.updateLeaderAndIsr(tableBucket, changed, coordinatorZkVersion); + fromCtx( + ctx -> { + ctx.putBucketLeaderAndIsr(tableBucket, changed); + return null; + }); + verifyIsr(tableBucket, newLeader, current.isr()); + } + + private void putBucketState( + TableBucket tableBucket, int leader, List isr, List assignment) + throws Exception { + fromCtx( + ctx -> { + ctx.updateBucketReplicaAssignment(tableBucket, assignment); + ctx.putBucketLeaderAndIsr( + tableBucket, + new LeaderAndIsr( + leader, + 1, + isr, + Collections.emptyList(), + ctx.getCoordinatorEpoch(), + 1)); + return null; + }); + } + + @Test + void testStaleNotifyLeaderAndIsrResponseCannotCompleteRebalance() throws Exception { + initCoordinatorChannel(); + + long tableId = createTableWithReplicas(Arrays.asList(0, 1, 2), 1); + TableBucket tableBucket = new TableBucket(tableId, 0); + + installBlockingNotifyGateways(pendingResponses); + RebalancePlanForBucket plan = + new RebalancePlanForBucket( + tableBucket, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 0, 2)); + eventProcessor + .getRebalanceManager() + .registerRebalance( + "stale-response-test", Collections.singletonMap(tableBucket, plan)); + retry(Duration.ofMinutes(1), () -> assertThat(pendingResponses).isNotEmpty()); + + LeaderAndIsr current = fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get()); + NotifyLeaderAndIsrRequestContext staleContext = + new NotifyLeaderAndIsrRequestContext( + eventProcessor.getCoordinatorEpoch(), + current.leader(), + current.leaderEpoch(), + current.bucketEpoch() - 1, + eventProcessor.getRebalanceManager().getExecutionKey(tableBucket)); + notifyLeaderResponse(tableBucket, current.leader(), staleContext); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()).isTrue(); + + NotifyLeaderAndIsrRequestContext oldAttemptContext = + new NotifyLeaderAndIsrRequestContext( + eventProcessor.getCoordinatorEpoch(), + current.leader(), + current.leaderEpoch(), + current.bucketEpoch(), + new RebalanceExecutionKey("old-rebalance", tableBucket)); + notifyLeaderResponse(tableBucket, current.leader(), oldAttemptContext); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()).isTrue(); + + NotifyLeaderAndIsrRequestContext currentContext = + new NotifyLeaderAndIsrRequestContext( + eventProcessor.getCoordinatorEpoch(), + current.leader(), + current.leaderEpoch(), + current.bucketEpoch(), + eventProcessor.getRebalanceManager().getExecutionKey(tableBucket)); + notifyLeaderResponse(tableBucket, current.leader(), currentContext); + retry( + Duration.ofMinutes(1), + () -> + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) + .isFalse()); + drainPendingNotifyTriggers(pendingResponses); + } + + private void notifyLeaderResponse( + TableBucket bucket, int server, NotifyLeaderAndIsrRequestContext context) + throws Exception { + eventProcessor + .getCoordinatorEventManager() + .put( + new NotifyLeaderAndIsrResponseReceivedEvent( + Collections.singletonList( + new NotifyLeaderAndIsrResultForBucket(bucket)), + server, + Collections.singletonMap(bucket, context))); + fromCtx(ctx -> null); + } + + private RebalanceStatus rebalanceStatus(TableBucket tableBucket) { + return eventProcessor + .getRebalanceManager() + .listRebalanceProgress(null) + .progressForBucketMap() + .get(tableBucket) + .status(); + } + + private long createTableWithReplicas(List replicas, int bucketCount) throws Exception { + Map assignments = new HashMap<>(); + for (int bucket = 0; bucket < bucketCount; bucket++) { + assignments.put(bucket, new BucketAssignment(replicas)); + } + long tableId = + metadataManager.createTable( + TablePath.of(defaultDatabase, "rebalance"), + remoteDataDir, + TEST_TABLE, + new TableAssignment(assignments), + false); + for (int bucket = 0; bucket < bucketCount; bucket++) { + verifyIsr(new TableBucket(tableId, bucket), replicas.get(0), replicas); + } + return tableId; + } + + private void registerRebalance(String id, Map plans) + throws Exception { + zookeeperClient.registerRebalanceTask( + new RebalanceTask(id, RebalanceStatus.NOT_STARTED, plans)); + fromCtx( + ctx -> { + eventProcessor.getRebalanceManager().registerRebalance(id, plans); + return null; + }); + } + + private void restartCoordinator(@Nullable ConcurrentLinkedDeque responses) + throws Exception { + eventProcessor.shutdown(); + testCoordinatorChannelManager.close(); + zkEpoch = zookeeperClient.fenceBecomeCoordinatorLeader("2"); + testCoordinatorChannelManager = new TestCoordinatorChannelManager(); + if (responses == null) { + initCoordinatorChannel(); + } else { + installBlockingNotifyGateways(responses); + } + eventProcessor = buildCoordinatorEventProcessor(testClock); + eventProcessor.startup(); + fromCtx(ctx -> null); + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestingControlledNotifyGateway.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestingControlledNotifyGateway.java new file mode 100644 index 00000000000..8e6a0487a10 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestingControlledNotifyGateway.java @@ -0,0 +1,77 @@ +/* + * 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.fluss.server.coordinator; + +import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrRequest; +import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrResponse; +import org.apache.fluss.server.tablet.TestTabletServerGateway; + +import java.util.Collections; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentLinkedDeque; + +/** A gateway that can hold notifications until a test releases their responses. */ +final class TestingControlledNotifyGateway extends TestTabletServerGateway { + private final int responseServerId; + private final ConcurrentLinkedDeque pendingTriggers; + + TestingControlledNotifyGateway( + int responseServerId, ConcurrentLinkedDeque pendingTriggers) { + super(false, Collections.emptySet()); + this.responseServerId = responseServerId; + this.pendingTriggers = pendingTriggers; + } + + @Override + public CompletableFuture notifyLeaderAndIsr( + NotifyLeaderAndIsrRequest request) { + NotifyLeaderAndIsrResponse response = super.notifyLeaderAndIsr(request).join(); + PendingNotify trigger = new PendingNotify(responseServerId, request); + pendingTriggers.add(trigger); + return trigger.getFuture().thenApply(ignored -> response); + } + + /** A captured request whose response is released explicitly by the test. */ + static final class PendingNotify { + private final int responseServerId; + private final NotifyLeaderAndIsrRequest request; + private final CompletableFuture future = new CompletableFuture<>(); + + PendingNotify(int responseServerId, NotifyLeaderAndIsrRequest request) { + this.responseServerId = responseServerId; + this.request = request; + } + + int getResponseServerId() { + return responseServerId; + } + + NotifyLeaderAndIsrRequest getRequest() { + return request; + } + + CompletableFuture getFuture() { + return future; + } + + void complete() { + future.complete(null); + } + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestingFailingNotifyGateway.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestingFailingNotifyGateway.java new file mode 100644 index 00000000000..02708480a8a --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestingFailingNotifyGateway.java @@ -0,0 +1,47 @@ +/* + * 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.fluss.server.coordinator; + +import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrRequest; +import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrResponse; +import org.apache.fluss.server.tablet.TestTabletServerGateway; + +import java.util.Collections; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; + +/** A failing gateway that records the number of notification attempts. */ +final class TestingFailingNotifyGateway extends TestTabletServerGateway { + private final AtomicInteger notifyLeaderAndIsrCount = new AtomicInteger(); + + TestingFailingNotifyGateway() { + super(true, Collections.emptySet()); + } + + int getNotifyLeaderAndIsrCount() { + return notifyLeaderAndIsrCount.get(); + } + + @Override + public CompletableFuture notifyLeaderAndIsr( + NotifyLeaderAndIsrRequest request) { + notifyLeaderAndIsrCount.incrementAndGet(); + return super.notifyLeaderAndIsr(request); + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerITCase.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerITCase.java index 810f9d30e7f..e0211840054 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerITCase.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerITCase.java @@ -18,7 +18,6 @@ package org.apache.fluss.server.coordinator.rebalance; import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket; -import org.apache.fluss.cluster.rebalance.RebalanceStatus; import org.apache.fluss.cluster.rebalance.ServerTag; import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; @@ -177,9 +176,7 @@ void testRebalanceWithRemoteLog() throws Exception { assertThat(newReplicaStates).allMatch(Objects::isNull); rebalanceManager.registerRebalance( - "test-rebalance-dsds", - Collections.singletonMap(tb, planForBucket), - RebalanceStatus.NOT_STARTED); + "test-rebalance-dsds", Collections.singletonMap(tb, planForBucket)); retry( Duration.ofMinutes(2), diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerTest.java index b88116d7e88..5009009c880 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerTest.java @@ -17,59 +17,78 @@ package org.apache.fluss.server.coordinator.rebalance; +import org.apache.fluss.cluster.Endpoint; +import org.apache.fluss.cluster.ServerType; import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket; import org.apache.fluss.cluster.rebalance.RebalanceResultForBucket; import org.apache.fluss.cluster.rebalance.RebalanceStatus; +import org.apache.fluss.cluster.rebalance.ServerTag; +import org.apache.fluss.config.ConfigOption; import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; import org.apache.fluss.metadata.TableBucket; -import org.apache.fluss.server.coordinator.AutoPartitionManager; +import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TablePath; import org.apache.fluss.server.coordinator.CoordinatorContext; -import org.apache.fluss.server.coordinator.CoordinatorEventProcessor; -import org.apache.fluss.server.coordinator.LakeCatalogDynamicLoader; -import org.apache.fluss.server.coordinator.LakeTableTieringManager; -import org.apache.fluss.server.coordinator.MetadataManager; -import org.apache.fluss.server.coordinator.ReplicaCapacityController; -import org.apache.fluss.server.coordinator.TestCoordinatorChannelManager; import org.apache.fluss.server.coordinator.event.CoordinatorEvent; import org.apache.fluss.server.coordinator.event.EventManager; +import org.apache.fluss.server.coordinator.event.FinalizeRebalanceEvent; import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent; +import org.apache.fluss.server.coordinator.event.ReconcileRebalanceTaskEvent; import org.apache.fluss.server.coordinator.event.RecoverRebalanceEvent; -import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager; -import org.apache.fluss.server.coordinator.remote.RemoteDirDynamicLoader; -import org.apache.fluss.server.metadata.CoordinatorMetadataCache; -import org.apache.fluss.server.metrics.group.TestingMetricGroups; +import org.apache.fluss.server.coordinator.rebalance.goal.ReplicaDistributionGoal; +import org.apache.fluss.server.metadata.ServerInfo; +import org.apache.fluss.server.zk.CuratorFrameworkWithUnhandledErrorListener; import org.apache.fluss.server.zk.NOPErrorHandler; import org.apache.fluss.server.zk.ZkEpoch; import org.apache.fluss.server.zk.ZooKeeperClient; import org.apache.fluss.server.zk.ZooKeeperExtension; +import org.apache.fluss.server.zk.data.LeaderAndIsr; import org.apache.fluss.server.zk.data.RebalanceTask; +import org.apache.fluss.shaded.curator5.org.apache.curator.retry.ExponentialBackoffRetry; import org.apache.fluss.testutils.common.AllCallbackWrapper; +import org.apache.fluss.testutils.common.ManuallyTriggeredScheduledExecutorService; import org.apache.fluss.utils.clock.ManualClock; -import org.apache.fluss.utils.clock.SystemClock; -import org.apache.fluss.utils.concurrent.ExecutorThreadFactory; -import org.apache.fluss.utils.concurrent.FlussScheduler; -import org.apache.fluss.utils.concurrent.Scheduler; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; -import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import static org.apache.fluss.cluster.rebalance.RebalanceStatus.CANCELED; import static org.apache.fluss.cluster.rebalance.RebalanceStatus.COMPLETED; +import static org.apache.fluss.cluster.rebalance.RebalanceStatus.FAILED; import static org.apache.fluss.cluster.rebalance.RebalanceStatus.NOT_STARTED; +import static org.apache.fluss.cluster.rebalance.RebalanceStatus.REBALANCING; import static org.apache.fluss.cluster.rebalance.RebalanceStatus.TIMEOUT; +import static org.apache.fluss.config.ConfigOptions.COORDINATOR_REBALANCE_MAX_TRACKED_TIMED_OUT_TASKS; +import static org.apache.fluss.config.ConfigOptions.COORDINATOR_REBALANCE_NO_PROGRESS_TIMEOUT; +import static org.apache.fluss.config.ConfigOptions.COORDINATOR_REBALANCE_TARGET_UNAVAILABLE_TIMEOUT; +import static org.apache.fluss.record.TestData.DATA1_TABLE_DESCRIPTOR; +import static org.apache.fluss.server.zk.ZooKeeperUtils.startZookeeperClient; +import static org.apache.fluss.shaded.curator5.org.apache.curator.framework.CuratorFrameworkFactory.Builder; +import static org.apache.fluss.shaded.curator5.org.apache.curator.framework.CuratorFrameworkFactory.builder; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Test for {@link RebalanceManager}. */ public class RebalanceManagerTest { @@ -79,18 +98,13 @@ public class RebalanceManagerTest { new AllCallbackWrapper<>(new ZooKeeperExtension()); private static ZooKeeperClient zookeeperClient; - private static MetadataManager metadataManager; private static ZkEpoch zkEpoch; - private CoordinatorMetadataCache serverMetadataCache; - private TestCoordinatorChannelManager testCoordinatorChannelManager; - private AutoPartitionManager autoPartitionManager; - private ReplicaCapacityController replicaCapacityController; - private LakeTableTieringManager lakeTableTieringManager; + private ManualClock clock; + private TestingRebalanceExecutor rebalanceExecutor; + private RecordingEventManager eventManager; private RebalanceManager rebalanceManager; - private KvSnapshotLeaseManager kvSnapshotLeaseManager; - private Scheduler scheduler; - private Configuration conf; + private ManuallyTriggeredScheduledExecutorService timeoutChecker; @BeforeAll static void baseBeforeAll() throws Exception { @@ -102,60 +116,18 @@ static void baseBeforeAll() throws Exception { } @BeforeEach - void beforeEach() { - serverMetadataCache = new CoordinatorMetadataCache(); - testCoordinatorChannelManager = new TestCoordinatorChannelManager(); - String remoteDataDir = "/tmp/fluss/remote-data"; - conf = new Configuration(); - conf.set(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); - - kvSnapshotLeaseManager = - new KvSnapshotLeaseManager( - Duration.ofMinutes(10).toMillis(), - zookeeperClient, - remoteDataDir, - SystemClock.getInstance(), - TestingMetricGroups.COORDINATOR_METRICS); - kvSnapshotLeaseManager.start(); - - scheduler = new FlussScheduler(1); - scheduler.startup(); - - replicaCapacityController = - new ReplicaCapacityController( - conf, serverMetadataCache, TestingMetricGroups.COORDINATOR_METRICS); - autoPartitionManager = - new AutoPartitionManager( - serverMetadataCache, - metadataManager, - new RemoteDirDynamicLoader(conf), - conf, - replicaCapacityController); - lakeTableTieringManager = - new LakeTableTieringManager(TestingMetricGroups.LAKE_TIERING_METRICS); - CoordinatorEventProcessor eventProcessor = buildCoordinatorEventProcessor(); - RecordingEventManager recordingEventManager = new RecordingEventManager(); - rebalanceManager = - new RebalanceManager( - eventProcessor, - zookeeperClient, - recordingEventManager, - SystemClock.getInstance()); - rebalanceManager.startup(); + void beforeEach() throws Exception { + zookeeperClient.deleteRebalanceTask(); + clock = new ManualClock(); + rebalanceExecutor = new TestingRebalanceExecutor(new CoordinatorContext(zkEpoch)); + eventManager = new RecordingEventManager(); + rebalanceManager = createManager(new Configuration()); } @AfterEach void afterEach() throws Exception { rebalanceManager.close(); - if (scheduler != null) { - scheduler.shutdown(); - } zookeeperClient.deleteRebalanceTask(); - metadataManager = - new MetadataManager( - zookeeperClient, - new Configuration(), - new LakeCatalogDynamicLoader(new Configuration(), null, true)); } @Test @@ -169,7 +141,7 @@ void testRebalanceWithoutTask() throws Exception { assertThat(zookeeperClient.getRebalanceTask()).hasValue(rebalanceTask); // register a rebalance task with empty plan. - rebalanceManager.registerRebalance(rebalanceId, new HashMap<>(), NOT_STARTED); + rebalanceManager.registerRebalance(rebalanceId, new HashMap<>()); assertThat(rebalanceManager.getRebalanceId()).isEqualTo(rebalanceId); RebalanceStatus status = rebalanceManager.getRebalanceStatus(); @@ -180,224 +152,768 @@ void testRebalanceWithoutTask() throws Exception { } @Test - void testStartupQueuesRecoverRebalanceEvent() throws Exception { - ManualClock clock = new ManualClock(0L); - RecordingEventManager eventManager = new RecordingEventManager(); - NoOpScheduledExecutor executor = new NoOpScheduledExecutor(); - CoordinatorEventProcessor eventProcessor = buildCoordinatorEventProcessor(); - - Map plan = createRebalancePlan(2); - RebalanceTask rebalanceTask = new RebalanceTask("recover-test", NOT_STARTED, plan); - zookeeperClient.registerRebalanceTask(rebalanceTask); + void testTimeoutEnqueuesEvent() throws Exception { + TableBucket tb1 = new TableBucket(1L, 0); + TableBucket tb2 = new TableBucket(1L, 1); + Map plan = plans(tb1, tb2); + rebalanceManager.registerRebalance("timeout-test", plan); + RebalanceExecutionKey executionKey = rebalanceManager.getExecutionKey(tb1); - RebalanceManager manager = - new RebalanceManager( - eventProcessor, zookeeperClient, eventManager, clock, executor); - // If startup() finds a pending rebalance task in ZooKeeper, it should enqueue a - // RecoverRebalanceEvent to be processed by the coordinator event thread, instead of - // calling registerRebalance() directly on the startup thread. - manager.startup(); + clock.advanceTime(Duration.ofMillis(100_000)); + rebalanceManager.checkTimeout(); + assertThat(eventManager.events).isEmpty(); - assertThat(eventManager.events).hasSize(1); - assertThat(eventManager.events.get(0)).isInstanceOf(RecoverRebalanceEvent.class); + clock.advanceTime(Duration.ofMillis(30_000)); + rebalanceManager.checkTimeout(); - RecoverRebalanceEvent recoverEvent = (RecoverRebalanceEvent) eventManager.events.get(0); - assertThat(recoverEvent.getRebalanceTask()).isEqualTo(rebalanceTask); + assertThat(eventManager.events).hasSize(1); + assertThat(eventManager.events.get(0)).isInstanceOf(RebalanceTaskTimeoutEvent.class); + RebalanceTaskTimeoutEvent timeoutEvent = + (RebalanceTaskTimeoutEvent) eventManager.events.get(0); + assertThat(timeoutEvent.getExecutionKey()).isEqualTo(executionKey); - manager.close(); + clock.advanceTime(Duration.ofMillis(30_000)); + rebalanceManager.checkTimeout(); + assertThat(eventManager.events).hasSize(1); } - private Map createRebalancePlan(int taskCount) { - Map plan = new HashMap<>(); - for (int i = 0; i < taskCount; i++) { - TableBucket tb = new TableBucket(1L, i); - plan.put( - tb, - new RebalancePlanForBucket( - tb, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2))); - } - return plan; + @Test + void testSoftTimeoutAdmitsNextTaskAndTracksLateCompletion() throws Exception { + TableBucket tb1 = new TableBucket(1L, 0); + TableBucket tb2 = new TableBucket(1L, 1); + rebalanceManager.registerRebalance("soft-timeout-test", plans(tb1, tb2)); + RebalanceExecutionKey firstAttempt = rebalanceManager.getExecutionKey(tb1); + assertThat(rebalanceExecutor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(tb1); + + clock.advanceTime(Duration.ofMillis(130_000)); + rebalanceManager.checkTimeout(); + RebalanceTaskTimeoutEvent timeoutEvent = + (RebalanceTaskTimeoutEvent) eventManager.events.get(0); + assertThat(rebalanceManager.timeoutRebalanceTask(timeoutEvent.getExecutionKey())).isTrue(); + + RebalanceExecutionKey secondAttempt = rebalanceManager.getExecutionKey(tb2); + assertThat(secondAttempt).isNotNull(); + assertThat(rebalanceExecutor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(tb1, tb2); + assertThat(rebalanceManager.listRebalanceProgress(null).status()).isEqualTo(REBALANCING); + assertThat( + rebalanceManager + .listRebalanceProgress(null) + .progressForBucketMap() + .get(tb1) + .status()) + .isEqualTo(TIMEOUT); + assertThat(eventManager.events.get(1)).isInstanceOf(ReconcileRebalanceTaskEvent.class); + + RebalancePlanForBucket retryPlan = rebalanceManager.getPlanForReconciliation(firstAttempt); + assertThat(retryPlan).isNotNull(); + assertThat(retryPlan.getTableBucket()).isEqualTo(tb1); + // the dispatched reconciliation backs off, so no event is enqueued right away. + rebalanceManager.checkTimeout(); + assertThat(eventManager.events).hasSize(2); + + clock.advanceTime(Duration.ofMillis(30_000)); + rebalanceManager.checkTimeout(); + assertThat(eventManager.events).hasSize(3); + assertThat(eventManager.events.get(2)).isInstanceOf(ReconcileRebalanceTaskEvent.class); + assertThat(((ReconcileRebalanceTaskEvent) eventManager.events.get(2)).getExecutionKey()) + .isEqualTo(firstAttempt); + + assertThat(rebalanceManager.timeoutRebalanceTask(firstAttempt)).isFalse(); + assertThat(rebalanceManager.finishRebalanceTask(firstAttempt, COMPLETED)).isTrue(); + assertThat(rebalanceManager.finishRebalanceTask(firstAttempt, COMPLETED)).isFalse(); + assertThat(rebalanceManager.finishRebalanceTask(secondAttempt, COMPLETED)).isTrue(); + + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(COMPLETED); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(COMPLETED); } @Test - void testTimeoutEnqueuesEvent() throws Exception { - ManualClock clock = new ManualClock(0L); - RecordingEventManager eventManager = new RecordingEventManager(); - NoOpScheduledExecutor executor = new NoOpScheduledExecutor(); - CoordinatorEventProcessor eventProcessor = buildCoordinatorEventProcessor(); + void testFailureIsAggregatedIntoOverallStatus() throws Exception { + TableBucket tb1 = new TableBucket(1L, 0); + TableBucket tb2 = new TableBucket(1L, 1); + rebalanceManager.registerRebalance("failed-test", plans(tb1, tb2)); + rebalanceManager.finishRebalanceTask(tb1, FAILED); + rebalanceManager.finishRebalanceTask(tb2, COMPLETED); - RebalanceManager manager = - new RebalanceManager( - eventProcessor, zookeeperClient, eventManager, clock, executor); - manager.startup(); + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(FAILED); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()).isEqualTo(FAILED); + } + @Test + void testCancelPersistsIntentAndDrainsOnlyAdmittedTasks() throws Exception { TableBucket tb1 = new TableBucket(1L, 0); TableBucket tb2 = new TableBucket(1L, 1); - Map plan = new HashMap<>(); - plan.put( - tb1, - new RebalancePlanForBucket( - tb1, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2))); - plan.put( - tb2, - new RebalancePlanForBucket( - tb2, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2))); - - zookeeperClient.registerRebalanceTask(new RebalanceTask("timeout-test", NOT_STARTED, plan)); - manager.registerRebalance("timeout-test", plan, NOT_STARTED); - - // Not yet timed out. - clock.advanceTime(Duration.ofMillis(100_000)); - manager.checkTimeout(); - assertThat(eventManager.events).isEmpty(); + // A fresh pending task has never executed, even if failover changed its origin state. + rebalanceManager.registerRebalance("cancel-test", plans(tb1, tb2)); + RebalanceExecutionKey runningAttempt = rebalanceManager.getExecutionKey(tb1); + + rebalanceManager.cancelRebalance("cancel-test"); + + RebalanceTask storedTask = zookeeperClient.getRebalanceTask().get(); + assertThat(storedTask.getRebalanceStatus()).isEqualTo(REBALANCING); + assertThat(storedTask.isCancelRequested()).isTrue(); + assertThat(storedTask.getCanceledBuckets()).containsExactly(tb2); + assertThat(rebalanceManager.isCancelRequested()).isTrue(); + assertThat( + rebalanceManager + .listRebalanceProgress(null) + .progressForBucketMap() + .get(tb2) + .status()) + .isEqualTo(CANCELED); + assertThat(rebalanceExecutor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(tb1); + + rebalanceManager.finishRebalanceTask(runningAttempt, COMPLETED); + + storedTask = zookeeperClient.getRebalanceTask().get(); + assertThat(storedTask.getRebalanceStatus()).isEqualTo(CANCELED); + assertThat(storedTask.isCancelRequested()).isTrue(); + assertThat(rebalanceManager.hasInProgressRebalance()).isFalse(); + } - // Cross the 2-minute boundary. + @Test + void testRecoverCancellationKeepsIntermediateBucketTracked() throws Exception { + TableBucket originBucket = new TableBucket(1L, 0); + TableBucket intermediateBucket = new TableBucket(1L, 1); + Map plans = plans(originBucket, intermediateBucket); + rebalanceExecutor.originBuckets.add(originBucket); + + rebalanceManager.recoverRebalance( + new RebalanceTask("recover-cancel-test", REBALANCING, plans, true)); + + Map statuses = statuses(rebalanceManager); + assertThat(statuses.get(originBucket)).isEqualTo(CANCELED); + assertThat(statuses.get(intermediateBucket)).isEqualTo(TIMEOUT); + RebalanceExecutionKey attempt = rebalanceManager.getExecutionKey(intermediateBucket); + rebalanceManager.finishRebalanceTask(attempt, COMPLETED); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(CANCELED); + } + + @Test + void testCanceledRebalanceReportsFailedWhenDrainFails() throws Exception { + configureManager( + new Configuration() + .set(COORDINATOR_REBALANCE_NO_PROGRESS_TIMEOUT, Duration.ofMillis(1))); + TableBucket tableBucket = new TableBucket(1L, 0); + rebalanceManager.registerRebalance("cancel-drain-failure", plans(tableBucket)); + RebalanceExecutionKey attempt = rebalanceManager.getExecutionKey(tableBucket); + assertThat(rebalanceManager.timeoutRebalanceTask(attempt)).isTrue(); + + rebalanceManager.cancelRebalance("cancel-drain-failure"); + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNotNull(); + clock.advanceTime(Duration.ofMillis(2)); + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNull(); + + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(FAILED); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()).isEqualTo(FAILED); + assertThat(zookeeperClient.getRebalanceTask().get().isCancelRequested()).isTrue(); + } + + @Test + void testRecoverPreservesExplicitlyCanceledBucket() throws Exception { + TableBucket canceledBucket = new TableBucket(1L, 0); + TableBucket intermediateBucket = new TableBucket(1L, 1); + Map plans = plans(canceledBucket, intermediateBucket); + + rebalanceManager.recoverRebalance( + new RebalanceTask( + "recover-explicit-cancel-test", + REBALANCING, + plans, + true, + Collections.singleton(canceledBucket))); + + Map statuses = statuses(rebalanceManager); + assertThat(statuses.get(canceledBucket)).isEqualTo(CANCELED); + assertThat(statuses.get(intermediateBucket)).isEqualTo(TIMEOUT); + } + + @Test + void testRecoveryTracksIntermediateTasksBeforeAdmittingNewWork() throws Exception { + Configuration conf = + new Configuration().set(COORDINATOR_REBALANCE_MAX_TRACKED_TIMED_OUT_TASKS, 1); + configureManager(conf); + + TableBucket[] buckets = new TableBucket[12]; + for (int i = 0; i < buckets.length; i++) { + buckets[i] = new TableBucket(i + 1L, 31 * (11 - i)); + } + RebalanceTask storedTask = + new RebalanceTask("recovery-admission-test", REBALANCING, plans(buckets)); + zookeeperClient.registerRebalanceTask(storedTask); + RebalanceTask recoveredTask = zookeeperClient.getRebalanceTask().get(); + List recoveredOrder = new ArrayList<>(recoveredTask.getExecutePlan().keySet()); + Set intermediateBuckets = + new HashSet<>( + recoveredOrder.subList(recoveredOrder.size() - 3, recoveredOrder.size())); + rebalanceExecutor.originBuckets.addAll(recoveredOrder); + rebalanceExecutor.originBuckets.removeAll(intermediateBuckets); + + rebalanceManager.recoverRebalance(recoveredTask); + + assertThat(rebalanceExecutor.executedPlans).isEmpty(); + assertThat(statuses(rebalanceManager)) + .extractingByKeys(intermediateBuckets.toArray(new TableBucket[0])) + .containsOnly(TIMEOUT); + + Set remainingIntermediateBuckets = new HashSet<>(intermediateBuckets); + while (!remainingIntermediateBuckets.isEmpty()) { + List activeBuckets = + remainingIntermediateBuckets.stream() + .filter(bucket -> rebalanceManager.getExecutionKey(bucket) != null) + .collect(Collectors.toList()); + assertThat(activeBuckets).hasSize(1); + TableBucket activeBucket = activeBuckets.get(0); + RebalanceExecutionKey attempt = rebalanceManager.getExecutionKey(activeBucket); + assertThat(rebalanceManager.getPlanForReconciliation(attempt)) + .isEqualTo(recoveredTask.getExecutePlan().get(activeBucket)); + assertThat(rebalanceManager.finishRebalanceTask(attempt, COMPLETED)).isTrue(); + remainingIntermediateBuckets.remove(activeBucket); + } + assertThat(rebalanceExecutor.executedPlans).hasSize(1); + } + + @Test + void testRecoverFinalTaskDoesNotExecuteAgain() { + TableBucket tableBucket = new TableBucket(1L, 0); + rebalanceManager.recoverRebalance( + new RebalanceTask("final-test", COMPLETED, plans(tableBucket))); + + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(COMPLETED); + assertThat(rebalanceExecutor.executedPlans).isEmpty(); + } + + @Test + void testReconciliationBacksOffBetweenRetries() { + TableBucket tableBucket = new TableBucket(1L, 0); + rebalanceManager.registerRebalance("backoff-test", plans(tableBucket)); + RebalanceExecutionKey attempt = rebalanceManager.getExecutionKey(tableBucket); + assertThat(rebalanceManager.timeoutRebalanceTask(attempt)).isTrue(); + assertThat(reconciliationsFor(eventManager, attempt)).isEqualTo(1); + + // first retry is dispatched at the base interval, the next one only after twice that. + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNotNull(); clock.advanceTime(Duration.ofMillis(30_000)); - manager.checkTimeout(); + rebalanceManager.checkTimeout(); + assertThat(reconciliationsFor(eventManager, attempt)).isEqualTo(2); - assertThat(eventManager.events).hasSize(1); - assertThat(eventManager.events.get(0)).isInstanceOf(RebalanceTaskTimeoutEvent.class); - RebalanceTaskTimeoutEvent timeoutEvent = - (RebalanceTaskTimeoutEvent) eventManager.events.get(0); - assertThat(timeoutEvent.getTableBucket()).isEqualTo(tb1); + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNotNull(); + clock.advanceTime(Duration.ofMillis(30_000)); + rebalanceManager.checkTimeout(); + assertThat(reconciliationsFor(eventManager, attempt)).isEqualTo(2); - // A second checkTimeout() should NOT enqueue another event because the - // inflight state was cleared after the first timeout. clock.advanceTime(Duration.ofMillis(30_000)); - manager.checkTimeout(); - assertThat(eventManager.events).hasSize(1); + rebalanceManager.checkTimeout(); + assertThat(reconciliationsFor(eventManager, attempt)).isEqualTo(3); + } - manager.close(); + @ParameterizedTest + @MethodSource("trackedTimedOutTaskLimits") + void testTrackedTimedOutTasksAreCapped(Configuration conf, int limit) { + configureManager(conf); + TableBucket[] tableBuckets = new TableBucket[limit + 2]; + for (int i = 0; i < tableBuckets.length; i++) { + tableBuckets[i] = new TableBucket(1L, i); + } + rebalanceManager.registerRebalance("cap-test", plans(tableBuckets)); + + // every timed-out task keeps being tracked, so admitting new work has to stop at the cap. + List timedOut = new ArrayList<>(); + for (int i = 0; i < limit; i++) { + TableBucket running = + rebalanceExecutor + .executedPlans + .get(rebalanceExecutor.executedPlans.size() - 1) + .getTableBucket(); + RebalanceExecutionKey attempt = rebalanceManager.getExecutionKey(running); + assertThat(rebalanceManager.timeoutRebalanceTask(attempt)).isTrue(); + timedOut.add(attempt); + } + assertThat(rebalanceExecutor.executedPlans).hasSize(limit); + + // once a tracked task reaches a final status the next pending task is admitted again. + assertThat(rebalanceManager.finishRebalanceTask(timedOut.get(0), COMPLETED)).isTrue(); + assertThat(rebalanceExecutor.executedPlans).hasSize(limit + 1); + } + + @ParameterizedTest + @MethodSource("targetUnavailableTimeouts") + void testTimedOutTaskFailsWhenTargetReplicasStayUnavailable( + Configuration conf, Duration timeout) throws Exception { + configureManager(conf); + TableBucket tableBucket = new TableBucket(1L, 0); + rebalanceManager.registerRebalance("give-up-test", plans(tableBucket)); + RebalanceExecutionKey attempt = rebalanceManager.getExecutionKey(tableBucket); + assertThat(rebalanceManager.timeoutRebalanceTask(attempt)).isTrue(); + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNotNull(); + + clock.advanceTime(timeout); + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNotNull(); + clock.advanceTime(Duration.ofMillis(1)); + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNull(); + + // the rebalance reaches a final status, so later rebalance requests are not blocked. + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(FAILED); + assertThat(rebalanceManager.hasInProgressRebalance()).isFalse(); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()).isEqualTo(FAILED); } @Test - void testTimeoutAfterCompletionIsNoOp() throws Exception { - ManualClock clock = new ManualClock(0L); - RecordingEventManager eventManager = new RecordingEventManager(); - NoOpScheduledExecutor executor = new NoOpScheduledExecutor(); - CoordinatorEventProcessor eventProcessor = buildCoordinatorEventProcessor(); + void testDeletedTaskCompletesBeforeReconciliationDeadlineFailure() throws Exception { + TableBucket tableBucket = new TableBucket(1L, 0); + rebalanceManager.registerRebalance("deleted-reconciliation-test", plans(tableBucket)); + RebalanceExecutionKey attempt = rebalanceManager.getExecutionKey(tableBucket); + assertThat(rebalanceManager.timeoutRebalanceTask(attempt)).isTrue(); + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNotNull(); + + clock.advanceTime(Duration.ofMinutes(30).plusMillis(1)); + rebalanceExecutor.deletedBuckets.add(tableBucket); + + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNull(); + assertThat(statuses(rebalanceManager).get(tableBucket)).isEqualTo(COMPLETED); + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(COMPLETED); + } - RebalanceManager manager = - new RebalanceManager( - eventProcessor, zookeeperClient, eventManager, clock, executor); - manager.startup(); + @ParameterizedTest + @MethodSource("noProgressTimeouts") + void testTimedOutTaskFailsWithoutProgressWhileTargetsAreLive( + Configuration conf, Duration timeout) throws Exception { + configureManager(conf); + CoordinatorContext context = rebalanceExecutor.getCoordinatorContext(); + for (int serverId : new int[] {1, 2, 3}) { + context.addLiveTabletServer(tabletServer(serverId)); + } + TableBucket tableBucket = new TableBucket(1L, 0); + rebalanceManager.registerRebalance("no-progress-test", plans(tableBucket)); + RebalanceExecutionKey attempt = rebalanceManager.getExecutionKey(tableBucket); + assertThat(rebalanceManager.timeoutRebalanceTask(attempt)).isTrue(); + + clock.advanceTime(timeout.minusMillis(1)); + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNotNull(); + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(REBALANCING); + assertThat(rebalanceManager.hasInProgressRebalance()).isTrue(); + clock.advanceTime(Duration.ofMillis(1)); + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNotNull(); + clock.advanceTime(Duration.ofMillis(1)); + assertThat(rebalanceManager.getPlanForReconciliation(attempt)).isNull(); + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(FAILED); + assertThat(rebalanceManager.hasInProgressRebalance()).isFalse(); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()).isEqualTo(FAILED); + + TableBucket nextBucket = new TableBucket(1L, 1); + Map nextPlan = plans(nextBucket); + rebalanceManager.registerRebalance("after-no-progress-timeout", nextPlan); + assertThat(rebalanceManager.getRebalanceId()).isEqualTo("after-no-progress-timeout"); + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(REBALANCING); + assertThat(rebalanceManager.hasInProgressRebalance()).isTrue(); + assertThat(rebalanceExecutor.executedPlans) + .hasSize(2) + .last() + .isEqualTo(nextPlan.get(nextBucket)); + } - TableBucket tb1 = new TableBucket(1L, 0); - Map plan = new HashMap<>(); - plan.put( - tb1, - new RebalancePlanForBucket( - tb1, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2))); + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testGenerateRebalanceFromIntermediateAssignment(boolean targetUnavailable) { + CoordinatorContext context = rebalanceExecutor.getCoordinatorContext(); + for (int serverId : new int[] {0, 1, 2, 3, 4}) { + if (!targetUnavailable || serverId != 3) { + context.addLiveTabletServer(tabletServer(serverId)); + } + } + context.putServerTag(0, ServerTag.PERMANENT_OFFLINE); + TableBucket bucket = new TableBucket(1L, 0); + putBucketForPlanning(context, bucket, Arrays.asList(0, 1, 2, 3)); + + RebalanceTask next = + rebalanceManager.generateRebalanceTask( + Collections.singletonList(new ReplicaDistributionGoal())); + RebalancePlanForBucket plan = next.getExecutePlan().get(bucket); + assertThat(plan).isNotNull(); + assertThat(plan.getOriginReplicas()).containsExactly(0, 1, 2, 3); + assertThat(plan.getNewReplicas()).hasSize(3).contains(1, 2).doesNotContain(0); + assertThat(context.liveTabletServerSet()).containsAll(plan.getNewReplicas()); + // Planning must leave the real assignment intact until the replacement task executes. + assertThat(context.getAssignment(bucket)).containsExactly(0, 1, 2, 3); + } + + @Test + void testGenerateRebalanceReplacesUnavailableReplica() { + CoordinatorContext context = rebalanceExecutor.getCoordinatorContext(); + for (int serverId : new int[] {0, 1, 4}) { + context.addLiveTabletServer(tabletServer(serverId)); + } + TableBucket bucket = new TableBucket(1L, 0); + putBucketForPlanning(context, bucket, Arrays.asList(0, 1, 2)); + + RebalanceTask task = + rebalanceManager.generateRebalanceTask( + Collections.singletonList(new ReplicaDistributionGoal())); + assertThat(task.getExecutePlan().get(bucket).getNewReplicas()) + .containsExactlyInAnyOrder(0, 1, 4); + } - zookeeperClient.registerRebalanceTask( - new RebalanceTask("completion-test", NOT_STARTED, plan)); - manager.registerRebalance("completion-test", plan, NOT_STARTED); + @Test + void testGenerateCleanupPlanWithoutOptimizationChanges() { + CoordinatorContext context = rebalanceExecutor.getCoordinatorContext(); + for (int serverId : new int[] {0, 1, 2, 3}) { + context.addLiveTabletServer(tabletServer(serverId)); + } + // The replica count comes from table metadata, even when the previous plan is gone. + TableBucket bucket = new TableBucket(1L, 10L, 0); + putBucketForPlanning(context, bucket, Arrays.asList(0, 1, 2, 3)); + + RebalanceTask task = + rebalanceManager.generateRebalanceTask( + Collections.singletonList(new ReplicaDistributionGoal())); + assertThat(task.getExecutePlan()).containsKey(bucket); + assertThat(task.getExecutePlan().get(bucket).getOriginReplicas()) + .containsExactly(0, 1, 2, 3); + assertThat(task.getExecutePlan().get(bucket).getNewReplicas()).containsExactly(0, 1, 2); + } - // The task completes normally before timeout. - manager.finishRebalanceTask(tb1, COMPLETED); + private static void putBucketForPlanning( + CoordinatorContext context, TableBucket bucket, List assignment) { + TableDescriptor descriptor = DATA1_TABLE_DESCRIPTOR.withReplicationFactor(3); + context.putTableInfo( + TableInfo.of( + TablePath.of("db", "table"), + bucket.getTableId(), + 1, + descriptor, + "file:///tmp/rebalance-planning", + 0L, + 0L)); + context.updateBucketReplicaAssignment(bucket, assignment); + context.putBucketLeaderAndIsr( + bucket, + new LeaderAndIsr( + 0, + 1, + Arrays.asList(0, 1, 2), + Collections.emptyList(), + context.getCoordinatorEpoch(), + 1)); + } - // Now the timeout fires, but the task is already done. - clock.advanceTime(Duration.ofMillis(130_000)); - manager.checkTimeout(); + @ParameterizedTest + @MethodSource("invalidTimeouts") + void testRejectsInvalidTimeout(ConfigOption option, Duration timeout) { + Configuration conf = new Configuration().set(option, timeout); + assertThatThrownBy(() -> createManager(conf)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining(option.key()) + .hasMessageContaining("must be between 1 ms"); + } - // No timeout event should be enqueued because inflightTaskStartMs was cleared. - assertThat(eventManager.events).isEmpty(); + @ParameterizedTest + @ValueSource(ints = {0, -1}) + void testRejectsInvalidTrackedTimedOutTaskLimit(int limit) { + Configuration conf = + new Configuration().set(COORDINATOR_REBALANCE_MAX_TRACKED_TIMED_OUT_TASKS, limit); + assertThatThrownBy(() -> createManager(conf)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining(COORDINATOR_REBALANCE_MAX_TRACKED_TIMED_OUT_TASKS.key()) + .hasMessageContaining("must be at least 1"); + } + + @ParameterizedTest + @ValueSource(longs = {1, Long.MAX_VALUE}) + void testAcceptsTimeoutConfigurationBounds(long timeoutMs) { + Configuration conf = + new Configuration() + .set( + COORDINATOR_REBALANCE_TARGET_UNAVAILABLE_TIMEOUT, + Duration.ofMillis(timeoutMs)) + .set( + COORDINATOR_REBALANCE_NO_PROGRESS_TIMEOUT, + Duration.ofMillis(timeoutMs)); + configureManager(conf); + } - manager.close(); + private static Stream targetUnavailableTimeouts() { + return Stream.of( + Arguments.of(new Configuration(), Duration.ofMinutes(30)), + Arguments.of( + new Configuration() + .set( + COORDINATOR_REBALANCE_TARGET_UNAVAILABLE_TIMEOUT, + Duration.ofMinutes(5)), + Duration.ofMinutes(5)), + Arguments.of( + new Configuration() + .set( + COORDINATOR_REBALANCE_TARGET_UNAVAILABLE_TIMEOUT, + Duration.ofHours(1)), + Duration.ofHours(1))); + } + + private static Stream noProgressTimeouts() { + return Stream.of( + Arguments.of(new Configuration(), Duration.ofHours(24)), + Arguments.of( + new Configuration() + .set( + COORDINATOR_REBALANCE_NO_PROGRESS_TIMEOUT, + Duration.ofMinutes(5)), + Duration.ofMinutes(5)), + Arguments.of( + new Configuration() + .set( + COORDINATOR_REBALANCE_NO_PROGRESS_TIMEOUT, + Duration.ofHours(48)), + Duration.ofHours(48))); + } + + private static Stream trackedTimedOutTaskLimits() { + return Stream.of( + Arguments.of(new Configuration(), 8), + Arguments.of( + new Configuration() + .set(COORDINATOR_REBALANCE_MAX_TRACKED_TIMED_OUT_TASKS, 1), + 1), + Arguments.of( + new Configuration() + .set(COORDINATOR_REBALANCE_MAX_TRACKED_TIMED_OUT_TASKS, 3), + 3)); + } + + private static Stream invalidTimeouts() { + return Stream.of( + COORDINATOR_REBALANCE_TARGET_UNAVAILABLE_TIMEOUT, + COORDINATOR_REBALANCE_NO_PROGRESS_TIMEOUT) + .flatMap( + option -> + Stream.of( + Duration.ZERO, + Duration.ofMillis(-1), + Duration.ofNanos(999_999), + Duration.ofMillis(Long.MAX_VALUE).plusMillis(1)) + .map(timeout -> Arguments.of(option, timeout))); } @Test - void testTimeoutTreatsTaskAsCompleted() throws Exception { - ManualClock clock = new ManualClock(0L); - RecordingEventManager eventManager = new RecordingEventManager(); - NoOpScheduledExecutor executor = new NoOpScheduledExecutor(); - CoordinatorEventProcessor eventProcessor = buildCoordinatorEventProcessor(); + void testCancelGivesUpImmediatelyOnAdmittedTaskStillAtOrigin() throws Exception { + TableBucket tb1 = new TableBucket(1L, 0); + TableBucket tb2 = new TableBucket(1L, 1); + rebalanceExecutor.originBuckets.add(tb1); + rebalanceManager.registerRebalance("cancel-at-origin-test", plans(tb1, tb2)); + + rebalanceManager.cancelRebalance("cancel-at-origin-test"); + + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(CANCELED); + assertThat(rebalanceManager.hasInProgressRebalance()).isFalse(); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(CANCELED); + } + @Test + void testScheduledRetryPersistsFinalStatusAfterTransientFailure() throws Exception { + rebalanceManager.close(); + try (TestingFinalizationZooKeeperClient testingClient = + new TestingFinalizationZooKeeperClient()) { + rebalanceManager = createManager(new Configuration(), testingClient); + TableBucket tableBucket = new TableBucket(1L, 0); + Map plan = plans(tableBucket); + String rebalanceId = "finalization-retry-test"; + testingClient.registerRebalanceTask(new RebalanceTask(rebalanceId, NOT_STARTED, plan)); + rebalanceManager.registerRebalance(rebalanceId, plan); + RebalanceExecutionKey attempt = rebalanceManager.getExecutionKey(tableBucket); + testingClient.failNextFinalWrite(); + + assertThat(rebalanceManager.finishRebalanceTask(attempt, COMPLETED)).isTrue(); + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(REBALANCING); + assertThat(rebalanceManager.hasInProgressRebalance()).isTrue(); + assertThat(testingClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(NOT_STARTED); + + timeoutChecker.triggerPeriodicScheduledTasks(); + assertThat(eventManager.events).hasSize(1); + FinalizeRebalanceEvent event = (FinalizeRebalanceEvent) eventManager.events.get(0); + rebalanceManager.retryFinalizeRebalance(event.getRebalanceId()); + + assertThat(testingClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(COMPLETED); + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(COMPLETED); + assertThat(rebalanceManager.hasInProgressRebalance()).isFalse(); + + TableBucket nextBucket = new TableBucket(2L, 0); + rebalanceManager.registerRebalance("next-rebalance", plans(nextBucket)); + assertThat(rebalanceManager.getExecutionKey(nextBucket)).isNotNull(); + } + } + + private RebalanceManager createManager(Configuration conf) { + return createManager(conf, zookeeperClient); + } + + private RebalanceManager createManager(Configuration conf, ZooKeeperClient client) { + timeoutChecker = new ManuallyTriggeredScheduledExecutorService(); RebalanceManager manager = new RebalanceManager( - eventProcessor, zookeeperClient, eventManager, clock, executor); + rebalanceExecutor, client, eventManager, clock, conf, timeoutChecker); manager.startup(); + return manager; + } - TableBucket tb1 = new TableBucket(1L, 0); - TableBucket tb2 = new TableBucket(1L, 1); - Map plan = new HashMap<>(); - plan.put( - tb1, - new RebalancePlanForBucket( - tb1, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2))); - plan.put( - tb2, - new RebalancePlanForBucket( - tb2, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2))); - - zookeeperClient.registerRebalanceTask( - new RebalanceTask("completed-test", NOT_STARTED, plan)); - manager.registerRebalance("completed-test", plan, NOT_STARTED); - - // Timeout fires. - clock.advanceTime(Duration.ofMillis(130_000)); - manager.checkTimeout(); + private void configureManager(Configuration conf) { + rebalanceManager.close(); + rebalanceManager = createManager(conf); + } + + private static int reconciliationsFor( + RecordingEventManager eventManager, RebalanceExecutionKey executionKey) { + int reconciliations = 0; + for (CoordinatorEvent event : eventManager.events) { + if (event instanceof ReconcileRebalanceTaskEvent + && ((ReconcileRebalanceTaskEvent) event) + .getExecutionKey() + .equals(executionKey)) { + reconciliations++; + } + } + return reconciliations; + } - // Simulate the coordinator event thread processing the timeout event. + private static ServerInfo tabletServer(int serverId) { + return new ServerInfo( + serverId, + "RACK" + serverId, + Endpoint.fromListenersString("CLIENT://host" + serverId + ":9124"), + ServerType.TABLET_SERVER); + } + + @Test + void testStartupFencesNewRebalanceUntilRecoveryEventRuns() throws Exception { + TableBucket tableBucket = new TableBucket(1L, 0); + RebalanceTask storedTask = + new RebalanceTask("startup-recovery-test", REBALANCING, plans(tableBucket)); + zookeeperClient.registerRebalanceTask(storedTask); + configureManager(new Configuration()); + + assertThat(rebalanceManager.hasInProgressRebalance()).isTrue(); + assertThat(rebalanceManager.getRebalanceId()).isNull(); + assertThat(timeoutChecker.getActivePeriodicScheduledTask()).hasSize(1); assertThat(eventManager.events).hasSize(1); - RebalanceTaskTimeoutEvent timeoutEvent = - (RebalanceTaskTimeoutEvent) eventManager.events.get(0); - manager.finishRebalanceTask(timeoutEvent.getTableBucket(), TIMEOUT); - - // The timed-out task should be in finishedRebalanceTasks as TIMEOUT. - assertThat(manager.hasInProgressRebalance()).isTrue(); - RebalanceResultForBucket result = - manager.listRebalanceProgress(null).progressForBucketMap().get(tb1); - assertThat(result.status()).isEqualTo(TIMEOUT); - - manager.close(); - } - - private CoordinatorEventProcessor buildCoordinatorEventProcessor() { - return new CoordinatorEventProcessor( - zookeeperClient, - serverMetadataCache, - testCoordinatorChannelManager, - new CoordinatorContext(zkEpoch), - replicaCapacityController, - autoPartitionManager, - lakeTableTieringManager, - TestingMetricGroups.COORDINATOR_METRICS, - conf, - Executors.newFixedThreadPool(1, new ExecutorThreadFactory("test-coordinator-io")), - metadataManager, - kvSnapshotLeaseManager, - scheduler, - SystemClock.getInstance()); + RecoverRebalanceEvent recoveryEvent = (RecoverRebalanceEvent) eventManager.events.get(0); + assertThat(recoveryEvent.getRebalanceTask()).isEqualTo(storedTask); + + rebalanceExecutor.originBuckets.add(tableBucket); + rebalanceManager.recoverRebalance(recoveryEvent.getRebalanceTask()); + assertThat(rebalanceManager.getRebalanceId()).isEqualTo("startup-recovery-test"); + assertThat(rebalanceExecutor.executedPlans).hasSize(1); } - /** Records events put into the coordinator event queue. */ - private static final class RecordingEventManager implements EventManager { - final List events = new ArrayList<>(); + private static Map plans(TableBucket... tableBuckets) { + Map plans = new LinkedHashMap<>(); + for (TableBucket tableBucket : tableBuckets) { + plans.put( + tableBucket, + new RebalancePlanForBucket( + tableBucket, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 2, 3))); + } + return plans; + } + + private static Map statuses(RebalanceManager manager) { + Map statuses = new HashMap<>(); + for (Map.Entry entry : + manager.listRebalanceProgress(null).progressForBucketMap().entrySet()) { + statuses.put(entry.getKey(), entry.getValue().status()); + } + return statuses; + } + + private static final class TestingRebalanceExecutor implements RebalanceExecutor { + private final CoordinatorContext coordinatorContext; + private final List executedPlans = new ArrayList<>(); + private final Set originBuckets = new HashSet<>(); + private final Set deletedBuckets = new HashSet<>(); + + private TestingRebalanceExecutor(CoordinatorContext coordinatorContext) { + this.coordinatorContext = coordinatorContext; + } @Override - public void put(CoordinatorEvent event) { - events.add(event); + public CoordinatorContext getCoordinatorContext() { + return coordinatorContext; + } + + @Override + public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) { + executedPlans.add(planForBucket); + } + + @Override + public boolean isRebalanceTaskAtOrigin(RebalancePlanForBucket planForBucket) { + return originBuckets.contains(planForBucket.getTableBucket()); + } + + @Override + public boolean isRebalanceTaskDeleted(TableBucket tableBucket) { + return deletedBuckets.contains(tableBucket); } } - /** - * A scheduled executor that never actually runs scheduled tasks, so tests retain full control - * over when {@link RebalanceManager#checkTimeout()} is invoked. - */ - private static final class NoOpScheduledExecutor extends ScheduledThreadPoolExecutor { + private static final class TestingFinalizationZooKeeperClient extends ZooKeeperClient { + + private boolean failNextFinalWrite; - NoOpScheduledExecutor() { - super(0); + private TestingFinalizationZooKeeperClient() { + super(createCuratorFrameworkWrapper(), createConfiguration()); + } + + private void failNextFinalWrite() { + failNextFinalWrite = true; } @Override - public java.util.concurrent.ScheduledFuture scheduleWithFixedDelay( - Runnable command, - long initialDelay, - long delay, - java.util.concurrent.TimeUnit unit) { - return null; + public void registerRebalanceTask(RebalanceTask rebalanceTask) throws Exception { + if (failNextFinalWrite + && RebalanceStatus.FINAL_STATUSES.contains( + rebalanceTask.getRebalanceStatus())) { + failNextFinalWrite = false; + throw new RuntimeException("Expected finalization failure"); + } + super.registerRebalanceTask(rebalanceTask); + } + + private static CuratorFrameworkWithUnhandledErrorListener createCuratorFrameworkWrapper() { + Builder builder = + builder() + .connectString( + ZOO_KEEPER_EXTENSION_WRAPPER + .getCustomExtension() + .getConnectString()) + .namespace("fluss") + .retryPolicy(new ExponentialBackoffRetry(1000, 3)); + return startZookeeperClient(builder, NOPErrorHandler.INSTANCE); + } + + private static Configuration createConfiguration() { + Configuration configuration = new Configuration(); + configuration.set( + ConfigOptions.REMOTE_DATA_DIR, + ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().getRemoteDataDir()); + return configuration; + } + } + + /** Records events put into the coordinator event queue. */ + private static final class RecordingEventManager implements EventManager { + final List events = new ArrayList<>(); + + @Override + public void put(CoordinatorEvent event) { + events.add(event); } } } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java index 0fa1a1ddde4..f28eb59b621 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java @@ -961,6 +961,11 @@ void testRebalancePlan() throws Exception { new RebalanceTask("rebalance-task-2", COMPLETED, bucketPlan)); assertThat(zookeeperClient.getRebalanceTask()) .hasValue(new RebalanceTask("rebalance-task-2", COMPLETED, bucketPlan)); + + RebalanceTask cancelRequestedTask = + new RebalanceTask("rebalance-task-2", NOT_STARTED, bucketPlan, true); + zookeeperClient.registerRebalanceTask(cancelRequestedTask); + assertThat(zookeeperClient.getRebalanceTask()).hasValue(cancelRequestedTask); } @Test diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/data/RebalanceTaskJsonSerdeTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/data/RebalanceTaskJsonSerdeTest.java index 5711f7da3f7..fb1f7764b93 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/zk/data/RebalanceTaskJsonSerdeTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/zk/data/RebalanceTaskJsonSerdeTest.java @@ -20,12 +20,19 @@ import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.utils.json.JsonSerdeTestBase; +import org.apache.fluss.utils.json.JsonSerdeUtils; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.Map; import static org.apache.fluss.cluster.rebalance.RebalanceStatus.NOT_STARTED; +import static org.assertj.core.api.Assertions.assertThat; /** Test for {@link RebalanceTaskJsonSerde}. */ public class RebalanceTaskJsonSerdeTest extends JsonSerdeTestBase { @@ -80,17 +87,22 @@ protected RebalanceTask[] createObjects() { Arrays.asList(0, 1, 2), Arrays.asList(3, 4, 5))); return new RebalanceTask[] { - new RebalanceTask("rebalance-task-21jd", NOT_STARTED, bucketPlan) + new RebalanceTask( + "rebalance-task-21jd", + NOT_STARTED, + bucketPlan, + true, + Collections.singleton(new TableBucket(0L, 0))) }; } @Override protected String[] expectedJsons() { return new String[] { - "{\"version\":1,\"rebalance_id\":\"rebalance-task-21jd\",\"rebalance_status\":0,\"rebalance_plan\":" + "{\"version\":2,\"rebalance_id\":\"rebalance-task-21jd\",\"rebalance_status\":0,\"cancel_requested\":true,\"rebalance_plan\":" + "[{\"table_id\":0,\"buckets\":" + "[{\"bucket_id\":1,\"original_leader\":1,\"new_leader\":1,\"origin_replicas\":[0,1,2],\"new_replicas\":[1,2,3]}," - + "{\"bucket_id\":0,\"original_leader\":0,\"new_leader\":3,\"origin_replicas\":[0,1,2],\"new_replicas\":[3,4,5]}]}," + + "{\"bucket_id\":0,\"canceled\":true,\"original_leader\":0,\"new_leader\":3,\"origin_replicas\":[0,1,2],\"new_replicas\":[3,4,5]}]}," + "{\"table_id\":1,\"partition_id\":0,\"buckets\":[" + "{\"bucket_id\":0,\"original_leader\":0,\"new_leader\":3,\"origin_replicas\":[0,1,2],\"new_replicas\":[3,4,5]}," + "{\"bucket_id\":1,\"original_leader\":1,\"new_leader\":1,\"origin_replicas\":[0,1,2],\"new_replicas\":[1,2,3]}]}," @@ -98,4 +110,22 @@ protected String[] expectedJsons() { + "{\"bucket_id\":0,\"original_leader\":0,\"new_leader\":3,\"origin_replicas\":[0,1,2],\"new_replicas\":[3,4,5]}]}]}" }; } + + @Test + void testDeserializeLegacyTaskWithoutCancelRequested() throws IOException { + String legacyJson = + expectedJsons()[0] + .replace("\"version\":2", "\"version\":1") + .replace(",\"cancel_requested\":true", "") + .replace(",\"canceled\":true", ""); + + RebalanceTask rebalanceTask = + JsonSerdeUtils.readValue( + legacyJson.getBytes(StandardCharsets.UTF_8), + RebalanceTaskJsonSerde.INSTANCE); + + assertThat(rebalanceTask.isCancelRequested()).isFalse(); + assertThat(rebalanceTask.getCanceledBuckets()).isEmpty(); + assertThat(rebalanceTask.getExecutePlan()).hasSize(5); + } } diff --git a/website/docs/maintenance/operations/rebalance.md b/website/docs/maintenance/operations/rebalance.md index 405bc413bbf..e6e77082af6 100644 --- a/website/docs/maintenance/operations/rebalance.md +++ b/website/docs/maintenance/operations/rebalance.md @@ -132,10 +132,26 @@ Optional latestProgress = admin.listRebalanceProgress(null).g Rebalance statuses: - **NOT_STARTED**: The rebalance has been created but not yet started - **REBALANCING**: The rebalance is currently in progress -- **COMPLETED**: The rebalance has successfully completed -- **FAILED**: The rebalance has failed +- **COMPLETED**: The rebalance has successfully completed, or its table/partition was deleted while the task was pending or being reconciled +- **FAILED**: The rebalance has failed, for example because a bucket migration was given up on after its target servers stayed unavailable for too long. The affected buckets may be left with an intermediate assignment. A subsequent rebalance uses the table's configured replication factor to plan the target assignment and removes excess intermediate replicas during execution. Unavailable servers are included as sources to evacuate, not as destinations for new replicas - **CANCELED**: The rebalance has been canceled -- **TIMEOUT**: The rebalance task timed out (e.g., ISR could not converge within the timeout period) +- **TIMEOUT**: The normal execution slot was released after the timeout, but the task remains non-final. The coordinator retries the current migration phase idempotently, with a growing backoff, until the task reaches `COMPLETED` or `FAILED`. Only a bounded number of timed-out tasks is tracked at the same time, so a rebalance stops admitting new bucket migrations while too many tasks are still being reconciled + +#### Timeout and Recovery Settings + +Configure the following coordinator options in `server.yaml` to match your cluster's recovery objectives and resource capacity. The values below are the defaults. Restart the coordinator servers after changing these settings. + +```yaml +coordinator.rebalance.target-unavailable-timeout: 30 min +coordinator.rebalance.no-progress-timeout: 24 h +coordinator.rebalance.max-tracked-timed-out-tasks: 8 +``` + +- `target-unavailable-timeout`: A timed-out bucket migration is marked `FAILED` after its target tablet servers remain unavailable without observed bucket-state progress for this duration. Progress or all target servers becoming live resets the timer. +- `no-progress-timeout`: A timed-out bucket migration is marked `FAILED` after this duration without observed bucket-state progress, even when its target servers are live. Progress resets the timer, so this is not a limit on the total migration duration. Either timeout can cause a migration to fail. +- `max-tracked-timed-out-tasks`: New bucket migrations stop being admitted when this many timed-out migrations remain non-final. Admission resumes when a tracked migration reaches a final status. A higher limit allows more migrations to remain active and increases coordinator and ZooKeeper work. + +Both durations must be at least `1 ms` and fit in a signed 64-bit millisecond value; the task limit must be at least `1`. Invalid values prevent coordinator initialization. Timeouts are evaluated during periodic reconciliation, so the transition to `FAILED` can occur after the configured duration has elapsed. ### 4. Cancel Rebalance (If Needed) @@ -152,6 +168,7 @@ admin.cancelRebalance(null).get(); **Important Notes:** - Only rebalance operations in `NOT_STARTED` or `REBALANCING` status can be canceled - Already completed bucket migrations will not be rolled back +- Bucket migrations that have not changed anything yet are canceled right away, migrations that are already under way are drained first so that no bucket is left half-migrated - After cancellation, the rebalance status will change to `CANCELED` ### 5. Remove Server Tags (After Completion)