[morazow]

Purpose

Linked issues: close #4622, close #3869

RebalanceITCase#testRebalanceForLogTable intermittently fails because one tablet server ends a completed rebalance with 0 leaders. The test exposes a real coordinator bug: a planned leadership move is dropped when the same bucket also has a replica move (#4622). As discussed with @swuferhong, this PR also makes that combined move wait for the planned leader's acknowledgement before the task completes (#3869).

Brief change log

The goal optimizer encodes the target leader as the first target replica (GoalOptimizerUtils#getDiff). A plan can move a replica and the leadership of the same bucket, for example:

originReplicas=[2, 0, 1], originalLeader=2
newReplicas   =[3, 2, 1], newLeader=3

Phase A - elect the planned leader. After the adding replicas catch up, CoordinatorEventProcessor#maybeReassignedBucketLeaderIfRequired re-elected a leader only when the current leader was not in the target replicas. Here the old leader 2 stays, so no election ran. Server 2 kept leadership and the task was still marked COMPLETED. The coordinator now re-elects whenever the current leader is not the first target replica.

Phase B - wait for the planned leader's ACK. The combined move sent asynchronous NotifyLeaderAndIsr requests and then marked the task COMPLETED in the same call. Only leader-only plans waited for an ACK, and that check looked only at the server id.

• After the final leader state is written (leader-only election, or the end of phase B), RebalanceManager records the bucket epoch the planned leader must acknowledge.
• CoordinatorRequestBatch passes the bucket epoch each request sent to the response event, because the response carries no epoch.
• A response completes the task only if it succeeded, came from the planned leader, answers a request with at least the recorded epoch, and ZooKeeper shows the planned leader. Responses to earlier requests do not count.
• While a bucket waits for its ACK, ISR changes do not run phase B again.
• A failed ACK leaves the task REBALANCING until the existing task timeout. After a coordinator restart, ZooKeeper already holds the target replicas, so the recovered plan runs the leader-only path, which also waits for the ACK.

Tests

CoordinatorEventProcessorTest:

• testBucketReassignmentMovesLeaderWhenOldLeaderStaysInReplicas: plan [0, 1, 3] -> [1, 0, 2] with new leader 1. It holds the planned leader's final ACK and releases all earlier responses, then checks that an ISR change sends no new requests. The task stays REBALANCING until the ACK is released, then completes with leader 1 and replicas [1, 0, 2].
• testBucketReassignmentWaitsForTimeoutWhenPlannedLeaderAckFails: a failed ACK keeps the task REBALANCING, and the task timeout then moves it to TIMEOUT.
• testBucketReassignmentWithLeaderChangeResumesAfterCoordinatorRestart: parameterized over four phase B boundaries (before leader election, after election, after the ISR shrink, after the assignment update). Each case completes only on the planned leader's ACK.

Red-green: reverting each part of the fix makes these tests fail. Phase B finishing immediately fails 5 tests, dropping the epoch check fails 2, and dropping the ISR-change guard fails 1.

Runs: CoordinatorEventProcessorTest (40), RebalanceManagerTest (5) and RebalanceManagerITCase (2) pass. RebalanceITCase passed 10 of 10 runs. Before the fix, the unmodified test failed 1 in 8 local runs.

API and Format

No.

Documentation

No.

🤖 Generated with Claude Code

[swuferhong]

Hi, @morazow thanks for your contribution, I left one comment:

[swuferhong]

This fixes the missing leader election, but combined replica-and-leader reassignment still completes before the planned leader acknowledges the final state.

onBucketReassignment(..., true) sends asynchronous NotifyLeaderAndIsr requests and then immediately calls finishRebalanceTask(..., COMPLETED). The existing ACK-based completion check only applies to leader-only plans.

I reproduced this on 9266dec: after Phase A, hold server 1's notification responses and submit the ISR catch-up for [0,1,3] -> [1,0,2]. The metadata leader becomes 1, but the task is already COMPLETED while those responses remain pending.

Could we complete this flow before merging, as required by #3869? Completion should wait for a successful ACK from the planned leader for the final assignment and epoch. Earlier-phase/stale responses must not satisfy that condition, and failures, timeout, or coordinator restart must leave the task recoverable.

Please extend the regression test to hold the final target-leader ACK and assert that the task remains REBALANCING until that ACK is released.

[morazow]

Hello @swuferhong ,

I see it now, I had to understand first on the phase B part of the rebalance.

You PR on #3860 already addressed this, right?

I'd suggest then you can add the part for if (currentLeader != targetReplicas.get(0)) {, then it will also fix the flakiness resulting from it. Please feel free to move the tests to your PR.

Then I will close this PR 🤝

[swuferhong]

Hi, @morazow thanks for your contribution, I left one comment:

[morazow]

Thanks @swuferhong for the review, checking and addressing today 🤝