KAFKA-13972; Ensure that replicas are stopped after cancelled reassignment

core/src/main/scala/kafka/controller/ReplicaStateMachine.scala

@@ -152,6 +152,9 @@ class ZkReplicaStateMachine(config: KafkaConfig,
    * ReplicaDeletionSuccessful -> NonExistentReplica
    * -- remove the replica from the in memory partition replica assignment cache
    *
+   * NewReplica,OnlineReplica,OfflineReplica -> UnassignedReplicas
+   * -- Following completion or cancellation of a reassignment, remove the replica state
+   *
    * @param replicaId The replica for which the state transition is invoked
    * @param replicas The partitions on this replica for which the state transition is invoked
    * @param targetState The end state that the replica should be moved to
@@ -218,22 +221,28 @@ class ZkReplicaStateMachine(config: KafkaConfig,
           controllerContext.putReplicaState(replica, OnlineReplica)
         }
       case OfflineReplica =>
-        validReplicas.foreach { replica =>
-          controllerBrokerRequestBatch.addStopReplicaRequestForBrokers(Seq(replicaId), replica.topicPartition, deletePartition = false)
-        }
         val (replicasWithLeadershipInfo, replicasWithoutLeadershipInfo) = validReplicas.partition { replica =>
           controllerContext.partitionLeadershipInfo(replica.topicPartition).isDefined
         }
         val updatedLeaderIsrAndControllerEpochs = removeReplicasFromIsr(replicaId, replicasWithLeadershipInfo.map(_.topicPartition))
         updatedLeaderIsrAndControllerEpochs.forKeyValue { (partition, leaderIsrAndControllerEpoch) =>
-          stateLogger.info(s"Partition $partition state changed to $leaderIsrAndControllerEpoch after removing replica $replicaId from the ISR as part of transition to $OfflineReplica")
+          stateLogger.info(s"Partition $partition state changed to $leaderIsrAndControllerEpoch " +
+            s"after transitioning replica $replicaId to $OfflineReplica")
           if (!controllerContext.isTopicQueuedUpForDeletion(partition.topic)) {
             val recipients = controllerContext.partitionReplicaAssignment(partition).filterNot(_ == replicaId)
             controllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(recipients,
               partition,
               leaderIsrAndControllerEpoch,
-              controllerContext.partitionFullReplicaAssignment(partition), isNew = false)
+              controllerContext.partitionFullReplicaAssignment(partition),
+              isNew = false
+            )
           }
+          controllerBrokerRequestBatch.addStopReplicaRequestForBrokers(
+            Seq(replicaId),
+            partition,
+            deletePartition = false
+          )
+
           val replica = PartitionAndReplica(partition, replicaId)
           val currentState = controllerContext.replicaState(replica)
           if (traceEnabled)
@@ -245,7 +254,15 @@ class ZkReplicaStateMachine(config: KafkaConfig,
           val currentState = controllerContext.replicaState(replica)
           if (traceEnabled)
             logSuccessfulTransition(stateLogger, replicaId, replica.topicPartition, currentState, OfflineReplica)
-          controllerBrokerRequestBatch.addUpdateMetadataRequestForBrokers(controllerContext.liveOrShuttingDownBrokerIds.toSeq, Set(replica.topicPartition))
+          controllerBrokerRequestBatch.addUpdateMetadataRequestForBrokers(
+            controllerContext.liveOrShuttingDownBrokerIds.toSeq,
+            Set(replica.topicPartition)
+          )
+          controllerBrokerRequestBatch.addStopReplicaRequestForBrokers(
+            Seq(replicaId),
+            replica.topicPartition,
+            deletePartition = false
+          )
           controllerContext.putReplicaState(replica, OfflineReplica)
         }
       case ReplicaDeletionStarted =>
@@ -282,6 +299,17 @@ class ZkReplicaStateMachine(config: KafkaConfig,
             logSuccessfulTransition(stateLogger, replicaId, replica.topicPartition, currentState, NonExistentReplica)
           controllerContext.removeReplicaState(replica)
         }
+      case UnassignedReplica =>
+        // For the case of a replica which has been removed following reassignment,
+        // replicas will be stopped and deleted as part of the final transition to `OnlinePartition`.
+        // The replica assignment also will be updated as part of the reassignment completion.
+        // So here we just need to remove the replica state.
+        validReplicas.foreach { replica =>
+          val currentState = controllerContext.replicaState(replica)
+          if (traceEnabled)
+            logSuccessfulTransition(stateLogger, replicaId, replica.topicPartition, currentState, UnassignedReplica)
+          controllerContext.removeReplicaState(replica)
+        }
     }
   }
 
