Merge pull request #883 from janstenpickle/no-offsets-transactional-p…

…roducer
fd4s · Mar 7, 2022 · 8f9c1e9 · 8f9c1e9
2 parents 4826432 + 0999df0
commit 8f9c1e9
Show file tree

Hide file tree

Showing 2 changed files with 216 additions and 128 deletions.
diff --git a/modules/core/src/main/scala/fs2/kafka/TransactionalKafkaProducer.scala b/modules/core/src/main/scala/fs2/kafka/TransactionalKafkaProducer.scala
@@ -7,7 +7,7 @@
 package fs2.kafka
 
 import cats.effect.syntax.all._
-import cats.effect.{Async, Resource, Outcome}
+import cats.effect.{Async, Outcome, Resource}
 import cats.syntax.all._
 import fs2.kafka.internal._
 import fs2.kafka.internal.converters.collection._
@@ -57,6 +57,22 @@ object TransactionalKafkaProducer {
     def metrics: F[Map[MetricName, Metric]]
   }
 
+  /**
+    * [[TransactionalKafkaProducer.WithoutOffsets]] extends [[TransactionalKafkaProducer.Metrics]]
+    * to allow producing of records without corresponding upstream offsets.
+    */
+  abstract class WithoutOffsets[F[_], K, V] extends Metrics[F, K, V] {
+
+    /**
+      * Produces the `ProducerRecord`s in the specified [[ProducerRecords]]
+      * in three steps: first a transaction is initialized, then the records are placed
+      * in the buffer of the producer, and lastly the transaction is committed. If errors
+      * or cancellation occurs, the transaction is aborted. The returned effect succeeds
+      * if the whole transaction completes successfully.
+      */
+    def produceWithoutOffsets[P](records: ProducerRecords[P, K, V]): F[ProducerResult[P, K, V]]
+  }
+
   /**
     * Creates a new [[TransactionalKafkaProducer]] in the `Resource` context,
     * using the specified [[TransactionalProducerSettings]]. Note that there
@@ -72,20 +88,20 @@ object TransactionalKafkaProducer {
   )(
     implicit F: Async[F],
     mk: MkProducer[F]
-  ): Resource[F, TransactionalKafkaProducer.Metrics[F, K, V]] =
+  ): Resource[F, TransactionalKafkaProducer.WithoutOffsets[F, K, V]] =
     (
       Resource.eval(settings.producerSettings.keySerializer),
       Resource.eval(settings.producerSettings.valueSerializer),
       WithProducer(mk, settings)
     ).mapN { (keySerializer, valueSerializer, withProducer) =>
-      new TransactionalKafkaProducer.Metrics[F, K, V] {
+      new TransactionalKafkaProducer.WithoutOffsets[F, K, V] {
         override def produce[P](
           records: TransactionalProducerRecords[F, P, K, V]
         ): F[ProducerResult[P, K, V]] =
-          produceTransaction(records)
+          produceTransactionWithOffsets(records)
             .map(ProducerResult(_, records.passthrough))
 
-        private[this] def produceTransaction[P](
+        private[this] def produceTransactionWithOffsets[P](
           records: TransactionalProducerRecords[F, P, K, V]
         ): F[Chunk[(ProducerRecord[K, V], RecordMetadata)]] =
           if (records.records.isEmpty) F.pure(Chunk.empty)
@@ -99,34 +115,50 @@ object TransactionalKafkaProducer {
               else F.pure(batch.consumerGroupIds.head)
 
             consumerGroupId.flatMap { groupId =>
-              withProducer { (producer, blocking) =>
-                blocking(producer.beginTransaction())
-                  .bracketCase { _ =>
-                    records.records
-                      .flatMap(_.records)
-                      .traverse(
-                        KafkaProducer
-                          .produceRecord(keySerializer, valueSerializer, producer, blocking)
-                      )
-                      .map(_.sequence)
-                      .flatTap { _ =>
-                        blocking {
-                          producer.sendOffsetsToTransaction(
-                            batch.offsets.asJava,
-                            groupId
-                          )
-                        }
-                      }
-                  } {
-                    case (_, Outcome.Succeeded(_)) =>
-                      blocking(producer.commitTransaction())
-                    case (_, Outcome.Canceled() | Outcome.Errored(_)) =>
-                      blocking(producer.abortTransaction())
-                  }
-              }.flatten
+              val sendOffsets: (KafkaByteProducer, Blocking[F]) => F[Unit] = (producer, blocking) =>
+                blocking {
+                  producer.sendOffsetsToTransaction(
+                    batch.offsets.asJava,
+                    groupId
+                  )
+                }
+
+              produceTransaction(records.records.flatMap(_.records), Some(sendOffsets))
             }
           }
 
+        override def produceWithoutOffsets[P](
+          records: ProducerRecords[P, K, V]
+        ): F[ProducerResult[P, K, V]] =
+          produceTransaction(records.records, None).map(ProducerResult(_, records.passthrough))
+
+        private[this] def produceTransaction[P](
+          records: Chunk[ProducerRecord[K, V]],
+          sendOffsets: Option[(KafkaByteProducer, Blocking[F]) => F[Unit]]
+        ): F[Chunk[(ProducerRecord[K, V], RecordMetadata)]] =
+          if (records.isEmpty) F.pure(Chunk.empty)
+          else {
+
+            withProducer { (producer, blocking) =>
+              blocking(producer.beginTransaction())
+                .bracketCase { _ =>
+                  val produce = records
+                    .traverse(
+                      KafkaProducer
+                        .produceRecord(keySerializer, valueSerializer, producer, blocking)
+                    )
+                    .map(_.sequence)
+
+                  sendOffsets.fold(produce)(f => produce.flatTap(_ => f(producer, blocking)))
+                } {
+                  case (_, Outcome.Succeeded(_)) =>
+                    blocking(producer.commitTransaction())
+                  case (_, Outcome.Canceled() | Outcome.Errored(_)) =>
+                    blocking(producer.abortTransaction())
+                }
+            }.flatten
+          }
+
         override def metrics: F[Map[MetricName, Metric]] =
           withProducer.blocking { _.metrics().asScala.toMap }
 
@@ -150,7 +182,7 @@ object TransactionalKafkaProducer {
   )(
     implicit F: Async[F],
     mk: MkProducer[F]
-  ): Stream[F, TransactionalKafkaProducer.Metrics[F, K, V]] =
+  ): Stream[F, TransactionalKafkaProducer.WithoutOffsets[F, K, V]] =
     Stream.resource(resource(settings)(F, mk))
 
   def apply[F[_]]: TransactionalProducerPartiallyApplied[F] =
@@ -172,7 +204,7 @@ object TransactionalKafkaProducer {
     def resource[K, V](settings: TransactionalProducerSettings[F, K, V])(
       implicit F: Async[F],
       mk: MkProducer[F]
-    ): Resource[F, TransactionalKafkaProducer.Metrics[F, K, V]] =
+    ): Resource[F, TransactionalKafkaProducer.WithoutOffsets[F, K, V]] =
       TransactionalKafkaProducer.resource(settings)(F, mk)
 
     /**
@@ -188,7 +220,7 @@ object TransactionalKafkaProducer {
     def stream[K, V](settings: TransactionalProducerSettings[F, K, V])(
       implicit F: Async[F],
       mk: MkProducer[F]
-    ): Stream[F, TransactionalKafkaProducer.Metrics[F, K, V]] =
+    ): Stream[F, TransactionalKafkaProducer.WithoutOffsets[F, K, V]] =
       TransactionalKafkaProducer.stream(settings)(F, mk)
 
     override def toString: String =