You are viewing a plain text version of this content. The canonical link for it is here.

Posted to jira@kafka.apache.org by "FrankYang0529 (via GitHub)" <gi...@apache.org> on 2024/04/07 14:01:22 UTC

[PR] MINOR: enable kraft test for ReassignPartitionsIntegrationTest [kafka]

FrankYang0529 opened a new pull request, #15675:
URL: https://github.com/apache/kafka/pull/15675

   Enable Kraft test for `ReassignPartitionsIntegrationTest`.
   
   All test cases can pass on my laptop.
   
   ### Committer Checklist (excluded from commit message)
   - [ ] Verify design and implementation 
   - [ ] Verify test coverage and CI build status
   - [ ] Verify documentation (including upgrade notes)
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: migrate ReassignPartitionsIntegrationTest to use ClusterTestExtensions [kafka]

Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.

FrankYang0529 commented on PR #15675:
URL: https://github.com/apache/kafka/pull/15675#issuecomment-2122582976

   Hi @chia7712, thanks for the review and suggestion. I addressed all comments. Thank you.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: enable kraft test for ReassignPartitionsIntegrationTest [kafka]

Posted by "chia7712 (via GitHub)" <gi...@apache.org>.

chia7712 commented on PR #15675:
URL: https://github.com/apache/kafka/pull/15675#issuecomment-2078363083

   btw, it would be great if you can rewrite it by new test infra.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: migrate ReassignPartitionsIntegrationTest to use ClusterTestExtensions [kafka]

Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.

FrankYang0529 commented on code in PR #15675:
URL: https://github.com/apache/kafka/pull/15675#discussion_r1616073769


##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,820 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.time.Duration;
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testHighWaterMarkAfterPartitionReassignment() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 100);
+
+        // Execute the assignment
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[3,1,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, -1L, -1L);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(foo0,
+                    new PartitionReassignmentState(asList(3, 1, 2), asList(3, 1, 2), true));
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            TestUtils.waitForCondition(() -> {
+                ListOffsetsResultInfo result = admin.listOffsets(Collections.singletonMap(foo0, new OffsetSpec.LatestSpec())).partitionResult(foo0).get();
+                return result.offset() == 100;
+            }, "Timeout for waiting offset");
+        }
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Execute the assignment with a low throttle
+            long initialThrottle = 1L;
+            runExecuteAssignment(false, assignment, initialThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), initialThrottle);
+
+            // Now update the throttle and verify the reassignment completes
+            long updatedThrottle = 300000L;
+            runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), updatedThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(admin, assignment, false));
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+
+            // Execute the assignment
+            long interBrokerThrottle = 300000L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Wait for the assignment to complete
+            TestUtils.waitForCondition(() -> {
+                // Check the reassignment status.
+                VerifyAssignmentResult result = runVerifyAssignment(admin, assignment, true);
+
+                if (!result.partsOngoing) {
+                    return true;
+                } else {
+                    assertFalse(
+                            result.partStates.values().stream().allMatch(state -> state.done),
+                            "Expected at least one partition reassignment to be ongoing when result = " + result
+                    );
+                    assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                    assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                    System.out.println("Current result: " + result);
+                    waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+                    return false;
+                }
+            }, "Expected reassignment to complete.");
+
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(finalAssignment));
+            // The throttles should still have been preserved, since we ran with --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            List<ConsumerRecord<byte[], byte[]>> allRecords = new ArrayList<>();
+            TestUtils.waitForCondition(() -> {
+                ConsumerRecords<byte[], byte[]> records = consumer.poll(Duration.ofMillis(100L));
+                System.out.println("records.count " + records.count());

Review Comment:
   Removed all `System.out.println`. Thank you.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: migrate ReassignPartitionsIntegrationTest to use ClusterTestExtensions [kafka]

Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.

FrankYang0529 commented on PR #15675:
URL: https://github.com/apache/kafka/pull/15675#issuecomment-2135496699

   The CI is finished and failed test cases are not related to this PR.
   
   OffsetsApiIntegrationTest.testAlterSinkConnectorOffsets: https://issues.apache.org/jira/browse/KAFKA-15892
   OffsetsApiIntegrationTest.testAlterSinkConnectorOffsetsZombieSinkTasks: https://issues.apache.org/jira/browse/KAFKA-15917
   TransactionsTest.testFencingOnCommit: https://issues.apache.org/jira/browse/KAFKA-12567
   ExactlyOnceSourceIntegrationTest.testSeparateOffsetsTopic: https://issues.apache.org/jira/browse/KAFKA-14089
   ConsumerBounceTest.testConsumptionWithBrokerFailures: https://issues.apache.org/jira/browse/KAFKA-15146
   ReplicationQuotasTest.shouldThrottleOldSegments: https://issues.apache.org/jira/browse/KAFKA-16635
   DelegationTokenEndToEndAuthorizationWithOwnerTest: https://issues.apache.org/jira/browse/KAFKA-15411
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: enable kraft test for ReassignPartitionsIntegrationTest [kafka]

Posted by "chia7712 (via GitHub)" <gi...@apache.org>.

chia7712 commented on PR #15675:
URL: https://github.com/apache/kafka/pull/15675#issuecomment-2041482458

   @FrankYang0529 Could we use `ClusterTestExtensions` to enable `kraft`? It seems to me that using the new test framework is the better way. Otherwise, we may rewrite it again in the future.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: migrate ReassignPartitionsIntegrationTest to use ClusterTestExtensions [kafka]

Posted by "chia7712 (via GitHub)" <gi...@apache.org>.

chia7712 commented on code in PR #15675:
URL: https://github.com/apache/kafka/pull/15675#discussion_r1606091199


##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,776 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import kafka.utils.TestUtils;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+import scala.collection.JavaConverters;
+import scala.collection.Seq;
+
+import java.util.AbstractMap;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Execute the assignment with a low throttle
+        long initialThrottle = 1L;
+        runExecuteAssignment(false, assignment, initialThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), initialThrottle);
+
+        // Now update the throttle and verify the reassignment completes
+        long updatedThrottle = 300000L;
+        runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), updatedThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(assignment, false));
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Execute the assignment
+        long interBrokerThrottle = 300000L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Wait for the assignment to complete
+        TestUtils.waitUntilTrue(

Review Comment:
   We don't need to create `Admin` for each loop. Also, we ought to apply java version. For example:
   ```java
           try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
               org.apache.kafka.test.TestUtils.waitForCondition(() -> {
                   Set<Integer> isr = admin.describeTopics(Collections.singleton(foo0.topic()))
                       .allTopicNames().get().get(foo0.topic()).partitions().stream()
                       .filter(p -> p.partition() == foo0.partition())
                       .flatMap(p -> p.isr().stream())
                       .map(Node::id).collect(Collectors.toSet());
                   return isr.containsAll(Arrays.asList(0, 1, 2, 3));
               }, "Timed out while waiting for replica 3 to join the ISR");
           }
   ```



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java:
##########
@@ -64,40 +48,23 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.Map.Entry;
-import java.util.Objects;
-import java.util.Optional;
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singleton;
-import static java.util.Collections.singletonList;
 import static java.util.Collections.singletonMap;
-import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
-import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
-import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
-import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
 import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
-import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
 import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
 import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
 import static org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS;
-import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
-import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
 import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
-import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
 import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
 import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
-@SuppressWarnings("ClassFanOutComplexity")
 @Timeout(300)
 public class ReassignPartitionsIntegrationTest extends QuorumTestHarness {

Review Comment:
   `unthrottledBrokerConfigs` is unused in this test now. could you please remove it? 



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,776 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import kafka.utils.TestUtils;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+import scala.collection.JavaConverters;
+import scala.collection.Seq;
+
+import java.util.AbstractMap;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Execute the assignment with a low throttle
+        long initialThrottle = 1L;
+        runExecuteAssignment(false, assignment, initialThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), initialThrottle);
+
+        // Now update the throttle and verify the reassignment completes
+        long updatedThrottle = 300000L;
+        runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), updatedThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(assignment, false));
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Execute the assignment
+        long interBrokerThrottle = 300000L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Wait for the assignment to complete
+        TestUtils.waitUntilTrue(
+                () -> {
+                    // Check the reassignment status.
+                    VerifyAssignmentResult result = runVerifyAssignment(assignment, true);
+
+                    if (!result.partsOngoing) {
+                        return true;
+                    } else {
+                        assertFalse(
+                                result.partStates.values().stream().allMatch(state -> state.done),
+                                "Expected at least one partition reassignment to be ongoing when result = " + result
+                        );
+                        assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                        assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                        System.out.println("Current result: " + result);
+                        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+                        return false;
+                    }
+                }, () -> "Expected reassignment to complete.", DEFAULT_MAX_WAIT_MS, 100L);
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(finalAssignment));
+        // The throttles should still have been preserved, since we ran with --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            TestUtils.pollUntilAtLeastNumRecords(consumer, 100, DEFAULT_MAX_WAIT_MS);
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        assertEquals(unthrottledBrokerConfigs,
+                describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+        long interBrokerThrottle = 1L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+        partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+        partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+        // Verify that the reassignment is running.  The very low throttle should keep it
+        // from completing before this runs.
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+        // Cancel the reassignment.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+        // Broker throttles are still active because we passed --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Cancelling the reassignment again should reveal nothing to cancel.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+        // This time, the broker throttles were removed.
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+        // Verify that there are no ongoing reassignments.
+        assertFalse(runVerifyAssignment(assignment, false).partsOngoing);
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        TestUtils.waitUntilTrue(
+                () -> {
+                    try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+                        return Objects.equals(TestUtils.currentIsr(admin, foo0), mutableSet(0, 1, 2, 3).toSet());
+                    }
+                },
+                () -> "Timed out while waiting for replica 3 to join the ISR",
+                DEFAULT_MAX_WAIT_MS, 100L
+        );
+
+        // Now cancel the assignment and verify that the partition is removed from cancelled replicas
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(singleton(foo0), Collections.emptySet()), runCancelAssignment(assignment, true));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 4));
+    }
+
+    /**
+     * Test moving partitions between directories.
+     */
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testLogDirReassignment() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        ReassignPartitionsIntegrationTest.LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move, but throttle it to be very slow so that it can't complete
+        // before our next checks happen.
+        long logDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, logDirThrottle);
+
+        // Check the output of --verify
+        waitForVerifyAssignment(reassignment.json, true,
+                new VerifyAssignmentResult(singletonMap(
+                        topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                ), false, singletonMap(
+                        new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                        new ActiveMoveState(reassignment.currentDir, reassignment.targetDir, reassignment.targetDir)
+                ), true));
+        waitForLogDirThrottle(singleton(0), logDirThrottle);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Remove the throttle
+            admin.incrementalAlterConfigs(singletonMap(
+                            new ConfigResource(ConfigResource.Type.BROKER, "0"),
+                            singletonList(new AlterConfigOp(
+                                    new ConfigEntry(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, ""), AlterConfigOp.OpType.DELETE))))
+                    .all().get();
+            waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+
+            // Wait for the directory movement to complete.
+            waitForVerifyAssignment(reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+
+            ReassignPartitionsIntegrationTest.BrokerDirs info1 = new ReassignPartitionsIntegrationTest.BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
+            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
+        }
+    }
+
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testAlterLogDirReassignmentThrottle() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        ReassignPartitionsIntegrationTest.LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move with a low throttle so it does not complete
+        long initialLogDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
+        waitForLogDirThrottle(new HashSet<>(singletonList(0)), initialLogDirThrottle);
+
+        // Now increase the throttle and verify that the log dir movement completes
+        long updatedLogDirThrottle = 3000000L;
+        runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
+        waitForLogDirThrottle(singleton(0), updatedLogDirThrottle);
+
+        waitForVerifyAssignment(reassignment.json, true,
+                new VerifyAssignmentResult(singletonMap(
+                        topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                ), false, singletonMap(
+                        new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
+                        new CompletedMoveState(reassignment.targetDir)
+                ), false));
+    }
+
+    private void createTopics() {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<Integer, List<Integer>> fooReplicasAssignments = new HashMap<>();
+            fooReplicasAssignments.put(0, asList(0, 1, 2));
+            fooReplicasAssignments.put(1, asList(1, 2, 3));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("foo", fooReplicasAssignments))).topicId("foo").get());
+
+            Map<Integer, List<Integer>> barReplicasAssignments = new HashMap<>();
+            barReplicasAssignments.put(0, asList(3, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("bar", barReplicasAssignments))).topicId("bar").get());
+
+            Map<Integer, List<Integer>> bazReplicasAssignments = new HashMap<>();
+            bazReplicasAssignments.put(0, asList(1, 0, 2));
+            bazReplicasAssignments.put(1, asList(2, 0, 1));
+            bazReplicasAssignments.put(2, asList(0, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("baz", bazReplicasAssignments))).topicId("baz").get());
+        }
+    }
+
+    private void produceMessages(String topic, int partition, int numMessages) {
+        Properties properties = new Properties();
+        properties.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        try (KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(properties, new ByteArraySerializer(), new ByteArraySerializer())) {
+            IntStream.range(0, numMessages).forEach(i -> {
+                ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topic, partition, null, new byte[10000]);
+                assertDoesNotThrow(() -> producer.send(record).get());
+            });
+        }
+    }
+
+    private void executeAndVerifyReassignment() throws ExecutionException, InterruptedException {
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3,2,0],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition bar0 = new TopicPartition("bar", 0);
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+
+        initialAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 3), true));
+        initialAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 0), true));
+
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(initialAssignment));
+
+        // Execute the assignment
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 3), asList(0, 1, 3), true));
+        finalAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 0), asList(3, 2, 0), true));
+
+        VerifyAssignmentResult verifyAssignmentResult = runVerifyAssignment(assignment, false);
+        assertFalse(verifyAssignmentResult.movesOngoing);
+
+        // Wait for the assignment to complete
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+
+        assertEquals(unthrottledBrokerConfigs,
+                describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Verify that partitions are removed from brokers no longer assigned
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 2));
+        verifyReplicaDeleted(new TopicPartitionReplica(bar0.topic(), bar0.partition(), 1));
+    }
+
+    private final Map<Integer, Map<String, Long>> unthrottledBrokerConfigs = new HashMap<>(); {
+        IntStream.range(0, 4).forEach(brokerId -> {
+            Map<String, Long> brokerConfig = new HashMap<>();
+
+            BROKER_LEVEL_THROTTLES.forEach(throttle -> brokerConfig.put(throttle, -1L));
+
+            unthrottledBrokerConfigs.put(brokerId, brokerConfig);
+        });
+    }
+
+    private void verifyReplicaDeleted(TopicPartitionReplica topicPartitionReplica) throws ExecutionException, InterruptedException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitUntilTrue(
+                    () -> {
+                        try {
+                            TopicDescription topicDescription = admin.describeTopics(singleton(topicPartitionReplica.topic())).topicNameValues().get(topicPartitionReplica.topic()).get();
+                            return topicDescription.partitions().stream().noneMatch(topicPartitionInfo -> {
+                                if (topicPartitionInfo.partition() != topicPartitionReplica.partition()) {
+                                    return false;
+                                }
+                                return topicPartitionInfo.replicas().stream().anyMatch(node -> node.id() == topicPartitionReplica.brokerId());
+                            });
+                        } catch (InterruptedException | ExecutionException e) {
+                            throw new RuntimeException(e);
+                        }
+                    },
+                    () -> "Timed out waiting for replica " + topicPartitionReplica.brokerId() + " of " + topicPartitionReplica + " to be deleted",
+                    DEFAULT_MAX_WAIT_MS,
+                    100L
+            );
+        }
+    }
+
+    private void waitForLogDirThrottle(Set<Integer> throttledBrokers, Long logDirThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, logDirThrottle);
+        waitForBrokerThrottles(throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForInterBrokerThrottle(List<Integer> throttledBrokers, Long interBrokerThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, -1L);
+        waitForBrokerThrottles(throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForBrokerThrottles(Collection<Integer> throttledBrokers, Map<String, Long> throttleConfig) {
+        Map<Integer, Map<String, Long>> throttledBrokerConfigs = new HashMap<>();
+        unthrottledBrokerConfigs.forEach((brokerId, unthrottledConfig) -> {
+            Map<String, Long> expectedThrottleConfig = throttledBrokers.contains(brokerId)
+                    ? throttleConfig
+                    : unthrottledConfig;
+            throttledBrokerConfigs.put(brokerId, expectedThrottleConfig);
+        });
+        waitForBrokerLevelThrottles(throttledBrokerConfigs);
+    }
+
+    private void waitForBrokerLevelThrottles(Map<Integer, Map<String, Long>> targetThrottles) {
+        AtomicReference<Map<Integer, Map<String, Long>>> curThrottles = new AtomicReference<>(new HashMap<>());
+        TestUtils.waitUntilTrue(() -> {
+            assertDoesNotThrow(() -> curThrottles.set(describeBrokerLevelThrottles(targetThrottles.keySet())));
+            return targetThrottles.equals(curThrottles.get());
+        }, () -> "timed out waiting for broker throttle to become " + targetThrottles + ".  " +
+                "Latest throttles were " + curThrottles.get(), DEFAULT_MAX_WAIT_MS, 25);
+    }
+
+    /**
+     * Describe the broker-level throttles in the cluster.
+     *
+     * @return                A map whose keys are broker IDs and whose values are throttle
+     *                        information.  The nested maps are keyed on throttle name.
+     */
+    private Map<Integer, Map<String, Long>> describeBrokerLevelThrottles(Collection<Integer> brokerIds) throws ExecutionException, InterruptedException {
+        Map<Integer, Map<String, Long>> results = new HashMap<>();
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            for (Integer brokerId : brokerIds) {
+                ConfigResource brokerResource = new ConfigResource(ConfigResource.Type.BROKER, brokerId.toString());
+                Config brokerConfigs = admin.describeConfigs(singleton(brokerResource)).values()
+                        .get(brokerResource)
+                        .get();
+
+                Map<String, Long> throttles = new HashMap<>();
+                BROKER_LEVEL_THROTTLES.forEach(throttleName -> {
+                    String configValue = Optional.ofNullable(brokerConfigs.get(throttleName)).map(ConfigEntry::value).orElse("-1");
+                    throttles.put(throttleName, Long.parseLong(configValue));
+                });
+                results.put(brokerId, throttles);
+            }
+        }
+        return results;
+    }
+
+    static class LogDirReassignment {

Review Comment:
   It seems this is unused.



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java:
##########
@@ -119,93 +86,6 @@ public void tearDown() {
         });
     }
 