@@ -299,14 +327,14 @@ class ZkReplicaStateMachine(config: KafkaConfig,
     var results = Map.empty[TopicPartition, LeaderIsrAndControllerEpoch]
     var remaining = partitions
     while (remaining.nonEmpty) {
-      val (finishedRemoval, removalsToRetry) = doRemoveReplicasFromIsr(replicaId, remaining)
+      val (finishedRemoval, removalsToRetry) = tryRemoveReplicasFromIsr(replicaId, remaining)
       remaining = removalsToRetry
 
       finishedRemoval.foreach {
         case (partition, Left(e)) =>
-            val replica = PartitionAndReplica(partition, replicaId)
-            val currentState = controllerContext.replicaState(replica)
-            logFailedStateChange(replica, currentState, OfflineReplica, e)
+          val replica = PartitionAndReplica(partition, replicaId)
+          val currentState = controllerContext.replicaState(replica)
+          logFailedStateChange(replica, currentState, OfflineReplica, e)
         case (partition, Right(leaderIsrAndEpoch)) =>
           results += partition -> leaderIsrAndEpoch
       }
@@ -327,7 +355,7 @@ class ZkReplicaStateMachine(config: KafkaConfig,
    *         2. The partitions that we should retry due to a zookeeper BADVERSION conflict. Version conflicts can occur if
    *         the partition leader updated partition state while the controller attempted to update partition state.
    */
-  private def doRemoveReplicasFromIsr(
+  private def tryRemoveReplicasFromIsr(
     replicaId: Int,
     partitions: Seq[TopicPartition]
   ): (Map[TopicPartition, Either[Exception, LeaderIsrAndControllerEpoch]], Seq[TopicPartition]) = {
@@ -489,3 +517,8 @@ case object NonExistentReplica extends ReplicaState {
   val state: Byte = 7
   val validPreviousStates: Set[ReplicaState] = Set(ReplicaDeletionSuccessful)
 }
+
+case object UnassignedReplica extends ReplicaState {
+  val state: Byte = 8
+  val validPreviousStates: Set[ReplicaState] = Set(NewReplica, OnlineReplica, OfflineReplica)
+}

core/src/main/scala/kafka/server/KafkaServer.scala

@@ -584,7 +584,6 @@ class KafkaServer(
       var shutdownSucceeded: Boolean = false
 
       try {
-
         var remainingRetries = retries
         var prevController: Node = null
         var ioException = false

core/src/test/scala/integration/kafka/admin/ReassignPartitionsIntegrationTest.scala

@@ -24,6 +24,7 @@ import kafka.utils.{TestInfoUtils, TestUtils}
 import org.apache.kafka.clients.admin._
 import org.apache.kafka.clients.producer.ProducerRecord
 import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.errors.InvalidReplicaAssignmentException
 import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.common.{TopicPartition, TopicPartitionReplica}
 import org.apache.kafka.server.common.MetadataVersion.IBP_2_7_IV1
@@ -33,7 +34,7 @@ import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.ValueSource
 
 import java.io.Closeable
-import java.util.{Collections, HashMap, List}
+import java.util.{Collections, HashMap, List, Optional}
 import scala.collection.{Map, Seq, mutable}
 import scala.jdk.CollectionConverters._
 
@@ -108,12 +109,13 @@ class ReassignPartitionsIntegrationTest extends QuorumTestHarness {
       """{"topic":"bar","partition":0,"replicas":[3,2,0],"log_dirs":["any","any","any"]}""" +
       """]}"""
 
+    val foo0 = new TopicPartition("foo", 0)
+    val bar0 = new TopicPartition("bar", 0)
+
     // Check that the assignment has not yet been started yet.
     val initialAssignment = Map(
-      new TopicPartition("foo", 0) ->
-        PartitionReassignmentState(Seq(0, 1, 2), Seq(0, 1, 3), true),
-      new TopicPartition("bar", 0) ->
-        PartitionReassignmentState(Seq(3, 2, 1), Seq(3, 2, 0), true)
+      foo0 -> PartitionReassignmentState(Seq(0, 1, 2), Seq(0, 1, 3), true),
+      bar0 -> PartitionReassignmentState(Seq(3, 2, 1), Seq(3, 2, 0), true)
     )
     waitForVerifyAssignment(cluster.adminClient, assignment, false,
       VerifyAssignmentResult(initialAssignment))
@@ -122,10 +124,8 @@ class ReassignPartitionsIntegrationTest extends QuorumTestHarness {
     runExecuteAssignment(cluster.adminClient, false, assignment, -1L, -1L)
     assertEquals(unthrottledBrokerConfigs, describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet.toSeq))
     val finalAssignment = Map(
-      new TopicPartition("foo", 0) ->
-        PartitionReassignmentState(Seq(0, 1, 3), Seq(0, 1, 3), true),
-      new TopicPartition("bar", 0) ->
-        PartitionReassignmentState(Seq(3, 2, 0), Seq(3, 2, 0), true)
+      foo0 -> PartitionReassignmentState(Seq(0, 1, 3), Seq(0, 1, 3), true),
+      bar0 -> PartitionReassignmentState(Seq(3, 2, 0), Seq(3, 2, 0), true)
     )
 
     val verifyAssignmentResult = runVerifyAssignment(cluster.adminClient, assignment, false)
@@ -137,6 +137,10 @@ class ReassignPartitionsIntegrationTest extends QuorumTestHarness {
 
     assertEquals(unthrottledBrokerConfigs,
       describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet.toSeq))
+
+    // Verify that partitions are removed from brokers no longer assigned
+    verifyReplicaDeleted(topicPartition = foo0, replicaId = 2)
+    verifyReplicaDeleted(topicPartition = bar0, replicaId = 1)
   }
 
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
@@ -296,10 +300,13 @@ class ReassignPartitionsIntegrationTest extends QuorumTestHarness {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
   def testCancellation(quorum: String): Unit = {
+    val foo0 = new TopicPartition("foo", 0)
+    val baz1 = new TopicPartition("baz", 1)
+
     cluster = new ReassignPartitionsTestCluster()
     cluster.setup()
-    cluster.produceMessages("foo", 0, 200)
-    cluster.produceMessages("baz", 1, 200)
+    cluster.produceMessages(foo0.topic, foo0.partition, 200)
+    cluster.produceMessages(baz1.topic, baz1.partition, 200)
     val 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"]}""" +
@@ -314,14 +321,11 @@ class ReassignPartitionsIntegrationTest extends QuorumTestHarness {
     // from completing before this runs.
     waitForVerifyAssignment(cluster.adminClient, assignment, true,
       VerifyAssignmentResult(Map(
-        new TopicPartition("foo", 0) -> PartitionReassignmentState(Seq(0, 1, 3, 2), Seq(0, 1, 3), false),
-        new TopicPartition("baz", 1) -> PartitionReassignmentState(Seq(0, 2, 3, 1), Seq(0, 2, 3), false)),
+        foo0 -> PartitionReassignmentState(Seq(0, 1, 3, 2), Seq(0, 1, 3), false),
+        baz1 -> PartitionReassignmentState(Seq(0, 2, 3, 1), Seq(0, 2, 3), false)),
         true, Map(), false))
     // Cancel the reassignment.
-    assertEquals((Set(
-      new TopicPartition("foo", 0),
-      new TopicPartition("baz", 1)
-    ), Set()), runCancelAssignment(cluster.adminClient, assignment, true))
+    assertEquals((Set(foo0, baz1), Set()), runCancelAssignment(cluster.adminClient, assignment, true))
     // Broker throttles are still active because we passed --preserve-throttles
     waitForInterBrokerThrottle(Set(0, 1, 2, 3), interBrokerThrottle)
     // Cancelling the reassignment again should reveal nothing to cancel.
@@ -330,6 +334,118 @@ class ReassignPartitionsIntegrationTest extends QuorumTestHarness {
     waitForBrokerLevelThrottles(unthrottledBrokerConfigs)
     // Verify that there are no ongoing reassignments.
     assertFalse(runVerifyAssignment(cluster.adminClient, assignment, false).partsOngoing)
+    // Verify that the partition is removed from cancelled replicas
+    verifyReplicaDeleted(topicPartition = foo0, replicaId = 3)
+    verifyReplicaDeleted(topicPartition = baz1, replicaId = 3)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
+  @ValueSource(strings = Array("zk", "kraft"))
+  def testCancellationWithAddingReplicaInIsr(quorum: String): Unit = {
+    val foo0 = new TopicPartition("foo", 0)
+
+    cluster = new ReassignPartitionsTestCluster()
+    cluster.setup()
+    cluster.produceMessages(foo0.topic, foo0.partition, 200)
+
+    // The reassignment will bring replicas 3 and 4 into the replica set and remove 1 and 2.
+    val 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
+    TestUtils.setReplicationThrottleForPartitions(
+      cluster.adminClient,
+      brokerIds = Seq(4),
+      partitions = Set(foo0),
+      throttleBytes = 1
+    )
+
+    // Execute the assignment and wait for replica 3 (only) to join the ISR
+    runExecuteAssignment(
+      cluster.adminClient,
+      additional = false,
+      reassignmentJson = assignment
+    )
+    TestUtils.waitUntilTrue(
+      () => TestUtils.currentIsr(cluster.adminClient, foo0) == Set(0, 1, 2, 3),
+      msg = "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((Set(foo0), Set()), runCancelAssignment(cluster.adminClient, assignment, preserveThrottles = true))
+    verifyReplicaDeleted(topicPartition = foo0, replicaId = 3)
+    verifyReplicaDeleted(topicPartition = foo0, replicaId = 4)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
+  @ValueSource(strings = Array("zk", "kraft"))
+  def testCancellationNotAllowedIfOnlyAddingReplicasInIsr(quorum: String): Unit = {
+    val boo0 = new TopicPartition("boo", 0)
+
+    cluster = new ReassignPartitionsTestCluster()
+    cluster.setup()
+    cluster.produceMessages(boo0.topic, boo0.partition, 200)
+
+    // The reassignment will add replicas 1 and 2
+    val assignment = """{"version":1,"partitions":""" +
+      """[{"topic":"boo","partition":0,"replicas":[0,1,2],"log_dirs":["any","any","any"]}""" +
+      """]}"""
+
+    // We will throttle replica 2 so that only replica 1 joins the ISR
+    TestUtils.setReplicationThrottleForPartitions(
+      cluster.adminClient,
+      brokerIds = Seq(2),
+      partitions = Set(boo0),
+      throttleBytes = 1
+    )
+
+    // Execute the assignment and wait for replica 1 (only) to join the ISR
+    runExecuteAssignment(
+      cluster.adminClient,
+      additional = false,
+      reassignmentJson = assignment
+    )
+    TestUtils.waitUntilTrue(
+      () => TestUtils.currentIsr(cluster.adminClient, boo0) == Set(0, 1),
+      msg = "Timed out while waiting for replica 3 to join the ISR"
+    )
+
+    // Now we shutdown broker 0 and wait for 1 to become the leader
+    cluster.servers(0).shutdown()
+    cluster.servers(0).awaitShutdown()
+
+    TestUtils.awaitLeaderChange(
+      cluster.servers,
+      boo0,
+      oldLeader = 0
+    )
+
+    TestUtils.waitUntilTrue(
+      () => TestUtils.currentIsr(cluster.adminClient, boo0) == Set(1),
+      msg = "Timed out while waiting for replica 3 to join the ISR"
+    )
+
+    // Reassignment cancellation should not be allowed since it would leave no remaining
+    // replicas in the ISR.
+    val future = cluster.adminClient.alterPartitionReassignments(
+      Map(boo0 -> Optional.empty[NewPartitionReassignment]()).asJava
+    ).all()
+    TestUtils.assertFutureExceptionTypeEquals(future, classOf[InvalidReplicaAssignmentException])
+  }
+
+  private def verifyReplicaDeleted(
+    topicPartition: TopicPartition,
+    replicaId: Int
+  ): Unit = {
+    def isReplicaStoppedAndDeleted(): Boolean = {
+      val server = cluster.servers(replicaId)
+      val partition = server.replicaManager.getPartition(topicPartition)
+      val log = server.logManager.getLog(topicPartition)
+      partition == HostedPartition.None && log.isEmpty
+    }
+    TestUtils.waitUntilTrue(isReplicaStoppedAndDeleted,
+      msg = s"Timed out waiting for replica $replicaId of $topicPartition to be deleted")
   }
 
   private def waitForLogDirThrottle(throttledBrokers: Set[Int], logDirThrottle: Long): Unit = {
@@ -541,8 +657,8 @@ class ReassignPartitionsIntegrationTest extends QuorumTestHarness {
   private def runExecuteAssignment(adminClient: Admin,
                                    additional: Boolean,
                                    reassignmentJson: String,
-                                   interBrokerThrottle: Long,
-                                   replicaAlterLogDirsThrottle: Long) = {
+                                   interBrokerThrottle: Long = -1,
+                                   replicaAlterLogDirsThrottle: Long = -1) = {
     println(s"==> executeAssignment(adminClient, additional=${additional}, " +
       s"reassignmentJson=${reassignmentJson}, " +
       s"interBrokerThrottle=${interBrokerThrottle}, " +
@@ -591,7 +707,8 @@ class ReassignPartitionsIntegrationTest extends QuorumTestHarness {
     val topics = Map(
       "foo" -> Seq(Seq(0, 1, 2), Seq(1, 2, 3)),
       "bar" -> Seq(Seq(3, 2, 1)),
-      "baz" -> Seq(Seq(1, 0, 2), Seq(2, 0, 1), Seq(0, 2, 1))
+      "baz" -> Seq(Seq(1, 0, 2), Seq(2, 0, 1), Seq(0, 2, 1)),
+      "boo" -> Seq(Seq(0))
     )
 
     val brokerConfigs = brokers.map {

core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala

@@ -1535,7 +1535,7 @@ class PlaintextAdminIntegrationTest extends BaseAdminIntegrationTest {
           s"Unexpected message: ${exception.getMessage}")
       } else {
         assertTrue(exception.getMessage.contains(
-          s"Failed to elect leader for partition $topicPartition under strategy PreferredReplicaPartitionLeaderElectionStrategy"),
+          s"Failed to elect preferred leader"),
           s"Unexpected message: ${exception.getMessage}")
       }
     }