KAFKA-14505; [6/N] Avoid recheduling callback in request thread

core/src/main/scala/kafka/coordinator/group/CoordinatorPartitionWriter.scala

@@ -197,8 +197,8 @@ class CoordinatorPartitionWriter[T](
       producerId = producerId,
       producerEpoch = producerEpoch,
       baseSequence = RecordBatch.NO_SEQUENCE,
-      requestLocal = RequestLocal.NoCaching,
-      callback = (error, _, verificationGuard) => {
+      callback = errorAndGuard => {
+        val (error, verificationGuard) = errorAndGuard
         if (error != Errors.NONE) {
           future.completeExceptionally(error.exception)
         } else {

core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala

@@ -916,10 +916,10 @@ private[group] class GroupCoordinator(
         val offsetTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, partitionFor(group.groupId))
 
         def postVerificationCallback(
-          error: Errors,
           newRequestLocal: RequestLocal,
-          verificationGuard: VerificationGuard
+          errorAndGuard: (Errors, VerificationGuard)
         ): Unit = {
+          val (error, verificationGuard) = errorAndGuard
           if (error != Errors.NONE) {
             val finalError = GroupMetadataManager.maybeConvertOffsetCommitError(error)
             responseCallback(offsetMetadata.map { case (k, _) => k -> finalError })
@@ -935,8 +935,13 @@ private[group] class GroupCoordinator(
           producerId,
           producerEpoch,
           RecordBatch.NO_SEQUENCE,
-          requestLocal,
-          postVerificationCallback
+          // Wrap the callback to be handled on an arbitrary request handler thread
+          // when transaction verification is complete. The request local passed in
+          // is only used when the callback is executed immediately.
+          KafkaRequestHandler.wrapAsyncCallback(
+            postVerificationCallback,
+            requestLocal
+          )
         )
     }
   }

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

@@ -800,6 +800,9 @@ class ReplicaManager(val config: KafkaConfig,
    * @param requestLocal                  container for the stateful instances scoped to this request -- this must correspond to the
    *                                      thread calling this method
    * @param actionQueue                   the action queue to use. ReplicaManager#defaultActionQueue is used by default.
+   *
+   * The responseCallback is wrapped so that it is scheduled on a request handler thread. There, it should be called with
+   * that request handler thread's thread local and not the one supplied to this method.
    */
   def handleProduceAppend(timeout: Long,
                           requiredAcks: Short,
@@ -813,19 +816,19 @@ class ReplicaManager(val config: KafkaConfig,
 
     val transactionalProducerInfo = mutable.HashSet[(Long, Short)]()
     val topicPartitionBatchInfo = mutable.Map[TopicPartition, Int]()
-    entriesPerPartition.foreach { case (topicPartition, records) =>
+    entriesPerPartition.forKeyValue { (topicPartition, records) =>
       // Produce requests (only requests that require verification) should only have one batch per partition in "batches" but check all just to be safe.
       val transactionalBatches = records.batches.asScala.filter(batch => batch.hasProducerId && batch.isTransactional)
       transactionalBatches.foreach(batch => transactionalProducerInfo.add(batch.producerId, batch.producerEpoch))
-      if (!transactionalBatches.isEmpty) topicPartitionBatchInfo.put(topicPartition, records.firstBatch.baseSequence)
+      if (transactionalBatches.nonEmpty) topicPartitionBatchInfo.put(topicPartition, records.firstBatch.baseSequence)
     }
     if (transactionalProducerInfo.size > 1) {
       throw new InvalidPidMappingException("Transactional records contained more than one producer ID")
     }
 
-    def postVerificationCallback(preAppendErrors: Map[TopicPartition, Errors],
-                                 newRequestLocal: RequestLocal,
-                                 verificationGuards: Map[TopicPartition, VerificationGuard]): Unit = {
+    def postVerificationCallback(newRequestLocal: RequestLocal,
+                                 results: (Map[TopicPartition, Errors], Map[TopicPartition, VerificationGuard])): Unit = {
+      val (preAppendErrors, verificationGuards) = results
       val errorResults = preAppendErrors.map {
         case (topicPartition, error) =>
           // translate transaction coordinator errors to known producer response errors
@@ -868,12 +871,26 @@ class ReplicaManager(val config: KafkaConfig,
     }
 
     if (transactionalProducerInfo.size < 1) {
-      postVerificationCallback(Map.empty[TopicPartition, Errors], requestLocal, Map.empty[TopicPartition, VerificationGuard])
+      postVerificationCallback(
+        requestLocal,
+        (Map.empty[TopicPartition, Errors], Map.empty[TopicPartition, VerificationGuard])
+      )
       return
     }
 
-    maybeStartTransactionVerificationForPartitions(topicPartitionBatchInfo, transactionalId,
-      transactionalProducerInfo.head._1, transactionalProducerInfo.head._2, requestLocal, postVerificationCallback)
+    maybeStartTransactionVerificationForPartitions(
+      topicPartitionBatchInfo,
+      transactionalId,
+      transactionalProducerInfo.head._1,
+      transactionalProducerInfo.head._2,
+      // Wrap the callback to be handled on an arbitrary request handler thread
+      // when transaction verification is complete. The request local passed in
+      // is only used when the callback is executed immediately.
+      KafkaRequestHandler.wrapAsyncCallback(
+        postVerificationCallback,
+        requestLocal
+      )
+    )
   }
 
   private def buildProducePartitionStatus(
@@ -968,8 +985,6 @@ class ReplicaManager(val config: KafkaConfig,
    * @param producerId      the producer id for the producer writing to the transaction
    * @param producerEpoch   the epoch of the producer writing to the transaction
    * @param baseSequence    the base sequence of the first record in the batch we are trying to append
-   * @param requestLocal    container for the stateful instances scoped to this request -- this must correspond to the
-   *                        thread calling this method
    * @param callback        the method to execute once the verification is either completed or returns an error
    *
    * When the verification returns, the callback will be supplied the error if it exists or Errors.NONE.
@@ -982,24 +997,21 @@ class ReplicaManager(val config: KafkaConfig,
     producerId: Long,
     producerEpoch: Short,
     baseSequence: Int,
-    requestLocal: RequestLocal,
-    callback: (Errors, RequestLocal, VerificationGuard) => Unit
+    callback: ((Errors, VerificationGuard)) => Unit
   ): Unit = {
-    def generalizedCallback(preAppendErrors: Map[TopicPartition, Errors],
-                            newRequestLocal: RequestLocal,
-                            verificationGuards: Map[TopicPartition, VerificationGuard]): Unit = {
-      callback(
+    def generalizedCallback(results: (Map[TopicPartition, Errors], Map[TopicPartition, VerificationGuard])): Unit = {
+      val (preAppendErrors, verificationGuards) = results
+      callback((
         preAppendErrors.getOrElse(topicPartition, Errors.NONE),
-        newRequestLocal,
-        verificationGuards.getOrElse(topicPartition, VerificationGuard.SENTINEL))
+        verificationGuards.getOrElse(topicPartition, VerificationGuard.SENTINEL)
+      ))
     }
 
     maybeStartTransactionVerificationForPartitions(
       Map(topicPartition -> baseSequence),
       transactionalId,
       producerId,
       producerEpoch,
-      requestLocal,
       generalizedCallback
     )
   }
@@ -1010,31 +1022,26 @@ class ReplicaManager(val config: KafkaConfig,
    * @param transactionalId          the transactional id for the transaction
    * @param producerId               the producer id for the producer writing to the transaction
    * @param producerEpoch            the epoch of the producer writing to the transaction
-   * @param requestLocal             container for the stateful instances scoped to this request -- this must correspond to the
-   *                                 thread calling this method
    * @param callback                 the method to execute once the verification is either completed or returns an error
    *
    * When the verification returns, the callback will be supplied the errors per topic partition if there were errors.
    * The callback will also be supplied the verification guards per partition if they exist. It is possible to have an
    * error and a verification guard for a topic partition if the topic partition was unable to be verified by the transaction
-   * coordinator. Transaction coordinator errors are mapped to append-friendly errors. The callback is wrapped so that it
-   * is scheduled on a request handler thread. There, it should be called with that request handler thread's thread local and
-   * not the one supplied to this method.
+   * coordinator. Transaction coordinator errors are mapped to append-friendly errors.
    */
   private def maybeStartTransactionVerificationForPartitions(
     topicPartitionBatchInfo: Map[TopicPartition, Int],
     transactionalId: String,
     producerId: Long,
     producerEpoch: Short,
-    requestLocal: RequestLocal,
-    callback: (Map[TopicPartition, Errors], RequestLocal, Map[TopicPartition, VerificationGuard]) => Unit
+    callback: ((Map[TopicPartition, Errors], Map[TopicPartition, VerificationGuard])) => Unit
   ): Unit = {
     // Skip verification if the request is not transactional or transaction verification is disabled.
     if (transactionalId == null ||
       !config.transactionPartitionVerificationEnable
       || addPartitionsToTxnManager.isEmpty
     ) {
-      callback(Map.empty[TopicPartition, Errors], requestLocal, Map.empty[TopicPartition, VerificationGuard])
+      callback((Map.empty[TopicPartition, Errors], Map.empty[TopicPartition, VerificationGuard]))
       return
     }
 
@@ -1058,29 +1065,22 @@ class ReplicaManager(val config: KafkaConfig,
 
     // No partitions require verification.
     if (verificationGuards.isEmpty) {
-      callback(errors.toMap, requestLocal, Map.empty[TopicPartition, VerificationGuard])
+      callback((errors.toMap, Map.empty[TopicPartition, VerificationGuard]))
       return
     }
 
     def invokeCallback(
-      requestLocal: RequestLocal,
       verificationErrors: Map[TopicPartition, Errors]
     ): Unit = {
-      callback(errors ++ verificationErrors, requestLocal, verificationGuards.toMap)
+      callback((errors ++ verificationErrors, verificationGuards.toMap))
     }
 
-    // Wrap the callback to be handled on an arbitrary request handler thread when transaction verification is complete.
-    val verificationCallback = KafkaRequestHandler.wrapAsyncCallback(
-      invokeCallback,
-      requestLocal
-    )
-
     addPartitionsToTxnManager.foreach(_.verifyTransaction(
       transactionalId = transactionalId,
       producerId = producerId,
       producerEpoch = producerEpoch,
       topicPartitions = verificationGuards.keys.toSeq,
-      callback = verificationCallback
+      callback = invokeCallback
     ))
 
   }

core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala

@@ -200,11 +200,10 @@ object AbstractCoordinatorConcurrencyTest {
       producerId: Long,
       producerEpoch: Short,
       baseSequence: Int,
-      requestLocal: RequestLocal,
-      callback: (Errors, RequestLocal, VerificationGuard) => Unit
+      callback: ((Errors, VerificationGuard)) => Unit
     ): Unit = {
       // Skip verification
-      callback(Errors.NONE, requestLocal, VerificationGuard.SENTINEL)
+      callback((Errors.NONE, VerificationGuard.SENTINEL))
     }
 
     override def tryCompleteActions(): Unit = watchKeys.map(producePurgatory.checkAndComplete)

core/src/test/scala/unit/kafka/coordinator/group/CoordinatorPartitionWriterTest.scala

@@ -16,7 +16,7 @@
  */
 package kafka.coordinator.group
 
-import kafka.server.{ReplicaManager, RequestLocal}
+import kafka.server.ReplicaManager
 import kafka.utils.TestUtils
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.config.TopicConfig
@@ -334,23 +334,21 @@ class CoordinatorPartitionWriterTest {
       VerificationGuard.SENTINEL
     }
 
-    val callbackCapture: ArgumentCaptor[(Errors, RequestLocal, VerificationGuard) => Unit] =
-      ArgumentCaptor.forClass(classOf[(Errors, RequestLocal, VerificationGuard) => Unit])
+    val callbackCapture: ArgumentCaptor[((Errors, VerificationGuard)) => Unit] =
+      ArgumentCaptor.forClass(classOf[((Errors, VerificationGuard)) => Unit])
 
     when(replicaManager.maybeStartTransactionVerificationForPartition(
       ArgumentMatchers.eq(tp),
       ArgumentMatchers.eq("transactional-id"),
       ArgumentMatchers.eq(10L),
       ArgumentMatchers.eq(5.toShort),
       ArgumentMatchers.eq(RecordBatch.NO_SEQUENCE),
-      ArgumentMatchers.eq(RequestLocal.NoCaching),
       callbackCapture.capture()
     )).thenAnswer(_ => {
-      callbackCapture.getValue.apply(
+      callbackCapture.getValue.apply((
         error,
-        RequestLocal.NoCaching,
         verificationGuard
-      )
+      ))
     })
 
     val future = partitionRecordWriter.maybeStartTransactionVerification(

core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorConcurrencyTest.scala

@@ -24,7 +24,7 @@ import kafka.common.OffsetAndMetadata
 import kafka.coordinator.AbstractCoordinatorConcurrencyTest
 import kafka.coordinator.AbstractCoordinatorConcurrencyTest._
 import kafka.coordinator.group.GroupCoordinatorConcurrencyTest._
-import kafka.server.{DelayedOperationPurgatory, KafkaConfig}
+import kafka.server.{DelayedOperationPurgatory, KafkaConfig, KafkaRequestHandler}
 import kafka.utils.CoreUtils
 import org.apache.kafka.common.{TopicIdPartition, TopicPartition, Uuid}
 import org.apache.kafka.common.internals.Topic
@@ -85,6 +85,10 @@ class GroupCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest
     groupCoordinator = GroupCoordinator(config, replicaManager, heartbeatPurgatory, rebalancePurgatory, timer.time, metrics)
     groupCoordinator.startup(() => zkClient.getTopicPartitionCount(Topic.GROUP_METADATA_TOPIC_NAME).getOrElse(config.offsetsTopicPartitions),
       false)
+
+    // Transactional appends attempt to schedule to the request handler thread using
+    // a non request handler thread. Set this to avoid error.
+    KafkaRequestHandler.setBypassThreadCheck(true)
   }
 
   @AfterEach

core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala

@@ -19,7 +19,7 @@ package kafka.coordinator.group
 
 import java.util.{Optional, OptionalInt}
 import kafka.common.OffsetAndMetadata
-import kafka.server.{ActionQueue, DelayedOperationPurgatory, HostedPartition, KafkaConfig, ReplicaManager, RequestLocal}
+import kafka.server.{ActionQueue, DelayedOperationPurgatory, HostedPartition, KafkaConfig, KafkaRequestHandler, ReplicaManager, RequestLocal}
 import kafka.utils._
 import org.apache.kafka.common.{TopicIdPartition, TopicPartition, Uuid}
 import org.apache.kafka.common.protocol.Errors
@@ -4111,7 +4111,7 @@ class GroupCoordinatorTest {
                                          memberId: String = JoinGroupRequest.UNKNOWN_MEMBER_ID,
                                          groupInstanceId: Option[String] = Option.empty,
                                          generationId: Int = JoinGroupRequest.UNKNOWN_GENERATION_ID,
-                                         verificationError: Errors = Errors.NONE) : CommitOffsetCallbackParams = {
+                                         verificationError: Errors = Errors.NONE): CommitOffsetCallbackParams = {
     val (responseFuture, responseCallback) = setupCommitOffsetsCallback
 
     val capturedArgument: ArgumentCaptor[scala.collection.Map[TopicPartition, PartitionResponse] => Unit] = ArgumentCaptor.forClass(classOf[scala.collection.Map[TopicPartition, PartitionResponse] => Unit])
@@ -4120,11 +4120,22 @@ class GroupCoordinatorTest {
     val transactionalId = "dummy-txn-id"
     val offsetTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupCoordinator.partitionFor(groupId))
 
-    val postVerificationCallback: ArgumentCaptor[(Errors, RequestLocal, VerificationGuard) => Unit] = ArgumentCaptor.forClass(classOf[(Errors, RequestLocal, VerificationGuard) => Unit])
+    val postVerificationCallback: ArgumentCaptor[((Errors, VerificationGuard)) => Unit] =
+      ArgumentCaptor.forClass(classOf[((Errors, VerificationGuard)) => Unit])
 
-    when(replicaManager.maybeStartTransactionVerificationForPartition(ArgumentMatchers.eq(offsetTopicPartition), ArgumentMatchers.eq(transactionalId),
-      ArgumentMatchers.eq(producerId), ArgumentMatchers.eq(producerEpoch), any(), any(), postVerificationCallback.capture())).thenAnswer(
-      _ => postVerificationCallback.getValue()(verificationError, RequestLocal.NoCaching, VerificationGuard.SENTINEL)
+    // Transactional appends attempt to schedule to the request handler thread using
+    // a non request handler thread. Set this to avoid error.
+    KafkaRequestHandler.setBypassThreadCheck(true)
+
+    when(replicaManager.maybeStartTransactionVerificationForPartition(
+      ArgumentMatchers.eq(offsetTopicPartition),
+      ArgumentMatchers.eq(transactionalId),
+      ArgumentMatchers.eq(producerId),
+      ArgumentMatchers.eq(producerEpoch),
+      any(),
+      postVerificationCallback.capture()
+    )).thenAnswer(
+      _ => postVerificationCallback.getValue()((verificationError, VerificationGuard.SENTINEL))
     )
     when(replicaManager.appendRecords(anyLong,
       anyShort(),

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

@@ -3071,9 +3071,8 @@ class ReplicaManagerTest {
                                                             producerEpoch: Short,
                                                             baseSequence: Int = 0): CallbackResult[Either[Errors, VerificationGuard]] = {
     val result = new CallbackResult[Either[Errors, VerificationGuard]]()
-    def postVerificationCallback(error: Errors,
-                                 requestLocal: RequestLocal,
-                                 verificationGuard: VerificationGuard): Unit = {
+    def postVerificationCallback(errorAndGuard: (Errors, VerificationGuard)): Unit = {
+      val (error, verificationGuard) = errorAndGuard
       val errorOrGuard = if (error != Errors.NONE) Left(error) else Right(verificationGuard)
       result.fire(errorOrGuard)
     }
@@ -3084,7 +3083,6 @@ class ReplicaManagerTest {
       producerId,
       producerEpoch,
       baseSequence,
-      RequestLocal.NoCaching,
       postVerificationCallback
     )
     result