-    @ParameterizedTest
-    @ValueSource(strings = {"zk", "kraft"})
-    public void testReassignment(String quorum) throws Exception {
-        cluster = new ReassignPartitionsTestCluster(Collections.emptyMap(), Collections.emptyMap());
-        cluster.setup();
-        executeAndVerifyReassignment();
-    }
-
-    @ParameterizedTest
-    @ValueSource(strings = "zk") // Note: KRaft requires AlterPartition
-    public void testReassignmentWithAlterPartitionDisabled(String quorum) throws Exception {
-        // Test reassignment when the IBP is on an older version which does not use
-        // the `AlterPartition` API. In this case, the controller will register individual
-        // watches for each reassigning partition so that the reassignment can be
-        // completed as soon as the ISR is expanded.
-        Map<String, String> configOverrides = singletonMap(INTER_BROKER_PROTOCOL_VERSION_CONFIG, IBP_2_7_IV1.version());
-        cluster = new ReassignPartitionsTestCluster(configOverrides, Collections.emptyMap());
-        cluster.setup();
-        executeAndVerifyReassignment();
-    }
-
-    @ParameterizedTest
-    @ValueSource(strings = "zk") // Note: KRaft requires AlterPartition
-    public void testReassignmentCompletionDuringPartialUpgrade(String quorum) throws Exception {
-        // Test reassignment during a partial upgrade when some brokers are relying on
-        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
-        // In this test case, broker 0 starts up first on the latest IBP and is typically
-        // elected as controller. The three remaining brokers start up on the older IBP.
-        // We want to ensure that reassignment can still complete through the ISR change
-        // notification path even though the controller expects `AlterPartition`.
-
-        // Override change notification settings so that test is not delayed by ISR
-        // change notification delay
-        ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
-
-        Map<String, String> oldIbpConfig = singletonMap(INTER_BROKER_PROTOCOL_VERSION_CONFIG, IBP_2_7_IV1.version());
-        Map<Integer, Map<String, String>> brokerConfigOverrides = new HashMap<>();
-        brokerConfigOverrides.put(1, oldIbpConfig);
-        brokerConfigOverrides.put(2, oldIbpConfig);
-        brokerConfigOverrides.put(3, oldIbpConfig);
-
-        cluster = new ReassignPartitionsTestCluster(Collections.emptyMap(), brokerConfigOverrides);
-        cluster.setup();
-
-        executeAndVerifyReassignment();
-    }
-
-    private void executeAndVerifyReassignment() throws ExecutionException, InterruptedException {
-        String assignment = "{\"version\":1,\"partitions\":" +
-            "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
-            "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3,2,0],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
-            "]}";
-
-        TopicPartition foo0 = new TopicPartition("foo", 0);
-        TopicPartition bar0 = new TopicPartition("bar", 0);
-
-        // Check that the assignment has not yet been started yet.
-        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
-
-        initialAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 3), true));
-        initialAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 0), true));
-
-        waitForVerifyAssignment(cluster.adminClient, assignment, false,
-            new VerifyAssignmentResult(initialAssignment));
-
-        // Execute the assignment
-        runExecuteAssignment(cluster.adminClient, false, assignment, -1L, -1L);
-        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
-        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
-        finalAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 3), asList(0, 1, 3), true));
-        finalAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 0), asList(3, 2, 0), true));
-
-        VerifyAssignmentResult verifyAssignmentResult = runVerifyAssignment(cluster.adminClient, assignment, false);
-        assertFalse(verifyAssignmentResult.movesOngoing);
-
-        // Wait for the assignment to complete
-        waitForVerifyAssignment(cluster.adminClient, assignment, false,
-            new VerifyAssignmentResult(finalAssignment));
-
-        assertEquals(unthrottledBrokerConfigs,
-            describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
-
-        // Verify that partitions are removed from brokers no longer assigned
-        verifyReplicaDeleted(foo0, 2);
-        verifyReplicaDeleted(bar0, 1);
-    }
-
     @ParameterizedTest
     @ValueSource(strings = {"zk", "kraft"})
     public void testHighWaterMarkAfterPartitionReassignment(String quorum) throws Exception {

Review Comment:
   It seems the check list of this case is the new leader replica has synced the HWM. Hence, we can rewrite it by following steps.
   
   1. create a partition with 2 replicas (for example: 0,1)
   2. add some records
   3. create new reassignment with a new replica leader (for example: 2,0,1)
   4. wait for the new leader to be on board
   5. check the HWD (by `Admin#listOffsets`)



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,776 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import kafka.utils.TestUtils;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+import scala.collection.JavaConverters;
+import scala.collection.Seq;
+
+import java.util.AbstractMap;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Execute the assignment with a low throttle
+        long initialThrottle = 1L;
+        runExecuteAssignment(false, assignment, initialThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), initialThrottle);
+
+        // Now update the throttle and verify the reassignment completes
+        long updatedThrottle = 300000L;
+        runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), updatedThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(assignment, false));
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Execute the assignment
+        long interBrokerThrottle = 300000L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Wait for the assignment to complete
+        TestUtils.waitUntilTrue(
+                () -> {
+                    // Check the reassignment status.
+                    VerifyAssignmentResult result = runVerifyAssignment(assignment, true);
+
+                    if (!result.partsOngoing) {
+                        return true;
+                    } else {
+                        assertFalse(
+                                result.partStates.values().stream().allMatch(state -> state.done),
+                                "Expected at least one partition reassignment to be ongoing when result = " + result
+                        );
+                        assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                        assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                        System.out.println("Current result: " + result);
+                        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+                        return false;
+                    }
+                }, () -> "Expected reassignment to complete.", DEFAULT_MAX_WAIT_MS, 100L);
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(finalAssignment));
+        // The throttles should still have been preserved, since we ran with --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            TestUtils.pollUntilAtLeastNumRecords(consumer, 100, DEFAULT_MAX_WAIT_MS);
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        assertEquals(unthrottledBrokerConfigs,
+                describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+        long interBrokerThrottle = 1L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+        partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+        partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+        // Verify that the reassignment is running.  The very low throttle should keep it
+        // from completing before this runs.
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+        // Cancel the reassignment.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+        // Broker throttles are still active because we passed --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Cancelling the reassignment again should reveal nothing to cancel.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+        // This time, the broker throttles were removed.
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+        // Verify that there are no ongoing reassignments.
+        assertFalse(runVerifyAssignment(assignment, false).partsOngoing);
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        TestUtils.waitUntilTrue(
+                () -> {
+                    try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+                        return Objects.equals(TestUtils.currentIsr(admin, foo0), mutableSet(0, 1, 2, 3).toSet());
+                    }
+                },
+                () -> "Timed out while waiting for replica 3 to join the ISR",
+                DEFAULT_MAX_WAIT_MS, 100L
+        );
+
+        // Now cancel the assignment and verify that the partition is removed from cancelled replicas
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(singleton(foo0), Collections.emptySet()), runCancelAssignment(assignment, true));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 4));
+    }
+
+    /**
+     * Test moving partitions between directories.
+     */
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testLogDirReassignment() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        ReassignPartitionsIntegrationTest.LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move, but throttle it to be very slow so that it can't complete
+        // before our next checks happen.
+        long logDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, logDirThrottle);
+
+        // Check the output of --verify
+        waitForVerifyAssignment(reassignment.json, true,
+                new VerifyAssignmentResult(singletonMap(
+                        topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                ), false, singletonMap(
+                        new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                        new ActiveMoveState(reassignment.currentDir, reassignment.targetDir, reassignment.targetDir)
+                ), true));
+        waitForLogDirThrottle(singleton(0), logDirThrottle);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Remove the throttle
+            admin.incrementalAlterConfigs(singletonMap(
+                            new ConfigResource(ConfigResource.Type.BROKER, "0"),
+                            singletonList(new AlterConfigOp(
+                                    new ConfigEntry(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, ""), AlterConfigOp.OpType.DELETE))))
+                    .all().get();
+            waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+
+            // Wait for the directory movement to complete.
+            waitForVerifyAssignment(reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+
+            ReassignPartitionsIntegrationTest.BrokerDirs info1 = new ReassignPartitionsIntegrationTest.BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
+            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
+        }
+    }
+
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testAlterLogDirReassignmentThrottle() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        ReassignPartitionsIntegrationTest.LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move with a low throttle so it does not complete
+        long initialLogDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
+        waitForLogDirThrottle(new HashSet<>(singletonList(0)), initialLogDirThrottle);
+
+        // Now increase the throttle and verify that the log dir movement completes
+        long updatedLogDirThrottle = 3000000L;
+        runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
+        waitForLogDirThrottle(singleton(0), updatedLogDirThrottle);
+
+        waitForVerifyAssignment(reassignment.json, true,
+                new VerifyAssignmentResult(singletonMap(
+                        topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                ), false, singletonMap(
+                        new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
+                        new CompletedMoveState(reassignment.targetDir)
+                ), false));
+    }
+
+    private void createTopics() {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<Integer, List<Integer>> fooReplicasAssignments = new HashMap<>();
+            fooReplicasAssignments.put(0, asList(0, 1, 2));
+            fooReplicasAssignments.put(1, asList(1, 2, 3));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("foo", fooReplicasAssignments))).topicId("foo").get());
+
+            Map<Integer, List<Integer>> barReplicasAssignments = new HashMap<>();
+            barReplicasAssignments.put(0, asList(3, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("bar", barReplicasAssignments))).topicId("bar").get());
+
+            Map<Integer, List<Integer>> bazReplicasAssignments = new HashMap<>();
+            bazReplicasAssignments.put(0, asList(1, 0, 2));
+            bazReplicasAssignments.put(1, asList(2, 0, 1));
+            bazReplicasAssignments.put(2, asList(0, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("baz", bazReplicasAssignments))).topicId("baz").get());
+        }
+    }
+
+    private void produceMessages(String topic, int partition, int numMessages) {
+        Properties properties = new Properties();
+        properties.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        try (KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(properties, new ByteArraySerializer(), new ByteArraySerializer())) {
+            IntStream.range(0, numMessages).forEach(i -> {
+                ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topic, partition, null, new byte[10000]);
+                assertDoesNotThrow(() -> producer.send(record).get());
+            });
+        }
+    }
+
+    private void executeAndVerifyReassignment() throws ExecutionException, InterruptedException {
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3,2,0],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition bar0 = new TopicPartition("bar", 0);
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+
+        initialAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 3), true));
+        initialAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 0), true));
+
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(initialAssignment));
+
+        // Execute the assignment
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 3), asList(0, 1, 3), true));
+        finalAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 0), asList(3, 2, 0), true));
+
+        VerifyAssignmentResult verifyAssignmentResult = runVerifyAssignment(assignment, false);
+        assertFalse(verifyAssignmentResult.movesOngoing);
+
+        // Wait for the assignment to complete
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+
+        assertEquals(unthrottledBrokerConfigs,
+                describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Verify that partitions are removed from brokers no longer assigned
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 2));
+        verifyReplicaDeleted(new TopicPartitionReplica(bar0.topic(), bar0.partition(), 1));
+    }
+
+    private final Map<Integer, Map<String, Long>> unthrottledBrokerConfigs = new HashMap<>(); {
+        IntStream.range(0, 4).forEach(brokerId -> {
+            Map<String, Long> brokerConfig = new HashMap<>();
+
+            BROKER_LEVEL_THROTTLES.forEach(throttle -> brokerConfig.put(throttle, -1L));
+
+            unthrottledBrokerConfigs.put(brokerId, brokerConfig);
+        });
+    }
+
+    private void verifyReplicaDeleted(TopicPartitionReplica topicPartitionReplica) throws ExecutionException, InterruptedException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitUntilTrue(

Review Comment:
   ditto



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,776 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import kafka.utils.TestUtils;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+import scala.collection.JavaConverters;
+import scala.collection.Seq;
+
+import java.util.AbstractMap;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Execute the assignment with a low throttle
+        long initialThrottle = 1L;
+        runExecuteAssignment(false, assignment, initialThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), initialThrottle);
+
+        // Now update the throttle and verify the reassignment completes
+        long updatedThrottle = 300000L;
+        runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), updatedThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(assignment, false));
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Execute the assignment
+        long interBrokerThrottle = 300000L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Wait for the assignment to complete
+        TestUtils.waitUntilTrue(
+                () -> {
+                    // Check the reassignment status.
+                    VerifyAssignmentResult result = runVerifyAssignment(assignment, true);
+
+                    if (!result.partsOngoing) {
+                        return true;
+                    } else {
+                        assertFalse(
+                                result.partStates.values().stream().allMatch(state -> state.done),
+                                "Expected at least one partition reassignment to be ongoing when result = " + result
+                        );
+                        assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                        assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                        System.out.println("Current result: " + result);
+                        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+                        return false;
+                    }
+                }, () -> "Expected reassignment to complete.", DEFAULT_MAX_WAIT_MS, 100L);
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(finalAssignment));
+        // The throttles should still have been preserved, since we ran with --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            TestUtils.pollUntilAtLeastNumRecords(consumer, 100, DEFAULT_MAX_WAIT_MS);
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        assertEquals(unthrottledBrokerConfigs,
+                describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+        long interBrokerThrottle = 1L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+        partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+        partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+        // Verify that the reassignment is running.  The very low throttle should keep it
+        // from completing before this runs.
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+        // Cancel the reassignment.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+        // Broker throttles are still active because we passed --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Cancelling the reassignment again should reveal nothing to cancel.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+        // This time, the broker throttles were removed.
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+        // Verify that there are no ongoing reassignments.
+        assertFalse(runVerifyAssignment(assignment, false).partsOngoing);
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        TestUtils.waitUntilTrue(

Review Comment:
   ditto. please consider using java version.



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,776 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import kafka.utils.TestUtils;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+import scala.collection.JavaConverters;
+import scala.collection.Seq;
+
+import java.util.AbstractMap;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Execute the assignment with a low throttle
+        long initialThrottle = 1L;
+        runExecuteAssignment(false, assignment, initialThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), initialThrottle);
+
+        // Now update the throttle and verify the reassignment completes
+        long updatedThrottle = 300000L;
+        runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), updatedThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(assignment, false));
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Execute the assignment
+        long interBrokerThrottle = 300000L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Wait for the assignment to complete
+        TestUtils.waitUntilTrue(
+                () -> {
+                    // Check the reassignment status.
+                    VerifyAssignmentResult result = runVerifyAssignment(assignment, true);
+
+                    if (!result.partsOngoing) {
+                        return true;
+                    } else {
+                        assertFalse(
+                                result.partStates.values().stream().allMatch(state -> state.done),
+                                "Expected at least one partition reassignment to be ongoing when result = " + result
+                        );
+                        assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                        assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                        System.out.println("Current result: " + result);
+                        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+                        return false;
+                    }
+                }, () -> "Expected reassignment to complete.", DEFAULT_MAX_WAIT_MS, 100L);
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(finalAssignment));
+        // The throttles should still have been preserved, since we ran with --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            TestUtils.pollUntilAtLeastNumRecords(consumer, 100, DEFAULT_MAX_WAIT_MS);
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        assertEquals(unthrottledBrokerConfigs,
+                describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+        long interBrokerThrottle = 1L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+        partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+        partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+        // Verify that the reassignment is running.  The very low throttle should keep it
+        // from completing before this runs.
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+        // Cancel the reassignment.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+        // Broker throttles are still active because we passed --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Cancelling the reassignment again should reveal nothing to cancel.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+        // This time, the broker throttles were removed.
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+        // Verify that there are no ongoing reassignments.
+        assertFalse(runVerifyAssignment(assignment, false).partsOngoing);
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        TestUtils.waitUntilTrue(
+                () -> {
+                    try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+                        return Objects.equals(TestUtils.currentIsr(admin, foo0), mutableSet(0, 1, 2, 3).toSet());
+                    }
+                },
+                () -> "Timed out while waiting for replica 3 to join the ISR",
+                DEFAULT_MAX_WAIT_MS, 100L
+        );
+
+        // Now cancel the assignment and verify that the partition is removed from cancelled replicas
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(singleton(foo0), Collections.emptySet()), runCancelAssignment(assignment, true));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 4));
+    }
+
+    /**
+     * Test moving partitions between directories.
+     */
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testLogDirReassignment() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        ReassignPartitionsIntegrationTest.LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move, but throttle it to be very slow so that it can't complete
+        // before our next checks happen.
+        long logDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, logDirThrottle);
+
+        // Check the output of --verify
+        waitForVerifyAssignment(reassignment.json, true,
+                new VerifyAssignmentResult(singletonMap(
+                        topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                ), false, singletonMap(
+                        new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                        new ActiveMoveState(reassignment.currentDir, reassignment.targetDir, reassignment.targetDir)
+                ), true));
+        waitForLogDirThrottle(singleton(0), logDirThrottle);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Remove the throttle
+            admin.incrementalAlterConfigs(singletonMap(
+                            new ConfigResource(ConfigResource.Type.BROKER, "0"),
+                            singletonList(new AlterConfigOp(
+                                    new ConfigEntry(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, ""), AlterConfigOp.OpType.DELETE))))
+                    .all().get();
+            waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+
+            // Wait for the directory movement to complete.
+            waitForVerifyAssignment(reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+
+            ReassignPartitionsIntegrationTest.BrokerDirs info1 = new ReassignPartitionsIntegrationTest.BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
+            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
+        }
+    }
+
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testAlterLogDirReassignmentThrottle() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        ReassignPartitionsIntegrationTest.LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move with a low throttle so it does not complete
+        long initialLogDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
+        waitForLogDirThrottle(new HashSet<>(singletonList(0)), initialLogDirThrottle);
+
+        // Now increase the throttle and verify that the log dir movement completes
+        long updatedLogDirThrottle = 3000000L;
+        runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
+        waitForLogDirThrottle(singleton(0), updatedLogDirThrottle);
+
+        waitForVerifyAssignment(reassignment.json, true,
+                new VerifyAssignmentResult(singletonMap(
+                        topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                ), false, singletonMap(
+                        new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
+                        new CompletedMoveState(reassignment.targetDir)
+                ), false));
+    }
+
+    private void createTopics() {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<Integer, List<Integer>> fooReplicasAssignments = new HashMap<>();
+            fooReplicasAssignments.put(0, asList(0, 1, 2));
+            fooReplicasAssignments.put(1, asList(1, 2, 3));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("foo", fooReplicasAssignments))).topicId("foo").get());
+
+            Map<Integer, List<Integer>> barReplicasAssignments = new HashMap<>();
+            barReplicasAssignments.put(0, asList(3, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("bar", barReplicasAssignments))).topicId("bar").get());
+
+            Map<Integer, List<Integer>> bazReplicasAssignments = new HashMap<>();
+            bazReplicasAssignments.put(0, asList(1, 0, 2));
+            bazReplicasAssignments.put(1, asList(2, 0, 1));
+            bazReplicasAssignments.put(2, asList(0, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("baz", bazReplicasAssignments))).topicId("baz").get());
+        }
+    }
+
+    private void produceMessages(String topic, int partition, int numMessages) {
+        Properties properties = new Properties();
+        properties.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        try (KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(properties, new ByteArraySerializer(), new ByteArraySerializer())) {
+            IntStream.range(0, numMessages).forEach(i -> {
+                ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topic, partition, null, new byte[10000]);
+                assertDoesNotThrow(() -> producer.send(record).get());
+            });
+        }
+    }
+
+    private void executeAndVerifyReassignment() throws ExecutionException, InterruptedException {
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3,2,0],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition bar0 = new TopicPartition("bar", 0);
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+
+        initialAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 3), true));
+        initialAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 0), true));
+
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(initialAssignment));
+
+        // Execute the assignment
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 3), asList(0, 1, 3), true));
+        finalAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 0), asList(3, 2, 0), true));
+
+        VerifyAssignmentResult verifyAssignmentResult = runVerifyAssignment(assignment, false);
+        assertFalse(verifyAssignmentResult.movesOngoing);
+
+        // Wait for the assignment to complete
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+
+        assertEquals(unthrottledBrokerConfigs,
+                describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Verify that partitions are removed from brokers no longer assigned
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 2));
+        verifyReplicaDeleted(new TopicPartitionReplica(bar0.topic(), bar0.partition(), 1));
+    }
+
+    private final Map<Integer, Map<String, Long>> unthrottledBrokerConfigs = new HashMap<>(); {
+        IntStream.range(0, 4).forEach(brokerId -> {
+            Map<String, Long> brokerConfig = new HashMap<>();
+
+            BROKER_LEVEL_THROTTLES.forEach(throttle -> brokerConfig.put(throttle, -1L));
+
+            unthrottledBrokerConfigs.put(brokerId, brokerConfig);
+        });
+    }
+
+    private void verifyReplicaDeleted(TopicPartitionReplica topicPartitionReplica) throws ExecutionException, InterruptedException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitUntilTrue(
+                    () -> {
+                        try {
+                            TopicDescription topicDescription = admin.describeTopics(singleton(topicPartitionReplica.topic())).topicNameValues().get(topicPartitionReplica.topic()).get();
+                            return topicDescription.partitions().stream().noneMatch(topicPartitionInfo -> {
+                                if (topicPartitionInfo.partition() != topicPartitionReplica.partition()) {
+                                    return false;
+                                }
+                                return topicPartitionInfo.replicas().stream().anyMatch(node -> node.id() == topicPartitionReplica.brokerId());
+                            });
+                        } catch (InterruptedException | ExecutionException e) {
+                            throw new RuntimeException(e);
+                        }
+                    },
+                    () -> "Timed out waiting for replica " + topicPartitionReplica.brokerId() + " of " + topicPartitionReplica + " to be deleted",
+                    DEFAULT_MAX_WAIT_MS,
+                    100L
+            );
+        }
+    }
+
+    private void waitForLogDirThrottle(Set<Integer> throttledBrokers, Long logDirThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, logDirThrottle);
+        waitForBrokerThrottles(throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForInterBrokerThrottle(List<Integer> throttledBrokers, Long interBrokerThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, -1L);
+        waitForBrokerThrottles(throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForBrokerThrottles(Collection<Integer> throttledBrokers, Map<String, Long> throttleConfig) {
+        Map<Integer, Map<String, Long>> throttledBrokerConfigs = new HashMap<>();
+        unthrottledBrokerConfigs.forEach((brokerId, unthrottledConfig) -> {
+            Map<String, Long> expectedThrottleConfig = throttledBrokers.contains(brokerId)
+                    ? throttleConfig
+                    : unthrottledConfig;
+            throttledBrokerConfigs.put(brokerId, expectedThrottleConfig);
+        });
+        waitForBrokerLevelThrottles(throttledBrokerConfigs);
+    }
+
+    private void waitForBrokerLevelThrottles(Map<Integer, Map<String, Long>> targetThrottles) {
+        AtomicReference<Map<Integer, Map<String, Long>>> curThrottles = new AtomicReference<>(new HashMap<>());
+        TestUtils.waitUntilTrue(() -> {
+            assertDoesNotThrow(() -> curThrottles.set(describeBrokerLevelThrottles(targetThrottles.keySet())));
+            return targetThrottles.equals(curThrottles.get());
+        }, () -> "timed out waiting for broker throttle to become " + targetThrottles + ".  " +
+                "Latest throttles were " + curThrottles.get(), DEFAULT_MAX_WAIT_MS, 25);
+    }
+
+    /**
+     * Describe the broker-level throttles in the cluster.
+     *
+     * @return                A map whose keys are broker IDs and whose values are throttle
+     *                        information.  The nested maps are keyed on throttle name.
+     */
+    private Map<Integer, Map<String, Long>> describeBrokerLevelThrottles(Collection<Integer> brokerIds) throws ExecutionException, InterruptedException {
+        Map<Integer, Map<String, Long>> results = new HashMap<>();
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            for (Integer brokerId : brokerIds) {
+                ConfigResource brokerResource = new ConfigResource(ConfigResource.Type.BROKER, brokerId.toString());
+                Config brokerConfigs = admin.describeConfigs(singleton(brokerResource)).values()
+                        .get(brokerResource)
+                        .get();
+
+                Map<String, Long> throttles = new HashMap<>();
+                BROKER_LEVEL_THROTTLES.forEach(throttleName -> {
+                    String configValue = Optional.ofNullable(brokerConfigs.get(throttleName)).map(ConfigEntry::value).orElse("-1");
+                    throttles.put(throttleName, Long.parseLong(configValue));
+                });
+                results.put(brokerId, throttles);
+            }
+        }
+        return results;
+    }
+
+    static class LogDirReassignment {
+        final String json;
+        final String currentDir;
+        final String targetDir;
+
+        public LogDirReassignment(String json, String currentDir, String targetDir) {
+            this.json = json;
+            this.currentDir = currentDir;
+            this.targetDir = targetDir;
+        }
+    }
+
+    private ReassignPartitionsIntegrationTest.LogDirReassignment buildLogDirReassignment(TopicPartition topicPartition,
+                                                                                         int brokerId,
+                                                                                         List<Integer> replicas) throws ExecutionException, InterruptedException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            DescribeLogDirsResult describeLogDirsResult = admin.describeLogDirs(
+                    IntStream.range(0, 4).boxed().collect(Collectors.toList()));
+
+            ReassignPartitionsIntegrationTest.BrokerDirs logDirInfo = new ReassignPartitionsIntegrationTest.BrokerDirs(describeLogDirsResult, brokerId);
+            assertTrue(logDirInfo.futureLogDirs.isEmpty());
+
+            String currentDir = logDirInfo.curLogDirs.get(topicPartition);
+            String newDir = logDirInfo.logDirs.stream().filter(dir -> !dir.equals(currentDir)).findFirst().get();
+
+            List<String> logDirs = replicas.stream().map(replicaId -> {
+                if (replicaId == brokerId)
+                    return "\"" + newDir + "\"";
+                else
+                    return "\"any\"";
+            }).collect(Collectors.toList());
+
+            String reassignmentJson =
+                    " { \"version\": 1," +
+                            "  \"partitions\": [" +
+                            "    {" +
+                            "     \"topic\": \"" + topicPartition.topic() + "\"," +
+                            "     \"partition\": " + topicPartition.partition() + "," +
+                            "     \"replicas\": [" + replicas.stream().map(Object::toString).collect(Collectors.joining(",")) + "]," +
+                            "     \"log_dirs\": [" + String.join(",", logDirs) + "]" +
+                            "    }" +
+                            "   ]" +
+                            "  }";
+
+            return new ReassignPartitionsIntegrationTest.LogDirReassignment(reassignmentJson, currentDir, newDir);
+        }
+    }
+
+    private VerifyAssignmentResult runVerifyAssignment(String jsonString,
+                                                       Boolean preserveThrottles) {
+        System.out.println("==> verifyAssignment(adminClient, jsonString=" + jsonString);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            return verifyAssignment(admin, jsonString, preserveThrottles);
+        } catch (ExecutionException | InterruptedException | JsonProcessingException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private void waitForVerifyAssignment(String jsonString,
+                                         Boolean preserveThrottles,
+                                         VerifyAssignmentResult expectedResult) {
+        final VerifyAssignmentResult[] latestResult = {null};
+        TestUtils.waitUntilTrue(
+                () -> {
+                    latestResult[0] = runVerifyAssignment(jsonString, preserveThrottles);
+                    return expectedResult.equals(latestResult[0]);
+                }, () -> "Timed out waiting for verifyAssignment result " + expectedResult.moveStates.get(new TopicPartitionReplica("foo", 0, 0)) + ".  " +
+                        "The latest result was " + latestResult[0].moveStates.get(new TopicPartitionReplica("foo", 0, 0)), DEFAULT_MAX_WAIT_MS, 10L);
+    }
+
+    private void runExecuteAssignment(Boolean additional,
+                                      String reassignmentJson,
+                                      Long interBrokerThrottle,
+                                      Long replicaAlterLogDirsThrottle) throws RuntimeException {
+        System.out.println("==> executeAssignment(adminClient, additional=" + additional + ", " +
+                "reassignmentJson=" + reassignmentJson + ", " +
+                "interBrokerThrottle=" + interBrokerThrottle + ", " +
+                "replicaAlterLogDirsThrottle=" + replicaAlterLogDirsThrottle + "))");
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            executeAssignment(admin, additional, reassignmentJson,
+                    interBrokerThrottle, replicaAlterLogDirsThrottle, 10000L, Time.SYSTEM);
+        } catch (ExecutionException | InterruptedException | JsonProcessingException | TerseException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private Map.Entry<Set<TopicPartition>, Set<TopicPartitionReplica>> runCancelAssignment(
+            String jsonString,
+            Boolean preserveThrottles
+    ) {
+        System.out.println("==> cancelAssignment(adminClient, jsonString=" + jsonString);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            return cancelAssignment(admin, jsonString, preserveThrottles, 10000L, Time.SYSTEM);
+        } catch (ExecutionException | InterruptedException | JsonProcessingException | TerseException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    static class BrokerDirs {

Review Comment:
   ditto.  unused



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,776 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import kafka.utils.TestUtils;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+import scala.collection.JavaConverters;
+import scala.collection.Seq;
+
+import java.util.AbstractMap;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Execute the assignment with a low throttle
+        long initialThrottle = 1L;
+        runExecuteAssignment(false, assignment, initialThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), initialThrottle);
+
+        // Now update the throttle and verify the reassignment completes
+        long updatedThrottle = 300000L;
+        runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), updatedThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(assignment, false));
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Execute the assignment
+        long interBrokerThrottle = 300000L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Wait for the assignment to complete
+        TestUtils.waitUntilTrue(
+                () -> {
+                    // Check the reassignment status.
+                    VerifyAssignmentResult result = runVerifyAssignment(assignment, true);
+
+                    if (!result.partsOngoing) {
+                        return true;
+                    } else {
+                        assertFalse(
+                                result.partStates.values().stream().allMatch(state -> state.done),
+                                "Expected at least one partition reassignment to be ongoing when result = " + result
+                        );
+                        assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                        assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                        System.out.println("Current result: " + result);
+                        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+                        return false;
+                    }
+                }, () -> "Expected reassignment to complete.", DEFAULT_MAX_WAIT_MS, 100L);
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(finalAssignment));
+        // The throttles should still have been preserved, since we ran with --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            TestUtils.pollUntilAtLeastNumRecords(consumer, 100, DEFAULT_MAX_WAIT_MS);
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        assertEquals(unthrottledBrokerConfigs,
+                describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+        long interBrokerThrottle = 1L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+        partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+        partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+        // Verify that the reassignment is running.  The very low throttle should keep it
+        // from completing before this runs.
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+        // Cancel the reassignment.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+        // Broker throttles are still active because we passed --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Cancelling the reassignment again should reveal nothing to cancel.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+        // This time, the broker throttles were removed.
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+        // Verify that there are no ongoing reassignments.
+        assertFalse(runVerifyAssignment(assignment, false).partsOngoing);
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        TestUtils.waitUntilTrue(
+                () -> {
+                    try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+                        return Objects.equals(TestUtils.currentIsr(admin, foo0), mutableSet(0, 1, 2, 3).toSet());
+                    }
+                },
+                () -> "Timed out while waiting for replica 3 to join the ISR",
+                DEFAULT_MAX_WAIT_MS, 100L
+        );
+
+        // Now cancel the assignment and verify that the partition is removed from cancelled replicas
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(singleton(foo0), Collections.emptySet()), runCancelAssignment(assignment, true));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 4));
+    }
+
+    /**
+     * Test moving partitions between directories.
+     */
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testLogDirReassignment() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        ReassignPartitionsIntegrationTest.LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move, but throttle it to be very slow so that it can't complete
+        // before our next checks happen.
+        long logDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, logDirThrottle);
+
+        // Check the output of --verify
+        waitForVerifyAssignment(reassignment.json, true,
+                new VerifyAssignmentResult(singletonMap(
+                        topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                ), false, singletonMap(
+                        new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                        new ActiveMoveState(reassignment.currentDir, reassignment.targetDir, reassignment.targetDir)
+                ), true));
+        waitForLogDirThrottle(singleton(0), logDirThrottle);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Remove the throttle
+            admin.incrementalAlterConfigs(singletonMap(
+                            new ConfigResource(ConfigResource.Type.BROKER, "0"),
+                            singletonList(new AlterConfigOp(
+                                    new ConfigEntry(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, ""), AlterConfigOp.OpType.DELETE))))
+                    .all().get();
+            waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+
+            // Wait for the directory movement to complete.
+            waitForVerifyAssignment(reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+
+            ReassignPartitionsIntegrationTest.BrokerDirs info1 = new ReassignPartitionsIntegrationTest.BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
+            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
+        }
+    }
+
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testAlterLogDirReassignmentThrottle() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        ReassignPartitionsIntegrationTest.LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move with a low throttle so it does not complete
+        long initialLogDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
+        waitForLogDirThrottle(new HashSet<>(singletonList(0)), initialLogDirThrottle);
+
+        // Now increase the throttle and verify that the log dir movement completes
+        long updatedLogDirThrottle = 3000000L;
+        runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
+        waitForLogDirThrottle(singleton(0), updatedLogDirThrottle);
+
+        waitForVerifyAssignment(reassignment.json, true,
+                new VerifyAssignmentResult(singletonMap(
+                        topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                ), false, singletonMap(
+                        new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
+                        new CompletedMoveState(reassignment.targetDir)
+                ), false));
+    }
+
+    private void createTopics() {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<Integer, List<Integer>> fooReplicasAssignments = new HashMap<>();
+            fooReplicasAssignments.put(0, asList(0, 1, 2));
+            fooReplicasAssignments.put(1, asList(1, 2, 3));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("foo", fooReplicasAssignments))).topicId("foo").get());
+
+            Map<Integer, List<Integer>> barReplicasAssignments = new HashMap<>();
+            barReplicasAssignments.put(0, asList(3, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("bar", barReplicasAssignments))).topicId("bar").get());
+
+            Map<Integer, List<Integer>> bazReplicasAssignments = new HashMap<>();
+            bazReplicasAssignments.put(0, asList(1, 0, 2));
+            bazReplicasAssignments.put(1, asList(2, 0, 1));
+            bazReplicasAssignments.put(2, asList(0, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("baz", bazReplicasAssignments))).topicId("baz").get());
+        }
+    }
+
+    private void produceMessages(String topic, int partition, int numMessages) {
+        Properties properties = new Properties();
+        properties.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        try (KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(properties, new ByteArraySerializer(), new ByteArraySerializer())) {
+            IntStream.range(0, numMessages).forEach(i -> {
+                ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topic, partition, null, new byte[10000]);
+                assertDoesNotThrow(() -> producer.send(record).get());
+            });
+        }
+    }
+
+    private void executeAndVerifyReassignment() throws ExecutionException, InterruptedException {
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3,2,0],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition bar0 = new TopicPartition("bar", 0);
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+
+        initialAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 3), true));
+        initialAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 0), true));
+
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(initialAssignment));
+
+        // Execute the assignment
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 3), asList(0, 1, 3), true));
+        finalAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 0), asList(3, 2, 0), true));
+
+        VerifyAssignmentResult verifyAssignmentResult = runVerifyAssignment(assignment, false);
+        assertFalse(verifyAssignmentResult.movesOngoing);
+
+        // Wait for the assignment to complete
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+
+        assertEquals(unthrottledBrokerConfigs,
+                describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Verify that partitions are removed from brokers no longer assigned
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 2));
+        verifyReplicaDeleted(new TopicPartitionReplica(bar0.topic(), bar0.partition(), 1));
+    }
+
+    private final Map<Integer, Map<String, Long>> unthrottledBrokerConfigs = new HashMap<>(); {
+        IntStream.range(0, 4).forEach(brokerId -> {
+            Map<String, Long> brokerConfig = new HashMap<>();
+
+            BROKER_LEVEL_THROTTLES.forEach(throttle -> brokerConfig.put(throttle, -1L));
+
+            unthrottledBrokerConfigs.put(brokerId, brokerConfig);
+        });
+    }
+
+    private void verifyReplicaDeleted(TopicPartitionReplica topicPartitionReplica) throws ExecutionException, InterruptedException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitUntilTrue(
+                    () -> {
+                        try {
+                            TopicDescription topicDescription = admin.describeTopics(singleton(topicPartitionReplica.topic())).topicNameValues().get(topicPartitionReplica.topic()).get();
+                            return topicDescription.partitions().stream().noneMatch(topicPartitionInfo -> {
+                                if (topicPartitionInfo.partition() != topicPartitionReplica.partition()) {
+                                    return false;
+                                }
+                                return topicPartitionInfo.replicas().stream().anyMatch(node -> node.id() == topicPartitionReplica.brokerId());
+                            });
+                        } catch (InterruptedException | ExecutionException e) {
+                            throw new RuntimeException(e);
+                        }
+                    },
+                    () -> "Timed out waiting for replica " + topicPartitionReplica.brokerId() + " of " + topicPartitionReplica + " to be deleted",
+                    DEFAULT_MAX_WAIT_MS,
+                    100L
+            );
+        }
+    }
+
+    private void waitForLogDirThrottle(Set<Integer> throttledBrokers, Long logDirThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, logDirThrottle);
+        waitForBrokerThrottles(throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForInterBrokerThrottle(List<Integer> throttledBrokers, Long interBrokerThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, -1L);
+        waitForBrokerThrottles(throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForBrokerThrottles(Collection<Integer> throttledBrokers, Map<String, Long> throttleConfig) {
+        Map<Integer, Map<String, Long>> throttledBrokerConfigs = new HashMap<>();
+        unthrottledBrokerConfigs.forEach((brokerId, unthrottledConfig) -> {
+            Map<String, Long> expectedThrottleConfig = throttledBrokers.contains(brokerId)
+                    ? throttleConfig
+                    : unthrottledConfig;
+            throttledBrokerConfigs.put(brokerId, expectedThrottleConfig);
+        });
+        waitForBrokerLevelThrottles(throttledBrokerConfigs);
+    }
+
+    private void waitForBrokerLevelThrottles(Map<Integer, Map<String, Long>> targetThrottles) {
+        AtomicReference<Map<Integer, Map<String, Long>>> curThrottles = new AtomicReference<>(new HashMap<>());
+        TestUtils.waitUntilTrue(() -> {

Review Comment:
   ditto



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,776 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import kafka.utils.TestUtils;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+import scala.collection.JavaConverters;
+import scala.collection.Seq;
+
+import java.util.AbstractMap;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Execute the assignment with a low throttle
+        long initialThrottle = 1L;
+        runExecuteAssignment(false, assignment, initialThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), initialThrottle);
+
+        // Now update the throttle and verify the reassignment completes
+        long updatedThrottle = 300000L;
+        runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), updatedThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(assignment, false));
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Execute the assignment
+        long interBrokerThrottle = 300000L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+        finalAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+        // Wait for the assignment to complete
+        TestUtils.waitUntilTrue(
+                () -> {
+                    // Check the reassignment status.
+                    VerifyAssignmentResult result = runVerifyAssignment(assignment, true);
+
+                    if (!result.partsOngoing) {
+                        return true;
+                    } else {
+                        assertFalse(
+                                result.partStates.values().stream().allMatch(state -> state.done),
+                                "Expected at least one partition reassignment to be ongoing when result = " + result
+                        );
+                        assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                        assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                        System.out.println("Current result: " + result);
+                        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+                        return false;
+                    }
+                }, () -> "Expected reassignment to complete.", DEFAULT_MAX_WAIT_MS, 100L);
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(finalAssignment));
+        // The throttles should still have been preserved, since we ran with --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Now remove the throttles.
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            TestUtils.pollUntilAtLeastNumRecords(consumer, 100, DEFAULT_MAX_WAIT_MS);
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        assertEquals(unthrottledBrokerConfigs,
+                describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+        long interBrokerThrottle = 1L;
+        runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+
+        Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+        partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+        partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+        // Verify that the reassignment is running.  The very low throttle should keep it
+        // from completing before this runs.
+        waitForVerifyAssignment(assignment, true,
+                new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+        // Cancel the reassignment.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+        // Broker throttles are still active because we passed --preserve-throttles
+        waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
+        // Cancelling the reassignment again should reveal nothing to cancel.
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+        // This time, the broker throttles were removed.
+        waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+        // Verify that there are no ongoing reassignments.
+        assertFalse(runVerifyAssignment(assignment, false).partsOngoing);
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        TestUtils.waitUntilTrue(
+                () -> {
+                    try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+                        return Objects.equals(TestUtils.currentIsr(admin, foo0), mutableSet(0, 1, 2, 3).toSet());
+                    }
+                },
+                () -> "Timed out while waiting for replica 3 to join the ISR",
+                DEFAULT_MAX_WAIT_MS, 100L
+        );
+
+        // Now cancel the assignment and verify that the partition is removed from cancelled replicas
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(singleton(foo0), Collections.emptySet()), runCancelAssignment(assignment, true));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 4));
+    }
+
+    /**
+     * Test moving partitions between directories.
+     */
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testLogDirReassignment() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        ReassignPartitionsIntegrationTest.LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move, but throttle it to be very slow so that it can't complete
+        // before our next checks happen.
+        long logDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, logDirThrottle);
+
+        // Check the output of --verify
+        waitForVerifyAssignment(reassignment.json, true,
+                new VerifyAssignmentResult(singletonMap(
+                        topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                ), false, singletonMap(
+                        new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                        new ActiveMoveState(reassignment.currentDir, reassignment.targetDir, reassignment.targetDir)
+                ), true));
+        waitForLogDirThrottle(singleton(0), logDirThrottle);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Remove the throttle
+            admin.incrementalAlterConfigs(singletonMap(
+                            new ConfigResource(ConfigResource.Type.BROKER, "0"),
+                            singletonList(new AlterConfigOp(
+                                    new ConfigEntry(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, ""), AlterConfigOp.OpType.DELETE))))
+                    .all().get();
+            waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
+
+            // Wait for the directory movement to complete.
+            waitForVerifyAssignment(reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+
+            ReassignPartitionsIntegrationTest.BrokerDirs info1 = new ReassignPartitionsIntegrationTest.BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
+            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
+        }
+    }
+
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testAlterLogDirReassignmentThrottle() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        ReassignPartitionsIntegrationTest.LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move with a low throttle so it does not complete
+        long initialLogDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
+        waitForLogDirThrottle(new HashSet<>(singletonList(0)), initialLogDirThrottle);
+
+        // Now increase the throttle and verify that the log dir movement completes
+        long updatedLogDirThrottle = 3000000L;
+        runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
+        waitForLogDirThrottle(singleton(0), updatedLogDirThrottle);
+
+        waitForVerifyAssignment(reassignment.json, true,
+                new VerifyAssignmentResult(singletonMap(
+                        topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                ), false, singletonMap(
+                        new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
+                        new CompletedMoveState(reassignment.targetDir)
+                ), false));
+    }
+
+    private void createTopics() {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<Integer, List<Integer>> fooReplicasAssignments = new HashMap<>();
+            fooReplicasAssignments.put(0, asList(0, 1, 2));
+            fooReplicasAssignments.put(1, asList(1, 2, 3));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("foo", fooReplicasAssignments))).topicId("foo").get());
+
+            Map<Integer, List<Integer>> barReplicasAssignments = new HashMap<>();
+            barReplicasAssignments.put(0, asList(3, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("bar", barReplicasAssignments))).topicId("bar").get());
+
+            Map<Integer, List<Integer>> bazReplicasAssignments = new HashMap<>();
+            bazReplicasAssignments.put(0, asList(1, 0, 2));
+            bazReplicasAssignments.put(1, asList(2, 0, 1));
+            bazReplicasAssignments.put(2, asList(0, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("baz", bazReplicasAssignments))).topicId("baz").get());
+        }
+    }
+
+    private void produceMessages(String topic, int partition, int numMessages) {
+        Properties properties = new Properties();
+        properties.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        try (KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(properties, new ByteArraySerializer(), new ByteArraySerializer())) {
+            IntStream.range(0, numMessages).forEach(i -> {
+                ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topic, partition, null, new byte[10000]);
+                assertDoesNotThrow(() -> producer.send(record).get());
+            });
+        }
+    }
+
+    private void executeAndVerifyReassignment() throws ExecutionException, InterruptedException {
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3,2,0],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition bar0 = new TopicPartition("bar", 0);
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+
+        initialAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 3), true));
+        initialAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 0), true));
+
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(initialAssignment));
+
+        // Execute the assignment
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+        finalAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 3), asList(0, 1, 3), true));
+        finalAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 0), asList(3, 2, 0), true));
+
+        VerifyAssignmentResult verifyAssignmentResult = runVerifyAssignment(assignment, false);
+        assertFalse(verifyAssignmentResult.movesOngoing);
+
+        // Wait for the assignment to complete
+        waitForVerifyAssignment(assignment, false,
+                new VerifyAssignmentResult(finalAssignment));
+
+        assertEquals(unthrottledBrokerConfigs,
+                describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet()));
+
+        // Verify that partitions are removed from brokers no longer assigned
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 2));
+        verifyReplicaDeleted(new TopicPartitionReplica(bar0.topic(), bar0.partition(), 1));
+    }
+
+    private final Map<Integer, Map<String, Long>> unthrottledBrokerConfigs = new HashMap<>(); {
+        IntStream.range(0, 4).forEach(brokerId -> {
+            Map<String, Long> brokerConfig = new HashMap<>();
+
+            BROKER_LEVEL_THROTTLES.forEach(throttle -> brokerConfig.put(throttle, -1L));
+
+            unthrottledBrokerConfigs.put(brokerId, brokerConfig);
+        });
+    }
+
+    private void verifyReplicaDeleted(TopicPartitionReplica topicPartitionReplica) throws ExecutionException, InterruptedException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitUntilTrue(
+                    () -> {
+                        try {
+                            TopicDescription topicDescription = admin.describeTopics(singleton(topicPartitionReplica.topic())).topicNameValues().get(topicPartitionReplica.topic()).get();
+                            return topicDescription.partitions().stream().noneMatch(topicPartitionInfo -> {
+                                if (topicPartitionInfo.partition() != topicPartitionReplica.partition()) {
+                                    return false;
+                                }
+                                return topicPartitionInfo.replicas().stream().anyMatch(node -> node.id() == topicPartitionReplica.brokerId());
+                            });
+                        } catch (InterruptedException | ExecutionException e) {
+                            throw new RuntimeException(e);
+                        }
+                    },
+                    () -> "Timed out waiting for replica " + topicPartitionReplica.brokerId() + " of " + topicPartitionReplica + " to be deleted",
+                    DEFAULT_MAX_WAIT_MS,
+                    100L
+            );
+        }
+    }
+
+    private void waitForLogDirThrottle(Set<Integer> throttledBrokers, Long logDirThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, logDirThrottle);
+        waitForBrokerThrottles(throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForInterBrokerThrottle(List<Integer> throttledBrokers, Long interBrokerThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, -1L);
+        waitForBrokerThrottles(throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForBrokerThrottles(Collection<Integer> throttledBrokers, Map<String, Long> throttleConfig) {
+        Map<Integer, Map<String, Long>> throttledBrokerConfigs = new HashMap<>();
+        unthrottledBrokerConfigs.forEach((brokerId, unthrottledConfig) -> {
+            Map<String, Long> expectedThrottleConfig = throttledBrokers.contains(brokerId)
+                    ? throttleConfig
+                    : unthrottledConfig;
+            throttledBrokerConfigs.put(brokerId, expectedThrottleConfig);
+        });
+        waitForBrokerLevelThrottles(throttledBrokerConfigs);
+    }
+
+    private void waitForBrokerLevelThrottles(Map<Integer, Map<String, Long>> targetThrottles) {
+        AtomicReference<Map<Integer, Map<String, Long>>> curThrottles = new AtomicReference<>(new HashMap<>());
+        TestUtils.waitUntilTrue(() -> {
+            assertDoesNotThrow(() -> curThrottles.set(describeBrokerLevelThrottles(targetThrottles.keySet())));
+            return targetThrottles.equals(curThrottles.get());
+        }, () -> "timed out waiting for broker throttle to become " + targetThrottles + ".  " +
+                "Latest throttles were " + curThrottles.get(), DEFAULT_MAX_WAIT_MS, 25);
+    }
+
+    /**
+     * Describe the broker-level throttles in the cluster.
+     *
+     * @return                A map whose keys are broker IDs and whose values are throttle
+     *                        information.  The nested maps are keyed on throttle name.
+     */
+    private Map<Integer, Map<String, Long>> describeBrokerLevelThrottles(Collection<Integer> brokerIds) throws ExecutionException, InterruptedException {
+        Map<Integer, Map<String, Long>> results = new HashMap<>();
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            for (Integer brokerId : brokerIds) {
+                ConfigResource brokerResource = new ConfigResource(ConfigResource.Type.BROKER, brokerId.toString());
+                Config brokerConfigs = admin.describeConfigs(singleton(brokerResource)).values()
+                        .get(brokerResource)
+                        .get();
+
+                Map<String, Long> throttles = new HashMap<>();
+                BROKER_LEVEL_THROTTLES.forEach(throttleName -> {
+                    String configValue = Optional.ofNullable(brokerConfigs.get(throttleName)).map(ConfigEntry::value).orElse("-1");
+                    throttles.put(throttleName, Long.parseLong(configValue));
+                });
+                results.put(brokerId, throttles);
+            }
+        }
+        return results;
+    }
+
+    static class LogDirReassignment {
+        final String json;
+        final String currentDir;
+        final String targetDir;
+
+        public LogDirReassignment(String json, String currentDir, String targetDir) {
+            this.json = json;
+            this.currentDir = currentDir;
+            this.targetDir = targetDir;
+        }
+    }
+
+    private ReassignPartitionsIntegrationTest.LogDirReassignment buildLogDirReassignment(TopicPartition topicPartition,
+                                                                                         int brokerId,
+                                                                                         List<Integer> replicas) throws ExecutionException, InterruptedException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            DescribeLogDirsResult describeLogDirsResult = admin.describeLogDirs(
+                    IntStream.range(0, 4).boxed().collect(Collectors.toList()));
+
+            ReassignPartitionsIntegrationTest.BrokerDirs logDirInfo = new ReassignPartitionsIntegrationTest.BrokerDirs(describeLogDirsResult, brokerId);
+            assertTrue(logDirInfo.futureLogDirs.isEmpty());
+
+            String currentDir = logDirInfo.curLogDirs.get(topicPartition);
+            String newDir = logDirInfo.logDirs.stream().filter(dir -> !dir.equals(currentDir)).findFirst().get();
+
+            List<String> logDirs = replicas.stream().map(replicaId -> {
+                if (replicaId == brokerId)
+                    return "\"" + newDir + "\"";
+                else
+                    return "\"any\"";
+            }).collect(Collectors.toList());
+
+            String reassignmentJson =
+                    " { \"version\": 1," +
+                            "  \"partitions\": [" +
+                            "    {" +
+                            "     \"topic\": \"" + topicPartition.topic() + "\"," +
+                            "     \"partition\": " + topicPartition.partition() + "," +
+                            "     \"replicas\": [" + replicas.stream().map(Object::toString).collect(Collectors.joining(",")) + "]," +
+                            "     \"log_dirs\": [" + String.join(",", logDirs) + "]" +
+                            "    }" +
+                            "   ]" +
+                            "  }";
+
+            return new ReassignPartitionsIntegrationTest.LogDirReassignment(reassignmentJson, currentDir, newDir);
+        }
+    }
+
+    private VerifyAssignmentResult runVerifyAssignment(String jsonString,
+                                                       Boolean preserveThrottles) {
+        System.out.println("==> verifyAssignment(adminClient, jsonString=" + jsonString);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            return verifyAssignment(admin, jsonString, preserveThrottles);
+        } catch (ExecutionException | InterruptedException | JsonProcessingException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private void waitForVerifyAssignment(String jsonString,
+                                         Boolean preserveThrottles,
+                                         VerifyAssignmentResult expectedResult) {
+        final VerifyAssignmentResult[] latestResult = {null};
+        TestUtils.waitUntilTrue(

Review Comment:
   ditto



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: migrate ReassignPartitionsIntegrationTest to use ClusterTestExtensions [kafka]

Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.

FrankYang0529 commented on PR #15675:
URL: https://github.com/apache/kafka/pull/15675#issuecomment-2146578561

   @chia7712 I address all comments. Thanks for the review.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: migrate ReassignPartitionsIntegrationTest to use ClusterTestExtensions [kafka]

Posted by "chia7712 (via GitHub)" <gi...@apache.org>.

chia7712 commented on code in PR #15675:
URL: https://github.com/apache/kafka/pull/15675#discussion_r1625194534


##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,811 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.time.Duration;
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testHighWaterMarkAfterPartitionReassignment() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 100);
+
+        // Execute the assignment
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[3,1,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, -1L, -1L);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(foo0,
+                    new PartitionReassignmentState(asList(3, 1, 2), asList(3, 1, 2), true));
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            TestUtils.waitForCondition(() -> {
+                ListOffsetsResultInfo result = admin.listOffsets(Collections.singletonMap(foo0, new OffsetSpec.LatestSpec())).partitionResult(foo0).get();
+                return result.offset() == 100;
+            }, "Timeout for waiting offset");
+        }
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Execute the assignment with a low throttle
+            long initialThrottle = 1L;
+            runExecuteAssignment(false, assignment, initialThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), initialThrottle);
+
+            // Now update the throttle and verify the reassignment completes
+            long updatedThrottle = 300000L;
+            runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), updatedThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(admin, assignment, false));
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+
+            // Execute the assignment
+            long interBrokerThrottle = 300000L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Wait for the assignment to complete
+            TestUtils.waitForCondition(() -> {
+                // Check the reassignment status.
+                VerifyAssignmentResult result = runVerifyAssignment(admin, assignment, true);
+
+                if (!result.partsOngoing) {
+                    return true;
+                } else {
+                    assertFalse(
+                            result.partStates.values().stream().allMatch(state -> state.done),
+                            "Expected at least one partition reassignment to be ongoing when result = " + result
+                    );
+                    assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                    assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                    waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+                    return false;
+                }
+            }, "Expected reassignment to complete.");
+
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(finalAssignment));
+            // The throttles should still have been preserved, since we ran with --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            List<ConsumerRecord<byte[], byte[]>> allRecords = new ArrayList<>();
+            TestUtils.waitForCondition(() -> {
+                ConsumerRecords<byte[], byte[]> records = consumer.poll(Duration.ofMillis(100L));
+                records.records("baz").forEach(allRecords::add);
+                return allRecords.size() >= 100;
+            }, "Timeout for waiting enough records");
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers())))  {
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+        }
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(unthrottledBrokerConfigs,
+                    describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+            long interBrokerThrottle = 1L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+            partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+            partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+            // Verify that the reassignment is running.  The very low throttle should keep it
+            // from completing before this runs.
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+            // Cancel the reassignment.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+            // Broker throttles are still active because we passed --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Cancelling the reassignment again should reveal nothing to cancel.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+            // This time, the broker throttles were removed.
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+            // Verify that there are no ongoing reassignments.
+            assertFalse(runVerifyAssignment(admin, assignment, false).partsOngoing);
+        }
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitForCondition(
+                    () -> {
+                        Set<Integer> isr = admin.describeTopics(Collections.singleton(foo0.topic()))
+                                .allTopicNames().get().get(foo0.topic()).partitions().stream()
+                                .filter(p -> p.partition() == foo0.partition())
+                                .flatMap(p -> p.isr().stream())
+                                .map(Node::id).collect(Collectors.toSet());
+                        return isr.containsAll(Arrays.asList(0, 1, 2, 3));
+                    },
+                    "Timed out while waiting for replica 3 to join the ISR"
+            );
+        }
+
+        // Now cancel the assignment and verify that the partition is removed from cancelled replicas
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(singleton(foo0), Collections.emptySet()), runCancelAssignment(assignment, true));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 4));
+    }
+
+    /**
+     * Test moving partitions between directories.
+     */
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testLogDirReassignment() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move, but throttle it to be very slow so that it can't complete
+        // before our next checks happen.
+        long logDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, logDirThrottle);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Check the output of --verify
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new ActiveMoveState(reassignment.currentDir, reassignment.targetDir, reassignment.targetDir)
+                    ), true));
+            waitForLogDirThrottle(admin, singleton(0), logDirThrottle);
+
+            // Remove the throttle
+            admin.incrementalAlterConfigs(singletonMap(
+                            new ConfigResource(ConfigResource.Type.BROKER, "0"),
+                            singletonList(new AlterConfigOp(
+                                    new ConfigEntry(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, ""), AlterConfigOp.OpType.DELETE))))
+                    .all().get();
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+
+            // Wait for the directory movement to complete.
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+
+            BrokerDirs info1 = new BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
+            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
+        }
+    }
+
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testAlterLogDirReassignmentThrottle() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move with a low throttle so it does not complete
+        long initialLogDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            waitForLogDirThrottle(admin, new HashSet<>(singletonList(0)), initialLogDirThrottle);
+
+            // Now increase the throttle and verify that the log dir movement completes
+            long updatedLogDirThrottle = 3000000L;
+            runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
+            waitForLogDirThrottle(admin, singleton(0), updatedLogDirThrottle);
+
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+        }
+    }
+
+    private void createTopics() {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<Integer, List<Integer>> fooReplicasAssignments = new HashMap<>();
+            fooReplicasAssignments.put(0, asList(0, 1, 2));
+            fooReplicasAssignments.put(1, asList(1, 2, 3));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("foo", fooReplicasAssignments))).topicId("foo").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("foo", fooReplicasAssignments.size()));
+
+            Map<Integer, List<Integer>> barReplicasAssignments = new HashMap<>();
+            barReplicasAssignments.put(0, asList(3, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("bar", barReplicasAssignments))).topicId("bar").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("bar", barReplicasAssignments.size()));
+
+            Map<Integer, List<Integer>> bazReplicasAssignments = new HashMap<>();
+            bazReplicasAssignments.put(0, asList(1, 0, 2));
+            bazReplicasAssignments.put(1, asList(2, 0, 1));
+            bazReplicasAssignments.put(2, asList(0, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("baz", bazReplicasAssignments))).topicId("baz").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("baz", bazReplicasAssignments.size()));
+        }
+    }
+
+    private void produceMessages(String topic, int partition, int numMessages) {
+        Properties properties = new Properties();
+        properties.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        try (KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(properties, new ByteArraySerializer(), new ByteArraySerializer())) {
+            IntStream.range(0, numMessages).forEach(i -> {
+                ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topic, partition, null, new byte[10000]);
+                assertDoesNotThrow(() -> producer.send(record).get());
+            });
+        }
+    }
+
+    private void executeAndVerifyReassignment() throws ExecutionException, InterruptedException {

Review Comment:
   `ExecutionException` is unused 



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,811 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.time.Duration;
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testHighWaterMarkAfterPartitionReassignment() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 100);
+
+        // Execute the assignment
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[3,1,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, -1L, -1L);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(foo0,
+                    new PartitionReassignmentState(asList(3, 1, 2), asList(3, 1, 2), true));
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            TestUtils.waitForCondition(() -> {
+                ListOffsetsResultInfo result = admin.listOffsets(Collections.singletonMap(foo0, new OffsetSpec.LatestSpec())).partitionResult(foo0).get();
+                return result.offset() == 100;
+            }, "Timeout for waiting offset");
+        }
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Execute the assignment with a low throttle
+            long initialThrottle = 1L;
+            runExecuteAssignment(false, assignment, initialThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), initialThrottle);
+
+            // Now update the throttle and verify the reassignment completes
+            long updatedThrottle = 300000L;
+            runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), updatedThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(admin, assignment, false));
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+
+            // Execute the assignment
+            long interBrokerThrottle = 300000L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Wait for the assignment to complete
+            TestUtils.waitForCondition(() -> {
+                // Check the reassignment status.
+                VerifyAssignmentResult result = runVerifyAssignment(admin, assignment, true);
+
+                if (!result.partsOngoing) {
+                    return true;
+                } else {
+                    assertFalse(
+                            result.partStates.values().stream().allMatch(state -> state.done),
+                            "Expected at least one partition reassignment to be ongoing when result = " + result
+                    );
+                    assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                    assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                    waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+                    return false;
+                }
+            }, "Expected reassignment to complete.");
+
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(finalAssignment));
+            // The throttles should still have been preserved, since we ran with --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            List<ConsumerRecord<byte[], byte[]>> allRecords = new ArrayList<>();
+            TestUtils.waitForCondition(() -> {
+                ConsumerRecords<byte[], byte[]> records = consumer.poll(Duration.ofMillis(100L));
+                records.records("baz").forEach(allRecords::add);
+                return allRecords.size() >= 100;
+            }, "Timeout for waiting enough records");
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers())))  {
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+        }
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(unthrottledBrokerConfigs,
+                    describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+            long interBrokerThrottle = 1L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+            partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+            partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+            // Verify that the reassignment is running.  The very low throttle should keep it
+            // from completing before this runs.
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+            // Cancel the reassignment.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+            // Broker throttles are still active because we passed --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Cancelling the reassignment again should reveal nothing to cancel.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+            // This time, the broker throttles were removed.
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+            // Verify that there are no ongoing reassignments.
+            assertFalse(runVerifyAssignment(admin, assignment, false).partsOngoing);
+        }
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitForCondition(
+                    () -> {
+                        Set<Integer> isr = admin.describeTopics(Collections.singleton(foo0.topic()))
+                                .allTopicNames().get().get(foo0.topic()).partitions().stream()
+                                .filter(p -> p.partition() == foo0.partition())
+                                .flatMap(p -> p.isr().stream())
+                                .map(Node::id).collect(Collectors.toSet());
+                        return isr.containsAll(Arrays.asList(0, 1, 2, 3));
+                    },
+                    "Timed out while waiting for replica 3 to join the ISR"
+            );
+        }
+
+        // Now cancel the assignment and verify that the partition is removed from cancelled replicas
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(singleton(foo0), Collections.emptySet()), runCancelAssignment(assignment, true));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 4));
+    }
+
+    /**
+     * Test moving partitions between directories.
+     */
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testLogDirReassignment() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move, but throttle it to be very slow so that it can't complete
+        // before our next checks happen.
+        long logDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, logDirThrottle);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Check the output of --verify
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new ActiveMoveState(reassignment.currentDir, reassignment.targetDir, reassignment.targetDir)
+                    ), true));
+            waitForLogDirThrottle(admin, singleton(0), logDirThrottle);
+
+            // Remove the throttle
+            admin.incrementalAlterConfigs(singletonMap(
+                            new ConfigResource(ConfigResource.Type.BROKER, "0"),
+                            singletonList(new AlterConfigOp(
+                                    new ConfigEntry(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, ""), AlterConfigOp.OpType.DELETE))))
+                    .all().get();
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+
+            // Wait for the directory movement to complete.
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+
+            BrokerDirs info1 = new BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
+            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
+        }
+    }
+
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testAlterLogDirReassignmentThrottle() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move with a low throttle so it does not complete
+        long initialLogDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            waitForLogDirThrottle(admin, new HashSet<>(singletonList(0)), initialLogDirThrottle);
+
+            // Now increase the throttle and verify that the log dir movement completes
+            long updatedLogDirThrottle = 3000000L;
+            runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
+            waitForLogDirThrottle(admin, singleton(0), updatedLogDirThrottle);
+
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+        }
+    }
+
+    private void createTopics() {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<Integer, List<Integer>> fooReplicasAssignments = new HashMap<>();
+            fooReplicasAssignments.put(0, asList(0, 1, 2));
+            fooReplicasAssignments.put(1, asList(1, 2, 3));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("foo", fooReplicasAssignments))).topicId("foo").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("foo", fooReplicasAssignments.size()));
+
+            Map<Integer, List<Integer>> barReplicasAssignments = new HashMap<>();
+            barReplicasAssignments.put(0, asList(3, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("bar", barReplicasAssignments))).topicId("bar").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("bar", barReplicasAssignments.size()));
+
+            Map<Integer, List<Integer>> bazReplicasAssignments = new HashMap<>();
+            bazReplicasAssignments.put(0, asList(1, 0, 2));
+            bazReplicasAssignments.put(1, asList(2, 0, 1));
+            bazReplicasAssignments.put(2, asList(0, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("baz", bazReplicasAssignments))).topicId("baz").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("baz", bazReplicasAssignments.size()));
+        }
+    }
+
+    private void produceMessages(String topic, int partition, int numMessages) {
+        Properties properties = new Properties();
+        properties.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        try (KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(properties, new ByteArraySerializer(), new ByteArraySerializer())) {
+            IntStream.range(0, numMessages).forEach(i -> {
+                ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topic, partition, null, new byte[10000]);
+                assertDoesNotThrow(() -> producer.send(record).get());
+            });
+        }
+    }
+
+    private void executeAndVerifyReassignment() throws ExecutionException, InterruptedException {
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3,2,0],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition bar0 = new TopicPartition("bar", 0);
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+
+        initialAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 3), true));
+        initialAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 0), true));
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(initialAssignment));
+
+            // Execute the assignment
+            runExecuteAssignment(false, assignment, -1L, -1L);
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 3), asList(0, 1, 3), true));
+            finalAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 0), asList(3, 2, 0), true));
+
+            VerifyAssignmentResult verifyAssignmentResult = runVerifyAssignment(admin, assignment, false);
+            assertFalse(verifyAssignmentResult.movesOngoing);
+
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+
+            assertEquals(unthrottledBrokerConfigs,
+                    describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+        }
+
+        // Verify that partitions are removed from brokers no longer assigned
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 2));
+        verifyReplicaDeleted(new TopicPartitionReplica(bar0.topic(), bar0.partition(), 1));
+    }
+
+    private final Map<Integer, Map<String, Long>> unthrottledBrokerConfigs = new HashMap<>(); {
+        IntStream.range(0, 4).forEach(brokerId -> {
+            Map<String, Long> brokerConfig = new HashMap<>();
+
+            BROKER_LEVEL_THROTTLES.forEach(throttle -> brokerConfig.put(throttle, -1L));
+
+            unthrottledBrokerConfigs.put(brokerId, brokerConfig);
+        });
+    }
+
+    private void verifyReplicaDeleted(TopicPartitionReplica topicPartitionReplica) throws InterruptedException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitForCondition(
+                    () -> {
+                        try {
+                            TopicDescription topicDescription = admin.describeTopics(singleton(topicPartitionReplica.topic())).topicNameValues().get(topicPartitionReplica.topic()).get();

Review Comment:
   We can use `assertDoesNotThrow` to convert the checked exception to unchecked version.



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,811 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.time.Duration;
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testHighWaterMarkAfterPartitionReassignment() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 100);
+
+        // Execute the assignment
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[3,1,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, -1L, -1L);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(foo0,
+                    new PartitionReassignmentState(asList(3, 1, 2), asList(3, 1, 2), true));
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            TestUtils.waitForCondition(() -> {
+                ListOffsetsResultInfo result = admin.listOffsets(Collections.singletonMap(foo0, new OffsetSpec.LatestSpec())).partitionResult(foo0).get();
+                return result.offset() == 100;
+            }, "Timeout for waiting offset");
+        }
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Execute the assignment with a low throttle
+            long initialThrottle = 1L;
+            runExecuteAssignment(false, assignment, initialThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), initialThrottle);
+
+            // Now update the throttle and verify the reassignment completes
+            long updatedThrottle = 300000L;
+            runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), updatedThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(admin, assignment, false));
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+
+            // Execute the assignment
+            long interBrokerThrottle = 300000L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Wait for the assignment to complete
+            TestUtils.waitForCondition(() -> {
+                // Check the reassignment status.
+                VerifyAssignmentResult result = runVerifyAssignment(admin, assignment, true);
+
+                if (!result.partsOngoing) {
+                    return true;
+                } else {
+                    assertFalse(
+                            result.partStates.values().stream().allMatch(state -> state.done),
+                            "Expected at least one partition reassignment to be ongoing when result = " + result
+                    );
+                    assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                    assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                    waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+                    return false;
+                }
+            }, "Expected reassignment to complete.");
+
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(finalAssignment));
+            // The throttles should still have been preserved, since we ran with --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            List<ConsumerRecord<byte[], byte[]>> allRecords = new ArrayList<>();
+            TestUtils.waitForCondition(() -> {
+                ConsumerRecords<byte[], byte[]> records = consumer.poll(Duration.ofMillis(100L));
+                records.records("baz").forEach(allRecords::add);
+                return allRecords.size() >= 100;
+            }, "Timeout for waiting enough records");
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers())))  {
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+        }
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(unthrottledBrokerConfigs,
+                    describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+            long interBrokerThrottle = 1L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+            partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+            partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+            // Verify that the reassignment is running.  The very low throttle should keep it
+            // from completing before this runs.
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+            // Cancel the reassignment.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+            // Broker throttles are still active because we passed --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Cancelling the reassignment again should reveal nothing to cancel.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+            // This time, the broker throttles were removed.
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+            // Verify that there are no ongoing reassignments.
+            assertFalse(runVerifyAssignment(admin, assignment, false).partsOngoing);
+        }
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitForCondition(
+                    () -> {
+                        Set<Integer> isr = admin.describeTopics(Collections.singleton(foo0.topic()))
+                                .allTopicNames().get().get(foo0.topic()).partitions().stream()
+                                .filter(p -> p.partition() == foo0.partition())
+                                .flatMap(p -> p.isr().stream())
+                                .map(Node::id).collect(Collectors.toSet());
+                        return isr.containsAll(Arrays.asList(0, 1, 2, 3));
+                    },
+                    "Timed out while waiting for replica 3 to join the ISR"
+            );
+        }
+
+        // Now cancel the assignment and verify that the partition is removed from cancelled replicas
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(singleton(foo0), Collections.emptySet()), runCancelAssignment(assignment, true));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 4));
+    }
+
+    /**
+     * Test moving partitions between directories.
+     */
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testLogDirReassignment() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move, but throttle it to be very slow so that it can't complete
+        // before our next checks happen.
+        long logDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, logDirThrottle);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Check the output of --verify
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new ActiveMoveState(reassignment.currentDir, reassignment.targetDir, reassignment.targetDir)
+                    ), true));
+            waitForLogDirThrottle(admin, singleton(0), logDirThrottle);
+
+            // Remove the throttle
+            admin.incrementalAlterConfigs(singletonMap(
+                            new ConfigResource(ConfigResource.Type.BROKER, "0"),
+                            singletonList(new AlterConfigOp(
+                                    new ConfigEntry(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, ""), AlterConfigOp.OpType.DELETE))))
+                    .all().get();
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+
+            // Wait for the directory movement to complete.
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+
+            BrokerDirs info1 = new BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
+            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
+        }
+    }
+
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testAlterLogDirReassignmentThrottle() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move with a low throttle so it does not complete
+        long initialLogDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            waitForLogDirThrottle(admin, new HashSet<>(singletonList(0)), initialLogDirThrottle);
+
+            // Now increase the throttle and verify that the log dir movement completes
+            long updatedLogDirThrottle = 3000000L;
+            runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
+            waitForLogDirThrottle(admin, singleton(0), updatedLogDirThrottle);
+
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+        }
+    }
+
+    private void createTopics() {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<Integer, List<Integer>> fooReplicasAssignments = new HashMap<>();
+            fooReplicasAssignments.put(0, asList(0, 1, 2));
+            fooReplicasAssignments.put(1, asList(1, 2, 3));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("foo", fooReplicasAssignments))).topicId("foo").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("foo", fooReplicasAssignments.size()));
+
+            Map<Integer, List<Integer>> barReplicasAssignments = new HashMap<>();
+            barReplicasAssignments.put(0, asList(3, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("bar", barReplicasAssignments))).topicId("bar").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("bar", barReplicasAssignments.size()));
+
+            Map<Integer, List<Integer>> bazReplicasAssignments = new HashMap<>();
+            bazReplicasAssignments.put(0, asList(1, 0, 2));
+            bazReplicasAssignments.put(1, asList(2, 0, 1));
+            bazReplicasAssignments.put(2, asList(0, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("baz", bazReplicasAssignments))).topicId("baz").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("baz", bazReplicasAssignments.size()));
+        }
+    }
+
+    private void produceMessages(String topic, int partition, int numMessages) {
+        Properties properties = new Properties();
+        properties.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        try (KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(properties, new ByteArraySerializer(), new ByteArraySerializer())) {
+            IntStream.range(0, numMessages).forEach(i -> {
+                ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topic, partition, null, new byte[10000]);
+                assertDoesNotThrow(() -> producer.send(record).get());
+            });
+        }
+    }
+
+    private void executeAndVerifyReassignment() throws ExecutionException, InterruptedException {
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3,2,0],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition bar0 = new TopicPartition("bar", 0);
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+
+        initialAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 3), true));
+        initialAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 0), true));
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(initialAssignment));
+
+            // Execute the assignment
+            runExecuteAssignment(false, assignment, -1L, -1L);
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 3), asList(0, 1, 3), true));
+            finalAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 0), asList(3, 2, 0), true));
+
+            VerifyAssignmentResult verifyAssignmentResult = runVerifyAssignment(admin, assignment, false);
+            assertFalse(verifyAssignmentResult.movesOngoing);
+
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+
+            assertEquals(unthrottledBrokerConfigs,
+                    describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+        }
+
+        // Verify that partitions are removed from brokers no longer assigned
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 2));
+        verifyReplicaDeleted(new TopicPartitionReplica(bar0.topic(), bar0.partition(), 1));
+    }
+
+    private final Map<Integer, Map<String, Long>> unthrottledBrokerConfigs = new HashMap<>(); {
+        IntStream.range(0, 4).forEach(brokerId -> {
+            Map<String, Long> brokerConfig = new HashMap<>();
+
+            BROKER_LEVEL_THROTTLES.forEach(throttle -> brokerConfig.put(throttle, -1L));
+
+            unthrottledBrokerConfigs.put(brokerId, brokerConfig);
+        });
+    }
+
+    private void verifyReplicaDeleted(TopicPartitionReplica topicPartitionReplica) throws InterruptedException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitForCondition(
+                    () -> {
+                        try {
+                            TopicDescription topicDescription = admin.describeTopics(singleton(topicPartitionReplica.topic())).topicNameValues().get(topicPartitionReplica.topic()).get();
+                            return topicDescription.partitions().stream().noneMatch(topicPartitionInfo -> {
+                                if (topicPartitionInfo.partition() != topicPartitionReplica.partition()) {
+                                    return false;
+                                }
+                                return topicPartitionInfo.replicas().stream().anyMatch(node -> node.id() == topicPartitionReplica.brokerId());
+                            });
+                        } catch (InterruptedException | ExecutionException e) {
+                            throw new RuntimeException(e);
+                        }
+                    },
+                    "Timed out waiting for replica " + topicPartitionReplica.brokerId() + " of " + topicPartitionReplica + " to be deleted"
+            );
+        }
+    }
+
+    private void waitForLogDirThrottle(Admin admin, Set<Integer> throttledBrokers, Long logDirThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, logDirThrottle);
+        waitForBrokerThrottles(admin, throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForInterBrokerThrottle(Admin admin, List<Integer> throttledBrokers, Long interBrokerThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, -1L);
+        waitForBrokerThrottles(admin, throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForBrokerThrottles(Admin admin, Collection<Integer> throttledBrokers, Map<String, Long> throttleConfig) {
+        Map<Integer, Map<String, Long>> throttledBrokerConfigs = new HashMap<>();
+        unthrottledBrokerConfigs.forEach((brokerId, unthrottledConfig) -> {
+            Map<String, Long> expectedThrottleConfig = throttledBrokers.contains(brokerId)
+                    ? throttleConfig
+                    : unthrottledConfig;
+            throttledBrokerConfigs.put(brokerId, expectedThrottleConfig);
+        });
+        Assertions.assertDoesNotThrow(() -> waitForBrokerLevelThrottles(admin, throttledBrokerConfigs));
+    }
+
+    private void waitForBrokerLevelThrottles(Admin admin, Map<Integer, Map<String, Long>> targetThrottles) throws InterruptedException {
+        AtomicReference<Map<Integer, Map<String, Long>>> curThrottles = new AtomicReference<>(new HashMap<>());
+        TestUtils.waitForCondition(() -> {
+            assertDoesNotThrow(() -> curThrottles.set(describeBrokerLevelThrottles(admin, targetThrottles.keySet())));
+            return targetThrottles.equals(curThrottles.get());
+        }, "timed out waiting for broker throttle to become " + targetThrottles + ".  " +
+                "Latest throttles were " + curThrottles.get());
+    }
+
+    /**
+     * Describe the broker-level throttles in the cluster.
+     *
+     * @return                A map whose keys are broker IDs and whose values are throttle
+     *                        information.  The nested maps are keyed on throttle name.
+     */
+    private Map<Integer, Map<String, Long>> describeBrokerLevelThrottles(Admin admin, Collection<Integer> brokerIds) throws ExecutionException, InterruptedException {

Review Comment:
   Could we simplify this function by following style?
   ```java
       private Map<Integer, Map<String, Long>> describeBrokerLevelThrottles(Admin admin, Collection<Integer> brokerIds) {
           return brokerIds.stream().collect(Collectors.toMap(Function.identity(), brokerId -> {
               ConfigResource brokerResource = new ConfigResource(ConfigResource.Type.BROKER, brokerId.toString());
               Config brokerConfigs = assertDoesNotThrow(() -> admin.describeConfigs(singleton(brokerResource)).values()
                   .get(brokerResource)
                   .get());
               return BROKER_LEVEL_THROTTLES.stream().collect(Collectors.toMap(Function.identity(),
                   name -> Optional.ofNullable(brokerConfigs.get(name)).map(e -> Long.parseLong(e.value())).orElse(-1L)));
           }));
       }
   ```



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,811 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.time.Duration;
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),

Review Comment:
   not sure which test cases require this config. Could you point out it for me?



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,811 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.time.Duration;
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testHighWaterMarkAfterPartitionReassignment() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 100);
+
+        // Execute the assignment
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[3,1,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, -1L, -1L);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(foo0,
+                    new PartitionReassignmentState(asList(3, 1, 2), asList(3, 1, 2), true));
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            TestUtils.waitForCondition(() -> {
+                ListOffsetsResultInfo result = admin.listOffsets(Collections.singletonMap(foo0, new OffsetSpec.LatestSpec())).partitionResult(foo0).get();
+                return result.offset() == 100;
+            }, "Timeout for waiting offset");
+        }
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Execute the assignment with a low throttle
+            long initialThrottle = 1L;
+            runExecuteAssignment(false, assignment, initialThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), initialThrottle);
+
+            // Now update the throttle and verify the reassignment completes
+            long updatedThrottle = 300000L;
+            runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), updatedThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(admin, assignment, false));
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+
+            // Execute the assignment
+            long interBrokerThrottle = 300000L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Wait for the assignment to complete
+            TestUtils.waitForCondition(() -> {
+                // Check the reassignment status.
+                VerifyAssignmentResult result = runVerifyAssignment(admin, assignment, true);
+
+                if (!result.partsOngoing) {
+                    return true;
+                } else {
+                    assertFalse(
+                            result.partStates.values().stream().allMatch(state -> state.done),
+                            "Expected at least one partition reassignment to be ongoing when result = " + result
+                    );
+                    assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                    assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                    waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+                    return false;
+                }
+            }, "Expected reassignment to complete.");
+
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(finalAssignment));
+            // The throttles should still have been preserved, since we ran with --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            List<ConsumerRecord<byte[], byte[]>> allRecords = new ArrayList<>();
+            TestUtils.waitForCondition(() -> {
+                ConsumerRecords<byte[], byte[]> records = consumer.poll(Duration.ofMillis(100L));
+                records.records("baz").forEach(allRecords::add);
+                return allRecords.size() >= 100;
+            }, "Timeout for waiting enough records");
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers())))  {
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+        }
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(unthrottledBrokerConfigs,
+                    describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+            long interBrokerThrottle = 1L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+            partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+            partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+            // Verify that the reassignment is running.  The very low throttle should keep it
+            // from completing before this runs.
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+            // Cancel the reassignment.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+            // Broker throttles are still active because we passed --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Cancelling the reassignment again should reveal nothing to cancel.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+            // This time, the broker throttles were removed.
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+            // Verify that there are no ongoing reassignments.
+            assertFalse(runVerifyAssignment(admin, assignment, false).partsOngoing);
+        }
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitForCondition(
+                    () -> {
+                        Set<Integer> isr = admin.describeTopics(Collections.singleton(foo0.topic()))
+                                .allTopicNames().get().get(foo0.topic()).partitions().stream()
+                                .filter(p -> p.partition() == foo0.partition())
+                                .flatMap(p -> p.isr().stream())
+                                .map(Node::id).collect(Collectors.toSet());
+                        return isr.containsAll(Arrays.asList(0, 1, 2, 3));
+                    },
+                    "Timed out while waiting for replica 3 to join the ISR"
+            );
+        }
+
+        // Now cancel the assignment and verify that the partition is removed from cancelled replicas
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(singleton(foo0), Collections.emptySet()), runCancelAssignment(assignment, true));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 4));
+    }
+
+    /**
+     * Test moving partitions between directories.
+     */
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testLogDirReassignment() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move, but throttle it to be very slow so that it can't complete
+        // before our next checks happen.
+        long logDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, logDirThrottle);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Check the output of --verify
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new ActiveMoveState(reassignment.currentDir, reassignment.targetDir, reassignment.targetDir)
+                    ), true));
+            waitForLogDirThrottle(admin, singleton(0), logDirThrottle);
+
+            // Remove the throttle
+            admin.incrementalAlterConfigs(singletonMap(
+                            new ConfigResource(ConfigResource.Type.BROKER, "0"),
+                            singletonList(new AlterConfigOp(
+                                    new ConfigEntry(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, ""), AlterConfigOp.OpType.DELETE))))
+                    .all().get();
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+
+            // Wait for the directory movement to complete.
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+
+            BrokerDirs info1 = new BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
+            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
+        }
+    }
+
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testAlterLogDirReassignmentThrottle() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move with a low throttle so it does not complete
+        long initialLogDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            waitForLogDirThrottle(admin, new HashSet<>(singletonList(0)), initialLogDirThrottle);
+
+            // Now increase the throttle and verify that the log dir movement completes
+            long updatedLogDirThrottle = 3000000L;
+            runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
+            waitForLogDirThrottle(admin, singleton(0), updatedLogDirThrottle);
+
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+        }
+    }
+
+    private void createTopics() {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<Integer, List<Integer>> fooReplicasAssignments = new HashMap<>();
+            fooReplicasAssignments.put(0, asList(0, 1, 2));
+            fooReplicasAssignments.put(1, asList(1, 2, 3));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("foo", fooReplicasAssignments))).topicId("foo").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("foo", fooReplicasAssignments.size()));
+
+            Map<Integer, List<Integer>> barReplicasAssignments = new HashMap<>();
+            barReplicasAssignments.put(0, asList(3, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("bar", barReplicasAssignments))).topicId("bar").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("bar", barReplicasAssignments.size()));
+
+            Map<Integer, List<Integer>> bazReplicasAssignments = new HashMap<>();
+            bazReplicasAssignments.put(0, asList(1, 0, 2));
+            bazReplicasAssignments.put(1, asList(2, 0, 1));
+            bazReplicasAssignments.put(2, asList(0, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("baz", bazReplicasAssignments))).topicId("baz").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("baz", bazReplicasAssignments.size()));
+        }
+    }
+
+    private void produceMessages(String topic, int partition, int numMessages) {
+        Properties properties = new Properties();
+        properties.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        try (KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(properties, new ByteArraySerializer(), new ByteArraySerializer())) {
+            IntStream.range(0, numMessages).forEach(i -> {
+                ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topic, partition, null, new byte[10000]);
+                assertDoesNotThrow(() -> producer.send(record).get());
+            });
+        }
+    }
+
+    private void executeAndVerifyReassignment() throws ExecutionException, InterruptedException {
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3,2,0],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition bar0 = new TopicPartition("bar", 0);
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+
+        initialAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 3), true));
+        initialAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 0), true));
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(initialAssignment));
+
+            // Execute the assignment
+            runExecuteAssignment(false, assignment, -1L, -1L);
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 3), asList(0, 1, 3), true));
+            finalAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 0), asList(3, 2, 0), true));
+
+            VerifyAssignmentResult verifyAssignmentResult = runVerifyAssignment(admin, assignment, false);
+            assertFalse(verifyAssignmentResult.movesOngoing);
+
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+
+            assertEquals(unthrottledBrokerConfigs,
+                    describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+        }
+
+        // Verify that partitions are removed from brokers no longer assigned
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 2));
+        verifyReplicaDeleted(new TopicPartitionReplica(bar0.topic(), bar0.partition(), 1));
+    }
+
+    private final Map<Integer, Map<String, Long>> unthrottledBrokerConfigs = new HashMap<>(); {

Review Comment:
   this code style is uncommon. Could we rewrite it by lambda and move it up to line#113?



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,811 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.time.Duration;
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testHighWaterMarkAfterPartitionReassignment() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 100);
+
+        // Execute the assignment
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[3,1,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, -1L, -1L);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(foo0,
+                    new PartitionReassignmentState(asList(3, 1, 2), asList(3, 1, 2), true));
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            TestUtils.waitForCondition(() -> {
+                ListOffsetsResultInfo result = admin.listOffsets(Collections.singletonMap(foo0, new OffsetSpec.LatestSpec())).partitionResult(foo0).get();
+                return result.offset() == 100;
+            }, "Timeout for waiting offset");
+        }
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Execute the assignment with a low throttle
+            long initialThrottle = 1L;
+            runExecuteAssignment(false, assignment, initialThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), initialThrottle);
+
+            // Now update the throttle and verify the reassignment completes
+            long updatedThrottle = 300000L;
+            runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), updatedThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(admin, assignment, false));
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+
+            // Execute the assignment
+            long interBrokerThrottle = 300000L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Wait for the assignment to complete
+            TestUtils.waitForCondition(() -> {
+                // Check the reassignment status.
+                VerifyAssignmentResult result = runVerifyAssignment(admin, assignment, true);
+
+                if (!result.partsOngoing) {
+                    return true;
+                } else {
+                    assertFalse(
+                            result.partStates.values().stream().allMatch(state -> state.done),
+                            "Expected at least one partition reassignment to be ongoing when result = " + result
+                    );
+                    assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                    assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                    waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+                    return false;
+                }
+            }, "Expected reassignment to complete.");
+
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(finalAssignment));
+            // The throttles should still have been preserved, since we ran with --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            List<ConsumerRecord<byte[], byte[]>> allRecords = new ArrayList<>();
+            TestUtils.waitForCondition(() -> {
+                ConsumerRecords<byte[], byte[]> records = consumer.poll(Duration.ofMillis(100L));
+                records.records("baz").forEach(allRecords::add);
+                return allRecords.size() >= 100;
+            }, "Timeout for waiting enough records");
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers())))  {
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+        }
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(unthrottledBrokerConfigs,
+                    describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+            long interBrokerThrottle = 1L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+            partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+            partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+            // Verify that the reassignment is running.  The very low throttle should keep it
+            // from completing before this runs.
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+            // Cancel the reassignment.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+            // Broker throttles are still active because we passed --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Cancelling the reassignment again should reveal nothing to cancel.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+            // This time, the broker throttles were removed.
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+            // Verify that there are no ongoing reassignments.
+            assertFalse(runVerifyAssignment(admin, assignment, false).partsOngoing);
+        }
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitForCondition(
+                    () -> {
+                        Set<Integer> isr = admin.describeTopics(Collections.singleton(foo0.topic()))
+                                .allTopicNames().get().get(foo0.topic()).partitions().stream()
+                                .filter(p -> p.partition() == foo0.partition())
+                                .flatMap(p -> p.isr().stream())
+                                .map(Node::id).collect(Collectors.toSet());
+                        return isr.containsAll(Arrays.asList(0, 1, 2, 3));
+                    },
+                    "Timed out while waiting for replica 3 to join the ISR"
+            );
+        }
+
+        // Now cancel the assignment and verify that the partition is removed from cancelled replicas
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(singleton(foo0), Collections.emptySet()), runCancelAssignment(assignment, true));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 4));
+    }
+
+    /**
+     * Test moving partitions between directories.
+     */
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft

