From 530d965d89efd0293274f3091a7a0f01bf3b7b67 Mon Sep 17 00:00:00 2001 From: yunhong <337361684@qq.com> Date: Mon, 3 Aug 2026 14:12:41 +0800 Subject: [PATCH 1/4] [server] Make timed-out rebalance tasks recoverable Keep timed-out bucket migrations under bounded reconciliation while admitting new work. Persist cancellation intent and final status, recover intermediate assignments, and require a current leader acknowledgement before completion. Fence stale responses and avoid repeated elections for applied leader moves. Expose recovery limits and replan failed migrations using the configured replica count. Document the state transitions and recovery settings, share test setup, and retain targeted failure and restart coverage. Exempt only ConfigOptions from the Checkstyle file length limit. Skip startup recovery on ZooKeeper read failure to allow new rebalance requests, and leave a TODO for recovery retries. Validation: 89 targeted and integration tests passed, including five rebalance integration cases. After the startup recovery adjustment, reran all 35 RebalanceManagerTest cases successfully. Checkstyle and Spotless checks passed. Co-Authored-By: Codex AI-Model: gpt-6 AI-Contributed/Feature: 7/7 AI-Contributed/UT: 0/0 --- .../cluster/rebalance/RebalanceStatus.java | 2 +- .../apache/fluss/config/ConfigOptions.java | 33 + .../CoordinatorEventProcessor.java | 433 ++++++--- .../coordinator/CoordinatorRequestBatch.java | 27 +- .../event/FinalizeRebalanceEvent.java | 36 + .../NotifyLeaderAndIsrRequestContext.java | 86 ++ ...tifyLeaderAndIsrResponseReceivedEvent.java | 18 + .../event/RebalanceTaskTimeoutEvent.java | 14 +- .../event/ReconcileRebalanceTaskEvent.java | 38 + .../event/RecoverRebalanceEvent.java | 5 +- .../rebalance/RebalanceExecutionKey.java | 70 ++ .../rebalance/RebalanceExecutor.java | 34 + .../rebalance/RebalanceManager.java | 823 ++++++++++++++---- .../fluss/server/zk/data/RebalanceTask.java | 27 +- .../zk/data/RebalanceTaskJsonSerde.java | 8 +- .../CoordinatorEventProcessorTest.java | 447 ++-------- .../CoordinatorEventProcessorTestBase.java | 154 +++- .../coordinator/CoordinatorRebalanceTest.java | 765 ++++++++++++++++ .../TestingControlledNotifyGateway.java | 77 ++ .../TestingFailingNotifyGateway.java | 47 + .../rebalance/RebalanceManagerITCase.java | 5 +- .../rebalance/RebalanceManagerTest.java | 808 +++++++++++------ .../fluss/server/zk/ZooKeeperClientTest.java | 5 + .../zk/data/RebalanceTaskJsonSerdeTest.java | 24 +- .../docs/maintenance/operations/rebalance.md | 23 +- 25 files changed, 3091 insertions(+), 918 deletions(-) create mode 100644 fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/FinalizeRebalanceEvent.java create mode 100644 fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/NotifyLeaderAndIsrRequestContext.java create mode 100644 fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/ReconcileRebalanceTaskEvent.java create mode 100644 fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceExecutionKey.java create mode 100644 fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceExecutor.java create mode 100644 fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorRebalanceTest.java create mode 100644 fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestingControlledNotifyGateway.java create mode 100644 fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestingFailingNotifyGateway.java 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..6d6b9c85557 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) { @@ -767,12 +767,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 +775,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 +1246,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 +1260,7 @@ private void processNotifyLeaderAndIsrResponseReceivedEvent( } offlineReplicas.add(new TableBucketReplica(tableBucket, serverId)); } else { - succeededBuckets.add(notifyLeaderAndIsrResultForBucket.getTableBucket()); + succeededBuckets.add(tableBucket); } } for (TableBucket tb : succeededBuckets) { @@ -1269,9 +1277,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 +1637,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 +1677,77 @@ 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()) { + resumeBucketReassignment(tableBucket, reassignment); + } 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(); @@ -1738,16 +1758,73 @@ private void tryToCompleteRebalanceTask( tableBucket, RebalanceStatus.COMPLETED); } } else if (notifyLeaderAndIsrResultForBucket.succeeded()) { - tryToCompleteReassignmentTask(tableBucket, reassignment); + if (isFinalReassignmentState(tableBucket, reassignment)) { + if (isSuccessfulReassignmentResponseFromLeader( + responseServerId, requestContext)) { + rebalanceManager.finishRebalanceTask( + tableBucket, RebalanceStatus.COMPLETED); + } + } else { + tryToCompleteReassignmentTask(tableBucket, reassignment); + } } } 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. + return leaderAndIsr.leader() == planForBucket.getOriginalLeader() + && new HashSet<>(planForBucket.getOriginReplicas()).containsAll(leaderAndIsr.isr()); + } + + private boolean completeRebalanceTaskIfDeleted(TableBucket tableBucket) { + if (coordinatorContext.getAllBuckets().contains(tableBucket) + && !coordinatorContext.isTableQueuedForDeletion(tableBucket.getTableId()) + && !coordinatorContext.isToBeDeleted(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); @@ -1761,17 +1838,29 @@ private void tryToCompleteRebalanceTaskOnLeaderAndIsrChange(TableBucket tableBuc try { tryToCompleteReassignmentTask(tableBucket, reassignment); } 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) { + if (isFinalReassignmentState(tableBucket, reassignment)) { + // 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()); + } else if (isReplicaAssignmentAtTarget(tableBucket, reassignment.getTargetReplicas())) { + // A successful Notify response alone must not start an immediate resend loop while a + // target is still offline. Leader/ISR changes may advance persisted phase B; periodic + // timeout reconciliation remains responsible for probing an unchanged state. + if (isReassignmentComplete(tableBucket, reassignment)) { + resumePersistedPhaseB(tableBucket, reassignment); + } + } else if (isReassignmentComplete(tableBucket, reassignment)) { LOG.info( "Target replicas {} have all caught up with the leader for reassigning bucket {}", reassignment.getTargetReplicas(), @@ -1780,13 +1869,9 @@ private void tryToCompleteReassignmentTask( } } - @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 +1918,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,9 +2004,85 @@ 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 void resumeBucketReassignment(TableBucket tableBucket, ReplicaReassignment reassignment) + throws Exception { + if (isFinalReassignmentState(tableBucket, reassignment)) { + LOG.info( + "Re-send final state for rebalance task of {} and wait for leader ack.", + tableBucket); + sendCurrentLeaderAndIsrRequest(tableBucket, reassignment.getTargetReplicas()); + } else if (isReplicaAssignmentAtTarget(tableBucket, reassignment.getTargetReplicas())) { + // 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, reassignment); + } else if (isReassignmentComplete(tableBucket, reassignment)) { + LOG.info("Resume phase B for rebalance task of {}.", tableBucket); + onBucketReassignment(tableBucket, reassignment, true); + } else 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); + } + } + + private void retryBucketReassignmentPhaseA( + TableBucket tableBucket, ReplicaReassignment reassignment) throws Exception { + // 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, ReplicaReassignment reassignment) + throws Exception { + List targetReplicas = reassignment.getTargetReplicas(); + if (!isReassignmentComplete(tableBucket, reassignment)) { + sendCurrentLeaderAndIsrRequest(tableBucket, targetReplicas); + return; } + + LeaderAndIsr leaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket).get(); + 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( @@ -1928,7 +2090,52 @@ private boolean isReassignmentComplete( LeaderAndIsr leaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket).get(); List isr = leaderAndIsr.isr(); List targetReplicas = reassignment.getTargetReplicas(); - return targetReplicas.isEmpty() || new HashSet<>(isr).containsAll(targetReplicas); + return coordinatorContext.liveTabletServerSet().containsAll(targetReplicas) + && (targetReplicas.isEmpty() || new HashSet<>(isr).containsAll(targetReplicas)); + } + + private boolean isFinalReassignmentState( + TableBucket tableBucket, ReplicaReassignment reassignment) throws Exception { + List targetReplicas = reassignment.getTargetReplicas(); + if (!coordinatorContext.getAllBuckets().contains(tableBucket) + || !coordinatorContext.liveTabletServerSet().containsAll(targetReplicas) + || !isReplicaAssignmentAtTarget(tableBucket, targetReplicas)) { + return false; + } + Optional leaderAndIsrOpt = zooKeeperClient.getLeaderAndIsr(tableBucket); + if (!leaderAndIsrOpt.isPresent()) { + return false; + } + LeaderAndIsr leaderAndIsr = leaderAndIsrOpt.get(); + 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 +2848,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 +2880,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..50cbb46b0d3 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceExecutor.java @@ -0,0 +1,34 @@ +/* + * 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.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); +} 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..ea7413ec090 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. New pending tasks are canceled directly; + * recovered or admitted tasks must still be at their origin to be canceled. Intermediate tasks + * continue to drain, including those queued again after coordinator recovery. Recovery resumes + * non-final plans from the current assignment and ISR, 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, a cancellation request + * yields CANCELED; otherwise any failed bucket yields FAILED, 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,65 @@ 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; - /** A queue of in progress table bucket to rebalance. */ - private final Queue inProgressRebalanceTasksQueue = new ArrayDeque<>(); + /** Tasks awaiting admission, including intermediate migrations loaded during recovery. */ + private final Queue pendingRebalanceTasks = new ArrayDeque<>(); /** A mapping from table bucket to rebalance status of pending and running 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 +197,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 +223,13 @@ 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() { + public synchronized void start() { timeoutChecker.scheduleWithFixedDelay( this::checkTimeoutSafely, TIMEOUT_CHECK_INTERVAL_MS, @@ -171,7 +241,7 @@ public void start() { TIMEOUT_CHECK_INTERVAL_MS); } - public @Nullable String getRebalanceId() { + public synchronized @Nullable String getRebalanceId() { return currentRebalanceId; } @@ -179,87 +249,119 @@ 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. + for (Map.Entry entry : rebalancePlan.entrySet()) { + addPendingTask(entry.getKey(), entry.getValue()); + } + + rebalanceStatus = REBALANCING; + processNewRebalanceTask(); + } + + /** Recovers a persisted task by comparing every bucket plan with current coordinator state. */ + public synchronized void recoverRebalance(RebalanceTask rebalanceTask) { + checkNotClosed(); + 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; + } + + for (Map.Entry entry : + rebalanceTask.getExecutePlan().entrySet()) { + addPendingTask(entry.getKey(), entry.getValue()); + } + if (cancelRequested) { + cancelPendingTasks(); + } + + if (inProgressRebalanceTasks.isEmpty()) { + finalizeRebalance(); + } else { rebalanceStatus = REBALANCING; processNewRebalanceTask(); - } else { - rebalanceStatus = newStatus; } } - public void finishRebalanceTask(TableBucket tableBucket, RebalanceStatus statusForBucket) { + 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(); - 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()); + checkArgument(statusForBucket != TIMEOUT, "Use timeoutRebalanceTask for soft timeouts."); + RebalanceTaskAttempt attempt = findActiveAttempt(executionKey); + if (attempt == null) { + LOG.debug("Ignore stale completion for {}.", executionKey); + return false; + } - if (inProgressRebalanceTasksQueue.isEmpty()) { - // All rebalance tasks are completed. - completeRebalance(); - } else { - // Trigger one rebalance task to execute. - processNewRebalanceTask(); - } + 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 @Nullable RebalanceProgress listRebalanceProgress(@Nullable String rebalanceId) { + public synchronized @Nullable RebalanceProgress listRebalanceProgress( + @Nullable String rebalanceId) { checkNotClosed(); if (rebalanceId != null && currentRebalanceId != null @@ -285,9 +387,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 +412,72 @@ public void cancelRebalance(@Nullable String rebalanceId) { return; } + Map executePlan = allRebalancePlans(); 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, executePlan, true)); } 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(); + + // 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()) { + RebalanceResultForBucket result = + inProgressRebalanceTasks.get(attempt.executionKey.getTableBucket()); + if (result != null && rebalanceExecutor.isRebalanceTaskAtOrigin(result.plan())) { + 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); + } + + private void cancelPendingTasks() { + Iterator pending = pendingRebalanceTasks.iterator(); + while (pending.hasNext()) { + TableBucket bucket = pending.next(); + RebalanceResultForBucket result = inProgressRebalanceTasks.get(bucket); + // Recovery can put an already started migration back into the pending queue. + if (!recoveredRebalance || rebalanceExecutor.isRebalanceTaskAtOrigin(result.plan())) { + pending.remove(); + inProgressRebalanceTasks.remove(bucket); + finishedRebalanceTasks.put( + bucket, RebalanceResultForBucket.of(result.plan(), CANCELED)); + } + } } - public boolean hasInProgressRebalance() { + 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 +485,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 +503,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 +517,272 @@ 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; + } + + // 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; + 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; + } + TableBucket tableBucket; + while ((tableBucket = pendingRebalanceTasks.poll()) != null) { RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.get(tableBucket); - RebalanceResultForBucket rebalanceResultForBucket = - RebalanceResultForBucket.of(resultForBucket.plan(), REBALANCING); - eventProcessor.tryToExecuteRebalanceTask(rebalanceResultForBucket.plan()); + 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 completeRebalance() { + private void finalizeRebalance() { + finalizationPending = true; + persistFinalStatus(); + } + + private void persistFinalStatus() { checkNotClosed(); + RebalanceStatus finalStatus = cancelRequested ? CANCELED : 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)); } 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(); + pendingRebalanceTasks.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(); - inProgressRebalanceTasksQueue.clear(); + pendingRebalanceTasks.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 @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; + } - // Here, it will not clear finishedRebalanceTasks, because it will be used by - // listRebalanceProgress. It will be cleared when next register. + /** + * 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)); + } + } + + 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; + } - LOG.info("Rebalance complete with {} ms.", System.currentTimeMillis() - registerTime); + private RebalanceStatus aggregateFinalStatus() { + for (RebalanceResultForBucket result : finishedRebalanceTasks.values()) { + if (result.status() == FAILED || result.status() == CANCELED) { + return FAILED; + } + } + return COMPLETED; } private ClusterModel buildClusterModel(CoordinatorContext coordinatorContext) { @@ -433,10 +801,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 +829,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 +839,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 +899,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..239402a0c61 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 @@ -45,6 +45,9 @@ 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; + /** A mapping from tableBucket to RebalancePlanForBuckets of none-partitioned table. */ private final Map> planForBuckets; @@ -56,8 +59,17 @@ 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 = rebalanceId; this.rebalanceStatus = rebalanceStatus; + this.cancelRequested = cancelRequested; this.planForBuckets = new HashMap<>(); this.planForBucketsOfPartitionedTable = new HashMap<>(); @@ -86,6 +98,10 @@ public RebalanceStatus getRebalanceStatus() { return rebalanceStatus; } + public boolean isCancelRequested() { + return cancelRequested; + } + public Map> getPlanForBuckets() { return planForBuckets; } @@ -121,6 +137,8 @@ public String toString() { + rebalanceId + ", rebalanceStatus=" + rebalanceStatus + + ", cancelRequested=" + + cancelRequested + ", planForBuckets=" + planForBuckets + ", planForBucketsOfPartitionedTable=" @@ -138,7 +156,8 @@ 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(planForBuckets, that.planForBuckets) && Objects.equals( @@ -148,6 +167,10 @@ public boolean equals(Object o) { @Override public int hashCode() { return Objects.hash( - rebalanceId, rebalanceStatus, planForBuckets, planForBucketsOfPartitionedTable); + rebalanceId, + rebalanceStatus, + cancelRequested, + 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..a410a7f6f20 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 @@ -42,6 +42,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"; @@ -54,7 +55,7 @@ public class RebalanceTaskJsonSerde 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 +63,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. @@ -102,6 +104,8 @@ 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<>(); for (JsonNode tablePartitionPlanNode : rebalancePlanNode) { @@ -140,7 +144,7 @@ public RebalanceTask deserialize(JsonNode node) { } } - return new RebalanceTask(rebalanceId, rebalanceStatus, planForBuckets); + return new RebalanceTask(rebalanceId, rebalanceStatus, planForBuckets, cancelRequested); } private void serializeRebalancePlanForBucket( 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..794fc9352be --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorRebalanceTest.java @@ -0,0 +1,765 @@ +/* + * 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.REBALANCING); + 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); + } + + @Test + void testRestartAfterCancellationDoesNotStartPendingBucketWithShrunkIsr() 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); + adjustRebalanceIsr(pending, Arrays.asList(0, 1)); + installBlockingNotifyGateways(pendingResponses); + Map plan = new LinkedHashMap<>(); + plan.put(running, new RebalancePlanForBucket(running, 0, 0, origin, target)); + plan.put(pending, 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, 0, Arrays.asList(0, 1)); + + 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(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 recoveredPending = + plans.keySet().stream() + .filter(bucket -> rebalanceStatus(bucket) == RebalanceStatus.NOT_STARTED) + .findFirst() + .get(); + assertThat( + zookeeperClient + .getTableAssignment(tableId) + .get() + .getBucketAssignment(recoveredPending.getBucket())) + .isEqualTo(new BucketAssignment(union)); + fromCtx( + ctx -> { + assertThat(eventProcessor.isRebalanceTaskAtOrigin(plans.get(recoveredPending))) + .isFalse(); + eventProcessor.getRebalanceManager().cancelRebalance(rebalanceId); + return null; + }); + assertThat(rebalanceStatus(recoveredPending)) + .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)).isFalse(); + } + + 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..08a01a580d3 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,70 @@ 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.config.ConfigOptions; +import org.apache.fluss.cluster.rebalance.ServerTag; +import org.apache.fluss.config.ConfigOption; 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.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.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.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.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.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Test for {@link RebalanceManager}. */ public class RebalanceManagerTest { @@ -79,18 +90,12 @@ 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; @BeforeAll static void baseBeforeAll() throws Exception { @@ -102,60 +107,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 +132,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 +143,579 @@ 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(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(); + } + + @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(REBALANCING); + RebalanceExecutionKey attempt = rebalanceManager.getExecutionKey(intermediateBucket); + rebalanceManager.finishRebalanceTask(attempt, COMPLETED); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(CANCELED); + } + + @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(); + } - // Cross the 2-minute boundary. + @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); + } + + @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); + } + + @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)); + } - manager.close(); + @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 testTimeoutAfterCompletionIsNoOp() throws Exception { - ManualClock clock = new ManualClock(0L); - RecordingEventManager eventManager = new RecordingEventManager(); - NoOpScheduledExecutor executor = new NoOpScheduledExecutor(); - CoordinatorEventProcessor eventProcessor = buildCoordinatorEventProcessor(); + 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); + } - RebalanceManager manager = - new RebalanceManager( - eventProcessor, zookeeperClient, eventManager, clock, executor); - manager.startup(); + @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); + } - 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))); + 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)); + } - zookeeperClient.registerRebalanceTask( - new RebalanceTask("completion-test", NOT_STARTED, plan)); - manager.registerRebalance("completion-test", plan, NOT_STARTED); + @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"); + } - // The task completes normally before timeout. - manager.finishRebalanceTask(tb1, COMPLETED); + @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"); + } - // Now the timeout fires, but the task is already done. - clock.advanceTime(Duration.ofMillis(130_000)); - manager.checkTimeout(); + @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); + } - // No timeout event should be enqueued because inflightTaskStartMs was cleared. - assertThat(eventManager.events).isEmpty(); + 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))); + } - manager.close(); + 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); + } + private RebalanceManager createManager(Configuration conf) { RebalanceManager manager = new RebalanceManager( - eventProcessor, zookeeperClient, eventManager, clock, executor); + rebalanceExecutor, + zookeeperClient, + eventManager, + clock, + conf, + new ManuallyTriggeredScheduledExecutorService()); 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); + } - // Simulate the coordinator event thread processing the timeout event. + 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; + } + + 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(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); + + 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; + } - @Override - public void put(CoordinatorEvent event) { - events.add(event); + 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; } - /** - * 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 TestingRebalanceExecutor implements RebalanceExecutor { + private final CoordinatorContext coordinatorContext; + private final List executedPlans = new ArrayList<>(); + private final Set originBuckets = new HashSet<>(); - NoOpScheduledExecutor() { - super(0); + private TestingRebalanceExecutor(CoordinatorContext coordinatorContext) { + this.coordinatorContext = coordinatorContext; } @Override - public java.util.concurrent.ScheduledFuture scheduleWithFixedDelay( - Runnable command, - long initialDelay, - long delay, - java.util.concurrent.TimeUnit unit) { - return null; + 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()); + } + } + + /** 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..4a464f10772 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,18 @@ 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.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 { @@ -87,7 +93,7 @@ protected RebalanceTask[] createObjects() { @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\":false,\"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]}]}," @@ -98,4 +104,20 @@ 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\":false", ""); + + RebalanceTask rebalanceTask = + JsonSerdeUtils.readValue( + legacyJson.getBytes(StandardCharsets.UTF_8), + RebalanceTaskJsonSerde.INSTANCE); + + assertThat(rebalanceTask.isCancelRequested()).isFalse(); + 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) From b9213972e517c91e390379a787613888280ef9f2 Mon Sep 17 00:00:00 2001 From: yunhong <337361684@qq.com> Date: Mon, 26 Oct 2026 19:59:09 +0800 Subject: [PATCH 2/4] address bing's comments Co-Authored-By: Codex AI-Model: gpt-5.6-sol AI-Contributed/Feature: 58/58 AI-Contributed/UT: 208/208 --- .../CoordinatorEventProcessor.java | 23 ++- .../rebalance/RebalanceExecutor.java | 4 + .../rebalance/RebalanceManager.java | 31 +++- ...oordinatorEventProcessorAdjustIsrTest.java | 6 +- .../coordinator/CoordinatorRebalanceTest.java | 46 +++++- .../rebalance/RebalanceManagerTest.java | 156 +++++++++++++++++- 6 files changed, 236 insertions(+), 30 deletions(-) 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 6d6b9c85557..eaf7533c396 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 @@ -1806,14 +1806,27 @@ public boolean isRebalanceTaskAtOrigin(RebalancePlanForBucket planForBucket) { 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. - return leaderAndIsr.leader() == planForBucket.getOriginalLeader() - && new HashSet<>(planForBucket.getOriginReplicas()).containsAll(leaderAndIsr.isr()); + 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. For a + // leader-only plan, however, the leader is the state being changed. + return reassignment.isBeingReassigned() + || leaderAndIsr.leader() == planForBucket.getOriginalLeader(); + } + + @Override + public boolean isRebalanceTaskDeleted(TableBucket tableBucket) { + return !coordinatorContext.getAllBuckets().contains(tableBucket) + || coordinatorContext.isTableQueuedForDeletion(tableBucket.getTableId()) + || coordinatorContext.isToBeDeleted(tableBucket); } private boolean completeRebalanceTaskIfDeleted(TableBucket tableBucket) { - if (coordinatorContext.getAllBuckets().contains(tableBucket) - && !coordinatorContext.isTableQueuedForDeletion(tableBucket.getTableId()) - && !coordinatorContext.isToBeDeleted(tableBucket)) { + if (!isRebalanceTaskDeleted(tableBucket)) { return false; } 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 index 50cbb46b0d3..e83732a3300 100644 --- 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 @@ -18,6 +18,7 @@ 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. */ @@ -31,4 +32,7 @@ public interface RebalanceExecutor { /** 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 ea7413ec090..746ccb0a9cd 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 @@ -105,9 +105,8 @@ *

    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. New pending tasks are canceled directly; - * recovered or admitted tasks must still be at their origin to be canceled. Intermediate tasks - * continue to drain, including those queued again after coordinator recovery. Recovery resumes - * non-final plans from the current assignment and ISR, and still requires a current leader + * recovered or admitted tasks must still be at their origin to be canceled. Recovery restores + * intermediate tasks directly into reconciliation, 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, @@ -142,7 +141,7 @@ public class RebalanceManager { private final long noProgressTimeoutMs; private final int maxTrackedTimedOutTasks; - /** Tasks awaiting admission, including intermediate migrations loaded during recovery. */ + /** Tasks awaiting admission. */ private final Queue pendingRebalanceTasks = new ArrayDeque<>(); /** A mapping from table bucket to rebalance status of pending and running tasks. */ @@ -300,7 +299,11 @@ public synchronized void recoverRebalance(RebalanceTask rebalanceTask) { for (Map.Entry entry : rebalanceTask.getExecutePlan().entrySet()) { - addPendingTask(entry.getKey(), entry.getValue()); + if (rebalanceExecutor.isRebalanceTaskAtOrigin(entry.getValue())) { + addPendingTask(entry.getKey(), entry.getValue()); + } else { + addRecoveredIntermediateTask(entry.getKey(), entry.getValue()); + } } if (cancelRequested) { cancelPendingTasks(); @@ -448,7 +451,7 @@ private void cancelPendingTasks() { while (pending.hasNext()) { TableBucket bucket = pending.next(); RebalanceResultForBucket result = inProgressRebalanceTasks.get(bucket); - // Recovery can put an already started migration back into the pending queue. + // A recovered pending task may have changed after it was classified at startup. if (!recoveredRebalance || rebalanceExecutor.isRebalanceTaskAtOrigin(result.plan())) { pending.remove(); inProgressRebalanceTasks.remove(bucket); @@ -570,6 +573,10 @@ public synchronized boolean timeoutRebalanceTask(RebalanceExecutionKey execution 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(); @@ -704,6 +711,18 @@ private void addPendingTask(TableBucket tableBucket, RebalancePlanForBucket plan inProgressRebalanceTasks.put(tableBucket, RebalanceResultForBucket.of(plan, NOT_STARTED)); } + private void addRecoveredIntermediateTask( + TableBucket tableBucket, RebalancePlanForBucket plan) { + long now = clock.milliseconds(); + RebalanceTaskAttempt attempt = + new RebalanceTaskAttempt( + new RebalanceExecutionKey(currentRebalanceId, tableBucket), now); + attempt.onTimedOut(now, observeBucketState(tableBucket)); + timedOutRebalanceTasks.put(tableBucket, attempt); + inProgressRebalanceTasks.put(tableBucket, RebalanceResultForBucket.of(plan, TIMEOUT)); + enqueueReconciliation(attempt); + } + private @Nullable RebalanceTaskAttempt findActiveAttempt(RebalanceExecutionKey executionKey) { RebalanceTaskAttempt attempt = runningRebalanceTasks.get(executionKey.getTableBucket()); if (attempt == null) { 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..a5bf6e7b99a 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 @@ -355,7 +355,7 @@ private AdjustIsrResponse adjustIsr(Map leaderAndIsrs return response.get(); } - private AdjustIsrResultForBucket submitAdjustIsr( + protected AdjustIsrResultForBucket submitAdjustIsr( TableBucket tableBucket, LeaderAndIsr newLeaderAndIsr) throws Exception { CompletableFuture response = new CompletableFuture<>(); eventProcessor @@ -395,7 +395,7 @@ private void assertAdjustedIsrInstalled( } } - private void initCoordinatorChannel() throws Exception { + protected void initCoordinatorChannel() throws Exception { makeSendLeaderAndStopRequestAlwaysSuccess( testCoordinatorChannelManager, Arrays.stream(zookeeperClient.getSortedTabletServerList()) @@ -404,7 +404,7 @@ private void initCoordinatorChannel() throws Exception { Collections.emptySet()); } - private T fromCtx(java.util.function.Function retrieveFunction) + protected T fromCtx(java.util.function.Function retrieveFunction) throws Exception { AccessContextEvent event = new AccessContextEvent<>(retrieveFunction); eventProcessor.getCoordinatorEventManager().put(event); 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 index 794fc9352be..04cec5a9ff1 100644 --- 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 @@ -358,7 +358,7 @@ void testRestartResumesRebalanceAndWaitsForLeaderAck(String phase) throws Except for (NotifyLeaderAndIsrRequestContext staleContext : Arrays.asList(oldContext, oldEpochContext)) { notifyLeaderResponse(tableBucket, current.leader(), staleContext); - assertThat(rebalanceStatus(tableBucket)).isEqualTo(RebalanceStatus.REBALANCING); + assertThat(rebalanceStatus(tableBucket)).isEqualTo(RebalanceStatus.TIMEOUT); assertThat(zookeeperClient.getRebalanceTask()).hasValue(persistedTask); } @@ -459,7 +459,7 @@ void testNewRebalanceConvergesAfterFailureAndRestart(boolean targetUnavailable) } @Test - void testRestartAfterCancellationDoesNotStartPendingBucketWithShrunkIsr() throws Exception { + void testRestartAfterCancellationDoesNotStartPendingBucketAfterLeaderChange() throws Exception { registerTabletServer(3); initCoordinatorChannel(); List origin = Arrays.asList(0, 1, 2); @@ -468,6 +468,7 @@ void testRestartAfterCancellationDoesNotStartPendingBucketWithShrunkIsr() throws TableBucket running = new TableBucket(tableId, 0); TableBucket pending = new TableBucket(tableId, 1); adjustRebalanceIsr(pending, Arrays.asList(0, 1)); + changeLeader(pending, 1); installBlockingNotifyGateways(pendingResponses); Map plan = new LinkedHashMap<>(); plan.put(running, new RebalancePlanForBucket(running, 0, 0, origin, target)); @@ -493,7 +494,7 @@ void testRestartAfterCancellationDoesNotStartPendingBucketWithShrunkIsr() throws assertThat(eventProcessor.getRebalanceManager().getExecutionKey(pending)).isNull(); assertThat(zookeeperClient.getTableAssignment(tableId).get().getBucketAssignment(1)) .isEqualTo(new BucketAssignment(origin)); - verifyIsr(pending, 0, Arrays.asList(0, 1)); + verifyIsr(pending, 1, Arrays.asList(0, 1)); adjustRebalanceIsr(running, Arrays.asList(0, 1, 2, 3)); retry( @@ -545,25 +546,27 @@ void testCancelAfterRestartDrainsEveryIntermediateBucket() throws Exception { () -> assertThat(eventProcessor.getRebalanceManager().getRebalanceId()) .isEqualTo(rebalanceId)); - TableBucket recoveredPending = + TableBucket recoveredIntermediate = plans.keySet().stream() - .filter(bucket -> rebalanceStatus(bucket) == RebalanceStatus.NOT_STARTED) + .filter(bucket -> rebalanceStatus(bucket) == RebalanceStatus.TIMEOUT) .findFirst() .get(); assertThat( zookeeperClient .getTableAssignment(tableId) .get() - .getBucketAssignment(recoveredPending.getBucket())) + .getBucketAssignment(recoveredIntermediate.getBucket())) .isEqualTo(new BucketAssignment(union)); fromCtx( ctx -> { - assertThat(eventProcessor.isRebalanceTaskAtOrigin(plans.get(recoveredPending))) + assertThat( + eventProcessor.isRebalanceTaskAtOrigin( + plans.get(recoveredIntermediate))) .isFalse(); eventProcessor.getRebalanceManager().cancelRebalance(rebalanceId); return null; }); - assertThat(rebalanceStatus(recoveredPending)) + assertThat(rebalanceStatus(recoveredIntermediate)) .as("A previously started migration still needs draining after recovery") .isNotEqualTo(RebalanceStatus.CANCELED); for (TableBucket bucket : plans.keySet()) { @@ -619,7 +622,32 @@ void testRebalanceAtOriginAllowsShrunkIsr() throws Exception { 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)).isFalse(); + assertThat(eventProcessor.isRebalanceTaskAtOrigin(plan)).isTrue(); + + RebalancePlanForBucket leaderOnlyPlan = + new RebalancePlanForBucket( + tableBucket, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2)); + assertThat(eventProcessor.isRebalanceTaskAtOrigin(leaderOnlyPlan)).isFalse(); + } + + 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( 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 08a01a580d3..ae6622419ff 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 @@ -24,6 +24,7 @@ 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.metadata.TableDescriptor; @@ -32,17 +33,20 @@ import org.apache.fluss.server.coordinator.CoordinatorContext; 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.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; @@ -79,6 +83,9 @@ 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; @@ -96,6 +103,7 @@ public class RebalanceManagerTest { private TestingRebalanceExecutor rebalanceExecutor; private RecordingEventManager eventManager; private RebalanceManager rebalanceManager; + private ManuallyTriggeredScheduledExecutorService timeoutChecker; @BeforeAll static void baseBeforeAll() throws Exception { @@ -280,13 +288,43 @@ void testRecoverCancellationKeepsIntermediateBucketTracked() throws Exception { Map statuses = statuses(rebalanceManager); assertThat(statuses.get(originBucket)).isEqualTo(CANCELED); - assertThat(statuses.get(intermediateBucket)).isEqualTo(REBALANCING); + assertThat(statuses.get(intermediateBucket)).isEqualTo(TIMEOUT); RebalanceExecutionKey attempt = rebalanceManager.getExecutionKey(intermediateBucket); rebalanceManager.finishRebalanceTask(attempt, COMPLETED); assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()) .isEqualTo(CANCELED); } + @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()); + TableBucket intermediateBucket = recoveredOrder.get(recoveredOrder.size() - 1); + rebalanceExecutor.originBuckets.addAll(recoveredOrder); + rebalanceExecutor.originBuckets.remove(intermediateBucket); + + rebalanceManager.recoverRebalance(recoveredTask); + + RebalanceExecutionKey intermediateAttempt = + rebalanceManager.getExecutionKey(intermediateBucket); + assertThat(intermediateAttempt).isNotNull(); + assertThat(rebalanceManager.getPlanForReconciliation(intermediateAttempt)) + .isEqualTo(recoveredTask.getExecutePlan().get(intermediateBucket)); + assertThat(rebalanceExecutor.executedPlans).isEmpty(); + assertThat(statuses(rebalanceManager).get(intermediateBucket)).isEqualTo(TIMEOUT); + } + @Test void testRecoverFinalTaskDoesNotExecuteAgain() { TableBucket tableBucket = new TableBucket(1L, 0); @@ -372,6 +410,22 @@ void testTimedOutTaskFailsWhenTargetReplicasStayUnavailable( assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()).isEqualTo(FAILED); } + @Test + 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); + } + @ParameterizedTest @MethodSource("noProgressTimeouts") void testTimedOutTaskFailsWithoutProgressWhileTargetsAreLive( @@ -605,15 +659,52 @@ void testCancelGivesUpImmediatelyOnAdmittedTaskStillAtOrigin() throws Exception .isEqualTo(CANCELED); } + @Test + void testScheduledRetryPersistsFinalStatusAfterTransientFailure() throws Exception { + rebalanceManager.close(); + try (TestingFinalizationZooKeeperClient testingClient = + new TestingFinalizationZooKeeperClient()) { + rebalanceManager = createManager(new Configuration(), testingClient); + rebalanceManager.start(); + 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( - rebalanceExecutor, - zookeeperClient, - eventManager, - clock, - conf, - new ManuallyTriggeredScheduledExecutorService()); + rebalanceExecutor, client, eventManager, clock, conf, timeoutChecker); manager.startup(); return manager; } @@ -659,6 +750,7 @@ void testStartupFencesNewRebalanceUntilRecoveryEventRuns() throws Exception { 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); @@ -688,6 +780,7 @@ 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; @@ -707,6 +800,55 @@ public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) { public boolean isRebalanceTaskAtOrigin(RebalancePlanForBucket planForBucket) { return originBuckets.contains(planForBucket.getTableBucket()); } + + @Override + public boolean isRebalanceTaskDeleted(TableBucket tableBucket) { + return deletedBuckets.contains(tableBucket); + } + } + + private static final class TestingFinalizationZooKeeperClient extends ZooKeeperClient { + + private boolean failNextFinalWrite; + + private TestingFinalizationZooKeeperClient() { + super(createCuratorFrameworkWrapper(), createConfiguration()); + } + + private void failNextFinalWrite() { + failNextFinalWrite = true; + } + + @Override + 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. */ From a2473967c9a501a49cdcb4e4e2dabd8329b51018 Mon Sep 17 00:00:00 2001 From: yunhong <337361684@qq.com> Date: Mon, 26 Oct 2026 21:02:10 +0800 Subject: [PATCH 3/4] code optimize Co-Authored-By: Codex Co-Authored-By: Qoder AI-Model: gpt-6 AI-Contributed/Feature: 149/149 AI-Contributed/UT: 29/29 Co-Authored-By: Codex AI-Model: gpt-5.6-sol AI-Contributed/Feature: 60/60 AI-Contributed/UT: 84/84 --- .../CoordinatorEventProcessor.java | 150 ++++++++---------- .../rebalance/RebalanceManager.java | 59 ++++--- ...oordinatorEventProcessorAdjustIsrTest.java | 30 ---- .../coordinator/CoordinatorRebalanceTest.java | 29 +++- .../rebalance/RebalanceManagerTest.java | 54 +++++-- 5 files changed, 174 insertions(+), 148 deletions(-) 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 eaf7533c396..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 @@ -340,7 +340,6 @@ public void startup() { // start rebalance manager. rebalanceManager.startup(); - rebalanceManager.start(); } public void shutdown() { @@ -1688,7 +1687,12 @@ public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) { try { if (reassignment.isBeingReassigned()) { - resumeBucketReassignment(tableBucket, reassignment); + 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()) { @@ -1757,16 +1761,11 @@ private void tryToCompleteRebalanceTask( rebalanceManager.finishRebalanceTask( tableBucket, RebalanceStatus.COMPLETED); } - } else if (notifyLeaderAndIsrResultForBucket.succeeded()) { - if (isFinalReassignmentState(tableBucket, reassignment)) { - if (isSuccessfulReassignmentResponseFromLeader( + } else if (notifyLeaderAndIsrResultForBucket.succeeded() + && advanceReassignment(tableBucket, reassignment, false) + && isSuccessfulReassignmentResponseFromLeader( responseServerId, requestContext)) { - rebalanceManager.finishRebalanceTask( - tableBucket, RebalanceStatus.COMPLETED); - } - } else { - tryToCompleteReassignmentTask(tableBucket, reassignment); - } + rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.COMPLETED); } } catch (Exception e) { LOG.warn( @@ -1812,15 +1811,16 @@ public boolean isRebalanceTaskAtOrigin(RebalancePlanForBucket planForBucket) { ReplicaReassignment reassignment = ReplicaReassignment.build( planForBucket.getOriginReplicas(), planForBucket.getNewReplicas()); - // An unrelated leader election is not evidence that a replica migration started. For a - // leader-only plan, however, the leader is the state being changed. + // 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.getOriginalLeader(); + || leaderAndIsr.leader() != planForBucket.getNewLeader(); } @Override public boolean isRebalanceTaskDeleted(TableBucket tableBucket) { - return !coordinatorContext.getAllBuckets().contains(tableBucket) + return coordinatorContext.getAssignment(tableBucket).isEmpty() || coordinatorContext.isTableQueuedForDeletion(tableBucket.getTableId()) || coordinatorContext.isToBeDeleted(tableBucket); } @@ -1849,7 +1849,12 @@ 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.warn( "Failed to complete the reassignment for table bucket {}. It remains " @@ -1860,26 +1865,52 @@ private void tryToCompleteRebalanceTaskOnLeaderAndIsrChange(TableBucket tableBuc } } - private void tryToCompleteReassignmentTask( - TableBucket tableBucket, ReplicaReassignment reassignment) throws Exception { - if (isFinalReassignmentState(tableBucket, reassignment)) { - // 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()); - } else if (isReplicaAssignmentAtTarget(tableBucket, reassignment.getTargetReplicas())) { - // A successful Notify response alone must not start an immediate resend loop while a - // target is still offline. Leader/ISR changes may advance persisted phase B; periodic - // timeout reconciliation remains responsible for probing an unchanged state. - if (isReassignmentComplete(tableBucket, reassignment)) { - resumePersistedPhaseB(tableBucket, reassignment); + /** + * 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 (isReassignmentComplete(tableBucket, reassignment)) { + } 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; } private static boolean isSuccessfulReassignmentResponseFromLeader( @@ -2023,31 +2054,6 @@ private void onBucketReassignment( } } - private void resumeBucketReassignment(TableBucket tableBucket, ReplicaReassignment reassignment) - throws Exception { - if (isFinalReassignmentState(tableBucket, reassignment)) { - LOG.info( - "Re-send final state for rebalance task of {} and wait for leader ack.", - tableBucket); - sendCurrentLeaderAndIsrRequest(tableBucket, reassignment.getTargetReplicas()); - } else if (isReplicaAssignmentAtTarget(tableBucket, reassignment.getTargetReplicas())) { - // 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, reassignment); - } else if (isReassignmentComplete(tableBucket, reassignment)) { - LOG.info("Resume phase B for rebalance task of {}.", tableBucket); - onBucketReassignment(tableBucket, reassignment, true); - } else 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); - } - } - private void retryBucketReassignmentPhaseA( TableBucket tableBucket, ReplicaReassignment reassignment) throws Exception { // Persist the idempotent union assignment again in case the previous attempt updated @@ -2067,15 +2073,10 @@ private void retryBucketReassignmentPhaseA( sendCurrentLeaderAndIsrRequest(tableBucket, reassignment.replicas); } - private void resumePersistedPhaseB(TableBucket tableBucket, ReplicaReassignment reassignment) + private void resumePersistedPhaseB( + TableBucket tableBucket, List targetReplicas, LeaderAndIsr currentLeaderAndIsr) throws Exception { - List targetReplicas = reassignment.getTargetReplicas(); - if (!isReassignmentComplete(tableBucket, reassignment)) { - sendCurrentLeaderAndIsrRequest(tableBucket, targetReplicas); - return; - } - - LeaderAndIsr leaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket).get(); + LeaderAndIsr leaderAndIsr = currentLeaderAndIsr; if (!targetReplicas.contains(leaderAndIsr.leader())) { tableBucketStateMachine.handleStateChange( Collections.singleton(tableBucket), @@ -2099,27 +2100,14 @@ private void resumePersistedPhaseB(TableBucket tableBucket, ReplicaReassignment } private boolean isReassignmentComplete( - TableBucket tableBucket, ReplicaReassignment reassignment) throws Exception { - LeaderAndIsr leaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket).get(); - List isr = leaderAndIsr.isr(); - List targetReplicas = reassignment.getTargetReplicas(); + LeaderAndIsr leaderAndIsr, List targetReplicas) { return coordinatorContext.liveTabletServerSet().containsAll(targetReplicas) - && (targetReplicas.isEmpty() || new HashSet<>(isr).containsAll(targetReplicas)); + && (targetReplicas.isEmpty() + || new HashSet<>(leaderAndIsr.isr()).containsAll(targetReplicas)); } - private boolean isFinalReassignmentState( - TableBucket tableBucket, ReplicaReassignment reassignment) throws Exception { - List targetReplicas = reassignment.getTargetReplicas(); - if (!coordinatorContext.getAllBuckets().contains(tableBucket) - || !coordinatorContext.liveTabletServerSet().containsAll(targetReplicas) - || !isReplicaAssignmentAtTarget(tableBucket, targetReplicas)) { - return false; - } - Optional leaderAndIsrOpt = zooKeeperClient.getLeaderAndIsr(tableBucket); - if (!leaderAndIsrOpt.isPresent()) { - return false; - } - LeaderAndIsr leaderAndIsr = leaderAndIsrOpt.get(); + private static boolean isFinalReassignmentState( + LeaderAndIsr leaderAndIsr, List targetReplicas) { return targetReplicas.contains(leaderAndIsr.leader()) && new HashSet<>(leaderAndIsr.isr()).equals(new HashSet<>(targetReplicas)); } 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 746ccb0a9cd..5e36c9efbfc 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 @@ -105,13 +105,13 @@ *

    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. New pending tasks are canceled directly; - * recovered or admitted tasks must still be at their origin to be canceled. Recovery restores - * intermediate tasks directly into reconciliation, and still requires a current leader - * acknowledgement even when the persisted state already matches the target. + * 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, a cancellation request - * yields CANCELED; otherwise any failed bucket yields FAILED, and all completed buckets yield + * 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. * @@ -144,7 +144,10 @@ public class RebalanceManager { /** Tasks awaiting admission. */ private final Queue pendingRebalanceTasks = new ArrayDeque<>(); - /** A mapping from table bucket to rebalance status of pending and running tasks. */ + /** Recovered intermediate tasks awaiting bounded reconciliation admission. */ + private final Queue pendingRecoveredTasks = new ArrayDeque<>(); + + /** A mapping from table bucket to rebalance status of all non-final tasks. */ private final Map inProgressRebalanceTasks = new HashMap<>(); @@ -224,11 +227,9 @@ public RebalanceManager( public synchronized void startup() { LOG.info("Start up rebalance manager."); - initialize(); - } - - /** Starts the periodic timeout checker. Call after {@link #startup()}. */ - public synchronized 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, @@ -238,6 +239,7 @@ public synchronized void start() { "RebalanceManager timeout checker started: timeoutMs={}, checkIntervalMs={}", REBALANCE_TASK_TIMEOUT_MS, TIMEOUT_CHECK_INTERVAL_MS); + initialize(); } public synchronized @Nullable String getRebalanceId() { @@ -619,6 +621,24 @@ public synchronized void retryFinalizeRebalance(String rebalanceId) { } private void processNewRebalanceTask() { + TableBucket tableBucket; + while (timedOutRebalanceTasks.size() < maxTrackedTimedOutTasks + && (tableBucket = pendingRecoveredTasks.poll()) != null) { + RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.get(tableBucket); + 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; } @@ -633,7 +653,6 @@ private void processNewRebalanceTask() { timedOutRebalanceTasks.size()); return; } - TableBucket tableBucket; while ((tableBucket = pendingRebalanceTasks.poll()) != null) { RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.get(tableBucket); if (resultForBucket == null || resultForBucket.status() != NOT_STARTED) { @@ -657,7 +676,7 @@ private void finalizeRebalance() { private void persistFinalStatus() { checkNotClosed(); - RebalanceStatus finalStatus = cancelRequested ? CANCELED : aggregateFinalStatus(); + RebalanceStatus finalStatus = aggregateFinalStatus(); try { zkClient.registerRebalanceTask( new RebalanceTask( @@ -677,6 +696,7 @@ private void persistFinalStatus() { finalizationEventQueued = false; inProgressRebalanceTasks.clear(); pendingRebalanceTasks.clear(); + pendingRecoveredTasks.clear(); runningRebalanceTasks.clear(); timedOutRebalanceTasks.clear(); queuedTimeoutEvents.clear(); @@ -699,6 +719,7 @@ private void resetRebalance(String rebalanceId, boolean cancelRequested) { finalizationEventQueued = false; inProgressRebalanceTasks.clear(); pendingRebalanceTasks.clear(); + pendingRecoveredTasks.clear(); runningRebalanceTasks.clear(); timedOutRebalanceTasks.clear(); finishedRebalanceTasks.clear(); @@ -713,14 +734,8 @@ private void addPendingTask(TableBucket tableBucket, RebalancePlanForBucket plan private void addRecoveredIntermediateTask( TableBucket tableBucket, RebalancePlanForBucket plan) { - long now = clock.milliseconds(); - RebalanceTaskAttempt attempt = - new RebalanceTaskAttempt( - new RebalanceExecutionKey(currentRebalanceId, tableBucket), now); - attempt.onTimedOut(now, observeBucketState(tableBucket)); - timedOutRebalanceTasks.put(tableBucket, attempt); + pendingRecoveredTasks.add(tableBucket); inProgressRebalanceTasks.put(tableBucket, RebalanceResultForBucket.of(plan, TIMEOUT)); - enqueueReconciliation(attempt); } private @Nullable RebalanceTaskAttempt findActiveAttempt(RebalanceExecutionKey executionKey) { @@ -797,11 +812,11 @@ private Map allRebalancePlans() { private RebalanceStatus aggregateFinalStatus() { for (RebalanceResultForBucket result : finishedRebalanceTasks.values()) { - if (result.status() == FAILED || result.status() == CANCELED) { + if (result.status() == FAILED) { return FAILED; } } - return COMPLETED; + return cancelRequested ? CANCELED : COMPLETED; } private ClusterModel buildClusterModel(CoordinatorContext coordinatorContext) { 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 a5bf6e7b99a..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(); } - 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); - } - 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( } } - protected void initCoordinatorChannel() throws Exception { - makeSendLeaderAndStopRequestAlwaysSuccess( - testCoordinatorChannelManager, - Arrays.stream(zookeeperClient.getSortedTabletServerList()) - .boxed() - .collect(Collectors.toSet()), - Collections.emptySet()); - } - - protected 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/CoordinatorRebalanceTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorRebalanceTest.java index 04cec5a9ff1..f1b6dccc3b9 100644 --- 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 @@ -458,8 +458,10 @@ void testNewRebalanceConvergesAfterFailureAndRestart(boolean targetUnavailable) verifyIsr(bucket, 0, origin); } - @Test - void testRestartAfterCancellationDoesNotStartPendingBucketAfterLeaderChange() throws Exception { + @ParameterizedTest(name = "leaderOnly={0}") + @ValueSource(booleans = {false, true}) + void testRestartAfterCancellationDoesNotStartPendingBucketAfterLeaderChange(boolean leaderOnly) + throws Exception { registerTabletServer(3); initCoordinatorChannel(); List origin = Arrays.asList(0, 1, 2); @@ -467,12 +469,22 @@ void testRestartAfterCancellationDoesNotStartPendingBucketAfterLeaderChange() th long tableId = createTableWithReplicas(origin, 2); TableBucket running = new TableBucket(tableId, 0); TableBucket pending = new TableBucket(tableId, 1); - adjustRebalanceIsr(pending, Arrays.asList(0, 1)); - changeLeader(pending, 1); + // An unrelated leader election must not make the pending plan look started. For a + // leader-only plan, the elected leader is neither the original nor the target leader. + int pendingLeader = leaderOnly ? 2 : 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, new RebalancePlanForBucket(pending, 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( @@ -494,7 +506,7 @@ void testRestartAfterCancellationDoesNotStartPendingBucketAfterLeaderChange() th assertThat(eventProcessor.getRebalanceManager().getExecutionKey(pending)).isNull(); assertThat(zookeeperClient.getTableAssignment(tableId).get().getBucketAssignment(1)) .isEqualTo(new BucketAssignment(origin)); - verifyIsr(pending, 1, Arrays.asList(0, 1)); + verifyIsr(pending, pendingLeader, pendingIsr); adjustRebalanceIsr(running, Arrays.asList(0, 1, 2, 3)); retry( @@ -508,6 +520,7 @@ void testRestartAfterCancellationDoesNotStartPendingBucketAfterLeaderChange() th .isEqualTo(RebalanceStatus.CANCELED); assertThat(zookeeperClient.getTableAssignment(tableId).get().getBucketAssignment(1)) .isEqualTo(new BucketAssignment(origin)); + verifyIsr(pending, pendingLeader, pendingIsr); verifyIsr(running, 0, target); } @@ -627,7 +640,11 @@ void testRebalanceAtOriginAllowsShrunkIsr() throws Exception { 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 { 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 ae6622419ff..7f8b7905d50 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 @@ -71,6 +71,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.stream.Collectors; import java.util.stream.Stream; import static org.apache.fluss.cluster.rebalance.RebalanceStatus.CANCELED; @@ -295,6 +296,26 @@ void testRecoverCancellationKeepsIntermediateBucketTracked() throws Exception { .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 testRecoveryTracksIntermediateTasksBeforeAdmittingNewWork() throws Exception { Configuration conf = @@ -310,19 +331,34 @@ void testRecoveryTracksIntermediateTasksBeforeAdmittingNewWork() throws Exceptio zookeeperClient.registerRebalanceTask(storedTask); RebalanceTask recoveredTask = zookeeperClient.getRebalanceTask().get(); List recoveredOrder = new ArrayList<>(recoveredTask.getExecutePlan().keySet()); - TableBucket intermediateBucket = recoveredOrder.get(recoveredOrder.size() - 1); + Set intermediateBuckets = + new HashSet<>( + recoveredOrder.subList(recoveredOrder.size() - 3, recoveredOrder.size())); rebalanceExecutor.originBuckets.addAll(recoveredOrder); - rebalanceExecutor.originBuckets.remove(intermediateBucket); + rebalanceExecutor.originBuckets.removeAll(intermediateBuckets); rebalanceManager.recoverRebalance(recoveredTask); - RebalanceExecutionKey intermediateAttempt = - rebalanceManager.getExecutionKey(intermediateBucket); - assertThat(intermediateAttempt).isNotNull(); - assertThat(rebalanceManager.getPlanForReconciliation(intermediateAttempt)) - .isEqualTo(recoveredTask.getExecutePlan().get(intermediateBucket)); assertThat(rebalanceExecutor.executedPlans).isEmpty(); - assertThat(statuses(rebalanceManager).get(intermediateBucket)).isEqualTo(TIMEOUT); + 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 @@ -665,7 +701,6 @@ void testScheduledRetryPersistsFinalStatusAfterTransientFailure() throws Excepti try (TestingFinalizationZooKeeperClient testingClient = new TestingFinalizationZooKeeperClient()) { rebalanceManager = createManager(new Configuration(), testingClient); - rebalanceManager.start(); TableBucket tableBucket = new TableBucket(1L, 0); Map plan = plans(tableBucket); String rebalanceId = "finalization-retry-test"; @@ -746,6 +781,7 @@ void testStartupFencesNewRebalanceUntilRecoveryEventRuns() throws Exception { assertThat(rebalanceManager.hasInProgressRebalance()).isTrue(); assertThat(rebalanceManager.getRebalanceId()).isNull(); + assertThat(timeoutChecker.getActivePeriodicScheduledTask()).hasSize(1); assertThat(eventManager.events).hasSize(1); RecoverRebalanceEvent recoveryEvent = (RecoverRebalanceEvent) eventManager.events.get(0); assertThat(recoveryEvent.getRebalanceTask()).isEqualTo(storedTask); From 26a1f8cd033facd8a244304db6097bdb62f12249 Mon Sep 17 00:00:00 2001 From: yunhong <337361684@qq.com> Date: Thu, 29 Oct 2026 11:26:52 +0800 Subject: [PATCH 4/4] addressing bing's comments2 --- .../rebalance/RebalanceManager.java | 77 +++++++++++++++---- .../fluss/server/zk/data/RebalanceTask.java | 24 ++++++ .../zk/data/RebalanceTaskJsonSerde.java | 26 ++++++- .../coordinator/CoordinatorRebalanceTest.java | 16 +++- .../rebalance/RebalanceManagerTest.java | 20 +++++ .../zk/data/RebalanceTaskJsonSerdeTest.java | 16 +++- 6 files changed, 150 insertions(+), 29 deletions(-) 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 5e36c9efbfc..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 @@ -104,10 +104,11 @@ * *

    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. 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. + * 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 @@ -301,14 +302,17 @@ public synchronized void recoverRebalance(RebalanceTask rebalanceTask) { for (Map.Entry entry : rebalanceTask.getExecutePlan().entrySet()) { - if (rebalanceExecutor.isRebalanceTaskAtOrigin(entry.getValue())) { + 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 { addRecoveredIntermediateTask(entry.getKey(), entry.getValue()); } } if (cancelRequested) { - cancelPendingTasks(); + cancelPendingTasks(findTasksToCancel()); } if (inProgressRebalanceTasks.isEmpty()) { @@ -417,24 +421,29 @@ public synchronized void cancelRebalance(@Nullable String rebalanceId) { return; } - Map executePlan = allRebalancePlans(); + Set tasksToCancel = findTasksToCancel(); + Set canceledBuckets = canceledBuckets(); + canceledBuckets.addAll(tasksToCancel); try { zkClient.registerRebalanceTask( - new RebalanceTask(currentRebalanceId, REBALANCING, executePlan, true)); + new RebalanceTask( + currentRebalanceId, + REBALANCING, + allRebalancePlans(), + true, + canceledBuckets)); } catch (Exception e) { throw new RebalanceFailureException( "Failed to persist rebalance cancellation request.", e); } cancelRequested = true; - cancelPendingTasks(); + 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()) { - RebalanceResultForBucket result = - inProgressRebalanceTasks.get(attempt.executionKey.getTableBucket()); - if (result != null && rebalanceExecutor.isRebalanceTaskAtOrigin(result.plan())) { + if (tasksToCancel.contains(attempt.executionKey.getTableBucket())) { finishRebalanceTask(attempt.executionKey, CANCELED); } } @@ -448,13 +457,32 @@ public synchronized void cancelRebalance(@Nullable String rebalanceId) { currentRebalanceId); } - private void cancelPendingTasks() { + 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(); - RebalanceResultForBucket result = inProgressRebalanceTasks.get(bucket); - // A recovered pending task may have changed after it was classified at startup. - if (!recoveredRebalance || rebalanceExecutor.isRebalanceTaskAtOrigin(result.plan())) { + if (tasksToCancel.contains(bucket)) { + RebalanceResultForBucket result = inProgressRebalanceTasks.get(bucket); pending.remove(); inProgressRebalanceTasks.remove(bucket); finishedRebalanceTasks.put( @@ -680,7 +708,11 @@ private void persistFinalStatus() { try { zkClient.registerRebalanceTask( new RebalanceTask( - currentRebalanceId, finalStatus, allRebalancePlans(), cancelRequested)); + currentRebalanceId, + finalStatus, + allRebalancePlans(), + cancelRequested, + canceledBuckets())); } catch (Exception e) { rebalanceStatus = REBALANCING; finalizationPending = true; @@ -810,6 +842,17 @@ private Map allRebalancePlans() { 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; + } + private RebalanceStatus aggregateFinalStatus() { for (RebalanceResultForBucket result : finishedRebalanceTasks.values()) { if (result.status() == FAILED) { 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 239402a0c61..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. @@ -48,6 +51,9 @@ public class RebalanceTask { /** 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; @@ -67,9 +73,19 @@ public RebalanceTask( 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<>(); @@ -102,6 +118,10 @@ public boolean isCancelRequested() { return cancelRequested; } + public Set getCanceledBuckets() { + return Collections.unmodifiableSet(canceledBuckets); + } + public Map> getPlanForBuckets() { return planForBuckets; } @@ -139,6 +159,8 @@ public String toString() { + rebalanceStatus + ", cancelRequested=" + cancelRequested + + ", canceledBuckets=" + + canceledBuckets + ", planForBuckets=" + planForBuckets + ", planForBucketsOfPartitionedTable=" @@ -159,6 +181,7 @@ public boolean equals(Object o) { 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); @@ -170,6 +193,7 @@ public int hashCode() { 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 a410a7f6f20..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 @@ -50,6 +52,7 @@ 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"; @@ -73,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(); @@ -87,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(); @@ -108,6 +117,7 @@ public RebalanceTask deserialize(JsonNode node) { 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(); @@ -120,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(); @@ -144,13 +157,18 @@ public RebalanceTask deserialize(JsonNode node) { } } - return new RebalanceTask(rebalanceId, rebalanceStatus, planForBuckets, cancelRequested); + 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/CoordinatorRebalanceTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorRebalanceTest.java index f1b6dccc3b9..b0e1eadb493 100644 --- 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 @@ -469,9 +469,9 @@ void testRestartAfterCancellationDoesNotStartPendingBucketAfterLeaderChange(bool long tableId = createTableWithReplicas(origin, 2); TableBucket running = new TableBucket(tableId, 0); TableBucket pending = new TableBucket(tableId, 1); - // An unrelated leader election must not make the pending plan look started. For a - // leader-only plan, the elected leader is neither the original nor the target leader. - int pendingLeader = leaderOnly ? 2 : 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); @@ -508,6 +508,14 @@ void testRestartAfterCancellationDoesNotStartPendingBucketAfterLeaderChange(bool .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), @@ -520,7 +528,7 @@ void testRestartAfterCancellationDoesNotStartPendingBucketAfterLeaderChange(bool .isEqualTo(RebalanceStatus.CANCELED); assertThat(zookeeperClient.getTableAssignment(tableId).get().getBucketAssignment(1)) .isEqualTo(new BucketAssignment(origin)); - verifyIsr(pending, pendingLeader, pendingIsr); + verifyIsr(pending, expectedFinalPendingLeader, pendingIsr); verifyIsr(running, 0, target); } 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 7f8b7905d50..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 @@ -257,6 +257,7 @@ void testCancelPersistsIntentAndDrainsOnlyAdmittedTasks() throws Exception { 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 @@ -316,6 +317,25 @@ void testCanceledRebalanceReportsFailedWhenDrainFails() throws Exception { 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 = 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 4a464f10772..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 @@ -27,6 +27,7 @@ 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; @@ -86,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\":2,\"rebalance_id\":\"rebalance-task-21jd\",\"rebalance_status\":0,\"cancel_requested\":false,\"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]}]}," @@ -110,7 +116,8 @@ void testDeserializeLegacyTaskWithoutCancelRequested() throws IOException { String legacyJson = expectedJsons()[0] .replace("\"version\":2", "\"version\":1") - .replace(",\"cancel_requested\":false", ""); + .replace(",\"cancel_requested\":true", "") + .replace(",\"canceled\":true", ""); RebalanceTask rebalanceTask = JsonSerdeUtils.readValue( @@ -118,6 +125,7 @@ void testDeserializeLegacyTaskWithoutCancelRequested() throws IOException { RebalanceTaskJsonSerde.INSTANCE); assertThat(rebalanceTask.isCancelRequested()).isFalse(); + assertThat(rebalanceTask.getCanceledBuckets()).isEmpty(); assertThat(rebalanceTask.getExecutePlan()).hasSize(5); } }