From c3a1a1320c179ce1b988a99c730116e8dceb3ed6 Mon Sep 17 00:00:00 2001
From: Justine <jolshan@confluent.io>
Date: Tue, 24 Oct 2023 15:21:56 -0700
Subject: [PATCH 01/11] pass requestLocal

---
 .../scala/kafka/network/RequestChannel.scala  |   4 +-
 .../kafka/server/KafkaRequestHandler.scala    |  10 +-
 .../scala/kafka/server/ReplicaManager.scala   | 195 ++++++++++--------
 .../server/KafkaRequestHandlerTest.scala      |  42 +++-
 4 files changed, 149 insertions(+), 102 deletions(-)

diff --git a/core/src/main/scala/kafka/network/RequestChannel.scala b/core/src/main/scala/kafka/network/RequestChannel.scala
index 2521e581aba2a..00f7b17757fca 100644
--- a/core/src/main/scala/kafka/network/RequestChannel.scala
+++ b/core/src/main/scala/kafka/network/RequestChannel.scala
@@ -24,7 +24,7 @@ import com.fasterxml.jackson.databind.JsonNode
 import com.typesafe.scalalogging.Logger
 import com.yammer.metrics.core.Meter
 import kafka.network
-import kafka.server.KafkaConfig
+import kafka.server.{KafkaConfig, RequestLocal}
 import kafka.utils.{Logging, NotNothing, Pool}
 import kafka.utils.Implicits._
 import org.apache.kafka.common.config.ConfigResource
@@ -80,7 +80,7 @@ object RequestChannel extends Logging {
     }
   }
 
