Only schedule AlterIsr thread when we have an ISR change

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

@@ -17,11 +17,10 @@
 package kafka.server
 
 import java.util
-import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
 import java.util.concurrent.{ConcurrentHashMap, TimeUnit}
-
 import kafka.api.LeaderAndIsr
 import kafka.metrics.KafkaMetricsGroup
+import kafka.utils.CoreUtils.{inReadLock, inWriteLock}
 import kafka.utils.{KafkaScheduler, Logging, Scheduler}
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.clients.ClientResponse
@@ -32,6 +31,7 @@ import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.{AlterIsrRequest, AlterIsrResponse}
 import org.apache.kafka.common.utils.Time
 
+import java.util.concurrent.locks.ReentrantReadWriteLock
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
 import scala.jdk.CollectionConverters._
@@ -116,63 +116,78 @@ class DefaultAlterIsrManager(
   private val unsentIsrUpdates: util.Map[TopicPartition, AlterIsrItem] = new ConcurrentHashMap[TopicPartition, AlterIsrItem]()
 
   // Used to allow only one in-flight request at a time
-  private val inflightRequest: AtomicBoolean = new AtomicBoolean(false)
+  @volatile
+  private var inflightRequest: Boolean = false
 
-  private val lastIsrPropagationMs = new AtomicLong(0)
+  // Protects the updates of the inflight flag and prevents new pending items from being submitted while we are
+  // preparing a request
+  private val inflightLock: ReentrantReadWriteLock = new ReentrantReadWriteLock()
 
   override def start(): Unit = {
     controllerChannelManager.start()
-    scheduler.schedule("send-alter-isr", propagateIsrChanges, 50, 50, TimeUnit.MILLISECONDS)
   }
 
   override def shutdown(): Unit = {
     controllerChannelManager.shutdown()
   }
 
   override def submit(alterIsrItem: AlterIsrItem): Boolean = {
-    unsentIsrUpdates.putIfAbsent(alterIsrItem.topicPartition, alterIsrItem) == null
+    val (didSubmit, needsPropagate) = inReadLock(inflightLock) {
+      if (unsentIsrUpdates.putIfAbsent(alterIsrItem.topicPartition, alterIsrItem) == null) {
+        (true, !inflightRequest)
+      } else {
+        (false, false)
+      }
+    }
+    if (needsPropagate) {
+      propagateIsrChanges(true)
+    }
+    didSubmit
   }
 
   override def clearPending(topicPartition: TopicPartition): Unit = {
     unsentIsrUpdates.remove(topicPartition)
   }
 
-  private def propagateIsrChanges(): Unit = {
-    if (!unsentIsrUpdates.isEmpty && inflightRequest.compareAndSet(false, true)) {
-      // Copy current unsent ISRs but don't remove from the map
-      val inflightAlterIsrItems = new ListBuffer[AlterIsrItem]()
-      unsentIsrUpdates.values().forEach(item => inflightAlterIsrItems.append(item))
-
-      val now = time.milliseconds()
-      lastIsrPropagationMs.set(now)
-      sendRequest(inflightAlterIsrItems.toSeq)
+  private def propagateIsrChanges(checkInflight: Boolean): Unit = inWriteLock(inflightLock) {
+    // If we're checking for unsent items in the response handler, we ignore the inflight flag and send a request if
+    // there are any unsent items. Otherwise, we got here from the submit(AlterIsrItem) method and need to check that
+    // there is not an inflight request.
+    if (!checkInflight || !inflightRequest) {
+      if (!unsentIsrUpdates.isEmpty) {
+        // Copy current unsent ISRs but don't remove from the map
+        val inflightAlterIsrItems = new ListBuffer[AlterIsrItem]()
+        unsentIsrUpdates.values().forEach(item => inflightAlterIsrItems.append(item))
+
+        sendRequest(inflightAlterIsrItems.toSeq)
+        inflightRequest = true
+      } else {
+        // No items were pending, so no request was sent -- clear the inflight flag
+        inflightRequest = false
+      }
     }
   }
 
   private def sendRequest(inflightAlterIsrItems: Seq[AlterIsrItem]): Unit = {
     val message = buildRequest(inflightAlterIsrItems)
 
-    def clearInflightRequests(): Unit = {
-      // Be sure to clear the in-flight flag to allow future AlterIsr requests
-      if (!inflightRequest.compareAndSet(true, false)) {
-        throw new IllegalStateException("AlterIsr response callback called when no requests were in flight")
-      }
-    }
-
     debug(s"Sending AlterIsr to controller $message")
 
     // We will not timeout AlterISR request, instead letting it retry indefinitely
     // until a response is received, or a new LeaderAndIsr overwrites the existing isrState
-    // which causes the inflight requests to be ignored.
+    // which causes the response for those partitions to be ignored.
     controllerChannelManager.sendRequest(new AlterIsrRequest.Builder(message),
       new ControllerRequestCompletionHandler {
         override def onComplete(response: ClientResponse): Unit = {
-          try {
-            debug(s"Received AlterIsr response $response")
-            val body = response.responseBody().asInstanceOf[AlterIsrResponse]
-            handleAlterIsrResponse(body, message.brokerEpoch, inflightAlterIsrItems)
-          } finally {
-            clearInflightRequests()
+          debug(s"Received AlterIsr response $response")
+          val body = response.responseBody().asInstanceOf[AlterIsrResponse]
+          handleAlterIsrResponse(body, message.brokerEpoch, inflightAlterIsrItems) match {
+            case Errors.NONE =>
+              // In the normal case, check for pending updates to send immediately
+              propagateIsrChanges(false)
+            case _ =>
+              // If we received a top-level error from the controller, retry the request in the near future
+              scheduler.schedule("send-alter-isr", () => propagateIsrChanges(false), 50, -1, TimeUnit.MILLISECONDS)
           }
         }
 
@@ -207,7 +222,7 @@ class DefaultAlterIsrManager(
 
   def handleAlterIsrResponse(alterIsrResponse: AlterIsrResponse,
                              sentBrokerEpoch: Long,
-                             inflightAlterIsrItems: Seq[AlterIsrItem]): Unit = {
+                             inflightAlterIsrItems: Seq[AlterIsrItem]): Errors = {
     val data: AlterIsrResponseData = alterIsrResponse.data
 
     Errors.forCode(data.errorCode) match {
@@ -254,5 +269,7 @@ class DefaultAlterIsrManager(
       case e: Errors =>
         warn(s"Controller returned an unexpected top-level error when handling AlterIsr request: $e")
     }
+
+    Errors.forCode(data.errorCode)
   }
 }

core/src/test/scala/unit/kafka/server/AlterIsrManagerTest.scala

@@ -108,10 +108,9 @@ class AlterIsrManagerTest {
     for (i <- 0 to 9) {
       alterIsrManager.submit(AlterIsrItem(new TopicPartition(topic, i),
         new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}, 0))
-      time.sleep(1)
     }
 
-    time.sleep(50)
+    time.sleep(1)
     scheduler.tick()
 
     // This should not be included in the batch
@@ -161,7 +160,7 @@ class AlterIsrManagerTest {
     alterIsrManager.start()
     isrs.foreach(alterIsrManager.submit)
 
-    time.sleep(100)
+    time.sleep(1)
     scheduler.tick()
 
     EasyMock.verify(brokerToController)
@@ -238,14 +237,14 @@ class AlterIsrManagerTest {
     alterIsrManager.start()
     alterIsrManager.submit(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}, 0))
 
-    time.sleep(100)
+    time.sleep(1)
     scheduler.tick() // Triggers a request
 
     // Enqueue more updates
     alterIsrManager.submit(AlterIsrItem(tp1, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}, 0))
     alterIsrManager.submit(AlterIsrItem(tp2, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}, 0))
 
-    time.sleep(100)
+    time.sleep(1)
     scheduler.tick() // Trigger the schedule again, but no request this time
 
     EasyMock.verify(brokerToController)
@@ -260,6 +259,10 @@ class AlterIsrManagerTest {
     EasyMock.expect(brokerToController.sendRequest(EasyMock.anyObject(), EasyMock.capture(callbackCapture))).once()
     EasyMock.replay(brokerToController)
 
+    // Need to re-enqueue again to trigger the thread to be scheduled
+    alterIsrManager.clearPending(tp2)
+    alterIsrManager.submit(AlterIsrItem(tp2, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}, 0))
+
     time.sleep(100)
     scheduler.tick()
     EasyMock.verify(brokerToController)