B9. Wait for the new leader to acknowledge the final state, then mark the ongoing
+ * rebalance task to finish.
*
*
* In general, there are two goals we want to aim for:
@@ -1918,9 +1933,38 @@ private void onBucketReassignment(
null,
null,
Collections.singleton(tableBucket));
- // B8. Mark the ongoing rebalance task to finish.
- rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.COMPLETED);
+ // B9. Finish the task only after the planned leader acknowledges this final state.
+ awaitPlannedLeaderAck(tableBucket);
+ }
+ }
+
+ private void awaitPlannedLeaderAck(TableBucket tableBucket) {
+ RebalancePlanForBucket planForBucket =
+ rebalanceManager.getRebalancePlanForBucket(tableBucket);
+ Optional leaderAndIsr = coordinatorContext.getBucketLeaderAndIsr(tableBucket);
+ if (planForBucket == null) {
+ return;
+ }
+ // Without a leader the task can neither be acknowledged nor safely re-run, so fail it.
+ if (!leaderAndIsr.isPresent()) {
+ LOG.warn(
+ "No leader for tableBucket {} after applying its rebalance task.", tableBucket);
+ rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED);
+ return;
+ }
+ // The election picks the first live in-sync target replica, which is not the planned
+ // leader when that one is offline. Its ACK can then never arrive, so fail the task now
+ // instead of blocking the next migrations until the task timeout.
+ if (leaderAndIsr.get().leader() != planForBucket.getNewLeader()) {
+ LOG.warn(
+ "Planned leader {} for tableBucket {} was not elected, current leader is {}.",
+ planForBucket.getNewLeader(),
+ tableBucket,
+ leaderAndIsr.get().leader());
+ rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED);
+ return;
}
+ rebalanceManager.awaitLeaderAck(tableBucket, leaderAndIsr.get().bucketEpoch());
}
private boolean isReassignmentComplete(
@@ -1935,7 +1979,9 @@ private void maybeReassignedBucketLeaderIfRequired(
TableBucket tableBucket, List targetReplicas) throws Exception {
LeaderAndIsr leaderAndIsr = coordinatorContext.getBucketLeaderAndIsr(tableBucket).get();
int currentLeader = leaderAndIsr.leader();
- if (!targetReplicas.contains(currentLeader)) {
+ // The first target replica is the planned leader (see GoalOptimizerUtils#getDiff), so
+ // re-elect also when the current leader stays in the target replicas but is not first.
+ if (currentLeader != targetReplicas.get(0)) {
LOG.info(
"Leader {} for tableBucket {} being reassigned. Re-electing leader to {}",
currentLeader,
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatch.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatch.java
index b4bf7e90d41..d35108204ce 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
@@ -460,8 +460,12 @@ private void sendNotifyLeaderAndIsrRequest(int coordinatorEpoch) {
// 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
Set addedToPendingLeaderActivation = new HashSet<>();
+ // The response carries no epoch, so remember what this request sent to tell a stale
+ // response apart from the one for the latest leader state.
+ Map requestBucketEpochs = new HashMap<>();
for (Map.Entry entry :
notifyRequestEntry.getValue().entrySet()) {
+ requestBucketEpochs.put(entry.getKey(), entry.getValue().getBucketEpoch());
int leader = entry.getValue().getLeader();
if (leader == serverId) {
coordinatorContext.addPendingLeaderActivation(entry.getKey());
@@ -504,7 +508,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,
+ requestBucketEpochs));
});
}
notifyLeaderAndIsrRequestMap.clear();
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..3f07482e1cf 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,12 @@
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.List;
+import java.util.Map;
/** An event for receive the response of {@link NotifyLeaderAndIsrRequest} from tablet server. */
public class NotifyLeaderAndIsrResponseReceivedEvent implements CoordinatorEvent {
@@ -30,11 +32,16 @@ public class NotifyLeaderAndIsrResponseReceivedEvent implements CoordinatorEvent
// the server id that return the response
private final int responseServerId;
+ // the bucket epoch sent in the request for each bucket, the response itself carries no epoch
+ private final Map requestBucketEpochs;
+
public NotifyLeaderAndIsrResponseReceivedEvent(
List notifyLeaderAndIsrResultForBuckets,
- int responseServerId) {
+ int responseServerId,
+ Map requestBucketEpochs) {
this.notifyLeaderAndIsrResultForBuckets = notifyLeaderAndIsrResultForBuckets;
this.responseServerId = responseServerId;
+ this.requestBucketEpochs = requestBucketEpochs;
}
public int getResponseServerId() {
@@ -44,4 +51,9 @@ public int getResponseServerId() {
public List getNotifyLeaderAndIsrResultForBuckets() {
return notifyLeaderAndIsrResultForBuckets;
}
+
+ /** Returns the bucket epoch the request sent for the bucket, or -1 if unknown. */
+ public int getRequestBucketEpoch(TableBucket tableBucket) {
+ return requestBucketEpochs.getOrDefault(tableBucket, -1);
+ }
}
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java
index ce95dd862a0..1ef5e2312d2 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
@@ -102,6 +102,13 @@ public class RebalanceManager {
private final Map finishedRebalanceTasks =
new ConcurrentHashMap<>();
+ /**
+ * A mapping from table bucket to the bucket epoch its planned leader must acknowledge before
+ * the in-flight task completes.
+ */
+ private final Map awaitedLeaderAckBucketEpochs =
+ new ConcurrentHashMap<>();
+
private final GoalOptimizer goalOptimizer;
private volatile long registerTime;
private volatile @Nullable RebalanceStatus rebalanceStatus;
@@ -199,6 +206,7 @@ public void registerRebalance(
inProgressRebalanceTasks.clear();
inProgressRebalanceTasksQueue.clear();
finishedRebalanceTasks.clear();
+ awaitedLeaderAckBucketEpochs.clear();
// Clear gate (bucket) first, then data (startMs).
inflightTaskBucket = null;
inflightTaskStartMs = -1;
@@ -236,6 +244,7 @@ public void finishRebalanceTask(TableBucket tableBucket, RebalanceStatus statusF
if (inProgressRebalanceTasksQueue.contains(tableBucket)) {
inProgressRebalanceTasksQueue.remove(tableBucket);
RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.remove(tableBucket);
+ awaitedLeaderAckBucketEpochs.remove(tableBucket);
checkNotNull(resultForBucket, "RebalanceResultForBucket is null.");
finishedRebalanceTasks.put(
tableBucket,
@@ -323,6 +332,7 @@ public void cancelRebalance(@Nullable String rebalanceId) {
rebalanceStatus = CANCELED;
inProgressRebalanceTasksQueue.clear();
inProgressRebalanceTasks.clear();
+ awaitedLeaderAckBucketEpochs.clear();
// Clear gate (bucket) first, then data (startMs).
inflightTaskBucket = null;
inflightTaskStartMs = -1;
@@ -375,6 +385,27 @@ public RebalanceTask generateRebalanceTask(List goalsByPriority) {
return null;
}
+ /**
+ * Marks that the coordinator applied the final state of the in-flight task for the bucket. The
+ * task completes once the planned leader acknowledges a request with at least the given bucket
+ * epoch.
+ */
+ public void awaitLeaderAck(TableBucket tableBucket, int bucketEpoch) {
+ checkNotClosed();
+ if (inProgressRebalanceTasks.containsKey(tableBucket)) {
+ awaitedLeaderAckBucketEpochs.put(tableBucket, bucketEpoch);
+ }
+ }
+
+ /**
+ * Returns the bucket epoch the planned leader must acknowledge, or null if the task for the
+ * bucket is not waiting for the leader acknowledgement.
+ */
+ public @Nullable Integer getAwaitedLeaderAckBucketEpoch(TableBucket tableBucket) {
+ checkNotClosed();
+ return awaitedLeaderAckBucketEpochs.get(tableBucket);
+ }
+
private void processNewRebalanceTask() {
TableBucket tableBucket = inProgressRebalanceTasksQueue.peek();
if (tableBucket != null && inProgressRebalanceTasks.containsKey(tableBucket)) {
@@ -410,6 +441,7 @@ private void completeRebalance() {
rebalanceStatus = COMPLETED;
inProgressRebalanceTasks.clear();
inProgressRebalanceTasksQueue.clear();
+ awaitedLeaderAckBucketEpochs.clear();
// Here, it will not clear finishedRebalanceTasks, because it will be used by
// listRebalanceProgress. It will be cleared when next register.
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java
index b26d8cabca2..fc6c6ff8016 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
@@ -60,6 +60,7 @@
import org.apache.fluss.server.coordinator.event.CommitRemoteLogManifestEvent;
import org.apache.fluss.server.coordinator.event.CoordinatorEventManager;
import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrResponseReceivedEvent;
+import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent;
import org.apache.fluss.server.coordinator.event.RetryOfflineLeaderEvent;
import org.apache.fluss.server.coordinator.remote.RemoteDirDynamicLoader;
import org.apache.fluss.server.coordinator.statemachine.BucketState;
@@ -83,6 +84,7 @@
import org.apache.fluss.server.zk.data.BucketAssignment;
import org.apache.fluss.server.zk.data.LeaderAndIsr;
import org.apache.fluss.server.zk.data.PartitionAssignment;
+import org.apache.fluss.server.zk.data.RebalanceTask;
import org.apache.fluss.server.zk.data.TableAssignment;
import org.apache.fluss.server.zk.data.TabletServerRegistration;
import org.apache.fluss.server.zk.data.ZkData;
@@ -95,6 +97,8 @@
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
import java.nio.file.Path;
import java.time.Duration;
@@ -1403,7 +1407,8 @@ void testDiskWriteLockedNotifyLeaderResponseMarksReplicaOffline() throws Excepti
new ApiError(
Errors.DISK_WRITE_LOCKED,
"disk write locked"))),
- leader));
+ leader,
+ Collections.emptyMap()));
fromCtx(
ctx -> {
@@ -1449,7 +1454,8 @@ void testRetryOfflineLeaderEventRetriesOfflineReplicaOnLiveServer() throws Excep
new ApiError(
Errors.DISK_WRITE_LOCKED,
"disk write locked"))),
- leader));
+ leader,
+ Collections.emptyMap()));
retryVerifyContext(
ctx -> {
@@ -1507,7 +1513,8 @@ void testRetryOfflineLeaderEventKeepsReplicaOfflineWhenRetryStillFails() throws
new ApiError(
Errors.DISK_WRITE_LOCKED,
"disk write locked"))),
- leader));
+ leader,
+ Collections.emptyMap()));
retryVerifyContext(
ctx -> {
@@ -1568,7 +1575,8 @@ void testRetryOfflineLeaderEventSkipsDeletedBucket() throws Exception {
new ApiError(
Errors.DISK_WRITE_LOCKED,
"disk write locked"))),
- leader));
+ leader,
+ Collections.emptyMap()));
retryVerifyContext(
ctx -> {
@@ -2088,6 +2096,352 @@ void testDoBucketReassignment() throws Exception {
ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(ZkData.ServerIdZNode.path(3));
}
+ @Test
+ void testBucketReassignmentMovesLeaderWhenOldLeaderStaysInReplicas() throws Exception {
+ registerTabletServer3();
+ try {
+ ConcurrentLinkedDeque pendingTriggers =
+ new ConcurrentLinkedDeque<>();
+ TableBucket tb0 =
+ reassignWithLeaderChangeUntilFinalLeaderAck(
+ "test_reassignment_with_leader_change", pendingTriggers);
+ ControlledNotifyTrigger finalLeaderAck = holdFinalLeaderAck(pendingTriggers);
+ assertRebalancing(tb0);
+
+ // An ISR change while the task waits for the ACK must not run phase B again, which
+ // would send new NotifyLeaderAndIsr requests.
+ LeaderAndIsr currentLeaderAndIsr = fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tb0)).get();
+ CompletableFuture respCallback = new CompletableFuture<>();
+ eventProcessor
+ .getCoordinatorEventManager()
+ .put(
+ new AdjustIsrReceivedEvent(
+ Collections.singletonMap(tb0, currentLeaderAndIsr),
+ respCallback));
+ respCallback.get();
+ assertThat(fromCtx(ctx -> ctx.getBucketLeaderAndIsr(tb0)).get().bucketEpoch())
+ .isEqualTo(currentLeaderAndIsr.bucketEpoch() + 1);
+ assertThat(pendingTriggers).isEmpty();
+ assertRebalancing(tb0);
+
+ finalLeaderAck.complete(null);
+ retry(
+ Duration.ofMinutes(1),
+ () -> assertRebalanceStatus(tb0, RebalanceStatus.COMPLETED));
+ verifyIsr(tb0, 1, Arrays.asList(1, 0, 2));
+ assertThat(
+ zookeeperClient
+ .getTableAssignment(tb0.getTableId())
+ .get()
+ .getBucketAssignment(0)
+ .getReplicas())
+ .containsExactly(1, 0, 2);
+ } finally {
+ cleanupTabletServer3();
+ }
+ }
+
+ @Test
+ void testBucketReassignmentWaitsForTimeoutWhenPlannedLeaderAckFails() throws Exception {
+ registerTabletServer3();
+ try {
+ ConcurrentLinkedDeque pendingTriggers =
+ new ConcurrentLinkedDeque<>();
+ TableBucket tb0 =
+ reassignWithLeaderChangeUntilFinalLeaderAck(
+ "test_reassignment_leader_ack_fails", pendingTriggers);
+ ControlledNotifyTrigger finalLeaderAck = holdFinalLeaderAck(pendingTriggers);
+
+ // A failed ACK must not record a false success.
+ finalLeaderAck.fail();
+ fromCtx(ctx -> null);
+ assertRebalancing(tb0);
+
+ // The task timeout releases the bucket, so it does not block the rebalance forever.
+ eventProcessor.getCoordinatorEventManager().put(new RebalanceTaskTimeoutEvent(tb0));
+ retry(Duration.ofMinutes(1), () -> assertRebalanceStatus(tb0, RebalanceStatus.TIMEOUT));
+ assertThat(eventProcessor.getRebalanceManager().getAwaitedLeaderAckBucketEpoch(tb0))
+ .isNull();
+ } finally {
+ cleanupTabletServer3();
+ }
+ }
+
+ @Test
+ void testBucketReassignmentFailsWhenPlannedLeaderIsNotElected() throws Exception {
+ registerTabletServer3();
+ try {
+ TableBucket tb0 = createBucketWithReplicas("test_reassignment_planned_leader_offline");
+ eventProcessor
+ .getRebalanceManager()
+ .registerRebalance(
+ "rebalance-planned-leader-offline",
+ reassignmentWithLeaderChangePlan(tb0),
+ RebalanceStatus.NOT_STARTED);
+ // The planned leader 1 goes offline, so phase B elects 0 instead.
+ fromCtx(
+ ctx -> {
+ ctx.addOfflineBucketInServer(tb0, 1);
+ return null;
+ });
+
+ // The adding replica 2 catches up with the leader, which runs phase B.
+ CompletableFuture respCallback = new CompletableFuture<>();
+ eventProcessor
+ .getCoordinatorEventManager()
+ .put(
+ new AdjustIsrReceivedEvent(
+ Collections.singletonMap(
+ tb0,
+ new LeaderAndIsr(
+ 0,
+ 0,
+ Arrays.asList(0, 1, 2, 3),
+ Collections.emptyList(),
+ 0,
+ 1)),
+ respCallback));
+ respCallback.get();
+
+ retry(Duration.ofMinutes(1), () -> assertRebalanceStatus(tb0, RebalanceStatus.FAILED));
+ assertThat(eventProcessor.getRebalanceManager().getAwaitedLeaderAckBucketEpoch(tb0))
+ .isNull();
+ } finally {
+ cleanupTabletServer3();
+ }
+ }
+
+ /**
+ * The ZooKeeper state a coordinator leaves behind when it crashes during phase B of reassigning
+ * [0, 1, 3] with leader 0 to [1, 0, 2] with leader 1.
+ */
+ private enum PhaseBBoundary {
+ // Phase A finished and all target replicas caught up.
+ BEFORE_LEADER_ELECTION(0, Arrays.asList(1, 0, 2, 3), Arrays.asList(0, 1, 2, 3)),
+ // B2 elected the planned leader.
+ AFTER_LEADER_ELECTION(1, Arrays.asList(1, 0, 2, 3), Arrays.asList(0, 1, 2, 3)),
+ // B5 removed the old replica from the ISR.
+ AFTER_ISR_SHRINK(1, Arrays.asList(1, 0, 2, 3), Arrays.asList(1, 0, 2)),
+ // B7 stored the target replicas, only the planned leader's ACK is missing.
+ AFTER_ASSIGNMENT_UPDATE(1, Arrays.asList(1, 0, 2), Arrays.asList(1, 0, 2));
+
+ private final int leader;
+ private final List replicas;
+ private final List isr;
+
+ PhaseBBoundary(int leader, List replicas, List isr) {
+ this.leader = leader;
+ this.replicas = replicas;
+ this.isr = isr;
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(PhaseBBoundary.class)
+ void testBucketReassignmentWithLeaderChangeResumesAfterCoordinatorRestart(
+ PhaseBBoundary boundary) throws Exception {
+ registerTabletServer3();
+ try {
+ ConcurrentLinkedDeque pendingTriggers =
+ new ConcurrentLinkedDeque<>();
+ List controlledGateways =
+ setControlledNotifyGateways(pendingTriggers);
+ TableBucket tb0 =
+ createBucketWithReplicas(
+ "test_reassignment_restart_" + boundary.name().toLowerCase());
+
+ eventProcessor.shutdown();
+ LeaderAndIsr leaderAndIsr = zookeeperClient.getLeaderAndIsr(tb0).get();
+ zookeeperClient.updateTableAssignment(
+ tb0.getTableId(),
+ new TableAssignment(
+ Collections.singletonMap(0, new BucketAssignment(boundary.replicas))),
+ ZkVersion.MATCH_ANY_VERSION.getVersion());
+ zookeeperClient.updateLeaderAndIsr(
+ tb0,
+ new LeaderAndIsr(
+ boundary.leader,
+ leaderAndIsr.leaderEpoch() + 1,
+ boundary.isr,
+ Collections.emptyList(),
+ leaderAndIsr.coordinatorEpoch(),
+ leaderAndIsr.bucketEpoch() + 1),
+ ZkVersion.MATCH_ANY_VERSION.getVersion());
+ zookeeperClient.registerRebalanceTask(
+ new RebalanceTask(
+ "rebalance-reassign-restart",
+ RebalanceStatus.REBALANCING,
+ reassignmentWithLeaderChangePlan(tb0)));
+
+ controlledGateways.forEach(ControlledNotifyGateway::enableControlMode);
+ eventProcessor = buildCoordinatorEventProcessor();
+ eventProcessor.startup();
+
+ // Responses from the other servers drive the recovered task to its final state, but
+ // only the planned leader's ACK completes it.
+ retry(
+ Duration.ofMinutes(1),
+ () -> {
+ completePendingNotifyTriggersExcept(pendingTriggers, 1);
+ assertThat(
+ eventProcessor
+ .getRebalanceManager()
+ .getAwaitedLeaderAckBucketEpoch(tb0))
+ .isNotNull();
+ });
+ assertRebalancing(tb0);
+ verifyIsr(tb0, 1, Arrays.asList(1, 0, 2));
+
+ drainPendingNotifyTriggers(pendingTriggers);
+ retry(
+ Duration.ofMinutes(1),
+ () -> assertRebalanceStatus(tb0, RebalanceStatus.COMPLETED));
+ verifyIsr(tb0, 1, Arrays.asList(1, 0, 2));
+ assertThat(
+ zookeeperClient
+ .getTableAssignment(tb0.getTableId())
+ .get()
+ .getBucketAssignment(0)
+ .getReplicas())
+ .containsExactly(1, 0, 2);
+ } finally {
+ zookeeperClient.deleteRebalanceTask();
+ cleanupTabletServer3();
+ }
+ }
+
+ private void registerTabletServer3() throws Exception {
+ zookeeperClient.registerTabletServer(
+ 3,
+ new TabletServerRegistration(
+ "rack3",
+ Collections.singletonList(
+ new Endpoint("host3", 1001, DEFAULT_LISTENER_NAME)),
+ System.currentTimeMillis()));
+ }
+
+ private static void cleanupTabletServer3() {
+ ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(ZkData.ServerIdZNode.path(3));
+ }
+
+ /** Sets gateways that hold NotifyLeaderAndIsr responses once in control mode. */
+ private List setControlledNotifyGateways(
+ ConcurrentLinkedDeque pendingTriggers) throws Exception {
+ Map gateways = new HashMap<>();
+ List controlledGateways = new ArrayList<>();
+ for (int server : zookeeperClient.getSortedTabletServerList()) {
+ ControlledNotifyGateway gateway = new ControlledNotifyGateway(server, pendingTriggers);
+ gateways.put(server, gateway);
+ controlledGateways.add(gateway);
+ }
+ testCoordinatorChannelManager.setGateways(gateways);
+ return controlledGateways;
+ }
+
+ /** Creates a table with one bucket on replicas [0, 1, 3] and waits for leader 0. */
+ private TableBucket createBucketWithReplicas(String tableName) throws Exception {
+ long tableId =
+ metadataManager.createTable(
+ TablePath.of(defaultDatabase, tableName),
+ remoteDataDir,
+ TEST_TABLE,
+ new TableAssignment(
+ Collections.singletonMap(0, BucketAssignment.of(0, 1, 3))),
+ false);
+ TableBucket tableBucket = new TableBucket(tableId, 0);
+ verifyIsr(tableBucket, 0, Arrays.asList(0, 1, 3));
+ return tableBucket;
+ }
+
+ /**
+ * Replica 3 moves to 2 and the leader moves from 0 to 1, while the old leader 0 stays in the
+ * target replicas. The plan puts the target leader first.
+ */
+ private static Map reassignmentWithLeaderChangePlan(
+ TableBucket tableBucket) {
+ return Collections.singletonMap(
+ tableBucket,
+ new RebalancePlanForBucket(
+ tableBucket, 0, 1, Arrays.asList(0, 1, 3), Arrays.asList(1, 0, 2)));
+ }
+
+ /**
+ * Runs the plan of {@link #reassignmentWithLeaderChangePlan} until the coordinator applied the
+ * final state, while holding every NotifyLeaderAndIsr response.
+ */
+ private TableBucket reassignWithLeaderChangeUntilFinalLeaderAck(
+ String tableName, ConcurrentLinkedDeque pendingTriggers)
+ throws Exception {
+ List controlledGateways =
+ setControlledNotifyGateways(pendingTriggers);
+ TableBucket tb0 = createBucketWithReplicas(tableName);
+
+ controlledGateways.forEach(ControlledNotifyGateway::enableControlMode);
+ pendingTriggers.clear();
+ eventProcessor
+ .getRebalanceManager()
+ .registerRebalance(
+ "rebalance-reassign-with-leader-change",
+ reassignmentWithLeaderChangePlan(tb0),
+ RebalanceStatus.NOT_STARTED);
+
+ // The adding replica 2 catches up with the leader, which runs phase B.
+ CompletableFuture respCallback = new CompletableFuture<>();
+ eventProcessor
+ .getCoordinatorEventManager()
+ .put(
+ new AdjustIsrReceivedEvent(
+ Collections.singletonMap(
+ tb0,
+ new LeaderAndIsr(
+ 0,
+ 0,
+ Arrays.asList(0, 1, 2, 3),
+ Collections.emptyList(),
+ 0,
+ 1)),
+ respCallback));
+ respCallback.get();
+ verifyIsr(tb0, 1, Arrays.asList(0, 1, 2));
+ return tb0;
+ }
+
+ /**
+ * Holds the planned leader's ACK of the final state (its latest request) and releases all other
+ * responses, including server 1's responses to the earlier requests.
+ */
+ private ControlledNotifyTrigger holdFinalLeaderAck(
+ ConcurrentLinkedDeque pendingTriggers) throws Exception {
+ ControlledNotifyTrigger finalLeaderAck = null;
+ for (ControlledNotifyTrigger trigger : pendingTriggers) {
+ if (trigger.getResponseServerId() == 1) {
+ finalLeaderAck = trigger;
+ }
+ }
+ assertThat(finalLeaderAck).isNotNull();
+ assertThat(pendingTriggers.remove(finalLeaderAck)).isTrue();
+ drainPendingNotifyTriggers(pendingTriggers);
+ fromCtx(ctx -> null);
+ return finalLeaderAck;
+ }
+
+ private void assertRebalancing(TableBucket tableBucket) {
+ assertThat(countInProgressRebalanceTasks(tableBucket)).isEqualTo(1);
+ assertThat(eventProcessor.getRebalanceManager().listRebalanceProgress(null).status())
+ .isEqualTo(RebalanceStatus.REBALANCING);
+ }
+
+ private void assertRebalanceStatus(TableBucket tableBucket, RebalanceStatus expectedStatus) {
+ assertThat(
+ eventProcessor
+ .getRebalanceManager()
+ .listRebalanceProgress(null)
+ .progressForBucketMap()
+ .get(tableBucket)
+ .status())
+ .isEqualTo(expectedStatus);
+ }
+
@Test
void testLeaderOnlyRebalanceExecutesSequentially() throws Exception {
// Set up controlled gateways that capture NotifyLeaderAndIsr calls.
@@ -2208,19 +2562,16 @@ void testLeaderOnlyRebalanceCompletionCheckRequiresSuccessfulResponseFromNewLead
tableBucket, new ApiError(Errors.UNKNOWN_SERVER_ERROR, "failed"));
assertThat(
- CoordinatorEventProcessor
- .isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader(
- successResult, 1, planForBucket))
+ CoordinatorEventProcessor.isSuccessfulRebalanceResponseFromNewLeader(
+ successResult, 1, planForBucket))
.isTrue();
assertThat(
- CoordinatorEventProcessor
- .isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader(
- successResult, 0, planForBucket))
+ CoordinatorEventProcessor.isSuccessfulRebalanceResponseFromNewLeader(
+ successResult, 0, planForBucket))
.isFalse();
assertThat(
- CoordinatorEventProcessor
- .isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader(
- failedResult, 1, planForBucket))
+ CoordinatorEventProcessor.isSuccessfulRebalanceResponseFromNewLeader(
+ failedResult, 1, planForBucket))
.isFalse();
}
@@ -2799,6 +3150,15 @@ private static void drainPendingNotifyTriggers(
}
}
+ private static void completePendingNotifyTriggersExcept(
+ ConcurrentLinkedDeque pendingTriggers, int heldServerId) {
+ for (ControlledNotifyTrigger trigger : pendingTriggers) {
+ if (trigger.getResponseServerId() != heldServerId && pendingTriggers.remove(trigger)) {
+ trigger.complete(null);
+ }
+ }
+ }
+
private static boolean hasPendingNotifyTrigger(
ConcurrentLinkedDeque pendingTriggers, int responseServerId) {
for (ControlledNotifyTrigger trigger : pendingTriggers) {
@@ -2885,13 +3245,29 @@ public CompletableFuture notifyLeaderAndIsr(
// 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);
+ return trigger.getFuture()
+ .thenApply(
+ v -> {
+ if (trigger.isFailed()) {
+ response.getNotifyBucketsLeaderRespsList()
+ .forEach(
+ resp ->
+ resp.setErrorCode(
+ Errors
+ .UNKNOWN_SERVER_ERROR
+ .code())
+ .setErrorMessage(
+ "mock failure"));
+ }
+ return response;
+ });
}
}
private static class ControlledNotifyTrigger {
private final int responseServerId;
private final CompletableFuture future = new CompletableFuture<>();
+ private volatile boolean failed;
ControlledNotifyTrigger(int responseServerId) {
this.responseServerId = responseServerId;
@@ -2908,6 +3284,16 @@ CompletableFuture getFuture() {
void complete(Void value) {
future.complete(value);
}
+
+ /** Releases the response with an error for every bucket. */
+ void fail() {
+ failed = true;
+ future.complete(null);
+ }
+
+ boolean isFailed() {
+ return failed;
+ }
}
private static class PartitionIdName {