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 3bb4279daff..f1962e968e1 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 @@ -251,6 +251,19 @@ public class ConfigOptions { + "conditions, such as disk write protection, become electable " + "again after recovery."); + public static final ConfigOption COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS = + key("coordinator.rebalance.max-inflight-tasks") + .intType() + .defaultValue(1) + .withDescription( + "The maximum number of bucket-level rebalance tasks that can be " + + "executed concurrently by the coordinator. A higher value " + + "can speed up rebalance, while a lower value reduces the " + + "number of simultaneous bucket movements. Setting it to 0 " + + "pauses scheduling new rebalance tasks; already in-flight " + + "tasks continue until they complete or time out. The value " + + "must be non-negative."); + public static final ConfigOption LOG_TABLE_ALLOW_CREATION = key("allow.create.log.tables") .booleanType() diff --git a/fluss-common/src/test/java/org/apache/fluss/cluster/rebalance/RebalanceStatusTest.java b/fluss-common/src/test/java/org/apache/fluss/cluster/rebalance/RebalanceStatusTest.java new file mode 100644 index 00000000000..9fa6d889690 --- /dev/null +++ b/fluss-common/src/test/java/org/apache/fluss/cluster/rebalance/RebalanceStatusTest.java @@ -0,0 +1,37 @@ +/* + * 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.cluster.rebalance; + +import org.junit.jupiter.api.Test; + +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.TIMEOUT; +import static org.assertj.core.api.Assertions.assertThat; + +/** Test for {@link RebalanceStatus}. */ +class RebalanceStatusTest { + + @Test + void testTimeoutIsRecoverableRatherThanFinal() { + assertThat(FINAL_STATUSES).containsExactlyInAnyOrder(COMPLETED, CANCELED, FAILED); + assertThat(FINAL_STATUSES).doesNotContain(TIMEOUT); + } +} diff --git a/fluss-server/src/main/java/org/apache/fluss/server/DynamicConfigManager.java b/fluss-server/src/main/java/org/apache/fluss/server/DynamicConfigManager.java index dec53519b2a..51fc9463e3e 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/DynamicConfigManager.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/DynamicConfigManager.java @@ -99,6 +99,21 @@ public void register(ServerReconfigurable serverReconfigurable) { dynamicServerConfig.register(serverReconfigurable); } + /** + * Register a ServerReconfigurable and immediately apply the current effective configuration. + * + *

Use this for components created after {@link #startup()}, when persisted dynamic + * configuration may already have been loaded. + */ + public void registerAndApplyCurrentConfig(ServerReconfigurable serverReconfigurable) { + dynamicServerConfig.registerAndApplyCurrentConfig(serverReconfigurable); + } + + /** Unregister a ServerReconfigurable that no longer listens to configuration changes. */ + public void unregister(ServerReconfigurable serverReconfigurable) { + dynamicServerConfig.unregister(serverReconfigurable); + } + /** * Register a ConfigValidator for stateless validation. * diff --git a/fluss-server/src/main/java/org/apache/fluss/server/DynamicServerConfig.java b/fluss-server/src/main/java/org/apache/fluss/server/DynamicServerConfig.java index 98d721e152e..5f45adace63 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/DynamicServerConfig.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/DynamicServerConfig.java @@ -45,6 +45,7 @@ import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; +import static org.apache.fluss.config.ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS; import static org.apache.fluss.config.ConfigOptions.DATALAKE_FORMAT; import static org.apache.fluss.config.ConfigOptions.KV_LEADER_REPLICA_MEMORY_RESERVED; import static org.apache.fluss.config.ConfigOptions.KV_SHARED_RATE_LIMITER_BYTES_PER_SEC; @@ -77,6 +78,7 @@ class DynamicServerConfig { new HashSet<>( Arrays.asList( DATALAKE_FORMAT.key(), + COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), LOG_RETENTION_ROLL_ACTIVE_SEGMENT_ENABLED.key(), LOG_REPLICA_MIN_IN_SYNC_REPLICAS_NUMBER.key(), KV_LEADER_REPLICA_MEMORY_RESERVED.key(), @@ -137,6 +139,24 @@ void register(ServerReconfigurable serverReconfigurable) { serverReconfigures.put(serverReconfigurable.getClass(), serverReconfigurable); } + void registerAndApplyCurrentConfig(ServerReconfigurable serverReconfigurable) { + inWriteLock( + lock, + () -> { + serverReconfigurable.validate(currentConfig); + serverReconfigurable.reconfigure(currentConfig); + serverReconfigures.put(serverReconfigurable.getClass(), serverReconfigurable); + }); + } + + void unregister(ServerReconfigurable serverReconfigurable) { + inWriteLock( + lock, + () -> + serverReconfigures.remove( + serverReconfigurable.getClass(), serverReconfigurable)); + } + /** * Register a ConfigValidator for stateless validation. * 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 96220fee120..e8bb75dafda 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 @@ -78,14 +78,19 @@ 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.RebalanceMaxInflightTasksChangedEvent; 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.RemoveServerTagEvent; import org.apache.fluss.server.coordinator.event.ResumeDropEvent; import org.apache.fluss.server.coordinator.event.RetryOfflineLeaderEvent; @@ -95,6 +100,8 @@ 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.RebalanceExecutionKey; +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; @@ -126,7 +133,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; @@ -167,7 +173,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); @@ -219,22 +225,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); @@ -271,8 +272,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) { @@ -743,6 +746,11 @@ public void process(CoordinatorEvent event) { RebalanceEvent rebalanceEvent = (RebalanceEvent) event; completeFromCallable( rebalanceEvent.getRespCallback(), () -> processRebalance(rebalanceEvent)); + } else if (event instanceof RecoverRebalanceEvent) { + rebalanceManager.recoverRebalance(((RecoverRebalanceEvent) event).getRebalanceTask()); + } else if (event instanceof RebalanceMaxInflightTasksChangedEvent) { + rebalanceManager.updateMaxInflightRebalanceTasks( + ((RebalanceMaxInflightTasksChangedEvent) event).getMaxInflightTasks()); } else if (event instanceof CancelRebalanceEvent) { CancelRebalanceEvent cancelRebalanceEvent = (CancelRebalanceEvent) event; completeFromCallable( @@ -750,11 +758,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) { + reconcileRebalanceTask(reconcileEvent.getExecutionKey(), 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 @@ -1190,13 +1208,15 @@ 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()) { - offlineReplicas.add( - new TableBucketReplica( - notifyLeaderAndIsrResultForBucket.getTableBucket(), serverId)); + offlineReplicas.add(new TableBucketReplica(tableBucket, serverId)); } else { - succeededBuckets.add(notifyLeaderAndIsrResultForBucket.getTableBucket()); + succeededBuckets.add(tableBucket); } } for (TableBucket tb : succeededBuckets) { @@ -1213,9 +1233,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())); } } @@ -1580,34 +1606,27 @@ 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.", + if (!allBuckets.contains(tableBucket) + || coordinatorContext.isTableQueuedForDeletion(tableBucket.getTableId()) + || coordinatorContext.isToBeDeleted(tableBucket)) { + LOG.info( + "Complete rebalance task of tableBucket {} since it was deleted or is being " + + "deleted.", tableBucket); - rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED); + rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.COMPLETED); return; } List newReplicas = planForBucket.getNewReplicas(); ReplicaReassignment reassignment = - ReplicaReassignment.build( - coordinatorContext.getAssignment(tableBucket), newReplicas); + ReplicaReassignment.build(planForBucket.getOriginReplicas(), 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). + // from the tablet server to confirm the leader change has been applied. The task keeps + // occupying a RebalanceManager execution slot while waiting, so concurrent leader + // migrations remain bounded by the configured rebalance limit. LOG.info("trigger leader election for tableBucket {}.", tableBucket); tableBucketStateMachine.handleStateChange( Collections.singleton(tableBucket), @@ -1619,10 +1638,19 @@ public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) { "Try to processing bucket reassignment for tableBucket {} with assignment: {}.", tableBucket, reassignment); - onBucketReassignment(tableBucket, reassignment, false); + if (!reassignment.isBeingReassigned()) { + rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.COMPLETED); + } else { + resumeBucketReassignment(tableBucket, reassignment); + } } catch (Exception e) { - LOG.error("Error when processing bucket reassignment.", e); - rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED); + // The coordinator or a tablet server may be temporarily unavailable. Keep the + // task active so that timeout reconciliation can retry from persisted state. + LOG.warn( + "Failed to process bucket reassignment for {}. It remains recoverable and " + + "will be retried after timeout.", + tableBucket, + e); } } } @@ -1630,11 +1658,30 @@ public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) { /** 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 + && !rebalanceManager + .getExecutionKey(tableBucket) + .equals(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()); @@ -1652,13 +1699,120 @@ 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 non-final persisted plan can be inferred complete after recovery. */ + public boolean isRebalanceTaskComplete(RebalancePlanForBucket planForBucket) { + TableBucket tableBucket = planForBucket.getTableBucket(); + List targetReplicas = planForBucket.getNewReplicas(); + // Recovery runs right after the coordinator context has been loaded from ZooKeeper, so the + // in-memory view is authoritative here and no extra ZooKeeper reads are needed. + if (!coordinatorContext.getAllBuckets().contains(tableBucket) + || !coordinatorContext.liveTabletServerSet().containsAll(targetReplicas) + || !coordinatorContext.getAssignment(tableBucket).equals(targetReplicas)) { + return false; + } + Optional leaderAndIsrOpt = + coordinatorContext.getBucketLeaderAndIsr(tableBucket); + if (!leaderAndIsrOpt.isPresent()) { + return false; + } + LeaderAndIsr leaderAndIsr = leaderAndIsrOpt.get(); + // Replaying a plan whose target state is already in place would elect the very same leader + // again and bump the leader epoch of an already migrated bucket, which churns the metadata + // of every client of that bucket. So only replay a plan that has not reached its target. + if (planForBucket.isLeaderChanged() + && leaderAndIsr.leader() != planForBucket.getNewLeader()) { + return false; + } + return new HashSet<>(leaderAndIsr.isr()).equals(new HashSet<>(targetReplicas)); + } + + /** 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(); + return leaderAndIsr.leader() == planForBucket.getOriginalLeader() + && new HashSet<>(leaderAndIsr.isr()) + .equals(new HashSet<>(planForBucket.getOriginReplicas())); + } + + private void reconcileRebalanceTask( + RebalanceExecutionKey executionKey, RebalancePlanForBucket planForBucket) { + TableBucket tableBucket = planForBucket.getTableBucket(); + if (!coordinatorContext.getAllBuckets().contains(tableBucket) + || coordinatorContext.isTableQueuedForDeletion(tableBucket.getTableId()) + || coordinatorContext.isToBeDeleted(tableBucket)) { + LOG.info( + "Complete rebalance task {} because its table bucket was deleted or is being " + + "deleted.", + executionKey); + rebalanceManager.finishRebalanceTask(executionKey, RebalanceStatus.COMPLETED); + return; + } + + ReplicaReassignment reassignment = + ReplicaReassignment.build( + planForBucket.getOriginReplicas(), planForBucket.getNewReplicas()); + try { + if (!reassignment.isBeingReassigned()) { + LOG.info("Retry timed-out leader reassignment task {}.", executionKey); + tableBucketStateMachine.handleStateChange( + Collections.singleton(tableBucket), + OnlineBucket, + new ReassignmentLeaderElection(planForBucket.getNewReplicas())); + } else { + resumeBucketReassignment(tableBucket, reassignment); } + } catch (Exception e) { + // A timeout retry may fail because ZooKeeper or a tablet server is temporarily + // unavailable. Keep the task recoverable and retry it on the next reconciliation + // interval instead of abandoning a possibly intermediate assignment. + LOG.warn( + "Failed to retry timed-out rebalance task {}. It remains timed out and will " + + "be retried.", + executionKey, + e); } } @@ -1675,17 +1829,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(), @@ -1703,6 +1869,11 @@ static boolean isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader( && responseServerId == planForBucket.getNewLeader(); } + private static boolean isSuccessfulReassignmentResponseFromLeader( + int responseServerId, @Nullable NotifyLeaderAndIsrRequestContext requestContext) { + return requestContext != null && responseServerId == requestContext.getLeader(); + } + /** * Reassigning replicas for a tableBucket goes through a few steps listed in the code. * @@ -1747,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: @@ -1832,17 +2004,138 @@ 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( TableBucket tableBucket, ReplicaReassignment reassignment) throws Exception { LeaderAndIsr leaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket).get(); List isr = leaderAndIsr.isr(); List targetReplicas = reassignment.getTargetReplicas(); - return targetReplicas.isEmpty() || new HashSet<>(isr).containsAll(targetReplicas); + 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( @@ -2539,6 +2832,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()); @@ -2548,22 +2864,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 539fb900160..25a55bd601f 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( @@ -415,6 +425,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 @@ -463,7 +486,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/CoordinatorServer.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorServer.java index a2f9b4fad97..4f89bbf214c 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorServer.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorServer.java @@ -360,6 +360,8 @@ protected void initCoordinatorLeader() throws Exception { kvSnapshotLeaseManager, scheduler, clock); + dynamicConfigManager.registerAndApplyCurrentConfig( + coordinatorEventProcessor.getRebalanceManager()); coordinatorEventProcessor.startup(); // As the active leader, this server is the sole writer of dynamic configs and holds the @@ -393,10 +395,7 @@ protected void cleanupCoordinatorLeader() { // Clean up leader-specific resources in reverse order of initialization try { - if (coordinatorEventProcessor != null) { - coordinatorEventProcessor.shutdown(); - coordinatorEventProcessor = null; - } + shutdownCoordinatorEventProcessor(); } catch (Throwable t) { LOG.warn("Failed to shutdown coordinator event processor", t); } @@ -627,9 +626,7 @@ CompletableFuture stopServices() { } try { - if (coordinatorEventProcessor != null) { - coordinatorEventProcessor.shutdown(); - } + shutdownCoordinatorEventProcessor(); } catch (Throwable t) { exception = ExceptionUtils.firstOrSuppressed(t, exception); } @@ -731,6 +728,14 @@ CompletableFuture stopServices() { } } + private void shutdownCoordinatorEventProcessor() { + if (coordinatorEventProcessor != null) { + dynamicConfigManager.unregister(coordinatorEventProcessor.getRebalanceManager()); + coordinatorEventProcessor.shutdown(); + coordinatorEventProcessor = null; + } + } + @Override protected CompletableFuture getTerminationFuture() { return terminationFuture; 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..2cc622b7ca9 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/NotifyLeaderAndIsrRequestContext.java @@ -0,0 +1,110 @@ +/* + * 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; + +import java.util.Objects; + +/** 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 boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + NotifyLeaderAndIsrRequestContext that = (NotifyLeaderAndIsrRequestContext) o; + return coordinatorEpoch == that.coordinatorEpoch + && leader == that.leader + && leaderEpoch == that.leaderEpoch + && bucketEpoch == that.bucketEpoch + && Objects.equals(rebalanceExecutionKey, that.rebalanceExecutionKey); + } + + @Override + public int hashCode() { + return Objects.hash( + coordinatorEpoch, leader, leaderEpoch, bucketEpoch, 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/RebalanceMaxInflightTasksChangedEvent.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RebalanceMaxInflightTasksChangedEvent.java new file mode 100644 index 00000000000..c05aa723217 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RebalanceMaxInflightTasksChangedEvent.java @@ -0,0 +1,37 @@ +/* + * 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 fired when the rebalance max in-flight task limit changes dynamically. */ +public final class RebalanceMaxInflightTasksChangedEvent implements CoordinatorEvent { + + private final int maxInflightTasks; + + public RebalanceMaxInflightTasksChangedEvent(int maxInflightTasks) { + this.maxInflightTasks = maxInflightTasks; + } + + public int getMaxInflightTasks() { + return maxInflightTasks; + } + + @Override + public String toString() { + return "RebalanceMaxInflightTasksChangedEvent{maxInflightTasks=" + maxInflightTasks + "}"; + } +} 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 new file mode 100644 index 00000000000..2d329ef6b4e --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RecoverRebalanceEvent.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.zk.data.RebalanceTask; + +/** An event that recovers and reconciles a persisted rebalance task. */ +public final class RecoverRebalanceEvent implements CoordinatorEvent { + private final RebalanceTask rebalanceTask; + + public RecoverRebalanceEvent(RebalanceTask rebalanceTask) { + this.rebalanceTask = rebalanceTask; + } + + public RebalanceTask getRebalanceTask() { + return rebalanceTask; + } + + @Override + public String toString() { + return "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..f8bbda5ae7d --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceExecutionKey.java @@ -0,0 +1,79 @@ +/* + * 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 one execution attempt of a bucket-level rebalance task. */ +public final class RebalanceExecutionKey { + private final String rebalanceId; + private final TableBucket tableBucket; + private final long attemptId; + + public RebalanceExecutionKey(String rebalanceId, TableBucket tableBucket, long attemptId) { + this.rebalanceId = rebalanceId; + this.tableBucket = tableBucket; + this.attemptId = attemptId; + } + + public String getRebalanceId() { + return rebalanceId; + } + + public TableBucket getTableBucket() { + return tableBucket; + } + + public long getAttemptId() { + return attemptId; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + RebalanceExecutionKey that = (RebalanceExecutionKey) o; + return attemptId == that.attemptId + && Objects.equals(rebalanceId, that.rebalanceId) + && Objects.equals(tableBucket, that.tableBucket); + } + + @Override + public int hashCode() { + return Objects.hash(rebalanceId, tableBucket, attemptId); + } + + @Override + public String toString() { + return "RebalanceExecutionKey{" + + "rebalanceId='" + + rebalanceId + + '\'' + + ", tableBucket=" + + tableBucket + + ", attemptId=" + + attemptId + + '}'; + } +} 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..17846fe7391 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceExecutor.java @@ -0,0 +1,37 @@ +/* + * 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 a non-final persisted plan can be treated as complete during recovery. */ + boolean isRebalanceTaskComplete(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 cc29b982b07..97ca67408cd 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,20 @@ 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.ConfigOptions; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.config.cluster.ServerReconfigurable; +import org.apache.fluss.exception.ConfigException; import org.apache.fluss.exception.NoRebalanceInProgressException; +import org.apache.fluss.exception.RebalanceFailureException; import org.apache.fluss.metadata.TableBucket; 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.RebalanceMaxInflightTasksChangedEvent; 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; import org.apache.fluss.server.coordinator.rebalance.model.ClusterModel; @@ -48,6 +56,7 @@ import javax.annotation.Nullable; import java.util.ArrayDeque; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -64,18 +73,20 @@ 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.utils.Preconditions.checkArgument; -import static org.apache.fluss.utils.Preconditions.checkNotNull; /** * A rebalance manager to generate rebalance plan, and execution rebalance plan. * - *

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

    This manager is used in the coordinator event loop as a single-threaded model. Non-event + * threads enqueue coordinator events instead of directly advancing rebalance state. */ -public class RebalanceManager { +public class RebalanceManager implements ServerReconfigurable { private static final Logger LOG = LoggerFactory.getLogger(RebalanceManager.class); /** Hardcoded timeout for an in-flight rebalance task: 2 minutes. */ @@ -84,53 +95,82 @@ 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; + + /** + * Hardcoded time after which a timed-out task is given up on while it cannot make progress at + * all, because a target replica is not hosted by a live tablet server: 30 minutes. + */ + private static final long TARGET_UNAVAILABLE_TIMEOUT_MS = 30 * 60 * 1000L; + + /** + * Hardcoded time after which a timed-out task is given up on even though its target replicas + * are live: 24 hours. This is only a safety net that keeps a rebalance from staying non-final + * forever, a single bucket that does not change any observable state for that long is broken. + */ + private static final long NO_PROGRESS_TIMEOUT_MS = 24 * 60 * 60 * 1000L; + + /** + * Hardcoded upper bound on the number of timed-out tasks tracked at the same time. Every + * tracked task keeps costing coordinator and ZooKeeper work on each reconciliation, and each + * admitted task adds one more concurrent replica migration. + */ + private static final int MAX_TRACKED_TIMED_OUT_TASKS = 8; + private final ZooKeeperClient zkClient; - private final CoordinatorEventProcessor eventProcessor; + private final RebalanceExecutor rebalanceExecutor; private final EventManager eventManager; private final Clock clock; private final ScheduledExecutorService timeoutChecker; - /** A queue of in progress table bucket to rebalance. */ - private final Queue inProgressRebalanceTasksQueue = new ArrayDeque<>(); + /** A queue of bucket tasks that have not started. */ + 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<>(); + /** Normally running tasks that occupy configured execution slots. */ + private final Map runningRebalanceTasks = + new ConcurrentHashMap<>(); + + /** Soft-timed-out tasks that no longer occupy the normal execution slot. */ + private final Map timedOutRebalanceTasks = + new ConcurrentHashMap<>(); + + private final Set queuedTimeoutEvents = ConcurrentHashMap.newKeySet(); + private final Set queuedReconcileEvents = ConcurrentHashMap.newKeySet(); + /** A mapping from table bucket to rebalance status of failed or completed tasks. */ private final Map finishedRebalanceTasks = new ConcurrentHashMap<>(); private final GoalOptimizer goalOptimizer; + private int maxInflightRebalanceTasks; + private int queuedMaxInflightRebalanceTasks; private volatile long registerTime; private volatile @Nullable RebalanceStatus rebalanceStatus; private volatile @Nullable String currentRebalanceId; + private volatile boolean recoveryPending; + private volatile boolean cancelRequested; + private volatile boolean finalizationPending; + private volatile boolean finalizationEventQueued; private volatile boolean isClosed = false; - - /** - * 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; - - /** The bucket of the current in-flight task, or null if idle. Acts as the "gate" variable. */ - private volatile @Nullable TableBucket inflightTaskBucket; + private long nextAttemptId; 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( @@ -139,17 +179,22 @@ public RebalanceManager( @VisibleForTesting RebalanceManager( - CoordinatorEventProcessor eventProcessor, + RebalanceExecutor rebalanceExecutor, ZooKeeperClient zkClient, EventManager eventManager, Clock clock, + Configuration conf, ScheduledExecutorService timeoutChecker) { - this.eventProcessor = eventProcessor; + this.rebalanceExecutor = rebalanceExecutor; this.zkClient = zkClient; this.eventManager = eventManager; this.clock = clock == null ? SystemClock.getInstance() : clock; this.timeoutChecker = timeoutChecker; this.goalOptimizer = new GoalOptimizer(); + validate(conf); + this.maxInflightRebalanceTasks = + conf.get(ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS); + this.queuedMaxInflightRebalanceTasks = maxInflightRebalanceTasks; } public void startup() { @@ -174,94 +219,181 @@ public void start() { return currentRebalanceId; } + @Override + public void validate(Configuration newConfig) throws ConfigException { + int newMaxInflightRebalanceTasks = + newConfig.get(ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS); + if (newMaxInflightRebalanceTasks < 0) { + throw new ConfigException( + String.format( + "Invalid %s: must be non-negative, but was %s", + ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), + newMaxInflightRebalanceTasks)); + } + } + + @Override + public synchronized void reconfigure(Configuration newConfig) { + int newMaxInflightRebalanceTasks = + newConfig.get(ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS); + if (newMaxInflightRebalanceTasks == queuedMaxInflightRebalanceTasks) { + LOG.debug( + "{} unchanged: {}", + ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), + newMaxInflightRebalanceTasks); + return; + } + + int oldQueuedMaxInflightRebalanceTasks = queuedMaxInflightRebalanceTasks; + queuedMaxInflightRebalanceTasks = newMaxInflightRebalanceTasks; + LOG.info( + "{} change queued: {} -> {}", + ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), + oldQueuedMaxInflightRebalanceTasks, + newMaxInflightRebalanceTasks); + if (!isClosed) { + eventManager.put( + new RebalanceMaxInflightTasksChangedEvent(newMaxInflightRebalanceTasks)); + } + } + private void initialize() { try { zkClient.getRebalanceTask() .ifPresent( - rebalancePlan -> - registerRebalance( - rebalancePlan.getRebalanceId(), - rebalancePlan.getExecutePlan(), - rebalancePlan.getRebalanceStatus())); + rebalanceTask -> { + recoveryPending = true; + eventManager.put(new RecoverRebalanceEvent(rebalanceTask)); + }); } catch (Exception e) { LOG.error( - "Failed to get rebalance plan from zookeeper, it will be treated as no" - + "rebalance tasks.", + "Failed to get rebalance plan from zookeeper. New rebalance requests will be " + + "rejected until the coordinator is restarted and recovery succeeds.", e); + recoveryPending = true; } } - public void registerRebalance( + public synchronized void registerRebalance( String rebalanceId, Map rebalancePlan, RebalanceStatus newStatus) { 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()) { + if (FINAL_STATUSES.contains(newStatus)) { + finishedRebalanceTasks.put( + entry.getKey(), RebalanceResultForBucket.of(entry.getValue(), newStatus)); + } else { + addPendingTask(entry.getKey(), entry.getValue()); + } + } + + if (!pendingRebalanceTasks.isEmpty()) { rebalanceStatus = REBALANCING; - processNewRebalanceTask(); + processNewRebalanceTasks(); } else { rebalanceStatus = newStatus; } } - public void finishRebalanceTask(TableBucket tableBucket, RebalanceStatus statusForBucket) { + /** Recovers a persisted task by comparing every bucket plan with current coordinator state. */ + public synchronized void recoverRebalance(RebalanceTask rebalanceTask) { checkNotClosed(); - if (inProgressRebalanceTasksQueue.contains(tableBucket)) { - inProgressRebalanceTasksQueue.remove(tableBucket); - RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.remove(tableBucket); - checkNotNull(resultForBucket, "RebalanceResultForBucket is null."); - finishedRebalanceTasks.put( - tableBucket, - RebalanceResultForBucket.of(resultForBucket.plan(), statusForBucket)); - // Clear gate (bucket) first, then data (startMs). - inflightTaskBucket = null; - inflightTaskStartMs = -1; - LOG.info( - "Rebalance task {} in progress: {} tasks pending, {} completed.", - currentRebalanceId, - inProgressRebalanceTasksQueue.size(), - finishedRebalanceTasks.size()); + if (FINAL_STATUSES.contains(rebalanceTask.getRebalanceStatus())) { + resetRebalance(rebalanceTask.getRebalanceId(), rebalanceTask.isCancelRequested()); + for (Map.Entry entry : + rebalanceTask.getExecutePlan().entrySet()) { + finishedRebalanceTasks.put( + entry.getKey(), + RebalanceResultForBucket.of( + entry.getValue(), rebalanceTask.getRebalanceStatus())); + } + rebalanceStatus = rebalanceTask.getRebalanceStatus(); + return; + } - if (inProgressRebalanceTasksQueue.isEmpty()) { - // All rebalance tasks are completed. - completeRebalance(); + boolean recoveringCancellation = + rebalanceTask.isCancelRequested() || rebalanceTask.getRebalanceStatus() == CANCELED; + resetRebalance(rebalanceTask.getRebalanceId(), recoveringCancellation); + + for (Map.Entry entry : + rebalanceTask.getExecutePlan().entrySet()) { + TableBucket tableBucket = entry.getKey(); + RebalancePlanForBucket plan = entry.getValue(); + if (rebalanceExecutor.isRebalanceTaskComplete(plan)) { + finishedRebalanceTasks.put( + tableBucket, RebalanceResultForBucket.of(plan, COMPLETED)); + } else if (recoveringCancellation && rebalanceExecutor.isRebalanceTaskAtOrigin(plan)) { + finishedRebalanceTasks.put( + tableBucket, RebalanceResultForBucket.of(plan, CANCELED)); } else { - // Trigger one rebalance task to execute. - processNewRebalanceTask(); + addPendingTask(tableBucket, plan); } } + + if (inProgressRebalanceTasks.isEmpty()) { + rebalanceStatus = recoveringCancellation ? CANCELED : aggregateFinalStatus(); + persistFinalStatus(); + } else { + rebalanceStatus = REBALANCING; + processNewRebalanceTasks(); + } + } + + public synchronized void finishRebalanceTask( + TableBucket tableBucket, RebalanceStatus statusForBucket) { + RebalanceExecutionKey executionKey = getExecutionKey(tableBucket); + if (executionKey != null) { + finishRebalanceTask(executionKey, statusForBucket); + } } - public @Nullable RebalanceProgress listRebalanceProgress(@Nullable String rebalanceId) { + public synchronized boolean finishRebalanceTask( + RebalanceExecutionKey executionKey, RebalanceStatus statusForBucket) { + checkNotClosed(); + checkArgument(statusForBucket != TIMEOUT, "Use timeoutRebalanceTask for soft timeouts."); + RebalanceTaskAttempt attempt = findActiveAttempt(executionKey); + if (attempt == null) { + LOG.debug("Ignore stale completion for {}.", executionKey); + return false; + } + + TableBucket tableBucket = executionKey.getTableBucket(); + runningRebalanceTasks.remove(tableBucket); + timedOutRebalanceTasks.remove(tableBucket); + queuedTimeoutEvents.remove(executionKey); + queuedReconcileEvents.remove(executionKey); + RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.remove(tableBucket); + if (resultForBucket == null) { + return false; + } + finishedRebalanceTasks.put( + tableBucket, RebalanceResultForBucket.of(resultForBucket.plan(), statusForBucket)); + LOG.info( + "Rebalance {} progress: {} pending, {} running, {} timed out and tracking, " + + "{} finished.", + currentRebalanceId, + pendingRebalanceTasks.size(), + runningRebalanceTasks.size(), + timedOutRebalanceTasks.size(), + finishedRebalanceTasks.size()); + + if (inProgressRebalanceTasks.isEmpty()) { + finalizeRebalance(); + } else { + processNewRebalanceTasks(); + } + return true; + } + + public synchronized @Nullable RebalanceProgress listRebalanceProgress( + @Nullable String rebalanceId) { checkNotClosed(); if (rebalanceId != null && currentRebalanceId != null @@ -287,9 +419,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)) { @@ -308,35 +444,47 @@ 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; + TableBucket pending; + while ((pending = pendingRebalanceTasks.poll()) != null) { + RebalanceResultForBucket result = inProgressRebalanceTasks.remove(pending); + if (result != null) { + finishedRebalanceTasks.put( + pending, RebalanceResultForBucket.of(result.plan(), CANCELED)); + } + } - LOG.info("Cancel rebalance task success."); + // 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); + } + } + + if (inProgressRebalanceTasks.isEmpty() && !FINAL_STATUSES.contains(rebalanceStatus)) { + finalizeRebalance(); + } + LOG.info( + "Accepted cancellation for rebalance {}. Running and timed-out tasks will be " + + "drained before the rebalance becomes canceled.", + currentRebalanceId); } - public boolean hasInProgressRebalance() { + public synchronized boolean hasInProgressRebalance() { checkNotClosed(); - return !inProgressRebalanceTasks.isEmpty() || !inProgressRebalanceTasksQueue.isEmpty(); + return recoveryPending || finalizationPending || !inProgressRebalanceTasks.isEmpty(); } public RebalanceTask generateRebalanceTask(List goalsByPriority) { @@ -346,7 +494,8 @@ public RebalanceTask generateRebalanceTask(List goalsByPriority) { try { // Generate the latest cluster model. long startTime = System.currentTimeMillis(); - ClusterModel clusterModel = buildClusterModel(eventProcessor.getCoordinatorContext()); + ClusterModel clusterModel = + buildClusterModel(rebalanceExecutor.getCoordinatorContext()); LOG.info( "Build cluster model for rebalance id {} with {} ms.", rebalanceId, @@ -368,8 +517,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(); @@ -377,46 +531,307 @@ public RebalanceTask generateRebalanceTask(List goalsByPriority) { return null; } - 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; + 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; + } + if (timedOutRebalanceTasks.size() >= MAX_TRACKED_TIMED_OUT_TASKS) { + // Keep the attempt in the normal running set until a reconciliation slot becomes + // available. The timeout checker will enqueue another timeout event on its next pass. + LOG.info( + "Keep timed-out rebalance task {} in the running set because {} other " + + "timed-out tasks are still being reconciled.", + executionKey, + timedOutRebalanceTasks.size()); + 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); + processNewRebalanceTasks(); + 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) > TARGET_UNAVAILABLE_TIMEOUT_MS + || now - attempt.lastProgressMs > NO_PROGRESS_TIMEOUT_MS) { + 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 processNewRebalanceTasks() { + if (timedOutRebalanceTasks.size() >= MAX_TRACKED_TIMED_OUT_TASKS) { + // Stop admitting work until some of the timed-out tasks reach a final status, so that + // a long cluster operation cannot grow the tracked set, and with it the reconciliation + // work and the number of concurrent replica migrations, without bound. + LOG.info( + "Hold back new tasks of rebalance {} because {} timed-out tasks are still " + + "being reconciled.", + currentRebalanceId, + timedOutRebalanceTasks.size()); + return; + } + + while (runningRebalanceTasks.size() < maxInflightRebalanceTasks) { + TableBucket tableBucket = pendingRebalanceTasks.poll(); + if (tableBucket == null) { + return; + } 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, ++nextAttemptId); + runningRebalanceTasks.put( + tableBucket, new RebalanceTaskAttempt(executionKey, clock.milliseconds())); + inProgressRebalanceTasks.put( + tableBucket, RebalanceResultForBucket.of(resultForBucket.plan(), REBALANCING)); + rebalanceExecutor.tryToExecuteRebalanceTask(resultForBucket.plan()); } } - private void completeRebalance() { + /** + * Applies a new rebalance concurrency limit on the coordinator event loop. + * + *

    Increasing the limit admits pending tasks immediately. Decreasing it does not cancel + * running tasks; new tasks are admitted only after the number of running tasks falls below the + * new limit. + */ + public synchronized void updateMaxInflightRebalanceTasks(int newMaxInflightRebalanceTasks) { + checkArgument( + newMaxInflightRebalanceTasks >= 0, + "%s must be non-negative.", + ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key()); + int oldMaxInflightRebalanceTasks = maxInflightRebalanceTasks; + if (newMaxInflightRebalanceTasks == oldMaxInflightRebalanceTasks) { + LOG.debug( + "{} unchanged: {}", + ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), + newMaxInflightRebalanceTasks); + return; + } + + maxInflightRebalanceTasks = newMaxInflightRebalanceTasks; + LOG.info( + "{} reconfigured: {} -> {}", + ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), + oldMaxInflightRebalanceTasks, + newMaxInflightRebalanceTasks); + if (!isClosed) { + processNewRebalanceTasks(); + } + } + + /** Returns the rebalance concurrency limit currently applied on the coordinator event loop. */ + @VisibleForTesting + public synchronized int getMaxInflightRebalanceTasks() { + return maxInflightRebalanceTasks; + } + + 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(); - inProgressRebalanceTasksQueue.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; + this.cancelRequested = cancelRequested; + finalizationPending = false; + finalizationEventQueued = false; + inProgressRebalanceTasks.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; + } + + /** + * 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 reconcileBackoffMs(int dispatchedAttempts) { + long backoff = TIMEOUT_CHECK_INTERVAL_MS << Math.min(dispatchedAttempts, 8); + return Math.min(backoff, MAX_RECONCILE_BACKOFF_MS); + } + + private void enqueueReconciliation(RebalanceTaskAttempt attempt) { + if (queuedReconcileEvents.add(attempt.executionKey)) { + eventManager.put(new ReconcileRebalanceTaskEvent(attempt.executionKey)); + } + } - // Here, it will not clear finishedRebalanceTasks, because it will be used by - // listRebalanceProgress. It will be cleared when next register. + private Map allRebalancePlans() { + Map plans = new HashMap<>(); + for (Map.Entry entry : + inProgressRebalanceTasks.entrySet()) { + plans.put(entry.getKey(), entry.getValue().plan()); + } + for (Map.Entry entry : + finishedRebalanceTasks.entrySet()) { + plans.put(entry.getKey(), entry.getValue().plan()); + } + return plans; + } - 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) { @@ -491,26 +906,32 @@ 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; + long now = clock.milliseconds(); + for (RebalanceTaskAttempt attempt : new HashMap<>(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 : new HashMap<>(timedOutRebalanceTasks).values()) { + if (now >= attempt.nextReconcileMs) { + enqueueReconciliation(attempt); + } + } + + String rebalanceId = currentRebalanceId; + if (finalizationPending && rebalanceId != null && !finalizationEventQueued) { + finalizationEventQueued = true; + eventManager.put(new FinalizeRebalanceEvent(rebalanceId)); } } @@ -525,7 +946,7 @@ public void close() { @VisibleForTesting public ClusterModel buildClusterModel() { - return buildClusterModel(eventProcessor.getCoordinatorContext()); + return buildClusterModel(rebalanceExecutor.getCoordinatorContext()); } @VisibleForTesting @@ -533,4 +954,69 @@ public ClusterModel buildClusterModel() { RebalanceStatus getRebalanceStatus() { return rebalanceStatus; } + + @VisibleForTesting + 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; + + /** Read by the timeout checker thread, written on the coordinator event loop. */ + private volatile 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/DynamicConfigChangeTest.java b/fluss-server/src/test/java/org/apache/fluss/server/DynamicConfigChangeTest.java index 5bc4a1fe8ce..67f815241c4 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/DynamicConfigChangeTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/DynamicConfigChangeTest.java @@ -569,6 +569,74 @@ public void reconfigure(Configuration newConfig) { assertThat(reconfiguredInterval.get()).isEqualTo(Duration.ofMinutes(5)); } + @Test + void testDynamicRebalanceMaxInflightTasksChange() throws Exception { + Configuration configuration = new Configuration(); + configuration.set(ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS, 1); + DynamicConfigManager dynamicConfigManager = createManager(configuration); + RebalanceConcurrencyConfigRecorder recorder = new RebalanceConcurrencyConfigRecorder(); + dynamicConfigManager.registerAndApplyCurrentConfig(recorder); + dynamicConfigManager.startup(); + + assertThat(recorder.value).isEqualTo(1); + dynamicConfigManager.alterConfigs( + Collections.singletonList( + new AlterConfig( + ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), + "5", + AlterConfigOpType.SET))); + + assertThat(zookeeperClient.fetchEntityConfig()) + .containsEntry(ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), "5"); + assertThat(recorder.value).isEqualTo(5); + + dynamicConfigManager.alterConfigs( + Collections.singletonList( + new AlterConfig( + ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), + "0", + AlterConfigOpType.SET))); + assertThat(recorder.value).isZero(); + + dynamicConfigManager.alterConfigs( + Collections.singletonList( + new AlterConfig( + ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), + null, + AlterConfigOpType.DELETE))); + assertThat(zookeeperClient.fetchEntityConfig()) + .doesNotContainKey(ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key()); + assertThat(recorder.value).isEqualTo(1); + } + + @Test + void testLateRegisterReceivesCurrentDynamicRebalanceMaxInflightTasks() throws Exception { + Configuration configuration = new Configuration(); + configuration.set(ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS, 1); + DynamicConfigManager dynamicConfigManager = createManager(configuration); + dynamicConfigManager.startup(); + dynamicConfigManager.alterConfigs( + Collections.singletonList( + new AlterConfig( + ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), + "4", + AlterConfigOpType.SET))); + + RebalanceConcurrencyConfigRecorder recorder = new RebalanceConcurrencyConfigRecorder(); + dynamicConfigManager.registerAndApplyCurrentConfig(recorder); + + assertThat(recorder.value).isEqualTo(4); + + dynamicConfigManager.unregister(recorder); + dynamicConfigManager.alterConfigs( + Collections.singletonList( + new AlterConfig( + ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(), + "6", + AlterConfigOpType.SET))); + assertThat(recorder.value).isEqualTo(4); + } + @Test void testDynamicKvLeaderReplicaMemoryReservedChange() throws Exception { Configuration configuration = new Configuration(); @@ -1216,4 +1284,16 @@ void testDescribeConfigsRedactsProviderResolvedValues() throws Exception { .get(); assertThat(entry.value()).isEqualTo("******"); } + + private static final class RebalanceConcurrencyConfigRecorder implements ServerReconfigurable { + private int value; + + @Override + public void validate(Configuration newConfig) throws ConfigException {} + + @Override + public void reconfigure(Configuration newConfig) { + value = newConfig.get(ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS); + } + } } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/ControlledNotifyGateway.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/ControlledNotifyGateway.java new file mode 100644 index 00000000000..06c9422c397 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/ControlledNotifyGateway.java @@ -0,0 +1,97 @@ +/* + * 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; +import java.util.concurrent.atomic.AtomicInteger; + +final 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); + } +} + +final class ControlledNotifyGateway extends TestTabletServerGateway { + private volatile boolean controlMode; + 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); + } + NotifyLeaderAndIsrResponse response = super.notifyLeaderAndIsr(request).join(); + ControlledNotifyTrigger trigger = new ControlledNotifyTrigger(responseServerId); + pendingTriggers.add(trigger); + return trigger.getFuture().thenApply(ignored -> response); + } +} + +final 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); + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorRebalanceTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorRebalanceTest.java new file mode 100644 index 00000000000..ccaabf1d856 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorRebalanceTest.java @@ -0,0 +1,975 @@ +/* + * 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.Endpoint; +import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket; +import org.apache.fluss.cluster.rebalance.RebalanceStatus; +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.metadata.DatabaseDescriptor; +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.gateway.TabletServerGateway; +import org.apache.fluss.rpc.messages.AdjustIsrResponse; +import org.apache.fluss.server.coordinator.event.AccessContextEvent; +import org.apache.fluss.server.coordinator.event.AdjustIsrReceivedEvent; +import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrRequestContext; +import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrResponseReceivedEvent; +import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent; +import org.apache.fluss.server.coordinator.event.ReconcileRebalanceTaskEvent; +import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager; +import org.apache.fluss.server.coordinator.rebalance.RebalanceExecutionKey; +import org.apache.fluss.server.coordinator.remote.RemoteDirDynamicLoader; +import org.apache.fluss.server.entity.NotifyLeaderAndIsrResultForBucket; +import org.apache.fluss.server.metadata.CoordinatorMetadataCache; +import org.apache.fluss.server.metrics.group.TestingMetricGroups; +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.BucketAssignment; +import org.apache.fluss.server.zk.data.CoordinatorAddress; +import org.apache.fluss.server.zk.data.LeaderAndIsr; +import org.apache.fluss.server.zk.data.TableAssignment; +import org.apache.fluss.server.zk.data.TabletServerRegistration; +import org.apache.fluss.server.zk.data.ZkData; +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.types.DataTypes; +import org.apache.fluss.utils.ExceptionUtils; +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 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.testutils.common.CommonTestUtils.retry; +import static org.apache.fluss.testutils.common.CommonTestUtils.waitValue; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests rebalance execution and recovery in {@link CoordinatorEventProcessor}. */ +class CoordinatorEventProcessorRebalanceTest { + + 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); + + @RegisterExtension + public static final AllCallbackWrapper ZOO_KEEPER_EXTENSION_WRAPPER = + new AllCallbackWrapper<>(new ZooKeeperExtension()); + + private static ZooKeeperClient zookeeperClient; + private static MetadataManager metadataManager; + private static ZkEpoch zkEpoch; + + private final String defaultDatabase = "db"; + + private CoordinatorEventProcessor eventProcessor; + private TestCoordinatorChannelManager testCoordinatorChannelManager; + private AutoPartitionManager autoPartitionManager; + private LakeTableTieringManager lakeTableTieringManager; + private CoordinatorMetadataCache serverMetadataCache; + private ReplicaCapacityController replicaCapacityController; + private KvSnapshotLeaseManager kvSnapshotLeaseManager; + private Scheduler scheduler; + private String remoteDataDir; + + @BeforeAll + static void baseBeforeAll() throws Exception { + zookeeperClient = + ZOO_KEEPER_EXTENSION_WRAPPER + .getCustomExtension() + .getZooKeeperClient(NOPErrorHandler.INSTANCE); + metadataManager = + new MetadataManager( + zookeeperClient, + new Configuration(), + new LakeCatalogDynamicLoader(new Configuration(), null, true)); + + zookeeperClient.registerCoordinatorLeader( + new CoordinatorAddress( + "2", Endpoint.fromListenersString("CLIENT://localhost:10012"))); + + zkEpoch = zookeeperClient.fenceBecomeCoordinatorLeader("2"); + for (int i = 0; i < 3; i++) { + zookeeperClient.registerTabletServer( + i, + new TabletServerRegistration( + "rack" + i, + Collections.singletonList( + new Endpoint("host" + i, 1000, DEFAULT_LISTENER_NAME)), + System.currentTimeMillis())); + } + } + + @BeforeEach + void beforeEach() { + serverMetadataCache = new CoordinatorMetadataCache(); + testCoordinatorChannelManager = new TestCoordinatorChannelManager(); + lakeTableTieringManager = + new LakeTableTieringManager(TestingMetricGroups.LAKE_TIERING_METRICS); + remoteDataDir = zookeeperClient.getDefaultRemoteDataDir(); + Configuration conf = new Configuration(); + conf.setString(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); + replicaCapacityController = new ReplicaCapacityController(conf, serverMetadataCache); + autoPartitionManager = + new AutoPartitionManager( + serverMetadataCache, + metadataManager, + new RemoteDirDynamicLoader(conf), + conf, + replicaCapacityController); + kvSnapshotLeaseManager = + new KvSnapshotLeaseManager( + Duration.ofMinutes(10).toMillis(), + zookeeperClient, + remoteDataDir, + SystemClock.getInstance(), + TestingMetricGroups.COORDINATOR_METRICS); + kvSnapshotLeaseManager.start(); + + scheduler = new FlussScheduler(1); + scheduler.startup(); + + eventProcessor = buildCoordinatorEventProcessor(); + eventProcessor.startup(); + metadataManager.createDatabase( + defaultDatabase, DatabaseDescriptor.builder().build(), false); + } + + @AfterEach + void afterEach() throws Exception { + if (eventProcessor != null) { + eventProcessor.shutdown(); + } + if (scheduler != null) { + scheduler.shutdown(); + } + metadataManager.dropDatabase(defaultDatabase, false, true); + ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(TableIdsZNode.path()); + ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(PartitionIdsZNode.path()); + } + + @Test + void testApplyRebalanceConcurrencyChangeOnCoordinatorEventThread() { + Configuration newConfig = new Configuration(); + newConfig.set(ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS, 3); + + eventProcessor.getRebalanceManager().reconfigure(newConfig); + + retry( + Duration.ofSeconds(10), + () -> + assertThat( + eventProcessor + .getRebalanceManager() + .getMaxInflightRebalanceTasks()) + .isEqualTo(3)); + } + + @Test + void testDoBucketReassignment() throws Exception { + registerTabletServer(3); + + 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, 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 testTimedOutReassignmentRetriesPhaseAIdempotently() throws Exception { + registerTabletServer(3); + + try { + initCoordinatorChannel(); + ConcurrentLinkedDeque pendingTriggers = + new ConcurrentLinkedDeque<>(); + + TablePath tablePath = + TablePath.of(defaultDatabase, "test_timed_out_reassignment_retry"); + Map bucketAssignments = new HashMap<>(); + bucketAssignments.put(0, BucketAssignment.of(0, 1, 3)); + long tableId = + metadataManager.createTable( + tablePath, + remoteDataDir, + TEST_TABLE, + new TableAssignment(bucketAssignments), + false); + TableBucket tableBucket = new TableBucket(tableId, 0); + verifyIsr(tableBucket, 0, Arrays.asList(0, 1, 3)); + installBlockingNotifyGateways(pendingTriggers); + + 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), + RebalanceStatus.NOT_STARTED); + RebalanceExecutionKey executionKey = + eventProcessor.getRebalanceManager().getExecutionKey(tableBucket); + retry(Duration.ofMinutes(1), () -> assertThat(pendingTriggers).isNotEmpty()); + int requestsAfterInitialPhaseA = pendingTriggers.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(pendingTriggers.size()) + .isGreaterThan(requestsAfterInitialPhaseA)); + int requestsAfterFirstRetry = pendingTriggers.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(pendingTriggers.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(pendingTriggers); + fromCtx(ctx -> null); + LeaderAndIsr current = fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get()); + CompletableFuture responseFuture = new CompletableFuture<>(); + eventProcessor + .getCoordinatorEventManager() + .put( + new AdjustIsrReceivedEvent( + Collections.singletonMap( + tableBucket, + new LeaderAndIsr( + current.leader(), + current.leaderEpoch(), + Arrays.asList(0, 1, 2, 3), + Collections.emptyList(), + current.coordinatorEpoch(), + current.bucketEpoch())), + responseFuture)); + responseFuture.get(); + fromCtx(ctx -> null); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()).isTrue(); + + ZOO_KEEPER_EXTENSION_WRAPPER + .getCustomExtension() + .cleanupPath(ZkData.ServerIdZNode.path(2)); + retryVerifyContext(ctx -> assertThat(ctx.liveTabletServerSet()).doesNotContain(2)); + drainPendingNotifyTriggers(pendingTriggers); + fromCtx(ctx -> null); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()).isTrue(); + assertThat(pendingTriggers).isEmpty(); + List assignmentWhileTargetOffline = + fromCtx(ctx -> ctx.getAssignment(tableBucket)); + assertThat(assignmentWhileTargetOffline).containsExactly(0, 1, 2); + + registerTabletServer(2); + retryVerifyContext(ctx -> assertThat(ctx.liveTabletServerSet()).contains(2)); + drainPendingNotifyTriggers(pendingTriggers); + fromCtx(ctx -> null); + current = fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get()); + responseFuture = new CompletableFuture<>(); + eventProcessor + .getCoordinatorEventManager() + .put( + new AdjustIsrReceivedEvent( + Collections.singletonMap( + tableBucket, + new LeaderAndIsr( + current.leader(), + current.leaderEpoch(), + Arrays.asList(0, 1, 2), + Collections.emptyList(), + current.coordinatorEpoch(), + current.bucketEpoch())), + responseFuture)); + responseFuture.get(); + drainPendingNotifyTriggers(pendingTriggers); + 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)); + } finally { + if (Arrays.stream(zookeeperClient.getSortedTabletServerList()) + .noneMatch(id -> id == 2)) { + registerTabletServer(2); + } + ZOO_KEEPER_EXTENSION_WRAPPER + .getCustomExtension() + .cleanupPath(ZkData.ServerIdZNode.path(3)); + } + } + + @Test + void testTimedOutRebalanceCompletesWhenTableIsBeingDeleted() throws Exception { + registerTabletServer(3); + initCoordinatorChannel(); + TablePath tablePath = TablePath.of(defaultDatabase, "test_rebalance_during_delete"); + Map assignments = new HashMap<>(); + assignments.put(0, BucketAssignment.of(0, 1, 2)); + long tableId = + metadataManager.createTable( + tablePath, + remoteDataDir, + TEST_TABLE, + new TableAssignment(assignments), + false); + TableBucket tableBucket = new TableBucket(tableId, 0); + ConcurrentLinkedDeque pendingTriggers = + new ConcurrentLinkedDeque<>(); + installBlockingNotifyGateways(pendingTriggers); + 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), + RebalanceStatus.NOT_STARTED); + RebalanceExecutionKey executionKey = + eventProcessor.getRebalanceManager().getExecutionKey(tableBucket); + + retry(Duration.ofMinutes(1), () -> assertThat(pendingTriggers).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); + ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(ZkData.ServerIdZNode.path(3)); + } + + @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), + RebalanceStatus.NOT_STARTED); + + retry( + Duration.ofMinutes(1), + () -> + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) + .isFalse()); + assertThat(rebalanceStatus(tableBucket)).isEqualTo(RebalanceStatus.COMPLETED); + } + + @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); + + // 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"); + Map bucketAssignments = new HashMap<>(); + bucketAssignments.put(0, BucketAssignment.of(0, 1, 2)); + bucketAssignments.put(1, BucketAssignment.of(0, 1, 2)); + bucketAssignments.put(2, 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); + TableBucket tb1 = new TableBucket(t1Id, 1); + TableBucket tb2 = new TableBucket(t1Id, 2); + + // Wait for initial leaders to be elected (all should be leader 0). + verifyIsr(tb0, 0, Arrays.asList(0, 1, 2)); + 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(); + + // Create leader-only rebalance plan (replicas stay the same, only leaders change): + // tb0: leader 0 -> 1 (newReplicas=[1,0,2] puts target leader first) + // tb1: leader 0 -> 2 (newReplicas=[2,0,1] puts target leader first) + // tb2: leader 0 -> 1 (newReplicas=[1,2,0] puts target leader first) + Map rebalancePlan = new HashMap<>(); + rebalancePlan.put( + tb0, + new RebalancePlanForBucket( + tb0, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 0, 2))); + rebalancePlan.put( + tb1, + new RebalancePlanForBucket( + tb1, 0, 2, Arrays.asList(0, 1, 2), Arrays.asList(2, 0, 1))); + rebalancePlan.put( + tb2, + new RebalancePlanForBucket( + tb2, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 2, 0))); + + // Register the rebalance. Only the FIRST task should trigger a leader election + // because subsequent tasks must wait for the NotifyLeaderAndIsr response. + eventProcessor + .getRebalanceManager() + .registerRebalance( + "rebalance-leader-sequential", rebalancePlan, RebalanceStatus.NOT_STARTED); + + // === Step 1: Verify only the first task started === + // registerRebalance() is synchronous, so after it returns, the first task's + // leader election has triggered NotifyLeaderAndIsr to replica servers. + // Other tasks must NOT have started because the first response is held. + assertThat(pendingTriggers).isNotEmpty(); + // All 3 tasks are still in progress (first executing, two waiting). + 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); + + // === Step 2: Wait for the second task to start === + // The event processor completes the first task via the response callback, + // then starts the second task which produces new pending triggers. + retry(Duration.ofMinutes(1), () -> assertThat(pendingTriggers).isNotEmpty()); + // First task completed, 2 tasks remaining. + assertThat(countInProgressRebalanceTasks(tb0, tb1, tb2)).isEqualTo(2); + drainPendingNotifyTriggers(pendingTriggers); + + // === Step 3: Wait for the third task to start === + retry(Duration.ofMinutes(1), () -> assertThat(pendingTriggers).isNotEmpty()); + // Two tasks completed, 1 task remaining. + assertThat(countInProgressRebalanceTasks(tb0, tb1, tb2)).isEqualTo(1); + drainPendingNotifyTriggers(pendingTriggers); + + // === Step 4: Wait for the rebalance to complete === + retry( + Duration.ofMinutes(1), + () -> + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) + .isFalse()); + + // Verify all leaders changed correctly. + verifyIsr(tb0, 1, Arrays.asList(0, 1, 2)); + verifyIsr(tb1, 2, Arrays.asList(0, 1, 2)); + verifyIsr(tb2, 1, Arrays.asList(0, 1, 2)); + } + + @Test + void testRebalanceRecoveryStateClassification() throws Exception { + // The classification only reads the coordinator state that recovery has just loaded. + TableBucket tableBucket = new TableBucket(987L, 0); + putBucketState(tableBucket, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2)); + + RebalancePlanForBucket plan = + new RebalancePlanForBucket( + tableBucket, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 0, 3)); + assertThat(eventProcessor.isRebalanceTaskComplete(plan)).isFalse(); + assertThat(eventProcessor.isRebalanceTaskAtOrigin(plan)).isTrue(); + + // Origin assignment with a leftover adding replica in ISR is not clean. + putBucketState(tableBucket, 0, Arrays.asList(0, 1, 2, 3), Arrays.asList(0, 1, 2)); + assertThat(eventProcessor.isRebalanceTaskComplete(plan)).isFalse(); + assertThat(eventProcessor.isRebalanceTaskAtOrigin(plan)).isFalse(); + + // The leader has not moved to the new leader of the plan yet, so recovery replays it. + putBucketState(tableBucket, 0, Arrays.asList(0, 1, 3), Arrays.asList(1, 0, 3)); + assertThat(eventProcessor.isRebalanceTaskComplete(plan)).isFalse(); + + // A plan whose target state is already in place must not be replayed, replaying it would + // elect the very same leader again and bump the epoch of an already migrated bucket. + RebalancePlanForBucket appliedPlan = + new RebalancePlanForBucket( + tableBucket, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 0, 2)); + putBucketState(tableBucket, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 0, 2)); + assertThat(eventProcessor.isRebalanceTaskComplete(appliedPlan)).isTrue(); + + // The plan targets replica 3, which is not hosted by a live server, so it stays replayable. + putBucketState(tableBucket, 1, Arrays.asList(0, 1, 3), Arrays.asList(1, 0, 3)); + assertThat(eventProcessor.isRebalanceTaskComplete(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 { + ConcurrentLinkedDeque pendingTriggers = + new ConcurrentLinkedDeque<>(); + int[] servers = zookeeperClient.getSortedTabletServerList(); + Map gateways = new HashMap<>(); + for (int server : servers) { + ControlledNotifyGateway gateway = new ControlledNotifyGateway(server, pendingTriggers); + gateways.put(server, gateway); + } + testCoordinatorChannelManager.setGateways(gateways); + + TablePath tablePath = TablePath.of(defaultDatabase, "test_stale_rebalance_response"); + Map bucketAssignments = new HashMap<>(); + bucketAssignments.put(0, BucketAssignment.of(0, 1, 2)); + long tableId = + metadataManager.createTable( + tablePath, + remoteDataDir, + TEST_TABLE, + new TableAssignment(bucketAssignments), + false); + TableBucket tableBucket = new TableBucket(tableId, 0); + verifyIsr(tableBucket, 0, Arrays.asList(0, 1, 2)); + + for (TabletServerGateway gateway : gateways.values()) { + ((ControlledNotifyGateway) gateway).enableControlMode(); + } + pendingTriggers.clear(); + 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), + RebalanceStatus.NOT_STARTED); + retry(Duration.ofMinutes(1), () -> assertThat(pendingTriggers).isNotEmpty()); + + LeaderAndIsr current = fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tableBucket).get()); + NotifyLeaderAndIsrResultForBucket success = + new NotifyLeaderAndIsrResultForBucket(tableBucket); + NotifyLeaderAndIsrRequestContext staleContext = + new NotifyLeaderAndIsrRequestContext( + eventProcessor.getCoordinatorEpoch(), + current.leader(), + current.leaderEpoch(), + current.bucketEpoch() - 1, + eventProcessor.getRebalanceManager().getExecutionKey(tableBucket)); + eventProcessor + .getCoordinatorEventManager() + .put( + new NotifyLeaderAndIsrResponseReceivedEvent( + Collections.singletonList(success), + current.leader(), + Collections.singletonMap(tableBucket, staleContext))); + fromCtx(ctx -> null); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()).isTrue(); + + NotifyLeaderAndIsrRequestContext oldAttemptContext = + new NotifyLeaderAndIsrRequestContext( + eventProcessor.getCoordinatorEpoch(), + current.leader(), + current.leaderEpoch(), + current.bucketEpoch(), + new RebalanceExecutionKey("old-rebalance", tableBucket, 1)); + eventProcessor + .getCoordinatorEventManager() + .put( + new NotifyLeaderAndIsrResponseReceivedEvent( + Collections.singletonList(success), + current.leader(), + Collections.singletonMap(tableBucket, oldAttemptContext))); + fromCtx(ctx -> null); + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()).isTrue(); + + NotifyLeaderAndIsrRequestContext currentContext = + new NotifyLeaderAndIsrRequestContext( + eventProcessor.getCoordinatorEpoch(), + current.leader(), + current.leaderEpoch(), + current.bucketEpoch(), + eventProcessor.getRebalanceManager().getExecutionKey(tableBucket)); + eventProcessor + .getCoordinatorEventManager() + .put( + new NotifyLeaderAndIsrResponseReceivedEvent( + Collections.singletonList(success), + current.leader(), + Collections.singletonMap(tableBucket, currentContext))); + retry( + Duration.ofMinutes(1), + () -> + assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) + .isFalse()); + drainPendingNotifyTriggers(pendingTriggers); + } + + @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 CoordinatorEventProcessor buildCoordinatorEventProcessor() { + Configuration conf = new Configuration(); + conf.set(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); + conf.set(ConfigOptions.COORDINATOR_OFFLINE_LEADER_RETRY_DELAY, Duration.ofDays(1)); + 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()); + } + + private void initCoordinatorChannel() throws Exception { + makeSendLeaderAndStopRequestAlwaysSuccess( + testCoordinatorChannelManager, + Arrays.stream(zookeeperClient.getSortedTabletServerList()) + .boxed() + .collect(Collectors.toSet()), + Collections.emptySet()); + } + + private 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())); + } + + private RebalanceStatus rebalanceStatus(TableBucket tableBucket) { + return eventProcessor + .getRebalanceManager() + .listRebalanceProgress(null) + .progressForBucketMap() + .get(tableBucket) + .status(); + } + + private void installBlockingNotifyGateways( + ConcurrentLinkedDeque pendingTriggers) throws Exception { + Map gateways = new HashMap<>(); + for (int server : zookeeperClient.getSortedTabletServerList()) { + ControlledNotifyGateway gateway = new ControlledNotifyGateway(server, pendingTriggers); + gateway.enableControlMode(); + gateways.put(server, gateway); + } + testCoordinatorChannelManager.setGateways(gateways); + } + + 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 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 void completePendingNotifyTrigger( + ConcurrentLinkedDeque pendingTriggers, int responseServerId) { + for (ControlledNotifyTrigger trigger : pendingTriggers) { + if (trigger.getResponseServerId() == responseServerId) { + assertThat(pendingTriggers.remove(trigger)).isTrue(); + trigger.complete(null); + return; + } + } + throw new AssertionError( + "No pending NotifyLeaderAndIsr response for server " + responseServerId); + } + + private int countInProgressRebalanceTasks(TableBucket... buckets) { + int count = 0; + for (TableBucket tableBucket : buckets) { + if (!RebalanceStatus.FINAL_STATUSES.contains( + eventProcessor + .getRebalanceManager() + .listRebalanceProgress(null) + .progressForBucketMap() + .get(tableBucket) + .status())) { + count++; + } + } + return count; + } +} 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 0bbb40b5c8d..a1988b4d8e7 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 @@ -19,8 +19,6 @@ import org.apache.fluss.cluster.Endpoint; import org.apache.fluss.cluster.TabletServerInfo; -import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket; -import org.apache.fluss.cluster.rebalance.RebalanceStatus; import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; import org.apache.fluss.exception.FencedLeaderEpochException; @@ -36,14 +34,11 @@ import org.apache.fluss.metadata.TableInfo; import org.apache.fluss.metadata.TablePartition; import org.apache.fluss.metadata.TablePath; -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.protocol.ApiError; @@ -54,6 +49,7 @@ 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.lease.KvSnapshotLeaseManager; @@ -114,10 +110,8 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import java.util.function.Function; import java.util.stream.Collectors; @@ -1359,10 +1353,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( @@ -1373,14 +1410,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 @@ -1894,283 +1925,6 @@ 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); - - // 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"); - Map bucketAssignments = new HashMap<>(); - bucketAssignments.put(0, BucketAssignment.of(0, 1, 2)); - bucketAssignments.put(1, BucketAssignment.of(0, 1, 2)); - bucketAssignments.put(2, 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); - TableBucket tb1 = new TableBucket(t1Id, 1); - TableBucket tb2 = new TableBucket(t1Id, 2); - - // Wait for initial leaders to be elected (all should be leader 0). - verifyIsr(tb0, 0, Arrays.asList(0, 1, 2)); - 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(); - - // Create leader-only rebalance plan (replicas stay the same, only leaders change): - // tb0: leader 0 -> 1 (newReplicas=[1,0,2] puts target leader first) - // tb1: leader 0 -> 2 (newReplicas=[2,0,1] puts target leader first) - // tb2: leader 0 -> 1 (newReplicas=[1,2,0] puts target leader first) - Map rebalancePlan = new HashMap<>(); - rebalancePlan.put( - tb0, - new RebalancePlanForBucket( - tb0, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 0, 2))); - rebalancePlan.put( - tb1, - new RebalancePlanForBucket( - tb1, 0, 2, Arrays.asList(0, 1, 2), Arrays.asList(2, 0, 1))); - rebalancePlan.put( - tb2, - new RebalancePlanForBucket( - tb2, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 2, 0))); - - // Register the rebalance. Only the FIRST task should trigger a leader election - // because subsequent tasks must wait for the NotifyLeaderAndIsr response. - eventProcessor - .getRebalanceManager() - .registerRebalance( - "rebalance-leader-sequential", rebalancePlan, RebalanceStatus.NOT_STARTED); - - // === Step 1: Verify only the first task started === - // registerRebalance() is synchronous, so after it returns, the first task's - // leader election has triggered NotifyLeaderAndIsr to replica servers. - // Other tasks must NOT have started because the first response is held. - assertThat(pendingTriggers).isNotEmpty(); - // All 3 tasks are still in progress (first executing, two waiting). - 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); - - // === Step 2: Wait for the second task to start === - // The event processor completes the first task via the response callback, - // then starts the second task which produces new pending triggers. - retry(Duration.ofMinutes(1), () -> assertThat(pendingTriggers).isNotEmpty()); - // First task completed, 2 tasks remaining. - assertThat(countInProgressRebalanceTasks(tb0, tb1, tb2)).isEqualTo(2); - drainPendingNotifyTriggers(pendingTriggers); - - // === Step 3: Wait for the third task to start === - retry(Duration.ofMinutes(1), () -> assertThat(pendingTriggers).isNotEmpty()); - // Two tasks completed, 1 task remaining. - assertThat(countInProgressRebalanceTasks(tb0, tb1, tb2)).isEqualTo(1); - drainPendingNotifyTriggers(pendingTriggers); - - // === Step 4: Wait for the rebalance to complete === - retry( - Duration.ofMinutes(1), - () -> - assertThat(eventProcessor.getRebalanceManager().hasInProgressRebalance()) - .isFalse()); - - // Verify all leaders changed correctly. - verifyIsr(tb0, 1, Arrays.asList(0, 1, 2)); - verifyIsr(tb1, 2, Arrays.asList(0, 1, 2)); - 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 CoordinatorEventProcessor buildCoordinatorEventProcessor() { Configuration conf = new Configuration(); conf.set(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); @@ -2648,125 +2402,6 @@ 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 void completePendingNotifyTrigger( - ConcurrentLinkedDeque pendingTriggers, int responseServerId) { - for (ControlledNotifyTrigger trigger : pendingTriggers) { - if (trigger.getResponseServerId() == responseServerId) { - assertThat(pendingTriggers.remove(trigger)).isTrue(); - trigger.complete(null); - return; - } - } - 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) { - 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; diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorHighAvailabilityITCase.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorHighAvailabilityITCase.java index 20f34e3ad24..c840c38d0d4 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorHighAvailabilityITCase.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorHighAvailabilityITCase.java @@ -246,7 +246,7 @@ void testStandbyTracksDynamicConfigAndNewLeaderUsesItAfterFailover() throws Exce assertThat(leader).isNotNull(); assertThat(standby).isNotNull(); - String configKey = ConfigOptions.LOG_REPLICA_MIN_IN_SYNC_REPLICAS_NUMBER.key(); + String configKey = ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS.key(); // Alter a dynamic config through the current leader. This persists it to ZK and inserts a // change notification that every server (including the standby) should react to. @@ -272,6 +272,10 @@ void testStandbyTracksDynamicConfigAndNewLeaderUsesItAfterFailover() throws Exce assertThat(dynamicConfigValue(standby, configKey)) .as("Newly promoted leader should use the dynamic config it tracked while standby") .isEqualTo("3"); + waitUntil( + () -> rebalanceMaxInflightTasks(standby) == 3, + Duration.ofSeconds(30), + "Newly promoted leader did not apply the dynamic rebalance concurrency"); // As the new leader it is now the sole writer: a further alter must take effect and stick. standby.getDynamicConfigManager() @@ -279,6 +283,20 @@ void testStandbyTracksDynamicConfigAndNewLeaderUsesItAfterFailover() throws Exce Collections.singletonList( new AlterConfig(configKey, "5", AlterConfigOpType.SET))); assertThat(dynamicConfigValue(standby, configKey)).isEqualTo("5"); + waitUntil( + () -> rebalanceMaxInflightTasks(standby) == 5, + Duration.ofSeconds(30), + "Leader did not apply the updated dynamic rebalance concurrency"); + } + + private static int rebalanceMaxInflightTasks(CoordinatorServer server) { + try { + return server.getCoordinatorEventProcessor() + .getRebalanceManager() + .getMaxInflightRebalanceTasks(); + } catch (IllegalStateException ignored) { + return -1; + } } /** Reads the effective value of a config key from a server's DynamicConfigManager. */ diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/RebalanceResponseCheckTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/RebalanceResponseCheckTest.java new file mode 100644 index 00000000000..3dad79012f6 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/RebalanceResponseCheckTest.java @@ -0,0 +1,63 @@ +/* + * 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.metadata.TableBucket; +import org.apache.fluss.rpc.protocol.ApiError; +import org.apache.fluss.rpc.protocol.Errors; +import org.apache.fluss.server.entity.NotifyLeaderAndIsrResultForBucket; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Test for the response checks that let {@link CoordinatorEventProcessor} complete a rebalance. */ +class RebalanceResponseCheckTest { + + @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(); + } +} 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 57731571a1e..d11714785e8 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,27 +17,22 @@ 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.config.Configuration; +import org.apache.fluss.exception.ConfigException; import org.apache.fluss.metadata.TableBucket; -import org.apache.fluss.server.coordinator.AutoPartitionManager; 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.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.event.ReconcileRebalanceTaskEvent; +import org.apache.fluss.server.coordinator.event.RecoverRebalanceEvent; +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; @@ -45,10 +40,6 @@ import org.apache.fluss.server.zk.data.RebalanceTask; import org.apache.fluss.testutils.common.AllCallbackWrapper; 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; @@ -60,15 +51,21 @@ import java.util.ArrayList; import java.util.Arrays; 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.Set; import java.util.concurrent.ScheduledThreadPoolExecutor; +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.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Test for {@link RebalanceManager}. */ public class RebalanceManagerTest { @@ -78,17 +75,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 TestingRebalanceExecutor rebalanceExecutor; + private RecordingEventManager eventManager; + private ManualClock clock; private RebalanceManager rebalanceManager; - private KvSnapshotLeaseManager kvSnapshotLeaseManager; - private Scheduler scheduler; @BeforeAll static void baseBeforeAll() throws Exception { @@ -100,60 +92,26 @@ static void baseBeforeAll() throws Exception { } @BeforeEach - void beforeEach() { - serverMetadataCache = new CoordinatorMetadataCache(); - testCoordinatorChannelManager = new TestCoordinatorChannelManager(); - String remoteDataDir = "/tmp/fluss/remote-data"; - Configuration 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(conf); - RecordingEventManager recordingEventManager = new RecordingEventManager(); + void beforeEach() throws Exception { + zookeeperClient.deleteRebalanceTask(); + rebalanceExecutor = new TestingRebalanceExecutor(new CoordinatorContext(zkEpoch)); + eventManager = new RecordingEventManager(); + clock = new ManualClock(); rebalanceManager = new RebalanceManager( - eventProcessor, + rebalanceExecutor, zookeeperClient, - recordingEventManager, - SystemClock.getInstance()); + eventManager, + clock, + new Configuration(), + new NoOpScheduledExecutor()); rebalanceManager.startup(); } @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 @@ -182,35 +140,26 @@ void testTimeoutEnqueuesEvent() throws Exception { ManualClock clock = new ManualClock(0L); RecordingEventManager eventManager = new RecordingEventManager(); NoOpScheduledExecutor executor = new NoOpScheduledExecutor(); - CoordinatorEventProcessor eventProcessor = - buildCoordinatorEventProcessor(new Configuration()); - RebalanceManager manager = new RebalanceManager( - eventProcessor, zookeeperClient, eventManager, clock, executor); + new TestingRebalanceExecutor(new CoordinatorContext(zkEpoch)), + zookeeperClient, + eventManager, + clock, + new Configuration(), + executor); manager.startup(); 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)); + Map plan = plans(tb1, tb2); manager.registerRebalance("timeout-test", plan, NOT_STARTED); + RebalanceExecutionKey executionKey = manager.getExecutionKey(tb1); - // Not yet timed out. clock.advanceTime(Duration.ofMillis(100_000)); manager.checkTimeout(); assertThat(eventManager.events).isEmpty(); - // Cross the 2-minute boundary. clock.advanceTime(Duration.ofMillis(30_000)); manager.checkTimeout(); @@ -218,10 +167,8 @@ void testTimeoutEnqueuesEvent() throws Exception { assertThat(eventManager.events.get(0)).isInstanceOf(RebalanceTaskTimeoutEvent.class); RebalanceTaskTimeoutEvent timeoutEvent = (RebalanceTaskTimeoutEvent) eventManager.events.get(0); - assertThat(timeoutEvent.getTableBucket()).isEqualTo(tb1); + assertThat(timeoutEvent.getExecutionKey()).isEqualTo(executionKey); - // 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); @@ -230,106 +177,575 @@ void testTimeoutEnqueuesEvent() throws Exception { } @Test - void testTimeoutAfterCompletionIsNoOp() throws Exception { + void testSoftTimeoutAdmitsNextTaskAndTracksLateCompletion() throws Exception { ManualClock clock = new ManualClock(0L); RecordingEventManager eventManager = new RecordingEventManager(); - NoOpScheduledExecutor executor = new NoOpScheduledExecutor(); - CoordinatorEventProcessor eventProcessor = - buildCoordinatorEventProcessor(new Configuration()); - + TestingRebalanceExecutor executor = + new TestingRebalanceExecutor(new CoordinatorContext(zkEpoch)); RebalanceManager manager = new RebalanceManager( - eventProcessor, zookeeperClient, eventManager, clock, executor); + executor, + zookeeperClient, + eventManager, + clock, + new Configuration(), + new NoOpScheduledExecutor()); manager.startup(); 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))); + TableBucket tb2 = new TableBucket(1L, 1); + manager.registerRebalance("soft-timeout-test", plans(tb1, tb2), NOT_STARTED); + RebalanceExecutionKey firstAttempt = manager.getExecutionKey(tb1); + assertThat(executor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(tb1); + + clock.advanceTime(Duration.ofMillis(130_000)); + manager.checkTimeout(); + RebalanceTaskTimeoutEvent timeoutEvent = + (RebalanceTaskTimeoutEvent) eventManager.events.get(0); + assertThat(manager.timeoutRebalanceTask(timeoutEvent.getExecutionKey())).isTrue(); + + RebalanceExecutionKey secondAttempt = manager.getExecutionKey(tb2); + assertThat(secondAttempt).isNotNull(); + assertThat(executor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(tb1, tb2); + assertThat(manager.listRebalanceProgress(null).status()).isEqualTo(REBALANCING); + assertThat(manager.listRebalanceProgress(null).progressForBucketMap().get(tb1).status()) + .isEqualTo(TIMEOUT); + assertThat(eventManager.events.get(1)).isInstanceOf(ReconcileRebalanceTaskEvent.class); + + RebalancePlanForBucket retryPlan = manager.getPlanForReconciliation(firstAttempt); + assertThat(retryPlan).isNotNull(); + assertThat(retryPlan.getTableBucket()).isEqualTo(tb1); + // the dispatched reconciliation backs off, so no event is enqueued right away. + manager.checkTimeout(); + assertThat(eventManager.events).hasSize(2); + + clock.advanceTime(Duration.ofMillis(30_000)); + manager.checkTimeout(); + assertThat(eventManager.events).hasSize(3); + assertThat(eventManager.events.get(2)).isInstanceOf(ReconcileRebalanceTaskEvent.class); + assertThat(((ReconcileRebalanceTaskEvent) eventManager.events.get(2)).getExecutionKey()) + .isEqualTo(firstAttempt); + + assertThat(manager.timeoutRebalanceTask(firstAttempt)).isFalse(); + assertThat(manager.finishRebalanceTask(firstAttempt, COMPLETED)).isTrue(); + assertThat(manager.finishRebalanceTask(firstAttempt, COMPLETED)).isFalse(); + assertThat(manager.finishRebalanceTask(secondAttempt, COMPLETED)).isTrue(); + + assertThat(manager.getRebalanceStatus()).isEqualTo(COMPLETED); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(COMPLETED); + + manager.close(); + } + + @Test + void testRegisterRebalanceRespectsMaxInflightTasks() { + rebalanceManager.updateMaxInflightRebalanceTasks(2); + TableBucket tb1 = new TableBucket(1L, 0); + TableBucket tb2 = new TableBucket(1L, 1); + TableBucket tb3 = new TableBucket(1L, 2); + TableBucket tb4 = new TableBucket(1L, 3); + + rebalanceManager.registerRebalance( + "concurrent-test", plans(tb1, tb2, tb3, tb4), NOT_STARTED); + + assertThat(rebalanceExecutor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(tb1, tb2); + assertThat(countStatus(rebalanceManager, REBALANCING)).isEqualTo(2); + assertThat(countStatus(rebalanceManager, NOT_STARTED)).isEqualTo(2); + + rebalanceManager.finishRebalanceTask(tb1, COMPLETED); + + assertThat(rebalanceExecutor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(tb1, tb2, tb3); + assertThat(countStatus(rebalanceManager, REBALANCING)).isEqualTo(2); + } + + @Test + void testIncreaseMaxInflightTasksStartsPendingTasks() { + TableBucket tb1 = new TableBucket(1L, 0); + TableBucket tb2 = new TableBucket(1L, 1); + TableBucket tb3 = new TableBucket(1L, 2); + TableBucket tb4 = new TableBucket(1L, 3); + rebalanceManager.registerRebalance("scale-up-test", plans(tb1, tb2, tb3, tb4), NOT_STARTED); + + assertThat(rebalanceManager.getMaxInflightRebalanceTasks()).isEqualTo(1); + assertThat(rebalanceExecutor.executedPlans).hasSize(1); - zookeeperClient.registerRebalanceTask( - new RebalanceTask("completion-test", NOT_STARTED, plan)); - manager.registerRebalance("completion-test", plan, NOT_STARTED); + rebalanceManager.updateMaxInflightRebalanceTasks(3); - // The task completes normally before timeout. - manager.finishRebalanceTask(tb1, COMPLETED); + assertThat(rebalanceManager.getMaxInflightRebalanceTasks()).isEqualTo(3); + assertThat(rebalanceExecutor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(tb1, tb2, tb3); + } + + @Test + void testZeroMaxInflightTasksPausesAndResumesScheduling() { + rebalanceManager.updateMaxInflightRebalanceTasks(0); + TableBucket tb1 = new TableBucket(1L, 0); + TableBucket tb2 = new TableBucket(1L, 1); + TableBucket tb3 = new TableBucket(1L, 2); + rebalanceManager.registerRebalance("pause-test", plans(tb1, tb2, tb3), NOT_STARTED); + + assertThat(rebalanceManager.getMaxInflightRebalanceTasks()).isZero(); + assertThat(rebalanceExecutor.executedPlans).isEmpty(); + assertThat(countStatus(rebalanceManager, NOT_STARTED)).isEqualTo(3); + assertThat(rebalanceManager.hasInProgressRebalance()).isTrue(); + + rebalanceManager.updateMaxInflightRebalanceTasks(2); + + assertThat(rebalanceExecutor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(tb1, tb2); + assertThat(countStatus(rebalanceManager, REBALANCING)).isEqualTo(2); + } + + @Test + void testDecreaseMaxInflightTasksDoesNotCancelRunningTasks() { + rebalanceManager.updateMaxInflightRebalanceTasks(3); + TableBucket tb1 = new TableBucket(1L, 0); + TableBucket tb2 = new TableBucket(1L, 1); + TableBucket tb3 = new TableBucket(1L, 2); + TableBucket tb4 = new TableBucket(1L, 3); + TableBucket tb5 = new TableBucket(1L, 4); + rebalanceManager.registerRebalance( + "scale-down-test", plans(tb1, tb2, tb3, tb4, tb5), NOT_STARTED); + + rebalanceManager.updateMaxInflightRebalanceTasks(1); + + assertThat(rebalanceExecutor.executedPlans).hasSize(3); + rebalanceManager.finishRebalanceTask(tb1, COMPLETED); + rebalanceManager.finishRebalanceTask(tb2, COMPLETED); + assertThat(rebalanceExecutor.executedPlans).hasSize(3); + + rebalanceManager.finishRebalanceTask(tb3, COMPLETED); + assertThat(rebalanceExecutor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(tb1, tb2, tb3, tb4); + } + + @Test + void testTimeoutEnqueuesEventsForAllInflightTasks() { + rebalanceManager.updateMaxInflightRebalanceTasks(2); + TableBucket tb1 = new TableBucket(1L, 0); + TableBucket tb2 = new TableBucket(1L, 1); + TableBucket tb3 = new TableBucket(1L, 2); + rebalanceManager.registerRebalance("timeout-all-test", plans(tb1, tb2, tb3), NOT_STARTED); - // Now the timeout fires, but the task is already done. clock.advanceTime(Duration.ofMillis(130_000)); + rebalanceManager.checkTimeout(); + + assertThat(eventManager.events) + .filteredOn(RebalanceTaskTimeoutEvent.class::isInstance) + .extracting( + event -> + ((RebalanceTaskTimeoutEvent) event) + .getExecutionKey() + .getTableBucket()) + .containsExactlyInAnyOrder(tb1, tb2); + RebalanceTaskTimeoutEvent firstTimeout = + (RebalanceTaskTimeoutEvent) eventManager.events.get(0); + RebalanceTaskTimeoutEvent secondTimeout = + (RebalanceTaskTimeoutEvent) eventManager.events.get(1); + + assertThat(rebalanceManager.timeoutRebalanceTask(firstTimeout.getExecutionKey())).isTrue(); + assertThat(rebalanceManager.timeoutRebalanceTask(secondTimeout.getExecutionKey())).isTrue(); + assertThat(rebalanceExecutor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(tb1, tb2, tb3); + } + + @Test + void testTimeoutTrackingCapIsPreservedWithConcurrentTasks() { + rebalanceManager.updateMaxInflightRebalanceTasks(10); + TableBucket[] tableBuckets = new TableBucket[10]; + for (int i = 0; i < tableBuckets.length; i++) { + tableBuckets[i] = new TableBucket(1L, i); + } + rebalanceManager.registerRebalance( + "concurrent-timeout-cap-test", plans(tableBuckets), NOT_STARTED); + List attempts = new ArrayList<>(); + for (TableBucket tableBucket : tableBuckets) { + attempts.add(rebalanceManager.getExecutionKey(tableBucket)); + } + + for (int i = 0; i < 8; i++) { + assertThat(rebalanceManager.timeoutRebalanceTask(attempts.get(i))).isTrue(); + } + assertThat(rebalanceManager.timeoutRebalanceTask(attempts.get(8))).isFalse(); + assertThat(countStatus(rebalanceManager, TIMEOUT)).isEqualTo(8); + assertThat(countStatus(rebalanceManager, REBALANCING)).isEqualTo(2); + + assertThat(rebalanceManager.finishRebalanceTask(attempts.get(0), COMPLETED)).isTrue(); + assertThat(rebalanceManager.timeoutRebalanceTask(attempts.get(8))).isTrue(); + assertThat(countStatus(rebalanceManager, TIMEOUT)).isEqualTo(8); + assertThat(countStatus(rebalanceManager, REBALANCING)).isEqualTo(1); + } + + @Test + void testRejectNegativeMaxInflightTasks() { + Configuration invalidConfig = new Configuration(); + invalidConfig.set(ConfigOptions.COORDINATOR_REBALANCE_MAX_INFLIGHT_TASKS, -1); + + assertThatThrownBy(() -> rebalanceManager.validate(invalidConfig)) + .isInstanceOf(ConfigException.class) + .hasMessageContaining("must be non-negative"); + } + + @Test + void testFailureIsAggregatedIntoOverallStatus() throws Exception { + TableBucket tb1 = new TableBucket(1L, 0); + TableBucket tb2 = new TableBucket(1L, 1); + rebalanceManager.registerRebalance("failed-test", plans(tb1, tb2), NOT_STARTED); + rebalanceManager.finishRebalanceTask(tb1, FAILED); + rebalanceManager.finishRebalanceTask(tb2, COMPLETED); + + 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); + rebalanceManager.registerRebalance("cancel-test", plans(tb1, tb2), NOT_STARTED); + 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 testRecoverReconcilesCompletedAndIntermediateBuckets() { + TableBucket completedBucket = new TableBucket(1L, 0); + TableBucket intermediateBucket = new TableBucket(1L, 1); + Map plans = plans(completedBucket, intermediateBucket); + rebalanceExecutor.completedBuckets.add(completedBucket); + + rebalanceManager.recoverRebalance(new RebalanceTask("recover-test", REBALANCING, plans)); + + Map statuses = statuses(rebalanceManager); + assertThat(statuses.get(completedBucket)).isEqualTo(COMPLETED); + assertThat(statuses.get(intermediateBucket)).isEqualTo(REBALANCING); + assertThat(rebalanceExecutor.executedPlans) + .extracting(RebalancePlanForBucket::getTableBucket) + .containsExactly(intermediateBucket); + } + + @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(); + } + + @Test + void testReconciliationBacksOffBetweenRetries() { + ManualClock clock = new ManualClock(0L); + RecordingEventManager eventManager = new RecordingEventManager(); + RebalanceManager manager = newManager(clock, eventManager, rebalanceExecutor); + + TableBucket tableBucket = new TableBucket(1L, 0); + manager.registerRebalance("backoff-test", plans(tableBucket), NOT_STARTED); + RebalanceExecutionKey attempt = manager.getExecutionKey(tableBucket); + assertThat(manager.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(manager.getPlanForReconciliation(attempt)).isNotNull(); + clock.advanceTime(Duration.ofMillis(30_000)); manager.checkTimeout(); + assertThat(reconciliationsFor(eventManager, attempt)).isEqualTo(2); - // No timeout event should be enqueued because inflightTaskStartMs was cleared. - assertThat(eventManager.events).isEmpty(); + assertThat(manager.getPlanForReconciliation(attempt)).isNotNull(); + clock.advanceTime(Duration.ofMillis(30_000)); + manager.checkTimeout(); + assertThat(reconciliationsFor(eventManager, attempt)).isEqualTo(2); + + clock.advanceTime(Duration.ofMillis(30_000)); + manager.checkTimeout(); + assertThat(reconciliationsFor(eventManager, attempt)).isEqualTo(3); manager.close(); } @Test - void testTimeoutTreatsTaskAsCompleted() throws Exception { + void testTrackedTimedOutTasksAreCapped() { ManualClock clock = new ManualClock(0L); - RecordingEventManager eventManager = new RecordingEventManager(); - NoOpScheduledExecutor executor = new NoOpScheduledExecutor(); - CoordinatorEventProcessor eventProcessor = - buildCoordinatorEventProcessor(new Configuration()); + TestingRebalanceExecutor executor = + new TestingRebalanceExecutor(new CoordinatorContext(zkEpoch)); + RebalanceManager manager = newManager(clock, new RecordingEventManager(), executor); + + TableBucket[] tableBuckets = new TableBucket[10]; + for (int i = 0; i < tableBuckets.length; i++) { + tableBuckets[i] = new TableBucket(1L, i); + } + manager.registerRebalance("cap-test", plans(tableBuckets), NOT_STARTED); + + // 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 < 8; i++) { + TableBucket running = + executor.executedPlans.get(executor.executedPlans.size() - 1).getTableBucket(); + RebalanceExecutionKey attempt = manager.getExecutionKey(running); + assertThat(manager.timeoutRebalanceTask(attempt)).isTrue(); + timedOut.add(attempt); + } + assertThat(executor.executedPlans).hasSize(8); + + // once a tracked task reaches a final status the next pending task is admitted again. + assertThat(manager.finishRebalanceTask(timedOut.get(0), COMPLETED)).isTrue(); + assertThat(executor.executedPlans).hasSize(9); + + manager.close(); + } + + @Test + void testTimedOutTaskFailsWhenTargetReplicasStayUnavailable() throws Exception { + ManualClock clock = new ManualClock(0L); + // no tablet server is live, so the target replicas can never catch up. + RebalanceManager manager = + newManager(clock, new RecordingEventManager(), rebalanceExecutor); + + TableBucket tableBucket = new TableBucket(1L, 0); + manager.registerRebalance("give-up-test", plans(tableBucket), NOT_STARTED); + RebalanceExecutionKey attempt = manager.getExecutionKey(tableBucket); + assertThat(manager.timeoutRebalanceTask(attempt)).isTrue(); + assertThat(manager.getPlanForReconciliation(attempt)).isNotNull(); + + clock.advanceTime(Duration.ofMinutes(31)); + assertThat(manager.getPlanForReconciliation(attempt)).isNull(); + + // the rebalance reaches a final status, so later rebalance requests are not blocked. + assertThat(manager.getRebalanceStatus()).isEqualTo(FAILED); + assertThat(manager.hasInProgressRebalance()).isFalse(); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()).isEqualTo(FAILED); + manager.close(); + } + + @Test + void testTimedOutTaskKeepsRetryingWhileTargetReplicasAreLive() { + ManualClock clock = new ManualClock(0L); + CoordinatorContext coordinatorContext = new CoordinatorContext(zkEpoch); + // the plans target replicas 1, 2 and 3, so the migration can still make progress. + for (int serverId : new int[] {1, 2, 3}) { + coordinatorContext.addLiveTabletServer(tabletServer(serverId)); + } RebalanceManager manager = - new RebalanceManager( - eventProcessor, zookeeperClient, eventManager, clock, executor); - manager.startup(); + newManager( + clock, + new RecordingEventManager(), + new TestingRebalanceExecutor(coordinatorContext)); + + TableBucket tableBucket = new TableBucket(1L, 0); + manager.registerRebalance("keep-retrying-test", plans(tableBucket), NOT_STARTED); + RebalanceExecutionKey attempt = manager.getExecutionKey(tableBucket); + assertThat(manager.timeoutRebalanceTask(attempt)).isTrue(); + + clock.advanceTime(Duration.ofMinutes(31)); + assertThat(manager.getPlanForReconciliation(attempt)).isNotNull(); + assertThat(manager.getRebalanceStatus()).isEqualTo(REBALANCING); + + manager.close(); + } + @Test + void testCancelGivesUpImmediatelyOnAdmittedTaskStillAtOrigin() 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("completed-test", NOT_STARTED, plan)); - manager.registerRebalance("completed-test", plan, NOT_STARTED); - - // Timeout fires. - clock.advanceTime(Duration.ofMillis(130_000)); - manager.checkTimeout(); + rebalanceExecutor.originBuckets.add(tb1); + rebalanceManager.registerRebalance("cancel-at-origin-test", plans(tb1, tb2), NOT_STARTED); - // Simulate the coordinator event thread processing the timeout event. - assertThat(eventManager.events).hasSize(1); - RebalanceTaskTimeoutEvent timeoutEvent = - (RebalanceTaskTimeoutEvent) eventManager.events.get(0); - manager.finishRebalanceTask(timeoutEvent.getTableBucket(), TIMEOUT); + rebalanceManager.cancelRebalance("cancel-at-origin-test"); - // 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); + assertThat(rebalanceManager.getRebalanceStatus()).isEqualTo(CANCELED); + assertThat(rebalanceManager.hasInProgressRebalance()).isFalse(); + assertThat(zookeeperClient.getRebalanceTask().get().getRebalanceStatus()) + .isEqualTo(CANCELED); + } - manager.close(); + private RebalanceManager newManager( + ManualClock clock, + RecordingEventManager eventManager, + TestingRebalanceExecutor executor) { + RebalanceManager manager = + new RebalanceManager( + executor, + zookeeperClient, + eventManager, + clock, + new Configuration(), + new NoOpScheduledExecutor()); + manager.startup(); + return manager; + } + + private static int countStatus(RebalanceManager manager, RebalanceStatus status) { + int count = 0; + for (RebalanceResultForBucket result : + manager.listRebalanceProgress(null).progressForBucketMap().values()) { + if (result.status() == status) { + count++; + } + } + return count; + } + + 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); } - private CoordinatorEventProcessor buildCoordinatorEventProcessor(Configuration conf) { - 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()); + @Test + void testStartupFencesNewRebalanceUntilRecoveryEventRuns() throws Exception { + TableBucket tableBucket = new TableBucket(1L, 0); + RebalanceTask storedTask = + new RebalanceTask("startup-recovery-test", REBALANCING, plans(tableBucket)); + zookeeperClient.registerRebalanceTask(storedTask); + TestingRebalanceExecutor executor = + new TestingRebalanceExecutor(new CoordinatorContext(zkEpoch)); + RecordingEventManager recordingEventManager = new RecordingEventManager(); + RebalanceManager recoveringManager = + new RebalanceManager( + executor, + zookeeperClient, + recordingEventManager, + new ManualClock(), + new Configuration(), + new NoOpScheduledExecutor()); + + recoveringManager.startup(); + + assertThat(recoveringManager.hasInProgressRebalance()).isTrue(); + assertThat(recoveringManager.getRebalanceId()).isNull(); + assertThat(recordingEventManager.events).hasSize(1); + RecoverRebalanceEvent recoveryEvent = + (RecoverRebalanceEvent) recordingEventManager.events.get(0); + assertThat(recoveryEvent.getRebalanceTask()).isEqualTo(storedTask); + + recoveringManager.recoverRebalance(recoveryEvent.getRebalanceTask()); + assertThat(recoveringManager.getRebalanceId()).isEqualTo("startup-recovery-test"); + assertThat(executor.executedPlans).hasSize(1); + recoveringManager.close(); + } + + private static Map plans(TableBucket... tableBuckets) { + Map plans = new LinkedHashMap<>(); + for (TableBucket tableBucket : tableBuckets) { + plans.put( + tableBucket, + new RebalancePlanForBucket( + tableBucket, 0, 1, Arrays.asList(0, 1, 2), Arrays.asList(1, 2, 3))); + } + return plans; + } + + private static Map statuses(RebalanceManager manager) { + Map statuses = new HashMap<>(); + for (Map.Entry entry : + manager.listRebalanceProgress(null).progressForBucketMap().entrySet()) { + statuses.put(entry.getKey(), entry.getValue().status()); + } + return statuses; + } + + private static final class TestingRebalanceExecutor implements RebalanceExecutor { + private final CoordinatorContext coordinatorContext; + private final List executedPlans = new ArrayList<>(); + private final Set completedBuckets = new HashSet<>(); + private final Set originBuckets = new HashSet<>(); + + private TestingRebalanceExecutor(CoordinatorContext coordinatorContext) { + this.coordinatorContext = coordinatorContext; + } + + @Override + public CoordinatorContext getCoordinatorContext() { + return coordinatorContext; + } + + @Override + public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) { + executedPlans.add(planForBucket); + } + + @Override + public boolean isRebalanceTaskComplete(RebalancePlanForBucket planForBucket) { + return completedBuckets.contains(planForBucket.getTableBucket()); + } + + @Override + public boolean isRebalanceTaskAtOrigin(RebalancePlanForBucket planForBucket) { + return originBuckets.contains(planForBucket.getTableBucket()); + } } /** Records events put into the coordinator event queue. */ 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 cbc0b85c6a2..899a32dbc4f 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 @@ -832,6 +832,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/configuration.md b/website/docs/maintenance/configuration.md index b9e44b84069..ae34f82b36b 100644 --- a/website/docs/maintenance/configuration.md +++ b/website/docs/maintenance/configuration.md @@ -21,8 +21,8 @@ auto-partition.check.interval: 5min ``` Server configuration refers to a set of configurations used to specify the running parameters of a server. -These settings can only be configured at the time of cluster startup and do not support dynamic modification -during the Fluss cluster working. +Most settings are parsed when the Fluss processes start and require restarting the relevant processes to take effect. +Some server configurations can be updated dynamically while the cluster is running. See [Updating Cluster Configs](operations/updating-configs.md#updating-cluster-configs) for the supported dynamic options. ## Common @@ -81,6 +81,7 @@ The logging-related environment options (`env.log.dir`, `env.log.level`, `env.lo | coordinator.lifecycle-throttler.inflight-timeout | Duration | 3min | The timeout for an in-flight drop event in the coordinator's TableLifecycleThrottler. If a drop event has been admitted but the corresponding completion callback has not arrived within this timeout, the throttler abandons tracking of that drop and continues admitting the next pending drop. | | coordinator.lifecycle-throttler.timeout-check-interval | Duration | 1min | The periodic interval at which the coordinator's TableLifecycleThrottler scans in-flight drops for timeouts. | | coordinator.offline-leader.retry-delay | Duration | 1min | The delay before the coordinator retries offline leaders on live tablet servers after they are marked offline. This lets a leader that was rejected because of temporary tablet-server conditions, such as disk write protection, become electable again after recovery. | +| coordinator.rebalance.max-inflight-tasks | Integer | 1 | The maximum number of bucket-level rebalance tasks that can be executed concurrently by the coordinator. A higher value can speed up rebalance, while a lower value reduces the number of simultaneous bucket movements. Setting it to 0 pauses scheduling new rebalance tasks; already in-flight tasks continue until they complete or time out. The value must be non-negative. | ## TabletServer diff --git a/website/docs/maintenance/operations/rebalance.md b/website/docs/maintenance/operations/rebalance.md index 418b7c585f2..db0e7228d3e 100644 --- a/website/docs/maintenance/operations/rebalance.md +++ b/website/docs/maintenance/operations/rebalance.md @@ -107,12 +107,20 @@ 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 that a new rebalance can move again - **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 -### 4. Cancel Rebalance (If Needed) +### 4. Control Rebalance Speed + +The coordinator controls rebalance parallelism with `coordinator.rebalance.max-inflight-tasks`, which limits how many bucket-level rebalance tasks can run at the same time. The default value is `1`, which keeps rebalance conservative. + +Increase the value to speed up rebalance when the cluster has enough network, disk, and TabletServer capacity. Decrease the value to reduce the number of simultaneous bucket movements. Setting it to `0` pauses scheduling new bucket-level rebalance tasks; tasks that are already in flight continue until they complete or time out. Set it back to a positive value to resume scheduling pending tasks. + +This option can be updated dynamically through cluster configuration updates. See [Updating Cluster Configs](updating-configs.md#updating-cluster-configs) for the Java and Flink SQL APIs. + +### 5. Cancel Rebalance (If Needed) Cancel an ongoing rebalance operation if necessary: @@ -127,9 +135,10 @@ 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) +### 6. Remove Server Tags (After Completion) After rebalance completes and maintenance is done, remove server tags to restore normal operation: @@ -272,11 +281,12 @@ If `RACK_AWARE` is not placed first, replica movements generated by earlier goal 1. **Plan Ahead**: Tag servers appropriately before triggering rebalance to guide the algorithm 2. **Monitor Progress**: Regularly check rebalance status to ensure smooth operation 3. **Off-Peak Hours**: Schedule rebalance operations during off-peak hours to minimize impact -4. **Single Rebalance**: Fluss supports only one active rebalance task at a time in the cluster -5. **Backup First**: For production environments, ensure data is backed up before major topology changes -6. **Goal Priority**: Order rebalance goals by priority - the system attempts to achieve them in order -7. **Server Tags**: Use `TEMPORARY_OFFLINE` for maintenance scenarios to allow buckets to return after maintenance -8. **Rack Awareness First**: In multi-rack deployments, always place `RACK_AWARE` as the first goal +4. **Single Rebalance**: Fluss supports only one active rebalance operation at a time in the cluster +5. **Rebalance Parallelism**: Tune `coordinator.rebalance.max-inflight-tasks` based on cluster capacity. Use `0` to pause scheduling new bucket-level rebalance tasks. +6. **Backup First**: For production environments, ensure data is backed up before major topology changes +7. **Goal Priority**: Order rebalance goals by priority - the system attempts to achieve them in order +8. **Server Tags**: Use `TEMPORARY_OFFLINE` for maintenance scenarios to allow buckets to return after maintenance +9. **Rack Awareness First**: In multi-rack deployments, always place `RACK_AWARE` as the first goal ## Troubleshooting @@ -293,6 +303,7 @@ If rebalance is taking longer than expected: - Check network bandwidth between TabletServers - Verify disk I/O performance on TabletServers - Monitor cluster load and resource utilization +- Consider increasing `coordinator.rebalance.max-inflight-tasks` if the cluster has enough capacity - Consider canceling and retrying with fewer goals ### Rebalance Stuck in REBALANCING Status diff --git a/website/docs/maintenance/operations/updating-configs.md b/website/docs/maintenance/operations/updating-configs.md index f6d3b12e38c..bb9a0f5b915 100644 --- a/website/docs/maintenance/operations/updating-configs.md +++ b/website/docs/maintenance/operations/updating-configs.md @@ -21,6 +21,7 @@ Currently, the supported dynamically updatable server configurations include: - `datalake.format`: Specify the lakehouse format, e.g., `paimon`, `iceberg`. When enabling lakehouse storage explicitly, use it together with `datalake.enabled = true`. - Options with prefix `datalake.${datalake.format}` - `kv.rocksdb.shared-rate-limiter.bytes-per-sec`: Control RocksDB flush and compaction write rate shared across all RocksDB instances on the TabletServer. The rate limiter is always enabled. Set to a lower value (e.g., 100MB) to limit the rate, or a very high value to effectively disable rate limiting. +- `coordinator.rebalance.max-inflight-tasks`: Control the maximum number of bucket-level rebalance tasks that can run concurrently. Set it to `0` to pause scheduling new rebalance tasks; already in-flight tasks continue until they complete or time out. You can update the configuration of a cluster with [Java client](#using-java-client) or [Flink SQL](#using-flink-sql).