-  case class CallbackRequest(fun: () => Unit,
+  case class CallbackRequest(fun: RequestLocal => Unit,
                              originalRequest: Request) extends BaseRequest
 
   class Request(val processor: Int,
diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
index 51e3aff9af866..266fff3ae0226 100755
--- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
+++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
@@ -55,23 +55,23 @@ object KafkaRequestHandler {
    * @param fun Callback function to execute
    * @return Wrapped callback that would execute `fun` on a request thread
    */
-  def wrap[T](fun: T => Unit): T => Unit = {
+  def wrap[T](fun: (RequestLocal, T) => Unit, requestLocal: RequestLocal): T => Unit = {
     val requestChannel = threadRequestChannel.get()
     val currentRequest = threadCurrentRequest.get()
     if (requestChannel == null || currentRequest == null) {
       if (!bypassThreadCheck)
         throw new IllegalStateException("Attempted to reschedule to request handler thread from non-request handler thread.")
-      T => fun(T)
+      T => fun(requestLocal, T)
     } else {
       T => {
         if (threadCurrentRequest.get() != null) {
           // If the callback is actually executed on a request thread, we can directly execute
           // it without re-scheduling it.
-          fun(T)
+          fun(requestLocal, T)
         } else {
           // The requestChannel and request are captured in this lambda, so when it's executed on the callback thread
           // we can re-schedule the original callback on a request thread and update the metrics accordingly.
-          requestChannel.sendCallbackRequest(RequestChannel.CallbackRequest(() => fun(T), currentRequest))
+          requestChannel.sendCallbackRequest(RequestChannel.CallbackRequest(newRequestLocal => fun(newRequestLocal, T), currentRequest))
         }
       }
     }
@@ -132,7 +132,7 @@ class KafkaRequestHandler(
             }
             
             threadCurrentRequest.set(originalRequest)
-            callback.fun()
+            callback.fun(requestLocal)
           } catch {
             case e: FatalExitError =>
               completeShutdown()
diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index 4eb9a6f9c75f1..bcfc8c0287fc2 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -747,96 +747,9 @@ class ReplicaManager(val config: KafkaConfig,
           (verifiedEntries.toMap, unverifiedEntries.toMap, errorEntries.toMap)
         }
 
-      def appendEntries(allEntries: Map[TopicPartition, MemoryRecords])(unverifiedEntries: Map[TopicPartition, Errors]): Unit = {
-        val verifiedEntries =
-          if (unverifiedEntries.isEmpty)
-            allEntries
-          else
-            allEntries.filter { case (tp, _) =>
-              !unverifiedEntries.contains(tp)
-            }
-
-        val localProduceResults = appendToLocalLog(internalTopicsAllowed = internalTopicsAllowed,
-          origin, verifiedEntries, requiredAcks, requestLocal, verificationGuards.toMap)
-        debug("Produce to local log in %d ms".format(time.milliseconds - sTime))
-
-        val errorResults = (unverifiedEntries ++ errorsPerPartition).map {
-          case (topicPartition, error) =>
-            // translate transaction coordinator errors to known producer response errors
-            val customException =
-              error match {
-                case Errors.INVALID_TXN_STATE => Some(error.exception("Partition was not added to the transaction"))
-                case Errors.CONCURRENT_TRANSACTIONS |
-                     Errors.COORDINATOR_LOAD_IN_PROGRESS |
-                     Errors.COORDINATOR_NOT_AVAILABLE |
-                     Errors.NOT_COORDINATOR => Some(new NotEnoughReplicasException(
-                         s"Unable to verify the partition has been added to the transaction. Underlying error: ${error.toString}"))
-                case _ => None
-            }
-            topicPartition -> LogAppendResult(
-              LogAppendInfo.UNKNOWN_LOG_APPEND_INFO,
-              Some(customException.getOrElse(error.exception)),
-              hasCustomErrorMessage = customException.isDefined
-            )
-        }
-
-        val allResults = localProduceResults ++ errorResults
-        val produceStatus = allResults.map { case (topicPartition, result) =>
-          topicPartition -> ProducePartitionStatus(
-            result.info.lastOffset + 1, // required offset
-            new PartitionResponse(
-              result.error,
-              result.info.firstOffset,
-              result.info.lastOffset,
-              result.info.logAppendTime,
-              result.info.logStartOffset,
-              result.info.recordErrors,
-              result.errorMessage
-            )
-          ) // response status
-        }
-
-        actionQueue.add {
-          () => allResults.foreach { case (topicPartition, result) =>
-            val requestKey = TopicPartitionOperationKey(topicPartition)
-            result.info.leaderHwChange match {
-              case LeaderHwChange.INCREASED =>
-                // some delayed operations may be unblocked after HW changed
-                delayedProducePurgatory.checkAndComplete(requestKey)
-                delayedFetchPurgatory.checkAndComplete(requestKey)
-                delayedDeleteRecordsPurgatory.checkAndComplete(requestKey)
-              case LeaderHwChange.SAME =>
-                // probably unblock some follower fetch requests since log end offset has been updated
-                delayedFetchPurgatory.checkAndComplete(requestKey)
-              case LeaderHwChange.NONE =>
-              // nothing
-            }
-          }
-        }
-
-        recordConversionStatsCallback(localProduceResults.map { case (k, v) => k -> v.info.recordConversionStats })
-
-        if (delayedProduceRequestRequired(requiredAcks, allEntries, allResults)) {
-          // create delayed produce operation
-          val produceMetadata = ProduceMetadata(requiredAcks, produceStatus)
-          val delayedProduce = new DelayedProduce(timeout, produceMetadata, this, responseCallback, delayedProduceLock)
-
-          // create a list of (topic, partition) pairs to use as keys for this delayed produce operation
-          val producerRequestKeys = allEntries.keys.map(TopicPartitionOperationKey(_)).toSeq
-
-          // try to complete the request immediately, otherwise put it into the purgatory
-          // this is because while the delayed produce operation is being created, new
-          // requests may arrive and hence make this operation completable.
-          delayedProducePurgatory.tryCompleteElseWatch(delayedProduce, producerRequestKeys)
-        } else {
-          // we can respond immediately
-          val produceResponseStatus = produceStatus.map { case (k, status) => k -> status.responseStatus }
-          responseCallback(produceResponseStatus)
-        }
-      }
-
       if (notYetVerifiedEntriesPerPartition.isEmpty || addPartitionsToTxnManager.isEmpty) {
-        appendEntries(verifiedEntriesPerPartition)(Map.empty)
+        appendEntries(verifiedEntriesPerPartition, internalTopicsAllowed, origin, requiredAcks, verificationGuards.toMap,
+          errorsPerPartition, sTime, recordConversionStatsCallback, timeout, responseCallback, delayedProduceLock)(requestLocal)(Map.empty)
       } else {
         // For unverified entries, send a request to verify. When verified, the append process will proceed via the callback.
         // We verify above that all partitions use the same producer ID.
@@ -846,7 +759,8 @@ class ReplicaManager(val config: KafkaConfig,
           producerId = batchInfo.producerId,
           producerEpoch = batchInfo.producerEpoch,
           topicPartitions = notYetVerifiedEntriesPerPartition.keySet.toSeq,
-          callback = KafkaRequestHandler.wrap(appendEntries(entriesPerPartition)(_))
+          callback = KafkaRequestHandler.wrap(appendEntries(entriesPerPartition, internalTopicsAllowed, origin, requiredAcks, verificationGuards.toMap,
+            errorsPerPartition, sTime, recordConversionStatsCallback, timeout, responseCallback, delayedProduceLock)(_)(_), requestLocal)
         ))
       }
     } else {
@@ -864,6 +778,107 @@ class ReplicaManager(val config: KafkaConfig,
     }
   }
 
+  private def appendEntries(allEntries: Map[TopicPartition, MemoryRecords],
+                            internalTopicsAllowed: Boolean,
+                            origin: AppendOrigin,
+                            requiredAcks: Short,
+                            verificationGuards: Map[TopicPartition, VerificationGuard],
+                            errorsPerPartition: Map[TopicPartition, Errors],
+                            sTime: Long,
+                            recordConversionStatsCallback: Map[TopicPartition, RecordConversionStats] => Unit,
+                            timeout: Long,
+                            responseCallback: Map[TopicPartition, PartitionResponse] => Unit,
+                            delayedProduceLock: Option[Lock])
+                           (requestLocal: RequestLocal)
+                           (unverifiedEntries: Map[TopicPartition, Errors]): Unit = {
+    val verifiedEntries =
+      if (unverifiedEntries.isEmpty)
+        allEntries
+      else
+        allEntries.filter { case (tp, _) =>
+          !unverifiedEntries.contains(tp)
+        }
+
+    val localProduceResults = appendToLocalLog(internalTopicsAllowed = internalTopicsAllowed,
+      origin, verifiedEntries, requiredAcks, requestLocal, verificationGuards.toMap)
+    debug("Produce to local log in %d ms".format(time.milliseconds - sTime))
+
+    val errorResults = (unverifiedEntries ++ errorsPerPartition).map {
+      case (topicPartition, error) =>
+        // translate transaction coordinator errors to known producer response errors
+        val customException =
+          error match {
+            case Errors.INVALID_TXN_STATE => Some(error.exception("Partition was not added to the transaction"))
+            case Errors.CONCURRENT_TRANSACTIONS |
+                 Errors.COORDINATOR_LOAD_IN_PROGRESS |
+                 Errors.COORDINATOR_NOT_AVAILABLE |
+                 Errors.NOT_COORDINATOR => Some(new NotEnoughReplicasException(
+              s"Unable to verify the partition has been added to the transaction. Underlying error: ${error.toString}"))
+            case _ => None
+          }
+        topicPartition -> LogAppendResult(
+          LogAppendInfo.UNKNOWN_LOG_APPEND_INFO,
+          Some(customException.getOrElse(error.exception)),
+          hasCustomErrorMessage = customException.isDefined
+        )
+    }
+
+    val allResults = localProduceResults ++ errorResults
+    val produceStatus = allResults.map { case (topicPartition, result) =>
+      topicPartition -> ProducePartitionStatus(
+        result.info.lastOffset + 1, // required offset
+        new PartitionResponse(
+          result.error,
+          result.info.firstOffset,
+          result.info.lastOffset,
+          result.info.logAppendTime,
+          result.info.logStartOffset,
+          result.info.recordErrors,
+          result.errorMessage
+        )
+      ) // response status
+    }
+
+    actionQueue.add {
+      () =>
+        allResults.foreach { case (topicPartition, result) =>
+          val requestKey = TopicPartitionOperationKey(topicPartition)
+          result.info.leaderHwChange match {
+            case LeaderHwChange.INCREASED =>
+              // some delayed operations may be unblocked after HW changed
+              delayedProducePurgatory.checkAndComplete(requestKey)
+              delayedFetchPurgatory.checkAndComplete(requestKey)
+              delayedDeleteRecordsPurgatory.checkAndComplete(requestKey)
+            case LeaderHwChange.SAME =>
+              // probably unblock some follower fetch requests since log end offset has been updated
+              delayedFetchPurgatory.checkAndComplete(requestKey)
+            case LeaderHwChange.NONE =>
+            // nothing
+          }
+        }
+    }
+
+    recordConversionStatsCallback(localProduceResults.map { case (k, v) => k -> v.info.recordConversionStats })
+
+    if (delayedProduceRequestRequired(requiredAcks, allEntries, allResults)) {
+      // create delayed produce operation
+      val produceMetadata = ProduceMetadata(requiredAcks, produceStatus)
+      val delayedProduce = new DelayedProduce(timeout, produceMetadata, this, responseCallback, delayedProduceLock)
+
+      // create a list of (topic, partition) pairs to use as keys for this delayed produce operation
+      val producerRequestKeys = allEntries.keys.map(TopicPartitionOperationKey(_)).toSeq
+
+      // try to complete the request immediately, otherwise put it into the purgatory
+      // this is because while the delayed produce operation is being created, new
+      // requests may arrive and hence make this operation completable.
+      delayedProducePurgatory.tryCompleteElseWatch(delayedProduce, producerRequestKeys)
+    } else {
+      // we can respond immediately
+      val produceResponseStatus = produceStatus.map { case (k, status) => k -> status.responseStatus }
+      responseCallback(produceResponseStatus)
+    }
+  }
+
   private def partitionEntriesForVerification(verificationGuards: mutable.Map[TopicPartition, VerificationGuard],
                                               entriesPerPartition: Map[TopicPartition, MemoryRecords],
                                               verifiedEntries: mutable.Map[TopicPartition, MemoryRecords],
diff --git a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
index 6016d7c99d395..15b211212869f 100644
--- a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
+++ b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
@@ -32,7 +32,7 @@ import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.ValueSource
 import org.mockito.ArgumentMatchers
 import org.mockito.ArgumentMatchers.any
-import org.mockito.Mockito.{mock, when}
+import org.mockito.Mockito.{mock, times, verify, when}
 
 import java.net.InetAddress
 import java.nio.ByteBuffer
@@ -57,10 +57,10 @@ class KafkaRequestHandlerTest {
       when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
         time.sleep(2)
         // Prepare the callback.
-        val callback = KafkaRequestHandler.wrap((ms: Int) => {
+        val callback = KafkaRequestHandler.wrap((reqLocal: RequestLocal, ms: Int) => {
           time.sleep(ms)
           handler.stop()
-        })
+        }, RequestLocal.NoCaching)
         // Execute the callback asynchronously.
         CompletableFuture.runAsync(() => callback(1))
         request.apiLocalCompleteTimeNanos = time.nanoseconds
@@ -94,9 +94,9 @@ class KafkaRequestHandlerTest {
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       handledCount = handledCount + 1
       // Prepare the callback.
-      val callback = KafkaRequestHandler.wrap((ms: Int) => {
+      val callback = KafkaRequestHandler.wrap((reqLocal: RequestLocal, ms: Int) => {
         handler.stop()
-      })
+      }, RequestLocal.NoCaching)
       // Execute the callback asynchronously.
       CompletableFuture.runAsync(() => callback(1))
     }
@@ -111,6 +111,38 @@ class KafkaRequestHandlerTest {
     assertEquals(1, tryCompleteActionCount)
   }
 
+  @Test
+  def testHandlingCallbackOnNewThread(): Unit = {
+    val time = new MockTime()
+    val metrics = mock(classOf[RequestChannel.Metrics])
+    val apiHandler = mock(classOf[ApiRequestHandler])
+    val requestChannel = new RequestChannel(10, "", time, metrics)
+    val handler = new KafkaRequestHandler(0, 0, mock(classOf[Meter]), new AtomicInteger(1), requestChannel, apiHandler, time)
+
+    val originalRequestLocal = mock(classOf[RequestLocal])
+
+    var handledCount = 0
+
+    val request = makeRequest(time, metrics)
+    requestChannel.sendRequest(request)
+
+    when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
+      // Prepare the callback.
+      val callback = KafkaRequestHandler.wrap((reqLocal: RequestLocal, ms: Int) => {
+        handler.stop()
+        reqLocal.bufferSupplier.close()
+        handledCount = handledCount + 1
+      }, originalRequestLocal)
+      // Execute the callback asynchronously.
+      CompletableFuture.runAsync(() => callback(1))
+    }
+
+    handler.run()
+    // Verify that we don't use the request local that we passed in.
+    verify(originalRequestLocal, times(0)).bufferSupplier
+    assertEquals(1, handledCount)
+  }
+
 
   @ParameterizedTest
   @ValueSource(booleans = Array(true, false))

From ddcf506aff798c1293965050b488ccade0d75db3 Mon Sep 17 00:00:00 2001
From: Justine <jolshan@confluent.io>
Date: Tue, 24 Oct 2023 15:57:16 -0700
Subject: [PATCH 02/11] Add handlingn for the callback request when the same
 request is being handled.

---
 .../kafka/server/KafkaRequestHandler.scala    |  4 +-
 .../server/KafkaRequestHandlerTest.scala      | 37 ++++++++++++++++++-
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
index 266fff3ae0226..af2fa93d9b2c3 100755
--- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
+++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
@@ -64,8 +64,8 @@ object KafkaRequestHandler {
       T => fun(requestLocal, T)
     } else {
       T => {
-        if (threadCurrentRequest.get() != null) {
-          // If the callback is actually executed on a request thread, we can directly execute
+        if (threadCurrentRequest.get() == currentRequest) {
+          // If the callback is actually executed on the same request thread, we can directly execute
           // it without re-scheduling it.
           fun(requestLocal, T)
         } else {
diff --git a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
index 15b211212869f..fd6880e7434e9 100644
--- a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
+++ b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
@@ -24,7 +24,7 @@ import org.apache.kafka.common.network.{ClientInformation, ListenerName}
 import org.apache.kafka.common.protocol.ApiKeys
 import org.apache.kafka.common.requests.{RequestContext, RequestHeader}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
-import org.apache.kafka.common.utils.{MockTime, Time}
+import org.apache.kafka.common.utils.{BufferSupplier, MockTime, Time}
 import org.apache.kafka.server.log.remote.storage.{RemoteLogManagerConfig, RemoteStorageMetrics}
 import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertTrue}
 import org.junit.jupiter.api.Test
@@ -129,9 +129,9 @@ class KafkaRequestHandlerTest {
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       // Prepare the callback.
       val callback = KafkaRequestHandler.wrap((reqLocal: RequestLocal, ms: Int) => {
-        handler.stop()
         reqLocal.bufferSupplier.close()
         handledCount = handledCount + 1
+        handler.stop()
       }, originalRequestLocal)
       // Execute the callback asynchronously.
       CompletableFuture.runAsync(() => callback(1))
@@ -143,6 +143,39 @@ class KafkaRequestHandlerTest {
     assertEquals(1, handledCount)
   }
 
+  @Test
+  def testCallbackOnSameThread(): Unit = {
+    val time = new MockTime()
+    val metrics = mock(classOf[RequestChannel.Metrics])
+    val apiHandler = mock(classOf[ApiRequestHandler])
+    val requestChannel = new RequestChannel(10, "", time, metrics)
+    val handler = new KafkaRequestHandler(0, 0, mock(classOf[Meter]), new AtomicInteger(1), requestChannel, apiHandler, time)
+
+    val originalRequestLocal = mock(classOf[RequestLocal])
+    when(originalRequestLocal.bufferSupplier).thenReturn(BufferSupplier.create())
+
+    var handledCount = 0
+
+    val request = makeRequest(time, metrics)
+    requestChannel.sendRequest(request)
+
+    when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
+      // Prepare the callback.
+      val callback = KafkaRequestHandler.wrap((reqLocal: RequestLocal, ms: Int) => {
+        reqLocal.bufferSupplier.close()
+        handledCount = handledCount + 1
+        handler.stop()
+      }, originalRequestLocal)
+      // Execute the callback before the request returns.
+      callback(1)
+    }
+
+    handler.run()
+    // Verify that we do use the request local that we passed in.
+    verify(originalRequestLocal, times(1)).bufferSupplier
+    assertEquals(1, handledCount)
+  }
+
 
   @ParameterizedTest
   @ValueSource(booleans = Array(true, false))

From 8b5fb69cc7d8d3e4401171b1199d963e2d71f088 Mon Sep 17 00:00:00 2001
From: Justine <jolshan@confluent.io>
Date: Wed, 25 Oct 2023 14:59:18 -0700
Subject: [PATCH 03/11] Add note about thread safety for appendEntries

---
 core/src/main/scala/kafka/server/ReplicaManager.scala | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index bcfc8c0287fc2..242b8abf9a12b 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -778,6 +778,10 @@ class ReplicaManager(val config: KafkaConfig,
     }
   }
 
+  /*
+   * Note: This method can be used as a callback in a different request thread. Ensure the that correct RequestLocal
+   * is passed when executing this method. Accessing non-thread-safe data structures should be avoided if possible.
+   */
   private def appendEntries(allEntries: Map[TopicPartition, MemoryRecords],
                             internalTopicsAllowed: Boolean,
                             origin: AppendOrigin,

From d99536227a5e936c2693b5be44a1ca90d463bdd1 Mon Sep 17 00:00:00 2001
From: Justine <jolshan@confluent.io>
Date: Wed, 25 Oct 2023 15:18:04 -0700
Subject: [PATCH 04/11] typo

---
 core/src/main/scala/kafka/server/ReplicaManager.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index 242b8abf9a12b..a200f8540cd0a 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -779,7 +779,7 @@ class ReplicaManager(val config: KafkaConfig,
   }
 
   /*
-   * Note: This method can be used as a callback in a different request thread. Ensure the that correct RequestLocal
+   * Note: This method can be used as a callback in a different request thread. Ensure that correct RequestLocal
    * is passed when executing this method. Accessing non-thread-safe data structures should be avoided if possible.
    */
   private def appendEntries(allEntries: Map[TopicPartition, MemoryRecords],

From 0e43aa93ff946c265bf52a1fabff6199264fee75 Mon Sep 17 00:00:00 2001
From: Justine <jolshan@confluent.io>
Date: Wed, 25 Oct 2023 17:00:33 -0700
Subject: [PATCH 05/11] add param documentation

---
 core/src/main/scala/kafka/server/KafkaRequestHandler.scala | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
index af2fa93d9b2c3..e6011427c0ba6 100755
--- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
+++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
@@ -53,6 +53,8 @@ object KafkaRequestHandler {
    * Wrap callback to schedule it on a request thread.
    * NOTE: this function must be called on a request thread.
    * @param fun Callback function to execute
+   * @param requestLocal The RequestLocal for the current request handler thread in case we need to call
+   *                     the callback function without queueing the callback request
    * @return Wrapped callback that would execute `fun` on a request thread
    */
   def wrap[T](fun: (RequestLocal, T) => Unit, requestLocal: RequestLocal): T => Unit = {

From 94c1d59c488ec7fa241630a1b5d91cda896d5da9 Mon Sep 17 00:00:00 2001
From: Justine <jolshan@confluent.io>
Date: Fri, 27 Oct 2023 14:15:43 -0700
Subject: [PATCH 06/11] Make thread safe callback clearer

---
 .../kafka/server/KafkaRequestHandler.scala    | 17 +++++---
 .../scala/kafka/server/ReplicaManager.scala   |  7 ++-
 .../server/KafkaRequestHandlerTest.scala      | 43 +++++++++++--------
 3 files changed, 43 insertions(+), 24 deletions(-)

diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
index e6011427c0ba6..e219e54fd36e5 100755
--- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
+++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
@@ -49,31 +49,38 @@ object KafkaRequestHandler {
     bypassThreadCheck = bypassCheck
   }
 
+  /**
+   * A method that expected to be executed once as part of a callback that can be performed in an abitrary request handler thread.
+   * The RequestLocal passed in must belong to the request handler thread that is executing the callback.
+   * The method should be thread-safe. If the method accesses any shared data structures, those accesses must also be thread-safe.
+   */
+  class ThreadSafeCallback[T](val fun: (RequestLocal, T) => Unit)
+
   /**
    * Wrap callback to schedule it on a request thread.
    * NOTE: this function must be called on a request thread.
-   * @param fun Callback function to execute
+   * @param threadSafeCallback thread-safe callback function to execute
    * @param requestLocal The RequestLocal for the current request handler thread in case we need to call
    *                     the callback function without queueing the callback request
    * @return Wrapped callback that would execute `fun` on a request thread
    */
-  def wrap[T](fun: (RequestLocal, T) => Unit, requestLocal: RequestLocal): T => Unit = {
+  def wrap[T](threadSafeCallback: ThreadSafeCallback[T], requestLocal: RequestLocal): T => Unit = {
     val requestChannel = threadRequestChannel.get()
     val currentRequest = threadCurrentRequest.get()
     if (requestChannel == null || currentRequest == null) {
       if (!bypassThreadCheck)
         throw new IllegalStateException("Attempted to reschedule to request handler thread from non-request handler thread.")
-      T => fun(requestLocal, T)
+      T => threadSafeCallback.fun(requestLocal, T)
     } else {
       T => {
         if (threadCurrentRequest.get() == currentRequest) {
           // If the callback is actually executed on the same request thread, we can directly execute
           // it without re-scheduling it.
-          fun(requestLocal, T)
+          threadSafeCallback.fun(requestLocal, T)
         } else {
           // The requestChannel and request are captured in this lambda, so when it's executed on the callback thread
           // we can re-schedule the original callback on a request thread and update the metrics accordingly.
-          requestChannel.sendCallbackRequest(RequestChannel.CallbackRequest(newRequestLocal => fun(newRequestLocal, T), currentRequest))
+          requestChannel.sendCallbackRequest(RequestChannel.CallbackRequest(newRequestLocal => threadSafeCallback.fun(newRequestLocal, T), currentRequest))
         }
       }
     }
diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index a200f8540cd0a..f05792ec8c570 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -23,6 +23,7 @@ import kafka.controller.{KafkaController, StateChangeLogger}
 import kafka.log.remote.RemoteLogManager
 import kafka.log.{LogManager, UnifiedLog}
 import kafka.server.HostedPartition.Online
+import kafka.server.KafkaRequestHandler.ThreadSafeCallback
 import kafka.server.QuotaFactory.QuotaManagers
 import kafka.server.ReplicaManager.{AtMinIsrPartitionCountMetricName, FailedIsrUpdatesPerSecMetricName, IsrExpandsPerSecMetricName, IsrShrinksPerSecMetricName, LeaderCountMetricName, OfflineReplicaCountMetricName, PartitionCountMetricName, PartitionsWithLateTransactionsCountMetricName, ProducerIdCountMetricName, ReassigningPartitionsMetricName, UnderMinIsrPartitionCountMetricName, UnderReplicatedPartitionsMetricName}
 import kafka.server.ReplicaManager.createLogReadResult
@@ -759,8 +760,10 @@ class ReplicaManager(val config: KafkaConfig,
           producerId = batchInfo.producerId,
           producerEpoch = batchInfo.producerEpoch,
           topicPartitions = notYetVerifiedEntriesPerPartition.keySet.toSeq,
-          callback = KafkaRequestHandler.wrap(appendEntries(entriesPerPartition, internalTopicsAllowed, origin, requiredAcks, verificationGuards.toMap,
-            errorsPerPartition, sTime, recordConversionStatsCallback, timeout, responseCallback, delayedProduceLock)(_)(_), requestLocal)
+          callback = KafkaRequestHandler.wrap(
+            new ThreadSafeCallback[Map[TopicPartition, Errors]](appendEntries(entriesPerPartition, internalTopicsAllowed, origin, requiredAcks, verificationGuards.toMap,
+            errorsPerPartition, sTime, recordConversionStatsCallback, timeout, responseCallback, delayedProduceLock)(_)(_)),
+            requestLocal)
         ))
       }
     } else {
diff --git a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
index fd6880e7434e9..3afaead7bbc80 100644
--- a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
+++ b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
@@ -19,6 +19,7 @@ package kafka.server
 
 import com.yammer.metrics.core.Meter
 import kafka.network.RequestChannel
+import kafka.server.KafkaRequestHandler.ThreadSafeCallback
 import org.apache.kafka.common.memory.MemoryPool
 import org.apache.kafka.common.network.{ClientInformation, ListenerName}
 import org.apache.kafka.common.protocol.ApiKeys
@@ -57,10 +58,12 @@ class KafkaRequestHandlerTest {
       when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
         time.sleep(2)
         // Prepare the callback.
-        val callback = KafkaRequestHandler.wrap((reqLocal: RequestLocal, ms: Int) => {
-          time.sleep(ms)
-          handler.stop()
-        }, RequestLocal.NoCaching)
+        val callback = KafkaRequestHandler.wrap(
+          new ThreadSafeCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+            time.sleep(ms)
+            handler.stop()
+          }),
+          RequestLocal.NoCaching)
         // Execute the callback asynchronously.
         CompletableFuture.runAsync(() => callback(1))
         request.apiLocalCompleteTimeNanos = time.nanoseconds
@@ -94,9 +97,11 @@ class KafkaRequestHandlerTest {
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       handledCount = handledCount + 1
       // Prepare the callback.
-      val callback = KafkaRequestHandler.wrap((reqLocal: RequestLocal, ms: Int) => {
-        handler.stop()
-      }, RequestLocal.NoCaching)
+      val callback = KafkaRequestHandler.wrap(
+        new ThreadSafeCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+          handler.stop()
+        }),
+        RequestLocal.NoCaching)
       // Execute the callback asynchronously.
       CompletableFuture.runAsync(() => callback(1))
     }
@@ -128,11 +133,13 @@ class KafkaRequestHandlerTest {
 
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       // Prepare the callback.
-      val callback = KafkaRequestHandler.wrap((reqLocal: RequestLocal, ms: Int) => {
-        reqLocal.bufferSupplier.close()
-        handledCount = handledCount + 1
-        handler.stop()
-      }, originalRequestLocal)
+      val callback = KafkaRequestHandler.wrap(
+        new ThreadSafeCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+          reqLocal.bufferSupplier.close()
+          handledCount = handledCount + 1
+          handler.stop()
+        }),
+        originalRequestLocal)
       // Execute the callback asynchronously.
       CompletableFuture.runAsync(() => callback(1))
     }
@@ -161,11 +168,13 @@ class KafkaRequestHandlerTest {
 
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       // Prepare the callback.
-      val callback = KafkaRequestHandler.wrap((reqLocal: RequestLocal, ms: Int) => {
-        reqLocal.bufferSupplier.close()
-        handledCount = handledCount + 1
-        handler.stop()
-      }, originalRequestLocal)
+      val callback = KafkaRequestHandler.wrap(
+        new ThreadSafeCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+          reqLocal.bufferSupplier.close()
+          handledCount = handledCount + 1
+          handler.stop()
+        }),
+        originalRequestLocal)
       // Execute the callback before the request returns.
       callback(1)
     }

From 7a0c6a7134157073496f83a45a209d37d8b77cd8 Mon Sep 17 00:00:00 2001
From: Justine <jolshan@confluent.io>
Date: Mon, 30 Oct 2023 15:26:42 -0700
Subject: [PATCH 07/11] remove the channel thread local

---
 core/src/main/scala/kafka/server/KafkaRequestHandler.scala | 1 +
 1 file changed, 1 insertion(+)

diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
index e219e54fd36e5..4a281c7064b9e 100755
--- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
+++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
@@ -183,6 +183,7 @@ class KafkaRequestHandler(
 
   private def completeShutdown(): Unit = {
     requestLocal.close()
+    threadRequestChannel.remove()
     shutdownComplete.countDown()
   }
 

From f53d5eb75b219cf378aabd717071ab4b153bc2e9 Mon Sep 17 00:00:00 2001
From: Justine <jolshan@confluent.io>
Date: Tue, 31 Oct 2023 11:25:09 -0700
Subject: [PATCH 08/11] Rename the callback and fix comments

---
 .../kafka/server/KafkaRequestHandler.scala    | 21 +++++++++----------
 .../scala/kafka/server/ReplicaManager.scala   |  4 ++--
 .../server/KafkaRequestHandlerTest.scala      | 10 ++++-----
 3 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
index 4a281c7064b9e..fc216d6a6db1c 100755
--- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
+++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
@@ -50,37 +50,36 @@ object KafkaRequestHandler {
   }
 
   /**
-   * A method that expected to be executed once as part of a callback that can be performed in an abitrary request handler thread.
+   * A callback method that expected to be executed once in an arbitrary request handler thread after an asynchronous action completes.
    * The RequestLocal passed in must belong to the request handler thread that is executing the callback.
-   * The method should be thread-safe. If the method accesses any shared data structures, those accesses must also be thread-safe.
    */
-  class ThreadSafeCallback[T](val fun: (RequestLocal, T) => Unit)
+  class AsynchronousCompletionCallback[T](val fun: (RequestLocal, T) => Unit)
 
   /**
-   * Wrap callback to schedule it on a request thread.
-   * NOTE: this function must be called on a request thread.
-   * @param threadSafeCallback thread-safe callback function to execute
+   * Wrap callback to schedule it on an abitrary request thread.
+   * NOTE: this function must be originally called from a request thread.
+   * @param asyncCompletionCallback a callback function to execute as the result of an asynchronous action completing
    * @param requestLocal The RequestLocal for the current request handler thread in case we need to call
    *                     the callback function without queueing the callback request
-   * @return Wrapped callback that would execute `fun` on a request thread
+   * @return Wrapped callback that would execute `asyncCompletionCallback` on an arbitrary request thread
    */
-  def wrap[T](threadSafeCallback: ThreadSafeCallback[T], requestLocal: RequestLocal): T => Unit = {
+  def wrap[T](asyncCompletionCallback: AsynchronousCompletionCallback[T], requestLocal: RequestLocal): T => Unit = {
     val requestChannel = threadRequestChannel.get()
     val currentRequest = threadCurrentRequest.get()
     if (requestChannel == null || currentRequest == null) {
       if (!bypassThreadCheck)
         throw new IllegalStateException("Attempted to reschedule to request handler thread from non-request handler thread.")
-      T => threadSafeCallback.fun(requestLocal, T)
+      T => asyncCompletionCallback.fun(requestLocal, T)
     } else {
       T => {
         if (threadCurrentRequest.get() == currentRequest) {
           // If the callback is actually executed on the same request thread, we can directly execute
           // it without re-scheduling it.
-          threadSafeCallback.fun(requestLocal, T)
+          asyncCompletionCallback.fun(requestLocal, T)
         } else {
           // The requestChannel and request are captured in this lambda, so when it's executed on the callback thread
           // we can re-schedule the original callback on a request thread and update the metrics accordingly.
-          requestChannel.sendCallbackRequest(RequestChannel.CallbackRequest(newRequestLocal => threadSafeCallback.fun(newRequestLocal, T), currentRequest))
+          requestChannel.sendCallbackRequest(RequestChannel.CallbackRequest(newRequestLocal => asyncCompletionCallback.fun(newRequestLocal, T), currentRequest))
         }
       }
     }
diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index f05792ec8c570..d90ad2aedd2d1 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -23,7 +23,7 @@ import kafka.controller.{KafkaController, StateChangeLogger}
 import kafka.log.remote.RemoteLogManager
 import kafka.log.{LogManager, UnifiedLog}
 import kafka.server.HostedPartition.Online
-import kafka.server.KafkaRequestHandler.ThreadSafeCallback
+import kafka.server.KafkaRequestHandler.AsynchronousCompletionCallback
 import kafka.server.QuotaFactory.QuotaManagers
 import kafka.server.ReplicaManager.{AtMinIsrPartitionCountMetricName, FailedIsrUpdatesPerSecMetricName, IsrExpandsPerSecMetricName, IsrShrinksPerSecMetricName, LeaderCountMetricName, OfflineReplicaCountMetricName, PartitionCountMetricName, PartitionsWithLateTransactionsCountMetricName, ProducerIdCountMetricName, ReassigningPartitionsMetricName, UnderMinIsrPartitionCountMetricName, UnderReplicatedPartitionsMetricName}
 import kafka.server.ReplicaManager.createLogReadResult
@@ -761,7 +761,7 @@ class ReplicaManager(val config: KafkaConfig,
           producerEpoch = batchInfo.producerEpoch,
           topicPartitions = notYetVerifiedEntriesPerPartition.keySet.toSeq,
           callback = KafkaRequestHandler.wrap(
-            new ThreadSafeCallback[Map[TopicPartition, Errors]](appendEntries(entriesPerPartition, internalTopicsAllowed, origin, requiredAcks, verificationGuards.toMap,
+            new AsynchronousCompletionCallback[Map[TopicPartition, Errors]](appendEntries(entriesPerPartition, internalTopicsAllowed, origin, requiredAcks, verificationGuards.toMap,
             errorsPerPartition, sTime, recordConversionStatsCallback, timeout, responseCallback, delayedProduceLock)(_)(_)),
             requestLocal)
         ))
diff --git a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
index 3afaead7bbc80..e2588069bbfcb 100644
--- a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
+++ b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
@@ -19,7 +19,7 @@ package kafka.server
 
 import com.yammer.metrics.core.Meter
 import kafka.network.RequestChannel
-import kafka.server.KafkaRequestHandler.ThreadSafeCallback
+import kafka.server.KafkaRequestHandler.AsynchronousCompletionCallback
 import org.apache.kafka.common.memory.MemoryPool
 import org.apache.kafka.common.network.{ClientInformation, ListenerName}
 import org.apache.kafka.common.protocol.ApiKeys
@@ -59,7 +59,7 @@ class KafkaRequestHandlerTest {
         time.sleep(2)
         // Prepare the callback.
         val callback = KafkaRequestHandler.wrap(
-          new ThreadSafeCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+          new AsynchronousCompletionCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
             time.sleep(ms)
             handler.stop()
           }),
@@ -98,7 +98,7 @@ class KafkaRequestHandlerTest {
       handledCount = handledCount + 1
       // Prepare the callback.
       val callback = KafkaRequestHandler.wrap(
-        new ThreadSafeCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+        new AsynchronousCompletionCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
           handler.stop()
         }),
         RequestLocal.NoCaching)
@@ -134,7 +134,7 @@ class KafkaRequestHandlerTest {
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       // Prepare the callback.
       val callback = KafkaRequestHandler.wrap(
-        new ThreadSafeCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+        new AsynchronousCompletionCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
           reqLocal.bufferSupplier.close()
           handledCount = handledCount + 1
           handler.stop()
@@ -169,7 +169,7 @@ class KafkaRequestHandlerTest {
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       // Prepare the callback.
       val callback = KafkaRequestHandler.wrap(
-        new ThreadSafeCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+        new AsynchronousCompletionCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
           reqLocal.bufferSupplier.close()
           handledCount = handledCount + 1
           handler.stop()

From a0e1cca4e47faf09607be26f7ac2732ad7a33950 Mon Sep 17 00:00:00 2001
From: Justine <jolshan@confluent.io>
Date: Tue, 31 Oct 2023 17:31:25 -0700
Subject: [PATCH 09/11] Update comments and clean up parameters

---
 .../kafka/server/KafkaRequestHandler.scala    |  4 ++--
 .../scala/kafka/server/ReplicaManager.scala   | 22 +++++++++++++------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
index fc216d6a6db1c..170eb12a60f74 100755
--- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
+++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
@@ -56,12 +56,12 @@ object KafkaRequestHandler {
   class AsynchronousCompletionCallback[T](val fun: (RequestLocal, T) => Unit)
 
   /**
-   * Wrap callback to schedule it on an abitrary request thread.
+   * Wrap callback to schedule it on an arbitrary request thread.
    * NOTE: this function must be originally called from a request thread.
    * @param asyncCompletionCallback a callback function to execute as the result of an asynchronous action completing
    * @param requestLocal The RequestLocal for the current request handler thread in case we need to call
    *                     the callback function without queueing the callback request
-   * @return Wrapped callback that would execute `asyncCompletionCallback` on an arbitrary request thread
+   * @return Wrapped callback that schedules `asyncCompletionCallback` on an arbitrary request thread
    */
   def wrap[T](asyncCompletionCallback: AsynchronousCompletionCallback[T], requestLocal: RequestLocal): T => Unit = {
     val requestChannel = threadRequestChannel.get()
diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index d90ad2aedd2d1..8089d14f579b1 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -734,7 +734,6 @@ class ReplicaManager(val config: KafkaConfig,
                     transactionalId: String = null,
                     actionQueue: ActionQueue = this.actionQueue): Unit = {
     if (isValidRequiredAcks(requiredAcks)) {
-      val sTime = time.milliseconds
 
       val verificationGuards: mutable.Map[TopicPartition, VerificationGuard] = mutable.Map[TopicPartition, VerificationGuard]()
       val (verifiedEntriesPerPartition, notYetVerifiedEntriesPerPartition, errorsPerPartition) =
@@ -750,7 +749,7 @@ class ReplicaManager(val config: KafkaConfig,
 
       if (notYetVerifiedEntriesPerPartition.isEmpty || addPartitionsToTxnManager.isEmpty) {
         appendEntries(verifiedEntriesPerPartition, internalTopicsAllowed, origin, requiredAcks, verificationGuards.toMap,
-          errorsPerPartition, sTime, recordConversionStatsCallback, timeout, responseCallback, delayedProduceLock)(requestLocal)(Map.empty)
+          errorsPerPartition, recordConversionStatsCallback, timeout, responseCallback, delayedProduceLock)(requestLocal, Map.empty)
       } else {
         // For unverified entries, send a request to verify. When verified, the append process will proceed via the callback.
         // We verify above that all partitions use the same producer ID.
@@ -761,8 +760,18 @@ class ReplicaManager(val config: KafkaConfig,
           producerEpoch = batchInfo.producerEpoch,
           topicPartitions = notYetVerifiedEntriesPerPartition.keySet.toSeq,
           callback = KafkaRequestHandler.wrap(
-            new AsynchronousCompletionCallback[Map[TopicPartition, Errors]](appendEntries(entriesPerPartition, internalTopicsAllowed, origin, requiredAcks, verificationGuards.toMap,
-            errorsPerPartition, sTime, recordConversionStatsCallback, timeout, responseCallback, delayedProduceLock)(_)(_)),
+            new AsynchronousCompletionCallback[Map[TopicPartition, Errors]](appendEntries(
+              entriesPerPartition,
+              internalTopicsAllowed,
+              origin,
+              requiredAcks,
+              verificationGuards.toMap,
+              errorsPerPartition,
+              recordConversionStatsCallback,
+              timeout,
+              responseCallback,
+              delayedProduceLock
+            )),
             requestLocal)
         ))
       }
@@ -791,13 +800,12 @@ class ReplicaManager(val config: KafkaConfig,
                             requiredAcks: Short,
                             verificationGuards: Map[TopicPartition, VerificationGuard],
                             errorsPerPartition: Map[TopicPartition, Errors],
-                            sTime: Long,
                             recordConversionStatsCallback: Map[TopicPartition, RecordConversionStats] => Unit,
                             timeout: Long,
                             responseCallback: Map[TopicPartition, PartitionResponse] => Unit,
                             delayedProduceLock: Option[Lock])
-                           (requestLocal: RequestLocal)
-                           (unverifiedEntries: Map[TopicPartition, Errors]): Unit = {
+                           (requestLocal: RequestLocal, unverifiedEntries: Map[TopicPartition, Errors]): Unit = {
+    val sTime = time.milliseconds
     val verifiedEntries =
       if (unverifiedEntries.isEmpty)
         allEntries

From 9cca64580dfad593d83a1a3c02a24d2463cafd06 Mon Sep 17 00:00:00 2001
From: Justine <jolshan@confluent.io>
Date: Fri, 3 Nov 2023 10:40:47 -0700
Subject: [PATCH 10/11] Simplify callback argument

---
 .../kafka/server/KafkaRequestHandler.scala    | 24 ++++++++----------
 .../scala/kafka/server/ReplicaManager.scala   |  7 +++---
 .../server/KafkaRequestHandlerTest.scala      | 25 +++++++++----------
 3 files changed, 25 insertions(+), 31 deletions(-)

diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
index 170eb12a60f74..3853a72d05341 100755
--- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
+++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
@@ -50,36 +50,32 @@ object KafkaRequestHandler {
   }
 
   /**
-   * A callback method that expected to be executed once in an arbitrary request handler thread after an asynchronous action completes.
-   * The RequestLocal passed in must belong to the request handler thread that is executing the callback.
-   */
-  class AsynchronousCompletionCallback[T](val fun: (RequestLocal, T) => Unit)
-
-  /**
-   * Wrap callback to schedule it on an arbitrary request thread.
+   * Execute or create a callback to be asynchronously scheduled on an arbitrary request thread
    * NOTE: this function must be originally called from a request thread.
-   * @param asyncCompletionCallback a callback function to execute as the result of an asynchronous action completing
-   * @param requestLocal The RequestLocal for the current request handler thread in case we need to call
-   *                     the callback function without queueing the callback request
+   * @param asyncCompletionCallback A callback method that is expected to be executed once in an arbitrary request
+   *                                handler thread after an asynchronous action completes. The RequestLocal passed in
+   *                                must belong to the request handler thread that is executing the callback.
+   * @param requestLocal The RequestLocal for the current request handler thread in case we need to execute the callback
+   *                     function immediately without queueing the callback request
    * @return Wrapped callback that schedules `asyncCompletionCallback` on an arbitrary request thread
    */
-  def wrap[T](asyncCompletionCallback: AsynchronousCompletionCallback[T], requestLocal: RequestLocal): T => Unit = {
+  def executeOrRegisterAsyncCallback[T](asyncCompletionCallback: (RequestLocal, T) => Unit, requestLocal: RequestLocal): T => Unit = {
     val requestChannel = threadRequestChannel.get()
     val currentRequest = threadCurrentRequest.get()
     if (requestChannel == null || currentRequest == null) {
       if (!bypassThreadCheck)
         throw new IllegalStateException("Attempted to reschedule to request handler thread from non-request handler thread.")
-      T => asyncCompletionCallback.fun(requestLocal, T)
+      T => asyncCompletionCallback(requestLocal, T)
     } else {
       T => {
         if (threadCurrentRequest.get() == currentRequest) {
           // If the callback is actually executed on the same request thread, we can directly execute
           // it without re-scheduling it.
-          asyncCompletionCallback.fun(requestLocal, T)
+          asyncCompletionCallback(requestLocal, T)
         } else {
           // The requestChannel and request are captured in this lambda, so when it's executed on the callback thread
           // we can re-schedule the original callback on a request thread and update the metrics accordingly.
-          requestChannel.sendCallbackRequest(RequestChannel.CallbackRequest(newRequestLocal => asyncCompletionCallback.fun(newRequestLocal, T), currentRequest))
+          requestChannel.sendCallbackRequest(RequestChannel.CallbackRequest(newRequestLocal => asyncCompletionCallback(newRequestLocal, T), currentRequest))
         }
       }
     }
diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index 8089d14f579b1..44ef1a8b6f64a 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -23,7 +23,6 @@ import kafka.controller.{KafkaController, StateChangeLogger}
 import kafka.log.remote.RemoteLogManager
 import kafka.log.{LogManager, UnifiedLog}
 import kafka.server.HostedPartition.Online
-import kafka.server.KafkaRequestHandler.AsynchronousCompletionCallback
 import kafka.server.QuotaFactory.QuotaManagers
 import kafka.server.ReplicaManager.{AtMinIsrPartitionCountMetricName, FailedIsrUpdatesPerSecMetricName, IsrExpandsPerSecMetricName, IsrShrinksPerSecMetricName, LeaderCountMetricName, OfflineReplicaCountMetricName, PartitionCountMetricName, PartitionsWithLateTransactionsCountMetricName, ProducerIdCountMetricName, ReassigningPartitionsMetricName, UnderMinIsrPartitionCountMetricName, UnderReplicatedPartitionsMetricName}
 import kafka.server.ReplicaManager.createLogReadResult
@@ -759,8 +758,8 @@ class ReplicaManager(val config: KafkaConfig,
           producerId = batchInfo.producerId,
           producerEpoch = batchInfo.producerEpoch,
           topicPartitions = notYetVerifiedEntriesPerPartition.keySet.toSeq,
-          callback = KafkaRequestHandler.wrap(
-            new AsynchronousCompletionCallback[Map[TopicPartition, Errors]](appendEntries(
+          callback = KafkaRequestHandler.executeOrRegisterAsyncCallback(
+            appendEntries(
               entriesPerPartition,
               internalTopicsAllowed,
               origin,
@@ -771,7 +770,7 @@ class ReplicaManager(val config: KafkaConfig,
               timeout,
               responseCallback,
               delayedProduceLock
-            )),
+            ),
             requestLocal)
         ))
       }
diff --git a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
index e2588069bbfcb..d897a70f3bdfe 100644
--- a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
+++ b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
@@ -19,7 +19,6 @@ package kafka.server
 
 import com.yammer.metrics.core.Meter
 import kafka.network.RequestChannel
-import kafka.server.KafkaRequestHandler.AsynchronousCompletionCallback
 import org.apache.kafka.common.memory.MemoryPool
 import org.apache.kafka.common.network.{ClientInformation, ListenerName}
 import org.apache.kafka.common.protocol.ApiKeys
@@ -58,11 +57,11 @@ class KafkaRequestHandlerTest {
       when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
         time.sleep(2)
         // Prepare the callback.
-        val callback = KafkaRequestHandler.wrap(
-          new AsynchronousCompletionCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+        val callback = KafkaRequestHandler.executeOrRegisterAsyncCallback(
+          (reqLocal: RequestLocal, ms: Int) => {
             time.sleep(ms)
             handler.stop()
-          }),
+          },
           RequestLocal.NoCaching)
         // Execute the callback asynchronously.
         CompletableFuture.runAsync(() => callback(1))
@@ -97,10 +96,10 @@ class KafkaRequestHandlerTest {
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       handledCount = handledCount + 1
       // Prepare the callback.
-      val callback = KafkaRequestHandler.wrap(
-        new AsynchronousCompletionCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+      val callback = KafkaRequestHandler.executeOrRegisterAsyncCallback(
+        (reqLocal: RequestLocal, ms: Int) => {
           handler.stop()
-        }),
+        },
         RequestLocal.NoCaching)
       // Execute the callback asynchronously.
       CompletableFuture.runAsync(() => callback(1))
@@ -133,12 +132,12 @@ class KafkaRequestHandlerTest {
 
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       // Prepare the callback.
-      val callback = KafkaRequestHandler.wrap(
-        new AsynchronousCompletionCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+      val callback = KafkaRequestHandler.executeOrRegisterAsyncCallback(
+        (reqLocal: RequestLocal, ms: Int) => {
           reqLocal.bufferSupplier.close()
           handledCount = handledCount + 1
           handler.stop()
-        }),
+        },
         originalRequestLocal)
       // Execute the callback asynchronously.
       CompletableFuture.runAsync(() => callback(1))
@@ -168,12 +167,12 @@ class KafkaRequestHandlerTest {
 
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       // Prepare the callback.
-      val callback = KafkaRequestHandler.wrap(
-        new AsynchronousCompletionCallback[Int]((reqLocal: RequestLocal, ms: Int) => {
+      val callback = KafkaRequestHandler.executeOrRegisterAsyncCallback(
+        (reqLocal: RequestLocal, ms: Int) => {
           reqLocal.bufferSupplier.close()
           handledCount = handledCount + 1
           handler.stop()
-        }),
+        },
         originalRequestLocal)
       // Execute the callback before the request returns.
       callback(1)

From 4f1b9deaa226ff1abd666abe0d1089c5102bef0a Mon Sep 17 00:00:00 2001
From: Justine <jolshan@confluent.io>
Date: Fri, 3 Nov 2023 15:34:05 -0700
Subject: [PATCH 11/11] Rename wrap again and update comments

---
 .../scala/kafka/server/KafkaRequestHandler.scala | 16 +++++++++-------
 .../main/scala/kafka/server/ReplicaManager.scala |  2 +-
 .../kafka/server/KafkaRequestHandlerTest.scala   |  8 ++++----
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
index 3853a72d05341..f5974b1e93f6e 100755
--- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
+++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
@@ -50,16 +50,18 @@ object KafkaRequestHandler {
   }
 
   /**
-   * Execute or create a callback to be asynchronously scheduled on an arbitrary request thread
+   * Creates a wrapped callback to be executed synchronously on the calling request thread or asynchronously
+   * on an arbitrary request thread.
    * NOTE: this function must be originally called from a request thread.
-   * @param asyncCompletionCallback A callback method that is expected to be executed once in an arbitrary request
-   *                                handler thread after an asynchronous action completes. The RequestLocal passed in
-   *                                must belong to the request handler thread that is executing the callback.
+   * @param asyncCompletionCallback A callback method that we intend to call from the current thread or in another
+   *                                thread after an asynchronous action completes. The RequestLocal passed in must
+   *                                belong to the request handler thread that is executing the callback.
    * @param requestLocal The RequestLocal for the current request handler thread in case we need to execute the callback
-   *                     function immediately without queueing the callback request
-   * @return Wrapped callback that schedules `asyncCompletionCallback` on an arbitrary request thread
+   *                     function synchronously from the calling thread.
+   * @return Wrapped callback will either immediately execute `asyncCompletionCallback` or schedule it on an arbitrary request thread
+   *         depending on where it is called
    */
-  def executeOrRegisterAsyncCallback[T](asyncCompletionCallback: (RequestLocal, T) => Unit, requestLocal: RequestLocal): T => Unit = {
+  def wrapAsyncCallback[T](asyncCompletionCallback: (RequestLocal, T) => Unit, requestLocal: RequestLocal): T => Unit = {
     val requestChannel = threadRequestChannel.get()
     val currentRequest = threadCurrentRequest.get()
     if (requestChannel == null || currentRequest == null) {
diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index 44ef1a8b6f64a..443f3c171cae3 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -758,7 +758,7 @@ class ReplicaManager(val config: KafkaConfig,
           producerId = batchInfo.producerId,
           producerEpoch = batchInfo.producerEpoch,
           topicPartitions = notYetVerifiedEntriesPerPartition.keySet.toSeq,
-          callback = KafkaRequestHandler.executeOrRegisterAsyncCallback(
+          callback = KafkaRequestHandler.wrapAsyncCallback(
             appendEntries(
               entriesPerPartition,
               internalTopicsAllowed,
diff --git a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
index d897a70f3bdfe..ac3d81f1a0acc 100644
--- a/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
+++ b/core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala
@@ -57,7 +57,7 @@ class KafkaRequestHandlerTest {
       when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
         time.sleep(2)
         // Prepare the callback.
-        val callback = KafkaRequestHandler.executeOrRegisterAsyncCallback(
+        val callback = KafkaRequestHandler.wrapAsyncCallback(
           (reqLocal: RequestLocal, ms: Int) => {
             time.sleep(ms)
             handler.stop()
@@ -96,7 +96,7 @@ class KafkaRequestHandlerTest {
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       handledCount = handledCount + 1
       // Prepare the callback.
-      val callback = KafkaRequestHandler.executeOrRegisterAsyncCallback(
+      val callback = KafkaRequestHandler.wrapAsyncCallback(
         (reqLocal: RequestLocal, ms: Int) => {
           handler.stop()
         },
@@ -132,7 +132,7 @@ class KafkaRequestHandlerTest {
 
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       // Prepare the callback.
-      val callback = KafkaRequestHandler.executeOrRegisterAsyncCallback(
+      val callback = KafkaRequestHandler.wrapAsyncCallback(
         (reqLocal: RequestLocal, ms: Int) => {
           reqLocal.bufferSupplier.close()
           handledCount = handledCount + 1
@@ -167,7 +167,7 @@ class KafkaRequestHandlerTest {
 
     when(apiHandler.handle(ArgumentMatchers.eq(request), any())).thenAnswer { _ =>
       // Prepare the callback.
-      val callback = KafkaRequestHandler.executeOrRegisterAsyncCallback(
+      val callback = KafkaRequestHandler.wrapAsyncCallback(
         (reqLocal: RequestLocal, ms: Int) => {
           reqLocal.bufferSupplier.close()
           handledCount = handledCount + 1