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.inLock
 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.ReentrantLock
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
 import scala.jdk.CollectionConverters._
@@ -116,63 +116,75 @@ 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)
+  // Protect updates of the inflight flag and prevent additional pending items from being submitted while we are
+  // preparing a request
+  private val inflightLock: ReentrantLock = new ReentrantLock()
 
   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
+    inLock(inflightLock) {
+      if (unsentIsrUpdates.putIfAbsent(alterIsrItem.topicPartition, alterIsrItem) == null) {
+        maybePropagateIsrChanges()
+        true
+      } else {
+        false
+      }
+    }
   }
 
   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
+  private[server] def maybePropagateIsrChanges(): Unit = inLock(inflightLock) {
+    // Send all pending items if there is not already a request in-flight.
+    if (!inflightRequest && !unsentIsrUpdates.isEmpty) {
+      // Copy current unsent ISRs but don't remove from the map, they get cleared in the response handler
       val inflightAlterIsrItems = new ListBuffer[AlterIsrItem]()
       unsentIsrUpdates.values().forEach(item => inflightAlterIsrItems.append(item))
-
-      val now = time.milliseconds()
-      lastIsrPropagationMs.set(now)
       sendRequest(inflightAlterIsrItems.toSeq)
+      inflightRequest = true
     }
   }
 
-  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")
-      }
+  private[server] def clearInFlightRequest(): Unit = inLock(inflightLock) {
+    if (!inflightRequest) {
+      warn("Attempting to clear AlterIsr in-flight flag when no apparent request is in-flight")
     }
+    inflightRequest = false
+  }
 
+  private def sendRequest(inflightAlterIsrItems: Seq[AlterIsrItem]): Unit = {
+    val message = buildRequest(inflightAlterIsrItems)
     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
+              clearInFlightRequest()
+              maybePropagateIsrChanges()
+            case _ =>
+              // If we received a top-level error from the controller, retry the request in the near future
+              clearInFlightRequest()
+              scheduler.schedule("send-alter-isr", () => maybePropagateIsrChanges(), 50, -1, TimeUnit.MILLISECONDS)
           }
         }
 
@@ -207,7 +219,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 +266,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

@@ -34,7 +34,6 @@ import org.junit.{Before, Test}
 import org.mockito.ArgumentMatchers.{any, anyString}
 import org.mockito.{ArgumentMatchers, Mockito}
 
-
 class AlterIsrManagerTest {
 
   val topic = "test-topic"
@@ -97,29 +96,37 @@ class AlterIsrManagerTest {
   @Test
   def testSingleBatch(): Unit = {
     val capture = EasyMock.newCapture[AbstractRequest.Builder[AlterIsrRequest]]()
+    val callbackCapture = EasyMock.newCapture[ControllerRequestCompletionHandler]()
+
     EasyMock.expect(brokerToController.start())
-    EasyMock.expect(brokerToController.sendRequest(EasyMock.capture(capture), EasyMock.anyObject())).once()
+    EasyMock.expect(brokerToController.sendRequest(EasyMock.capture(capture), EasyMock.capture(callbackCapture))).times(2)
     EasyMock.replay(brokerToController)
 
     val scheduler = new MockScheduler(time)
     val alterIsrManager = new DefaultAlterIsrManager(brokerToController, scheduler, time, brokerId, () => 2)
     alterIsrManager.start()
 
-    for (i <- 0 to 9) {
+    // First request will send batch of one
+    alterIsrManager.submit(AlterIsrItem(new TopicPartition(topic, 0),
+      new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}, 0))
+
+    // Other submissions will queue up until a response
+    for (i <- 1 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)
-    scheduler.tick()
+    // Simulate response, omitting partition 0 will allow it to stay in unsent queue
+    val alterIsrResp = new AlterIsrResponse(new AlterIsrResponseData())
+    val resp = new ClientResponse(null, null, "", 0L, 0L,
+      false, null, null, alterIsrResp)
 
-    // This should not be included in the batch
-    alterIsrManager.submit(AlterIsrItem(new TopicPartition(topic, 10),
-      new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}, 0))
+    // On the callback, we check for unsent items and send another request
+    callbackCapture.getValue.onComplete(resp)
 
     EasyMock.verify(brokerToController)
 
+    // Verify the last request sent had all 10 items
     val request = capture.getValue.build()
     assertEquals(request.data().topics().size(), 1)
     assertEquals(request.data().topics().get(0).partitions().size(), 10)
@@ -161,7 +168,7 @@ class AlterIsrManagerTest {
     alterIsrManager.start()
     isrs.foreach(alterIsrManager.submit)
 
-    time.sleep(100)
+    time.sleep(1)
     scheduler.tick()
 
     EasyMock.verify(brokerToController)
@@ -179,7 +186,7 @@ class AlterIsrManagerTest {
     errors.foreach(error => {
       val alterIsrManager = testPartitionError(tp0, error)
       // Any partition-level error should clear the item from the pending queue allowing for future updates
-      assertTrue(alterIsrManager.submit(AlterIsrItem(tp0, null, _ => { }, 0)))
+      assertTrue(alterIsrManager.submit(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => { }, 0)))
     })
   }
 
@@ -204,10 +211,8 @@ class AlterIsrManagerTest {
 
     alterIsrManager.submit(AlterIsrItem(tp, new LeaderAndIsr(1, 1, List(1,2,3), 10), callback, 0))
 
-    time.sleep(100)
-    scheduler.tick()
-
     EasyMock.verify(brokerToController)
+    EasyMock.reset(brokerToController)
 
     val alterIsrResp = new AlterIsrResponse(new AlterIsrResponseData()
       .setTopics(Collections.singletonList(
@@ -236,32 +241,24 @@ class AlterIsrManagerTest {
     val scheduler = new MockScheduler(time)
     val alterIsrManager = new DefaultAlterIsrManager(brokerToController, scheduler, time, brokerId, () => 2)
     alterIsrManager.start()
-    alterIsrManager.submit(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}, 0))
 
-    time.sleep(100)
-    scheduler.tick() // Triggers a request
+    // First submit will send the request
+    alterIsrManager.submit(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}, 0))
 
-    // Enqueue more updates
+    // These will become pending unsent items
     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)
-    scheduler.tick() // Trigger the schedule again, but no request this time
-
     EasyMock.verify(brokerToController)
 
-    // Even an empty response will clear the in-flight
+    // Once the callback runs, another request will be sent
+    EasyMock.reset(brokerToController)
+    EasyMock.expect(brokerToController.sendRequest(EasyMock.anyObject(), EasyMock.capture(callbackCapture))).once()
+    EasyMock.replay(brokerToController)
     val alterIsrResp = new AlterIsrResponse(new AlterIsrResponseData())
     val resp = new ClientResponse(null, null, "", 0L, 0L,
       false, null, null, alterIsrResp)
     callbackCapture.getValue.onComplete(resp)
-
-    EasyMock.reset(brokerToController)
-    EasyMock.expect(brokerToController.sendRequest(EasyMock.anyObject(), EasyMock.capture(callbackCapture))).once()
-    EasyMock.replay(brokerToController)
-
-    time.sleep(100)
-    scheduler.tick()
     EasyMock.verify(brokerToController)
   }