Review Comment:
   `BOD not yet implemented for KRaft` is this comment a bit awkward ?



##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,811 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.time.Duration;
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testHighWaterMarkAfterPartitionReassignment() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 100);
+
+        // Execute the assignment
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[3,1,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, -1L, -1L);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(foo0,
+                    new PartitionReassignmentState(asList(3, 1, 2), asList(3, 1, 2), true));
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            TestUtils.waitForCondition(() -> {
+                ListOffsetsResultInfo result = admin.listOffsets(Collections.singletonMap(foo0, new OffsetSpec.LatestSpec())).partitionResult(foo0).get();
+                return result.offset() == 100;
+            }, "Timeout for waiting offset");
+        }
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Execute the assignment with a low throttle
+            long initialThrottle = 1L;
+            runExecuteAssignment(false, assignment, initialThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), initialThrottle);
+
+            // Now update the throttle and verify the reassignment completes
+            long updatedThrottle = 300000L;
+            runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), updatedThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(admin, assignment, false));
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+
+            // Execute the assignment
+            long interBrokerThrottle = 300000L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Wait for the assignment to complete
+            TestUtils.waitForCondition(() -> {
+                // Check the reassignment status.
+                VerifyAssignmentResult result = runVerifyAssignment(admin, assignment, true);
+
+                if (!result.partsOngoing) {
+                    return true;
+                } else {
+                    assertFalse(
+                            result.partStates.values().stream().allMatch(state -> state.done),
+                            "Expected at least one partition reassignment to be ongoing when result = " + result
+                    );
+                    assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                    assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                    waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+                    return false;
+                }
+            }, "Expected reassignment to complete.");
+
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(finalAssignment));
+            // The throttles should still have been preserved, since we ran with --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            List<ConsumerRecord<byte[], byte[]>> allRecords = new ArrayList<>();
+            TestUtils.waitForCondition(() -> {
+                ConsumerRecords<byte[], byte[]> records = consumer.poll(Duration.ofMillis(100L));
+                records.records("baz").forEach(allRecords::add);
+                return allRecords.size() >= 100;
+            }, "Timeout for waiting enough records");
+        }
+        removeReplicationThrottleForPartitions(part);
+        Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(part,
+                new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers())))  {
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+        }
+    }
+
+    /**
+     * Test running a reassignment and then cancelling it.
+     */
+    @ClusterTest
+    public void testCancellation() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+        produceMessages(baz1.topic(), baz1.partition(), 200);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":1,\"replicas\":[0,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(unthrottledBrokerConfigs,
+                    describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+            long interBrokerThrottle = 1L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> partStates = new HashMap<>();
+
+            partStates.put(foo0, new PartitionReassignmentState(asList(0, 1, 3, 2), asList(0, 1, 3), false));
+            partStates.put(baz1, new PartitionReassignmentState(asList(0, 2, 3, 1), asList(0, 2, 3), false));
+
+            // Verify that the reassignment is running.  The very low throttle should keep it
+            // from completing before this runs.
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
+            // Cancel the reassignment.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(assignment, true));
+            // Broker throttles are still active because we passed --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Cancelling the reassignment again should reveal nothing to cancel.
+            assertEquals(new AbstractMap.SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(assignment, false));
+            // This time, the broker throttles were removed.
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+            // Verify that there are no ongoing reassignments.
+            assertFalse(runVerifyAssignment(admin, assignment, false).partsOngoing);
+        }
+        // Verify that the partition is removed from cancelled replicas
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
+    }
+
+    @ClusterTest
+    public void testCancellationWithAddingReplicaInIsr() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 200);
+
+        // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,4],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // We will throttle replica 4 so that only replica 3 joins the ISR
+        setReplicationThrottleForPartitions(foo0);
+
+        // Execute the assignment and wait for replica 3 (only) to join the ISR
+        runExecuteAssignment(false, assignment, -1L, -1L);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitForCondition(
+                    () -> {
+                        Set<Integer> isr = admin.describeTopics(Collections.singleton(foo0.topic()))
+                                .allTopicNames().get().get(foo0.topic()).partitions().stream()
+                                .filter(p -> p.partition() == foo0.partition())
+                                .flatMap(p -> p.isr().stream())
+                                .map(Node::id).collect(Collectors.toSet());
+                        return isr.containsAll(Arrays.asList(0, 1, 2, 3));
+                    },
+                    "Timed out while waiting for replica 3 to join the ISR"
+            );
+        }
+
+        // Now cancel the assignment and verify that the partition is removed from cancelled replicas
+        assertEquals(new AbstractMap.SimpleImmutableEntry<>(singleton(foo0), Collections.emptySet()), runCancelAssignment(assignment, true));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 4));
+    }
+
+    /**
+     * Test moving partitions between directories.
+     */
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testLogDirReassignment() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move, but throttle it to be very slow so that it can't complete
+        // before our next checks happen.
+        long logDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, logDirThrottle);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Check the output of --verify
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new ActiveMoveState(reassignment.currentDir, reassignment.targetDir, reassignment.targetDir)
+                    ), true));
+            waitForLogDirThrottle(admin, singleton(0), logDirThrottle);
+
+            // Remove the throttle
+            admin.incrementalAlterConfigs(singletonMap(
+                            new ConfigResource(ConfigResource.Type.BROKER, "0"),
+                            singletonList(new AlterConfigOp(
+                                    new ConfigEntry(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, ""), AlterConfigOp.OpType.DELETE))))
+                    .all().get();
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+
+            // Wait for the directory movement to complete.
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+
+            BrokerDirs info1 = new BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
+            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
+        }
+    }
+
+    @ClusterTest(types = {Type.ZK, Type.KRAFT}) // JBOD not yet implemented for KRaft
+    public void testAlterLogDirReassignmentThrottle() throws Exception {
+        createTopics();
+        TopicPartition topicPartition = new TopicPartition("foo", 0);
+        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);
+
+        int targetBrokerId = 0;
+        List<Integer> replicas = asList(0, 1, 2);
+        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);
+
+        // Start the replica move with a low throttle so it does not complete
+        long initialLogDirThrottle = 1L;
+        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            waitForLogDirThrottle(admin, new HashSet<>(singletonList(0)), initialLogDirThrottle);
+
+            // Now increase the throttle and verify that the log dir movement completes
+            long updatedLogDirThrottle = 3000000L;
+            runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
+            waitForLogDirThrottle(admin, singleton(0), updatedLogDirThrottle);
+
+            waitForVerifyAssignment(admin, reassignment.json, true,
+                    new VerifyAssignmentResult(singletonMap(
+                            topicPartition, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 2), true)
+                    ), false, singletonMap(
+                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
+                            new CompletedMoveState(reassignment.targetDir)
+                    ), false));
+        }
+    }
+
+    private void createTopics() {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<Integer, List<Integer>> fooReplicasAssignments = new HashMap<>();
+            fooReplicasAssignments.put(0, asList(0, 1, 2));
+            fooReplicasAssignments.put(1, asList(1, 2, 3));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("foo", fooReplicasAssignments))).topicId("foo").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("foo", fooReplicasAssignments.size()));
+
+            Map<Integer, List<Integer>> barReplicasAssignments = new HashMap<>();
+            barReplicasAssignments.put(0, asList(3, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("bar", barReplicasAssignments))).topicId("bar").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("bar", barReplicasAssignments.size()));
+
+            Map<Integer, List<Integer>> bazReplicasAssignments = new HashMap<>();
+            bazReplicasAssignments.put(0, asList(1, 0, 2));
+            bazReplicasAssignments.put(1, asList(2, 0, 1));
+            bazReplicasAssignments.put(2, asList(0, 2, 1));
+            Assertions.assertDoesNotThrow(() -> admin.createTopics(Collections.singletonList(new NewTopic("baz", bazReplicasAssignments))).topicId("baz").get());
+            Assertions.assertDoesNotThrow(() -> clusterInstance.waitForTopic("baz", bazReplicasAssignments.size()));
+        }
+    }
+
+    private void produceMessages(String topic, int partition, int numMessages) {
+        Properties properties = new Properties();
+        properties.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        try (KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(properties, new ByteArraySerializer(), new ByteArraySerializer())) {
+            IntStream.range(0, numMessages).forEach(i -> {
+                ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topic, partition, null, new byte[10000]);
+                assertDoesNotThrow(() -> producer.send(record).get());
+            });
+        }
+    }
+
+    private void executeAndVerifyReassignment() throws ExecutionException, InterruptedException {
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3,2,0],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition bar0 = new TopicPartition("bar", 0);
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+
+        initialAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 2), asList(0, 1, 3), true));
+        initialAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 0), true));
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(initialAssignment));
+
+            // Execute the assignment
+            runExecuteAssignment(false, assignment, -1L, -1L);
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(foo0, new PartitionReassignmentState(asList(0, 1, 3), asList(0, 1, 3), true));
+            finalAssignment.put(bar0, new PartitionReassignmentState(asList(3, 2, 0), asList(3, 2, 0), true));
+
+            VerifyAssignmentResult verifyAssignmentResult = runVerifyAssignment(admin, assignment, false);
+            assertFalse(verifyAssignmentResult.movesOngoing);
+
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+
+            assertEquals(unthrottledBrokerConfigs,
+                    describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+        }
+
+        // Verify that partitions are removed from brokers no longer assigned
+        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 2));
+        verifyReplicaDeleted(new TopicPartitionReplica(bar0.topic(), bar0.partition(), 1));
+    }
+
+    private final Map<Integer, Map<String, Long>> unthrottledBrokerConfigs = new HashMap<>(); {
+        IntStream.range(0, 4).forEach(brokerId -> {
+            Map<String, Long> brokerConfig = new HashMap<>();
+
+            BROKER_LEVEL_THROTTLES.forEach(throttle -> brokerConfig.put(throttle, -1L));
+
+            unthrottledBrokerConfigs.put(brokerId, brokerConfig);
+        });
+    }
+
+    private void verifyReplicaDeleted(TopicPartitionReplica topicPartitionReplica) throws InterruptedException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            TestUtils.waitForCondition(
+                    () -> {
+                        try {
+                            TopicDescription topicDescription = admin.describeTopics(singleton(topicPartitionReplica.topic())).topicNameValues().get(topicPartitionReplica.topic()).get();
+                            return topicDescription.partitions().stream().noneMatch(topicPartitionInfo -> {
+                                if (topicPartitionInfo.partition() != topicPartitionReplica.partition()) {
+                                    return false;
+                                }
+                                return topicPartitionInfo.replicas().stream().anyMatch(node -> node.id() == topicPartitionReplica.brokerId());
+                            });
+                        } catch (InterruptedException | ExecutionException e) {
+                            throw new RuntimeException(e);
+                        }
+                    },
+                    "Timed out waiting for replica " + topicPartitionReplica.brokerId() + " of " + topicPartitionReplica + " to be deleted"
+            );
+        }
+    }
+
+    private void waitForLogDirThrottle(Admin admin, Set<Integer> throttledBrokers, Long logDirThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, -1L);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, logDirThrottle);
+        waitForBrokerThrottles(admin, throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForInterBrokerThrottle(Admin admin, List<Integer> throttledBrokers, Long interBrokerThrottle) {
+        Map<String, Long> throttledConfigMap = new HashMap<>();
+        throttledConfigMap.put(QuotaConfigs.LEADER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_RATE_CONFIG, interBrokerThrottle);
+        throttledConfigMap.put(QuotaConfigs.REPLICA_ALTER_LOG_DIRS_IO_MAX_BYTES_PER_SECOND_CONFIG, -1L);
+        waitForBrokerThrottles(admin, throttledBrokers, throttledConfigMap);
+    }
+
+    private void waitForBrokerThrottles(Admin admin, Collection<Integer> throttledBrokers, Map<String, Long> throttleConfig) {
+        Map<Integer, Map<String, Long>> throttledBrokerConfigs = new HashMap<>();
+        unthrottledBrokerConfigs.forEach((brokerId, unthrottledConfig) -> {
+            Map<String, Long> expectedThrottleConfig = throttledBrokers.contains(brokerId)
+                    ? throttleConfig
+                    : unthrottledConfig;
+            throttledBrokerConfigs.put(brokerId, expectedThrottleConfig);
+        });
+        Assertions.assertDoesNotThrow(() -> waitForBrokerLevelThrottles(admin, throttledBrokerConfigs));
+    }
+
+    private void waitForBrokerLevelThrottles(Admin admin, Map<Integer, Map<String, Long>> targetThrottles) throws InterruptedException {
+        AtomicReference<Map<Integer, Map<String, Long>>> curThrottles = new AtomicReference<>(new HashMap<>());
+        TestUtils.waitForCondition(() -> {
+            assertDoesNotThrow(() -> curThrottles.set(describeBrokerLevelThrottles(admin, targetThrottles.keySet())));
+            return targetThrottles.equals(curThrottles.get());
+        }, "timed out waiting for broker throttle to become " + targetThrottles + ".  " +
+                "Latest throttles were " + curThrottles.get());
+    }
+
+    /**
+     * Describe the broker-level throttles in the cluster.
+     *
+     * @return                A map whose keys are broker IDs and whose values are throttle
+     *                        information.  The nested maps are keyed on throttle name.
+     */
+    private Map<Integer, Map<String, Long>> describeBrokerLevelThrottles(Admin admin, Collection<Integer> brokerIds) throws ExecutionException, InterruptedException {
+        Map<Integer, Map<String, Long>> results = new HashMap<>();
+        for (Integer brokerId : brokerIds) {
+            ConfigResource brokerResource = new ConfigResource(ConfigResource.Type.BROKER, brokerId.toString());
+            Config brokerConfigs = admin.describeConfigs(singleton(brokerResource)).values()
+                    .get(brokerResource)
+                    .get();
+
+            Map<String, Long> throttles = new HashMap<>();
+            BROKER_LEVEL_THROTTLES.forEach(throttleName -> {
+                String configValue = Optional.ofNullable(brokerConfigs.get(throttleName)).map(ConfigEntry::value).orElse("-1");
+                throttles.put(throttleName, Long.parseLong(configValue));
+            });
+            results.put(brokerId, throttles);
+        }
+        return results;
+    }
+
+    static class LogDirReassignment {
+        final String json;
+        final String currentDir;
+        final String targetDir;
+
+        public LogDirReassignment(String json, String currentDir, String targetDir) {
+            this.json = json;
+            this.currentDir = currentDir;
+            this.targetDir = targetDir;
+        }
+    }
+
+    private LogDirReassignment buildLogDirReassignment(TopicPartition topicPartition,
+                                                       int brokerId,
+                                                       List<Integer> replicas) throws ExecutionException, InterruptedException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            DescribeLogDirsResult describeLogDirsResult = admin.describeLogDirs(
+                    IntStream.range(0, 4).boxed().collect(Collectors.toList()));
+
+            BrokerDirs logDirInfo = new BrokerDirs(describeLogDirsResult, brokerId);
+            assertTrue(logDirInfo.futureLogDirs.isEmpty());
+
+            String currentDir = logDirInfo.curLogDirs.get(topicPartition);
+            String newDir = logDirInfo.logDirs.stream().filter(dir -> !dir.equals(currentDir)).findFirst().get();
+
+            List<String> logDirs = replicas.stream().map(replicaId -> {
+                if (replicaId == brokerId)
+                    return "\"" + newDir + "\"";
+                else
+                    return "\"any\"";
+            }).collect(Collectors.toList());
+
+            String reassignmentJson =
+                    " { \"version\": 1," +
+                            "  \"partitions\": [" +
+                            "    {" +
+                            "     \"topic\": \"" + topicPartition.topic() + "\"," +
+                            "     \"partition\": " + topicPartition.partition() + "," +
+                            "     \"replicas\": [" + replicas.stream().map(Object::toString).collect(Collectors.joining(",")) + "]," +
+                            "     \"log_dirs\": [" + String.join(",", logDirs) + "]" +
+                            "    }" +
+                            "   ]" +
+                            "  }";
+
+            return new LogDirReassignment(reassignmentJson, currentDir, newDir);
+        }
+    }
+
+    private VerifyAssignmentResult runVerifyAssignment(Admin admin,
+                                                       String jsonString,
+                                                       Boolean preserveThrottles) {
+        try {
+            return verifyAssignment(admin, jsonString, preserveThrottles);
+        } catch (ExecutionException | InterruptedException | JsonProcessingException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private void waitForVerifyAssignment(Admin admin,
+                                         String jsonString,
+                                         Boolean preserveThrottles,
+                                         VerifyAssignmentResult expectedResult) throws InterruptedException {
+        final VerifyAssignmentResult[] latestResult = {null};
+        TestUtils.waitForCondition(
+                () -> {
+                    latestResult[0] = runVerifyAssignment(admin, jsonString, preserveThrottles);
+                    return expectedResult.equals(latestResult[0]);
+                },
+                "Timed out waiting for verifyAssignment result " + expectedResult
+        );
+    }
+
+    private void runExecuteAssignment(Boolean additional,
+                                      String reassignmentJson,
+                                      Long interBrokerThrottle,
+                                      Long replicaAlterLogDirsThrottle) throws RuntimeException {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            executeAssignment(admin, additional, reassignmentJson,
+                    interBrokerThrottle, replicaAlterLogDirsThrottle, 10000L, Time.SYSTEM);
+        } catch (ExecutionException | InterruptedException | JsonProcessingException | TerseException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private Map.Entry<Set<TopicPartition>, Set<TopicPartitionReplica>> runCancelAssignment(
+            String jsonString,
+            Boolean preserveThrottles
+    ) {
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            return cancelAssignment(admin, jsonString, preserveThrottles, 10000L, Time.SYSTEM);
+        } catch (ExecutionException | InterruptedException | JsonProcessingException | TerseException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    static class BrokerDirs {

Review Comment:
   those static class can be private



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: enable kraft test for ReassignPartitionsIntegrationTest [kafka]

Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.

FrankYang0529 commented on PR #15675:
URL: https://github.com/apache/kafka/pull/15675#issuecomment-2078515889

   > btw, it would be great if you can rewrite it by new test infra.
   
   Yeah, but this class need to assign different config to different broker. I will wait for https://github.com/apache/kafka/pull/15761.
   
   https://github.com/apache/kafka/blob/025f9816f1a15d14aab25c9e8e5b03a87f0cefe2/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java#L754-L808


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: migrate ReassignPartitionsIntegrationTest to use ClusterTestExtensions [kafka]

Posted by "chia7712 (via GitHub)" <gi...@apache.org>.

chia7712 merged PR #15675:
URL: https://github.com/apache/kafka/pull/15675


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: enable kraft test for ReassignPartitionsIntegrationTest [kafka]

Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.

FrankYang0529 commented on PR #15675:
URL: https://github.com/apache/kafka/pull/15675#issuecomment-2041482907

   > @FrankYang0529 Could we use `ClusterTestExtensions` to enable `kraft`? It seems to me that using the new test framework is the better way. Otherwise, we may rewrite it again in the future.
   
   Of course. I will dig into it tomorrow. Thanks for the suggestion.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: migrate ReassignPartitionsIntegrationTest to use ClusterTestExtensions [kafka]

Posted by "chia7712 (via GitHub)" <gi...@apache.org>.

chia7712 commented on code in PR #15675:
URL: https://github.com/apache/kafka/pull/15675#discussion_r1615769708


##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,820 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.time.Duration;
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testHighWaterMarkAfterPartitionReassignment() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 100);
+
+        // Execute the assignment
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[3,1,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, -1L, -1L);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(foo0,
+                    new PartitionReassignmentState(asList(3, 1, 2), asList(3, 1, 2), true));
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            TestUtils.waitForCondition(() -> {
+                ListOffsetsResultInfo result = admin.listOffsets(Collections.singletonMap(foo0, new OffsetSpec.LatestSpec())).partitionResult(foo0).get();
+                return result.offset() == 100;
+            }, "Timeout for waiting offset");
+        }
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Execute the assignment with a low throttle
+            long initialThrottle = 1L;
+            runExecuteAssignment(false, assignment, initialThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), initialThrottle);
+
+            // Now update the throttle and verify the reassignment completes
+            long updatedThrottle = 300000L;
+            runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), updatedThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(admin, assignment, false));
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+
+            // Execute the assignment
+            long interBrokerThrottle = 300000L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Wait for the assignment to complete
+            TestUtils.waitForCondition(() -> {
+                // Check the reassignment status.
+                VerifyAssignmentResult result = runVerifyAssignment(admin, assignment, true);
+
+                if (!result.partsOngoing) {
+                    return true;
+                } else {
+                    assertFalse(
+                            result.partStates.values().stream().allMatch(state -> state.done),
+                            "Expected at least one partition reassignment to be ongoing when result = " + result
+                    );
+                    assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                    assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                    System.out.println("Current result: " + result);
+                    waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+                    return false;
+                }
+            }, "Expected reassignment to complete.");
+
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(finalAssignment));
+            // The throttles should still have been preserved, since we ran with --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            List<ConsumerRecord<byte[], byte[]>> allRecords = new ArrayList<>();
+            TestUtils.waitForCondition(() -> {
+                ConsumerRecords<byte[], byte[]> records = consumer.poll(Duration.ofMillis(100L));
+                System.out.println("records.count " + records.count());

Review Comment:
   It seems there are many similar `System.out.println` in this test



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] MINOR: migrate ReassignPartitionsIntegrationTest to use ClusterTestExtensions [kafka]

Posted by "chia7712 (via GitHub)" <gi...@apache.org>.

chia7712 commented on code in PR #15675:
URL: https://github.com/apache/kafka/pull/15675#discussion_r1615769194


##########
tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java:
##########
@@ -0,0 +1,820 @@
+/*
+ * 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.kafka.tools.reassign;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import kafka.test.ClusterInstance;
+import kafka.test.annotation.ClusterTest;
+import kafka.test.annotation.ClusterConfigProperty;
+import kafka.test.annotation.ClusterTestDefaults;
+import kafka.test.annotation.ClusterTests;
+import kafka.test.annotation.Type;
+import kafka.test.junit.ClusterTestExtensions;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.DescribeLogDirsResult;
+import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.config.QuotaConfigs;
+import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.TerseException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.time.Duration;
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.config.ConfigResource.Type.TOPIC;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1;
+import static org.apache.kafka.server.common.MetadataVersion.IBP_3_3_IV0;
+import static org.apache.kafka.server.config.QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.QuotaConfigs.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.INTER_BROKER_PROTOCOL_VERSION_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_FETCH_BACKOFF_MS_CONFIG;
+import static org.apache.kafka.server.config.ReplicationConfigs.REPLICA_LAG_TIME_MAX_MS_CONFIG;
+import static org.apache.kafka.tools.ToolsTestUtils.assignThrottledPartitionReplicas;
+import static org.apache.kafka.tools.ToolsTestUtils.throttleAllBrokersReplication;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.BROKER_LEVEL_THROTTLES;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.cancelAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.executeAssignment;
+import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.verifyAssignment;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ClusterTestDefaults(brokers = 5, disksPerBroker = 3, serverProperties = {
+        @ClusterConfigProperty(key = REPLICA_FETCH_BACKOFF_MS_CONFIG, value = "100"),
+        @ClusterConfigProperty(key = AUTO_LEADER_REBALANCE_ENABLE_CONFIG, value = "false"),
+        @ClusterConfigProperty(key = REPLICA_LAG_TIME_MAX_MS_CONFIG, value = "1000"),
+        @ClusterConfigProperty(id = 0, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 1, key = "broker.rack", value = "rack0"),
+        @ClusterConfigProperty(id = 2, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 3, key = "broker.rack", value = "rack1"),
+        @ClusterConfigProperty(id = 4, key = "broker.rack", value = "rack1"),
+})
+@ExtendWith(ClusterTestExtensions.class)
+public class ReassignPartitionsCommandTest {
+    private final ClusterInstance clusterInstance;
+
+    ReassignPartitionsCommandTest(ClusterInstance clusterInstance) {
+        this.clusterInstance = clusterInstance;
+    }
+
+    @ClusterTest
+    public void testReassignment() throws Exception {
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, metadataVersion = IBP_2_7_IV1),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, metadataVersion = IBP_3_3_IV0)
+    })
+    public void testReassignmentWithAlterPartitionDisabled() throws Exception {
+        // Test reassignment when the IBP is on an older version which does not use
+        // the `AlterPartition` API. In this case, the controller will register individual
+        // watches for each reassigning partition so that the reassignment can be
+        // completed as soon as the ISR is expanded.
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTests({
+            @ClusterTest(types = {Type.ZK}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "2.7-IV1"),
+            }),
+            @ClusterTest(types = {Type.KRAFT, Type.CO_KRAFT}, serverProperties = {
+                    @ClusterConfigProperty(id = 1, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 2, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+                    @ClusterConfigProperty(id = 3, key = INTER_BROKER_PROTOCOL_VERSION_CONFIG, value = "3.3-IV0"),
+            })
+    })
+    public void testReassignmentCompletionDuringPartialUpgrade() throws Exception {
+        // Test reassignment during a partial upgrade when some brokers are relying on
+        // `AlterPartition` and some rely on the old notification logic through Zookeeper.
+        // In this test case, broker 0 starts up first on the latest IBP and is typically
+        // elected as controller. The three remaining brokers start up on the older IBP.
+        // We want to ensure that reassignment can still complete through the ISR change
+        // notification path even though the controller expects `AlterPartition`.
+
+        // Override change notification settings so that test is not delayed by ISR
+        // change notification delay
+        // ZkAlterPartitionManager.DefaultIsrPropagationConfig_$eq(new IsrChangePropagationConfig(500, 100, 500));
+
+        createTopics();
+        executeAndVerifyReassignment();
+    }
+
+    @ClusterTest
+    public void testHighWaterMarkAfterPartitionReassignment() throws Exception {
+        createTopics();
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        produceMessages(foo0.topic(), foo0.partition(), 100);
+
+        // Execute the assignment
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[3,1,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, -1L, -1L);
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = singletonMap(foo0,
+                    new PartitionReassignmentState(asList(3, 1, 2), asList(3, 1, 2), true));
+            // Wait for the assignment to complete
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            TestUtils.waitForCondition(() -> {
+                ListOffsetsResultInfo result = admin.listOffsets(Collections.singletonMap(foo0, new OffsetSpec.LatestSpec())).partitionResult(foo0).get();
+                return result.offset() == 100;
+            }, "Timeout for waiting offset");
+        }
+    }
+
+    @ClusterTest
+    public void testAlterReassignmentThrottle() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            // Execute the assignment with a low throttle
+            long initialThrottle = 1L;
+            runExecuteAssignment(false, assignment, initialThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), initialThrottle);
+
+            // Now update the throttle and verify the reassignment completes
+            long updatedThrottle = 300000L;
+            runExecuteAssignment(true, assignment, updatedThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), updatedThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    /**
+     * Test running a reassignment with the interBrokerThrottle set.
+     */
+    @ClusterTest
+    public void testThrottledReassignment() throws Exception {
+        createTopics();
+        produceMessages("foo", 0, 50);
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,3,2],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
+                "{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+
+        // Check that the assignment has not yet been started yet.
+        Map<TopicPartition, PartitionReassignmentState> initialAssignment = new HashMap<>();
+        initialAssignment.put(new TopicPartition("foo", 0),
+                new PartitionReassignmentState(asList(0, 1, 2), asList(0, 3, 2), true));
+        initialAssignment.put(new TopicPartition("baz", 2),
+                new PartitionReassignmentState(asList(0, 2, 1), asList(3, 2, 1), true));
+        try (Admin admin = Admin.create(Collections.singletonMap(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
+            assertEquals(new VerifyAssignmentResult(initialAssignment), runVerifyAssignment(admin, assignment, false));
+            assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(admin, unthrottledBrokerConfigs.keySet()));
+
+            // Execute the assignment
+            long interBrokerThrottle = 300000L;
+            runExecuteAssignment(false, assignment, interBrokerThrottle, -1L);
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+
+            Map<TopicPartition, PartitionReassignmentState> finalAssignment = new HashMap<>();
+            finalAssignment.put(new TopicPartition("foo", 0),
+                    new PartitionReassignmentState(asList(0, 3, 2), asList(0, 3, 2), true));
+            finalAssignment.put(new TopicPartition("baz", 2),
+                    new PartitionReassignmentState(asList(3, 2, 1), asList(3, 2, 1), true));
+
+            // Wait for the assignment to complete
+            TestUtils.waitForCondition(() -> {
+                // Check the reassignment status.
+                VerifyAssignmentResult result = runVerifyAssignment(admin, assignment, true);
+
+                if (!result.partsOngoing) {
+                    return true;
+                } else {
+                    assertFalse(
+                            result.partStates.values().stream().allMatch(state -> state.done),
+                            "Expected at least one partition reassignment to be ongoing when result = " + result
+                    );
+                    assertEquals(asList(0, 3, 2), result.partStates.get(new TopicPartition("foo", 0)).targetReplicas);
+                    assertEquals(asList(3, 2, 1), result.partStates.get(new TopicPartition("baz", 2)).targetReplicas);
+                    System.out.println("Current result: " + result);
+                    waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+                    return false;
+                }
+            }, "Expected reassignment to complete.");
+
+            waitForVerifyAssignment(admin, assignment, true,
+                    new VerifyAssignmentResult(finalAssignment));
+            // The throttles should still have been preserved, since we ran with --preserve-throttles
+            waitForInterBrokerThrottle(admin, asList(0, 1, 2, 3), interBrokerThrottle);
+            // Now remove the throttles.
+            waitForVerifyAssignment(admin, assignment, false,
+                    new VerifyAssignmentResult(finalAssignment));
+            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
+        }
+    }
+
+    @ClusterTest
+    public void testProduceAndConsumeWithReassignmentInProgress() throws Exception {
+        createTopics();
+        produceMessages("baz", 2, 60);
+        String assignment = "{\"version\":1,\"partitions\":" +
+                "[{\"topic\":\"baz\",\"partition\":2,\"replicas\":[3,2,1],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
+                "]}";
+        runExecuteAssignment(false, assignment, 300L, -1L);
+        produceMessages("baz", 2, 100);
+
+        Properties consumerProps = new Properties();
+        consumerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers());
+        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        TopicPartition part = new TopicPartition("baz", 2);
+        try (Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(consumerProps, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+            consumer.assign(singleton(part));
+            List<ConsumerRecord<byte[], byte[]>> allRecords = new ArrayList<>();
+            TestUtils.waitForCondition(() -> {
+                ConsumerRecords<byte[], byte[]> records = consumer.poll(Duration.ofMillis(100L));
+                System.out.println("records.count " + records.count());

Review Comment:
   Could we remove those debug messages?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org