KAFKA-14505; [4/N] Wire transaction verification

build.gradle

@@ -1314,6 +1314,7 @@ project(':group-coordinator') {
     implementation project(':server-common')
     implementation project(':clients')
     implementation project(':metadata')
+    implementation project(':storage')
     implementation libs.slf4jApi
     implementation libs.metrics
 

checkstyle/import-control.xml

@@ -248,6 +248,7 @@
       <allow pkg="org.apache.kafka.server.common"/>
       <allow pkg="org.apache.kafka.server.record"/>
       <allow pkg="org.apache.kafka.server.util"/>
+      <allow pkg="org.apache.kafka.storage.internals.log"/>
       <allow pkg="org.apache.kafka.test" />
       <allow pkg="org.apache.kafka.timeline" />
       <subpackage name="metrics">

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

@@ -17,7 +17,7 @@
 package kafka.coordinator.group
 
 import kafka.cluster.PartitionListener
-import kafka.server.{ActionQueue, ReplicaManager}
+import kafka.server.{ActionQueue, ReplicaManager, RequestLocal}
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.RecordTooLargeException
 import org.apache.kafka.common.protocol.Errors
@@ -27,9 +27,10 @@ import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
 import org.apache.kafka.common.requests.TransactionResult
 import org.apache.kafka.common.utils.{BufferSupplier, Time}
 import org.apache.kafka.coordinator.group.runtime.PartitionWriter
-import org.apache.kafka.storage.internals.log.AppendOrigin
+import org.apache.kafka.storage.internals.log.{VerificationGuard}
 
 import java.util
+import java.util.concurrent.CompletableFuture
 import scala.collection.Map
 
 /**
@@ -121,6 +122,7 @@ class CoordinatorPartitionWriter[T](
     tp: TopicPartition,
     producerId: Long,
     producerEpoch: Short,
+    verificationGuard: VerificationGuard,
     records: util.List[T]
   ): Long = {
     if (records.isEmpty) throw new IllegalStateException("records must be non-empty.")
@@ -161,7 +163,7 @@ class CoordinatorPartitionWriter[T](
               s"in append to partition $tp which exceeds the maximum configured size of $maxBatchSize.")
           }
 
-          internalAppend(tp, recordsBuilder.build())
+          internalAppend(tp, recordsBuilder.build(), verificationGuard)
         } finally {
           bufferSupplier.release(buffer)
         }
@@ -201,18 +203,55 @@ class CoordinatorPartitionWriter[T](
     ))
   }
 
+  /**
+   * Verify the transaction.
+   *
+   * @param tp              The partition to write records to.
+   * @param transactionalId The transactional id.
+   * @param producerId      The producer id.
+   * @param producerEpoch   The producer epoch.
+   * @return A future containing the {@link VerificationGuard} or an exception.
+   * @throws KafkaException Any KafkaException caught during the operation.
+   */
+  override def maybeStartTransactionVerification(
+    tp: TopicPartition,
+    transactionalId: String,
+    producerId: Long,
+    producerEpoch: Short
+  ): CompletableFuture[VerificationGuard] = {
+    val future = new CompletableFuture[VerificationGuard]()
+    replicaManager.maybeStartTransactionVerificationForPartition(
+      topicPartition = tp,
+      transactionalId = transactionalId,
+      producerId = producerId,
+      producerEpoch = producerEpoch,
+      baseSequence = RecordBatch.NO_SEQUENCE,
+      requestLocal = RequestLocal.NoCaching,
+      callback = (error, _, verificationGuard) => {
+        if (error != Errors.NONE) {
+          future.completeExceptionally(error.exception)
+        } else {
+          future.complete(verificationGuard)
+        }
+      }
+    )
+    future
+  }
+
   private def internalAppend(
     tp: TopicPartition,
-    memoryRecords: MemoryRecords
+    memoryRecords: MemoryRecords,
+    verificationGuard: VerificationGuard = VerificationGuard.SENTINEL
   ): Long = {
     var appendResults: Map[TopicPartition, PartitionResponse] = Map.empty
-    replicaManager.appendRecords(
+    replicaManager.appendForGroup(
       timeout = 0L,
       requiredAcks = 1,
-      internalTopicsAllowed = true,
-      origin = AppendOrigin.COORDINATOR,
       entriesPerPartition = Map(tp -> memoryRecords),
       responseCallback = results => appendResults = results,
+      requestLocal = RequestLocal.NoCaching,
+      verificationGuards = Map(tp -> verificationGuard),
+      delayedProduceLock = None,
       // We can directly complete the purgatories here because we don't hold
       // any conflicting locks.
       actionQueue = directActionQueue

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

@@ -989,6 +989,7 @@ 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 verificationGuards            the mapping from topic partition to verification guards if transaction verification is used
+   * @param actionQueue                   the action queue to use
    */
   def appendForGroup(
     timeout: Long,
@@ -997,7 +998,8 @@ class ReplicaManager(val config: KafkaConfig,
     responseCallback: Map[TopicPartition, PartitionResponse] => Unit,
     delayedProduceLock: Option[Lock],
     requestLocal: RequestLocal,
-    verificationGuards: Map[TopicPartition, VerificationGuard]
+    verificationGuards: Map[TopicPartition, VerificationGuard],
+    actionQueue: ActionQueue = this.defaultActionQueue
   ): Unit = {
     if (!isValidRequiredAcks(requiredAcks)) {
       sendInvalidRequiredAcksResponse(entriesPerPartition, responseCallback)
@@ -1025,7 +1027,7 @@ class ReplicaManager(val config: KafkaConfig,
     val allResults = localProduceResults
     val produceStatus = buildProducePartitionStatus(allResults)
 
-    addCompletePurgatoryAction(defaultActionQueue, allResults)
+    addCompletePurgatoryAction(actionQueue, allResults)
 
     maybeAddDelayedProduce(
       requiredAcks,

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

@@ -215,7 +215,8 @@ object AbstractCoordinatorConcurrencyTest {
                                 responseCallback: Map[TopicPartition, PartitionResponse] => Unit,
                                 delayedProduceLock: Option[Lock] = None,
                                 requestLocal: RequestLocal = RequestLocal.NoCaching,
-                                verificationGuards: Map[TopicPartition, VerificationGuard] = Map.empty): Unit = {
+                                verificationGuards: Map[TopicPartition, VerificationGuard] = Map.empty,
+                                actionQueue: ActionQueue = null): Unit = {
       appendRecords(timeout, requiredAcks, true, AppendOrigin.COORDINATOR, entriesPerPartition, responseCallback,
         delayedProduceLock, requestLocal = requestLocal)
     }

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

@@ -16,7 +16,7 @@
  */
 package kafka.coordinator.group
 
-import kafka.server.ReplicaManager
+import kafka.server.{ReplicaManager, RequestLocal}
 import kafka.utils.TestUtils
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.config.TopicConfig
@@ -27,7 +27,8 @@ import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
 import org.apache.kafka.common.requests.TransactionResult
 import org.apache.kafka.common.utils.{MockTime, Time}
 import org.apache.kafka.coordinator.group.runtime.PartitionWriter
-import org.apache.kafka.storage.internals.log.{AppendOrigin, LogConfig}
+import org.apache.kafka.storage.internals.log.{LogConfig, VerificationGuard}
+import org.apache.kafka.test.TestUtils.assertFutureThrows
 import org.junit.jupiter.api.Assertions.{assertEquals, assertThrows, assertTrue}
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.params.ParameterizedTest
@@ -104,17 +105,14 @@ class CoordinatorPartitionWriterTest {
     val callbackCapture: ArgumentCaptor[Map[TopicPartition, PartitionResponse] => Unit] =
       ArgumentCaptor.forClass(classOf[Map[TopicPartition, PartitionResponse] => Unit])
 
-    when(replicaManager.appendRecords(
+    when(replicaManager.appendForGroup(
       ArgumentMatchers.eq(0L),
       ArgumentMatchers.eq(1.toShort),
-      ArgumentMatchers.eq(true),
-      ArgumentMatchers.eq(AppendOrigin.COORDINATOR),
       recordsCapture.capture(),
       callbackCapture.capture(),
       ArgumentMatchers.any(),
       ArgumentMatchers.any(),
-      ArgumentMatchers.any(),
-      ArgumentMatchers.any(),
+      ArgumentMatchers.eq(Map(tp -> VerificationGuard.SENTINEL)),
       ArgumentMatchers.any()
     )).thenAnswer( _ => {
       callbackCapture.getValue.apply(Map(
@@ -140,6 +138,7 @@ class CoordinatorPartitionWriterTest {
       tp,
       RecordBatch.NO_PRODUCER_ID,
       RecordBatch.NO_PRODUCER_EPOCH,
+      VerificationGuard.SENTINEL,
       records.asJava
     ))
 
@@ -168,6 +167,7 @@ class CoordinatorPartitionWriterTest {
       CompressionType.NONE,
       time
     )
+    val verificationGuard = new VerificationGuard()
 
     when(replicaManager.getLogConfig(tp)).thenReturn(Some(LogConfig.fromProps(
       Collections.emptyMap(),
@@ -179,17 +179,14 @@ class CoordinatorPartitionWriterTest {
     val callbackCapture: ArgumentCaptor[Map[TopicPartition, PartitionResponse] => Unit] =
       ArgumentCaptor.forClass(classOf[Map[TopicPartition, PartitionResponse] => Unit])
 
-    when(replicaManager.appendRecords(
+    when(replicaManager.appendForGroup(
       ArgumentMatchers.eq(0L),
       ArgumentMatchers.eq(1.toShort),
-      ArgumentMatchers.eq(true),
-      ArgumentMatchers.eq(AppendOrigin.COORDINATOR),
       recordsCapture.capture(),
       callbackCapture.capture(),
       ArgumentMatchers.any(),
       ArgumentMatchers.any(),
-      ArgumentMatchers.any(),
-      ArgumentMatchers.any(),
+      ArgumentMatchers.eq(Map(tp -> verificationGuard)),
       ArgumentMatchers.any()
     )).thenAnswer(_ => {
       callbackCapture.getValue.apply(Map(
@@ -215,6 +212,7 @@ class CoordinatorPartitionWriterTest {
       tp,
       100L,
       50.toShort,
+      verificationGuard,
       records.asJava
     ))
 
@@ -260,17 +258,14 @@ class CoordinatorPartitionWriterTest {
     val callbackCapture: ArgumentCaptor[Map[TopicPartition, PartitionResponse] => Unit] =
       ArgumentCaptor.forClass(classOf[Map[TopicPartition, PartitionResponse] => Unit])
 
-    when(replicaManager.appendRecords(
+    when(replicaManager.appendForGroup(
       ArgumentMatchers.eq(0L),
       ArgumentMatchers.eq(1.toShort),
-      ArgumentMatchers.eq(true),
-      ArgumentMatchers.eq(AppendOrigin.COORDINATOR),
       recordsCapture.capture(),
       callbackCapture.capture(),
       ArgumentMatchers.any(),
       ArgumentMatchers.any(),
-      ArgumentMatchers.any(),
-      ArgumentMatchers.any(),
+      ArgumentMatchers.eq(Map(tp -> VerificationGuard.SENTINEL)),
       ArgumentMatchers.any()
     )).thenAnswer(_ => {
       callbackCapture.getValue.apply(Map(
@@ -311,6 +306,58 @@ class CoordinatorPartitionWriterTest {
     assertEquals(List(controlRecordType), receivedRecords)
   }
 
+  @ParameterizedTest
+  @EnumSource(value = classOf[Errors], names = Array("NONE", "NOT_ENOUGH_REPLICAS"))
+  def testMaybeStartTransactionVerification(error: Errors): Unit = {
+    val tp = new TopicPartition("foo", 0)
+    val replicaManager = mock(classOf[ReplicaManager])
+    val time = new MockTime()
+    val partitionRecordWriter = new CoordinatorPartitionWriter(
+      replicaManager,
+      new StringKeyValueSerializer(),
+      CompressionType.NONE,
+      time
+    )
+
+    val verificationGuard = if (error == Errors.NONE) {
+      new VerificationGuard()
+    } else {
+      VerificationGuard.SENTINEL
+    }
+
+    val callbackCapture: ArgumentCaptor[(Errors, RequestLocal, VerificationGuard) => Unit] =
+      ArgumentCaptor.forClass(classOf[(Errors, RequestLocal, 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(
+        error,
+        RequestLocal.NoCaching,
+        verificationGuard
+      )
+    })
+
+    val future = partitionRecordWriter.maybeStartTransactionVerification(
+      tp,
+      "transactional-id",
+      10L,
+      5.toShort
+    )
+
+    if (error == Errors.NONE) {
+      assertEquals(verificationGuard, future.get)
+    } else {
+      assertFutureThrows(future, error.exception.getClass)
+    }
+  }
+
   @Test
   def testWriteRecordsWithFailure(): Unit = {
     val tp = new TopicPartition("foo", 0)
@@ -333,17 +380,14 @@ class CoordinatorPartitionWriterTest {
     val callbackCapture: ArgumentCaptor[Map[TopicPartition, PartitionResponse] => Unit] =
       ArgumentCaptor.forClass(classOf[Map[TopicPartition, PartitionResponse] => Unit])
 
-    when(replicaManager.appendRecords(
+    when(replicaManager.appendForGroup(
       ArgumentMatchers.eq(0L),
       ArgumentMatchers.eq(1.toShort),
-      ArgumentMatchers.eq(true),
-      ArgumentMatchers.eq(AppendOrigin.COORDINATOR),
       recordsCapture.capture(),
       callbackCapture.capture(),
       ArgumentMatchers.any(),
       ArgumentMatchers.any(),
-      ArgumentMatchers.any(),
-      ArgumentMatchers.any(),
+      ArgumentMatchers.eq(Map(tp -> VerificationGuard.SENTINEL)),
       ArgumentMatchers.any()
     )).thenAnswer(_ => {
       callbackCapture.getValue.apply(Map(
@@ -361,8 +405,9 @@ class CoordinatorPartitionWriterTest {
       tp,
       RecordBatch.NO_PRODUCER_ID,
       RecordBatch.NO_PRODUCER_EPOCH,
-      records.asJava)
-    )
+      VerificationGuard.SENTINEL,
+      records.asJava
+    ))
   }
 
   @Test
@@ -394,8 +439,9 @@ class CoordinatorPartitionWriterTest {
       tp,
       RecordBatch.NO_PRODUCER_ID,
       RecordBatch.NO_PRODUCER_EPOCH,
-      records.asJava)
-    )
+      VerificationGuard.SENTINEL,
+      records.asJava
+    ))
   }
 
   @Test
@@ -418,8 +464,9 @@ class CoordinatorPartitionWriterTest {
       tp,
       RecordBatch.NO_PRODUCER_ID,
       RecordBatch.NO_PRODUCER_EPOCH,
-      List.empty.asJava)
-    )
+      VerificationGuard.SENTINEL,
+      List.empty.asJava
+    ))
   }
 
   @Test
@@ -445,7 +492,8 @@ class CoordinatorPartitionWriterTest {
       tp,
       RecordBatch.NO_PRODUCER_ID,
       RecordBatch.NO_PRODUCER_EPOCH,
-      records.asJava)
-    )
+      VerificationGuard.SENTINEL,
+      records.asJava
+    ))
   }
 }