Skip to content

[server] Batch LeaderAndIsr updates for controlled shutdown #3654

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
swuferhong merged 5 commits into from
Jul 19, 2026
Merged

[server] Batch LeaderAndIsr updates for controlled shutdown #3654

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a42ad56
[server] Batch LeaderAndIsr updates for controlled shutdown
swuferhong Jul 14, 2026
6c5b079
address copilot's comments
swuferhong Jul 15, 2026
475c731
address gujian's comments
swuferhong Jul 15, 2026
8d33183
address gujian's comments
swuferhong Jul 15, 2026
3bb20b8
address anton's comments2
swuferhong Jul 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 24 additions & 2 deletions ...s-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2323,6 +2323,7 @@ private void handleCommitLakeTableSnapshotV2(

private ControlledShutdownResponse tryProcessControlledShutdown(
ControlledShutdownEvent controlledShutdownEvent) {
long startTimeMs = System.currentTimeMillis();
ControlledShutdownResponse response = new ControlledShutdownResponse();

// TODO here we need to check tabletServerEpoch, avoid to receive controlled shutdown
Expand All @@ -2346,8 +2347,10 @@ private ControlledShutdownResponse tryProcessControlledShutdown(
coordinatorContext.shuttingDownTabletServers());
LOG.debug("All live tabletServers: {}", coordinatorContext.liveTabletServerSet());

Set<TableBucketReplica> replicasOnTabletServer =
coordinatorContext.replicasOnTabletServer(tabletServerId);
List<TableBucketReplica> replicasToActOn =
coordinatorContext.replicasOnTabletServer(tabletServerId).stream()
replicasOnTabletServer.stream()
.filter(
replica -> {
TableBucket tableBucket = replica.getTableBucket();
Expand All @@ -2371,20 +2374,39 @@ private ControlledShutdownResponse tryProcessControlledShutdown(
}
}

long leaderMigrationStartTimeMs = System.currentTimeMillis();
tableBucketStateMachine.handleStateChange(
bucketsLedByServer, OnlineBucket, new ControlledShutdownLeaderElection());
LOG.info(
"Processed controlled shutdown leader state changes for {} buckets on tabletServer {} in {} ms.",
bucketsLedByServer.size(),
tabletServerId,
System.currentTimeMillis() - leaderMigrationStartTimeMs);

// TODO need send stop request to the leader?

// If the tabletServer is a follower, updates the isr in ZK and notifies the current leader.
long followerOfflineStartTimeMs = System.currentTimeMillis();
replicaStateMachine.handleStateChanges(replicasFollowedByServer, OfflineReplica);
LOG.info(
"Processed {} follower replica offline transitions for controlled shutdown of tabletServer {} in {} ms.",
replicasFollowedByServer.size(),
tabletServerId,
System.currentTimeMillis() - followerOfflineStartTimeMs);

// Return the list of buckets that are still being managed by the controlled shutdown
// tabletServer after leader migration.
Set<TableBucket> remainingLeaderBuckets =
coordinatorContext.getBucketsWithLeaderIn(tabletServerId);
response.addAllRemainingLeaderBuckets(
coordinatorContext.getBucketsWithLeaderIn(tabletServerId).stream()
remainingLeaderBuckets.stream()
.map(ServerRpcMessageUtils::fromTableBucket)
.collect(Collectors.toList()));
LOG.info(
"Completed controlled shutdown for tabletServer {} in {} ms; {} leader buckets remain.",
tabletServerId,
System.currentTimeMillis() - startTimeMs,
remainingLeaderBuckets.size());
return response;
}

Expand Down
217 changes: 199 additions & 18 deletions ...c/main/java/org/apache/fluss/server/coordinator/statemachine/TableBucketStateMachine.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,15 +31,19 @@
import org.apache.fluss.server.zk.ZooKeeperClient;
import org.apache.fluss.server.zk.data.LeaderAndIsr;
import org.apache.fluss.shaded.guava32.com.google.common.collect.Sets;
import org.apache.fluss.shaded.zookeeper3.org.apache.zookeeper.KeeperException;
import org.apache.fluss.utils.ExceptionUtils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
Expand Down Expand Up @@ -128,9 +132,15 @@ public void handleStateChange(
BucketState targetState,
ReplicaLeaderElection replicaLeaderElection) {
try {
boolean isControlledShutdown =
replicaLeaderElection instanceof ControlledShutdownLeaderElection
&& targetState == BucketState.OnlineBucket;
coordinatorRequestBatch.newBatch();

if (checkIfCreateTablePartitionRequest(tableBuckets, targetState)) {
if (isControlledShutdown) {
batchHandleControlledShutdown(
tableBuckets, (ControlledShutdownLeaderElection) replicaLeaderElection);
} else if (checkIfCreateTablePartitionRequest(tableBuckets, targetState)) {
// batch register table bucket lead and isr
batchHandleOnlineChangeAndInitLeader(tableBuckets);
} else {
Expand All @@ -145,6 +155,165 @@ public void handleStateChange(
}
}

private void batchHandleControlledShutdown(
Set<TableBucket> tableBuckets,
ControlledShutdownLeaderElection controlledShutdownLeaderElection) {
Map<TableBucket, LeaderAndIsr> currentLeaderAndIsrs;
long batchReadStartTimeMs = System.currentTimeMillis();
try {
currentLeaderAndIsrs = zooKeeperClient.getLeaderAndIsrs(tableBuckets);
} catch (Exception e) {
LOG.warn(
"Failed to batch read LeaderAndIsr for {} buckets during controlled shutdown. Falling back to individual updates.",
tableBuckets.size(),
e);
for (TableBucket tableBucket : tableBuckets) {
doHandleStateChange(
tableBucket, BucketState.OnlineBucket, controlledShutdownLeaderElection);
}
return;
}
LOG.info(
"Batch read LeaderAndIsr for {} of {} controlled shutdown buckets in {} ms.",
currentLeaderAndIsrs.size(),
tableBuckets.size(),
System.currentTimeMillis() - batchReadStartTimeMs);

long electionStartTimeMs = System.currentTimeMillis();
Map<TableBucket, ElectionResult> electionResults = new HashMap<>();
Map<TableBucket, String> partitionNames = new HashMap<>();
Set<TableBucket> bucketsToRetryIndividually = new HashSet<>();
for (TableBucket tableBucket : tableBuckets) {
coordinatorContext.putBucketStateIfNotExists(
tableBucket, BucketState.NonExistentBucket);
if (!checkValidTableBucketStateChange(tableBucket, BucketState.OnlineBucket)) {
continue;
}

BucketState currentState = coordinatorContext.getBucketState(tableBucket);
if (currentState == BucketState.NewBucket) {
bucketsToRetryIndividually.add(tableBucket);
Comment thread
swuferhong marked this conversation as resolved.
continue;
}

String partitionName = null;
if (tableBucket.getPartitionId() != null) {
partitionName = coordinatorContext.getPartitionName(tableBucket.getPartitionId());
if (partitionName == null) {
logFailedStateChange(
tableBucket,
currentState,
BucketState.OnlineBucket,
String.format(
"Can't find partition name for partition: %s.",
tableBucket.getPartitionId()));
continue;
}
}

LeaderAndIsr leaderAndIsr = currentLeaderAndIsrs.get(tableBucket);
if (leaderAndIsr == null) {
bucketsToRetryIndividually.add(tableBucket);
continue;
}

Optional<ElectionResult> optionalElectionResult =
electNewLeaderForTableBucket(
tableBucket, leaderAndIsr, controlledShutdownLeaderElection);
if (!optionalElectionResult.isPresent()) {
logFailedStateChange(
tableBucket,
currentState,
BucketState.OnlineBucket,
"Elect result is empty.");
continue;
}
electionResults.put(tableBucket, optionalElectionResult.get());
partitionNames.put(tableBucket, partitionName);
}
LOG.info(
"Prepared {} controlled shutdown leader elections in {} ms; {} buckets require individual retry.",
electionResults.size(),
System.currentTimeMillis() - electionStartTimeMs,
bucketsToRetryIndividually.size());

Map<TableBucket, LeaderAndIsr> leaderAndIsrsToUpdate = new HashMap<>();
electionResults.forEach(
(tableBucket, electionResult) ->
leaderAndIsrsToUpdate.put(tableBucket, electionResult.leaderAndIsr));
Set<TableBucket> updatedBuckets =
batchUpdateLeaderAndIsrWithFallback(leaderAndIsrsToUpdate);

@loserwang1024 loserwang1024 Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If BadVersionException occurs, need to stop here directly, on need to bucketsToRetryIndividually later.

@swuferhong swuferhong Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A simple return may be unsafe because earlier ZooKeeper transactions or individual retries may have already succeeded. Returning immediately would skip updating the coordinator cache for those buckets. Since subsequent fenced writes will fail safely, I prefer keeping the current behavior.


for (TableBucket tableBucket : updatedBuckets) {
ElectionResult electionResult = electionResults.get(tableBucket);
coordinatorContext.putBucketLeaderAndIsr(tableBucket, electionResult.leaderAndIsr);
doStateChange(tableBucket, BucketState.OnlineBucket);
coordinatorRequestBatch.addNotifyLeaderRequestForTabletServers(
new HashSet<>(electionResult.liveReplicas),
PhysicalTablePath.of(
coordinatorContext.getTablePathById(tableBucket.getTableId()),
partitionNames.get(tableBucket)),
tableBucket,
coordinatorContext.getAssignment(tableBucket),
electionResult.leaderAndIsr);
}

for (TableBucket tableBucket : bucketsToRetryIndividually) {
doHandleStateChange(
tableBucket, BucketState.OnlineBucket, controlledShutdownLeaderElection);
}
}

private Set<TableBucket> batchUpdateLeaderAndIsrWithFallback(
Map<TableBucket, LeaderAndIsr> leaderAndIsrs) {
Set<TableBucket> updatedBuckets = new HashSet<>();
if (leaderAndIsrs.isEmpty()) {
return updatedBuckets;
}

try {
zooKeeperClient.batchUpdateLeaderAndIsr(

@loserwang1024 loserwang1024 Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

batchUpdateLeaderAndIsr() splits the input into multiple ZooKeeper transactions of at most MAX_BATCH_SIZE, so an exception may be thrown after some earlier transactions have already committed. The current fallback retries every bucket individually, including buckets from transactions that were already confirmed successful.
Rewriting them should still succeed because the version check is against the coordinator epoch znode, which is not changed by LeaderAndIsr updates. Therefore, this appears correct but results in redundant ZooKeeper writes. Would it be worth propagating the set of definitely committed buckets from batchUpdateLeaderAndIsr() so that only failed or uncertain transactions need to be retried?

@swuferhong swuferhong Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think these retries write the same LeaderAndIsr values and validate against the coordinator epoch znode rather than the LeaderAndIsr node version, so they are idempotent and do not cause a correctness issue. The downside is limited to redundant ZooKeeper writes on the exceptional path.

Propagating the definitely committed buckets would require changing the ZooKeeper client API or introducing a partial-progress exception/result type. Given the additional complexity and the goal of keeping this PR minimally invasive, I suggest leaving it unchanged for now and considering it as a separate follow-up optimization.

leaderAndIsrs, coordinatorContext.getCoordinatorZkVersion());
updatedBuckets.addAll(leaderAndIsrs.keySet());
return updatedBuckets;
} catch (Exception e) {
if (ExceptionUtils.findThrowable(e, KeeperException.BadVersionException.class)
.isPresent()) {
LOG.error(
"Aborted batch LeaderAndIsr update during controlled shutdown because the coordinator was fenced.",
e);
return updatedBuckets;
}
LOG.warn(
"Failed to batch update LeaderAndIsr for {} buckets during controlled shutdown. Falling back to individual updates.",
leaderAndIsrs.size(),
e);
}

for (Map.Entry<TableBucket, LeaderAndIsr> entry : leaderAndIsrs.entrySet()) {
try {
zooKeeperClient.updateLeaderAndIsr(
entry.getKey(),
entry.getValue(),
coordinatorContext.getCoordinatorZkVersion());
updatedBuckets.add(entry.getKey());
} catch (Exception e) {
if (ExceptionUtils.findThrowable(e, KeeperException.BadVersionException.class)
.isPresent()) {
LOG.error(
"Stopped individual LeaderAndIsr updates during controlled shutdown because the coordinator was fenced.",
e);
break;
}
LOG.error(
"Failed to update bucket LeaderAndIsr for table bucket {} during controlled shutdown.",
stringifyBucket(entry.getKey()),
e);
}
}
return updatedBuckets;
}

/**
* Handle the state change of TableBucket. It's the core state transition logic of the state
* machine. It ensures that every state transition happens from a legal previous state to the
Expand Down Expand Up @@ -214,7 +383,7 @@ private void doHandleStateChange(
targetState,
String.format(
"Can't find partition name for partition: %s.",
tableBucket.getBucket()));
tableBucket.getPartitionId()));
return;
}
}
Expand Down Expand Up @@ -347,7 +516,7 @@ public void batchHandleOnlineChangeAndInitLeader(Set<TableBucket> tableBuckets)
BucketState.OnlineBucket,
String.format(
"Can't find partition name for partition: %s.",
tableBucket.getBucket()));
tableBucket.getPartitionId()));
continue;
}
}
Expand Down Expand Up @@ -486,6 +655,32 @@ private Optional<ElectionResult> electNewLeaderForTableBuckets(
LOG.error("Can't get state for table bucket {}.", stringifyBucket(tableBucket), e);
return Optional.empty();
}
Optional<ElectionResult> optionalElectionResult =
electNewLeaderForTableBucket(tableBucket, leaderAndIsr, electionStrategy);
if (!optionalElectionResult.isPresent()) {
return Optional.empty();
}
ElectionResult electionResult = optionalElectionResult.get();
try {
zooKeeperClient.updateLeaderAndIsr(
tableBucket,
electionResult.leaderAndIsr,
coordinatorContext.getCoordinatorZkVersion());
} catch (Exception e) {
LOG.error(
"Fail to update bucket LeaderAndIsr for table bucket {}.",
stringifyBucket(tableBucket),
e);
return Optional.empty();
}
coordinatorContext.putBucketLeaderAndIsr(tableBucket, electionResult.leaderAndIsr);
return optionalElectionResult;
}

private Optional<ElectionResult> electNewLeaderForTableBucket(
TableBucket tableBucket,
LeaderAndIsr leaderAndIsr,
ReplicaLeaderElection electionStrategy) {
if (leaderAndIsr.coordinatorEpoch() > coordinatorContext.getCoordinatorEpoch()) {
LOG.error(
"Aborted leader election for table bucket {} since the bucket state path was "
Expand All @@ -505,21 +700,7 @@ private Optional<ElectionResult> electNewLeaderForTableBuckets(
stringifyBucket(tableBucket));
return Optional.empty();
}
ElectionResult electionResult = optionalElectionResult.get();
try {
zooKeeperClient.updateLeaderAndIsr(
tableBucket,
electionResult.leaderAndIsr,
coordinatorContext.getCoordinatorZkVersion());
} catch (Exception e) {
LOG.error(
"Fail to update bucket LeaderAndIsr for table bucket {}.",
stringifyBucket(tableBucket),
e);
return Optional.empty();
}
coordinatorContext.putBucketLeaderAndIsr(tableBucket, electionResult.leaderAndIsr);
return Optional.of(electionResult);
return optionalElectionResult;
}

private boolean checkValidTableBucketStateChange(
Expand Down
6 changes: 6 additions & 0 deletions fluss-server/src/main/java/org/apache/fluss/server/tablet/TabletServer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -542,6 +542,7 @@ CompletableFuture<Void> stopServices() {
}

private void controlledShutDown() {
long startTime = System.currentTimeMillis();
LOG.info("Starting controlled shutdown.");

// We request the CoordinatorServer to do a controlled shutdown. On failure, we backoff for
Expand Down Expand Up @@ -597,6 +598,11 @@ private void controlledShutDown() {
LOG.warn(
"Proceeding to do an unclean shutdown as all the controlled shutdown attempts failed.");
}

LOG.info(
"Controlled shutdown attempts finished in {} ms, succeeded: {}.",
System.currentTimeMillis() - startTime,
shutdownSucceeded);
}

@Override
Expand Down
11 changes: 10 additions & 1 deletion fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -612,26 +612,35 @@ public void batchUpdateLeaderAndIsr(
return;
}

long startTimeMs = System.currentTimeMillis();
int transactionCount = 0;
List<CuratorOp> ops = new ArrayList<>(leaderAndIsrList.size());
for (Map.Entry<TableBucket, LeaderAndIsr> entry : leaderAndIsrList.entrySet()) {
TableBucket tableBucket = entry.getKey();
LeaderAndIsr leaderAndIsr = entry.getValue();

LOG.info("Batch Update {} for bucket {} in Zookeeper.", leaderAndIsr, tableBucket);
LOG.debug("Batch update {} for bucket {} in ZooKeeper.", leaderAndIsr, tableBucket);
String path = LeaderAndIsrZNode.path(tableBucket);
byte[] data = LeaderAndIsrZNode.encode(leaderAndIsr);
CuratorOp updateOp = zkClient.transactionOp().setData().forPath(path, data);
ops.add(updateOp);
if (ops.size() == MAX_BATCH_SIZE) {
List<CuratorOp> wrapOps = wrapRequestsWithEpochCheck(ops, expectedZkVersion);
zkClient.transaction().forOperations(wrapOps);
transactionCount++;
ops.clear();
}
}
if (!ops.isEmpty()) {
List<CuratorOp> wrapOps = wrapRequestsWithEpochCheck(ops, expectedZkVersion);
zkClient.transaction().forOperations(wrapOps);
transactionCount++;
}
LOG.info(
"Batch updated LeaderAndIsr for {} buckets in {} ZooKeeper transactions in {} ms.",
leaderAndIsrList.size(),
transactionCount,
System.currentTimeMillis() - startTimeMs);
}

protected void deleteLeaderAndIsr(TableBucket tableBucket) throws Exception {
Expand Down
Loading
Loading