Add a commit drain & stage to Kafka (#383)

Author: adamw

core/src/main/scala/ox/channels/actor.scala

@@ -15,15 +15,19 @@ object Actor:
     * cause the actor's channel to be closed with an error, and are propagated to the enclosing scope.
     *
     * The actor's mailbox (incoming channel) will have a capacity as specified by the [[BufferCapacity]] in scope.
+    *
+    * @param close
+    *   An optional callback that will be called uninterruptedly before the actor closes.
     */
   def create[T](logic: T, close: Option[T => Unit] = None)(using ox: Ox, sc: BufferCapacity): ActorRef[T] =
     val c = BufferCapacity.newChannel[T => Unit]
     val ref = ActorRef(c)
     forkDiscard {
       try
         forever {
-          val m = c.receive()
-          try m(logic)
+          try
+            val m = c.receive()
+            m(logic)
           catch
             case t: Throwable =>
               c.error(t)

core/src/main/scala/ox/flow/FlowOps.scala

@@ -28,6 +28,7 @@ import ox.unsupervised
 import java.util.concurrent.Semaphore
 import scala.concurrent.duration.DurationLong
 import scala.concurrent.duration.FiniteDuration
+import scala.util.control.ControlThrowable
 
 class FlowOps[+T]:
   outer: Flow[T] =>
@@ -271,7 +272,7 @@ class FlowOps[+T]:
 
       FlowEmit.channelToEmit(results, emit)
 
-  private val abortTake = new Exception("abort take")
+  private val abortTake = new ControlThrowable("abort take") {}
 
   /** Takes the first `n` elements from this flow and emits them. If the flow completes before emitting `n` elements, the returned flow
     * completes as well.

core/src/test/scala/ox/channels/ActorTest.scala

@@ -80,4 +80,15 @@ class ActorTest extends AnyFlatSpec with Matchers:
 
     thrown.getMessage shouldBe "boom"
   }
+
+  it should "throw a channel closed exception when the actor's scope becomes closed" in {
+    val actor = supervised:
+      val logic = new Test1:
+        override def f(x: Int): Long = 10
+
+      Actor.create(logic)
+      // when the scope ends, the actor's fork is interrupted
+
+    an[ChannelClosedException] should be thrownBy actor.ask(_.f(5))
+  }
 end ActorTest

doc/integrations/kafka.md

@@ -7,9 +7,16 @@ Dependency:
 ```
 
 `Flow`s which read from a Kafka topic, mapping stages and drains which publish to Kafka topics are available through
-the `KafkaFlow`, `KafkaStage` and `KafkaDrain` objects. In all cases either a manually constructed instance of a
-`KafkaProducer` / `KafkaConsumer` is needed, or `ProducerSettings` / `ConsumerSetttings` need to be provided with the
-bootstrap servers, consumer group id, key / value serializers, etc.
+the `KafkaFlow`, `KafkaStage` and `KafkaDrain` objects. 
+
+In all cases kafka producers and consumers can be provided:
+* by manually creating (and closing) an instance of a `KafkaProducer` / `KafkaConsumer`
+* through a `ProducerSettings` / `ConsumerSettings`, with the bootstrap servers, consumer group id, key/value
+  serializers, etc. The lifetime is then managed by the flow operators.
+* through a thread-safe wrapper on a consumer (`ActorRef[KafkaConsumerWrapper[K, V]]`), for which the lifetime is bound
+  to the current concurrency scope
+
+## Reading from Kafka
 
 To read from a Kafka topic, use:
 
@@ -25,6 +32,8 @@ val source = KafkaFlow.subscribe(settings, topic)
   .runForeach { (msg: ReceivedMessage[String, String]) => ??? }
 ```
 
+## Publishing to Kafka
+
 To publish data to a Kafka topic:
 
 ```scala mdoc:compile-only
@@ -40,6 +49,25 @@ Flow
   .pipe(KafkaDrain.runPublish(settings))
 ```
 
+To publish data as a mapping stage:
+
+```scala mdoc:compile-only
+import ox.flow.Flow
+import ox.kafka.ProducerSettings
+import ox.kafka.KafkaStage.*
+import org.apache.kafka.clients.producer.{ProducerRecord, RecordMetadata}
+
+val settings = ProducerSettings.default.bootstrapServers("localhost:9092")
+val metadatas: Flow[RecordMetadata] = Flow
+  .fromIterable(List("a", "b", "c"))
+  .map(msg => ProducerRecord[String, String]("my_topic", msg))
+  .mapPublish(settings)
+
+// process & run the metadatas flow further
+```
+
+## Reading & publishing to Kafka with offset commits
+
 Quite often data to be published to a topic (`topic1`) is computed basing on data received from another topic 
 (`topic2`). In such a case, it's possible to commit messages from `topic2`, after the messages to `topic1` are 
 successfully published. 
@@ -63,7 +91,7 @@ computed. For example:
 ```scala mdoc:compile-only
 import ox.kafka.{ConsumerSettings, KafkaDrain, KafkaFlow, ProducerSettings, SendPacket}
 import ox.kafka.ConsumerSettings.AutoOffsetReset
-import ox.pipe
+import ox.*
 import org.apache.kafka.clients.producer.ProducerRecord
 
 val consumerSettings = ConsumerSettings.default("my_group")
@@ -72,29 +100,43 @@ val producerSettings = ProducerSettings.default.bootstrapServers("localhost:9092
 val sourceTopic = "source_topic"
 val destTopic = "dest_topic"
 
-KafkaFlow
-  .subscribe(consumerSettings, sourceTopic)
-  .map(in => (in.value.toLong * 2, in))
-  .map((value, original) => 
-    SendPacket(ProducerRecord[String, String](destTopic, value.toString), original))
-  .pipe(KafkaDrain.runPublishAndCommit(producerSettings))
+supervised:
+  // the consumer is shared between the subscribe & offset stages
+  // its lifetime is bound to the current concurrency scope
+  val consumer = consumerSettings.toThreadSafeConsumerWrapper
+  KafkaFlow
+    .subscribe(consumer, sourceTopic)
+    .map(in => (in.value.toLong * 2, in))
+    .map((value, original) => 
+      SendPacket(ProducerRecord[String, String](destTopic, value.toString), original))
+    .pipe(KafkaDrain.runPublishAndCommit(producerSettings, consumer))
 ```
 
 The offsets are committed every second in a background process.
 
-To publish data as a mapping stage:
+## Reading from Kafka, processing data & committing offsets
+
+Offsets can also be committed after the data has been processed, without producing any records to write to a topic.
+For that, we can use the `runCommit` drain, or the `mapCommit` stage, both of which work with a `Flow[CommitPacket]`:
 
 ```scala mdoc:compile-only
-import ox.flow.Flow
-import ox.kafka.ProducerSettings
-import ox.kafka.KafkaStage.*
-import org.apache.kafka.clients.producer.{ProducerRecord, RecordMetadata}
+import ox.kafka.{ConsumerSettings, KafkaDrain, KafkaFlow, CommitPacket}
+import ox.kafka.ConsumerSettings.AutoOffsetReset
+import ox.*
 
-val settings = ProducerSettings.default.bootstrapServers("localhost:9092")
-val metadatas: Flow[RecordMetadata] = Flow
-  .fromIterable(List("a", "b", "c"))
-  .map(msg => ProducerRecord[String, String]("my_topic", msg))
-  .mapPublish(settings)
+val consumerSettings = ConsumerSettings.default("my_group")
+  .bootstrapServers("localhost:9092").autoOffsetReset(AutoOffsetReset.Earliest)
+val sourceTopic = "source_topic"
 
-// process & run the metadatas flow further
-```
+supervised:
+  // the consumer is shared between the subscribe & offset stages
+  // its lifetime is bound to the current concurrency scope
+  val consumer = consumerSettings.toThreadSafeConsumerWrapper
+  KafkaFlow
+    .subscribe(consumer, sourceTopic)
+    .mapPar(10) { in => 
+      // process the message, e.g. send an HTTP request
+      CommitPacket(in)
+    }
+    .pipe(KafkaDrain.runCommit(consumer))
+```

kafka/src/main/scala/ox/kafka/ConsumerSettings.scala

@@ -1,7 +1,11 @@
 package ox.kafka
 
-import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
-import org.apache.kafka.common.serialization.{Deserializer, StringDeserializer}
+import org.apache.kafka.clients.consumer.ConsumerConfig
+import org.apache.kafka.clients.consumer.KafkaConsumer
+import org.apache.kafka.common.serialization.Deserializer
+import org.apache.kafka.common.serialization.StringDeserializer
+import ox.Ox
+import ox.channels.ActorRef
 import ox.kafka.ConsumerSettings.AutoOffsetReset
 
 import java.util.Properties
@@ -32,7 +36,16 @@ case class ConsumerSettings[K, V](
     props
   end toProperties
 
+  /** Using these settings, create a new open [[KafkaConsumer]]. The consumer is not thread-safe, and should not be used concurrently, and
+    * has to be closed manually.
+    */
   def toConsumer: KafkaConsumer[K, V] = KafkaConsumer(toProperties, keyDeserializer, valueDeserializer)
+
+  /** Using these settings, create a thread-safe wrapper on top of a new open [[KafkaConsumer]]. The wrapper serializes calls using an
+    * actor. The actor is created within the current concurrency scope and will be closed (along with the consumer) when the scope ends.
+    */
+  def toThreadSafeConsumerWrapper(using Ox): ActorRef[KafkaConsumerWrapper[K, V]] =
+    KafkaConsumerWrapper(toConsumer, closeWhenComplete = true)
 end ConsumerSettings
 
 object ConsumerSettings:

kafka/src/main/scala/ox/kafka/KafkaConsumerWrapper.scala

@@ -16,6 +16,9 @@ trait KafkaConsumerWrapper[K, V]:
 object KafkaConsumerWrapper:
   private val logger = LoggerFactory.getLogger(classOf[KafkaConsumerWrapper.type])
 
+  /** Create a thread-safe wrapper on top of a [[KafkaConsumer]], which serializes calls using an actor. The actor is created within the
+    * current concurrency scope and will be closed when the scope ends if the `closeWhenComplete` flag is `true`.
+    */
   def apply[K, V](consumer: KafkaConsumer[K, V], closeWhenComplete: Boolean)(using Ox): ActorRef[KafkaConsumerWrapper[K, V]] =
     val logic = new KafkaConsumerWrapper[K, V]:
       override def subscribe(topics: Seq[String]): Unit =
@@ -41,7 +44,7 @@ object KafkaConsumerWrapper:
 
     def close(wrapper: KafkaConsumerWrapper[K, V]): Unit = if closeWhenComplete then
       logger.debug("Closing the Kafka consumer")
-      uninterruptible(consumer.close())
+      consumer.close()
 
     Actor.create(logic, Some(close))
   end apply

kafka/src/main/scala/ox/kafka/KafkaDrain.scala

@@ -52,20 +52,32 @@ object KafkaDrain:
 
   /** @return
     *   A drain, which consumes all packets emitted by the provided [[Flow]]. For each packet, first all `send` messages (producer records)
-    *   are sent. Then, all `commit` messages (consumer records) up to their offsets are committed.
+    *   are sent, using a producer created with the given `producerSettings`. Then, all `commit` messages (consumer records) up to their
+    *   offsets are committed, using the given `consumer`.
     */
-  def runPublishAndCommit[K, V](producerSettings: ProducerSettings[K, V])(using BufferCapacity): Flow[SendPacket[K, V]] => Unit =
-    flow => runPublishAndCommit(producerSettings.toProducer, closeWhenComplete = true)(flow)
+  def runPublishAndCommit[K, V](producerSettings: ProducerSettings[K, V], consumer: ActorRef[KafkaConsumerWrapper[K, V]])(using
+      BufferCapacity
+  ): Flow[SendPacket[K, V]] => Unit =
+    flow => runPublishAndCommit(producerSettings.toProducer, consumer, closeWhenComplete = true)(flow)
 
   /** @param producer
     *   The producer that is used to send messages.
     * @return
     *   A drain, which consumes all packets emitted by the provided [[Flow]]. For each packet, first all `send` messages (producer records)
-    *   are sent. Then, all `commit` messages (consumer records) up to their offsets are committed.
+    *   are sent, using the given `producer`. Then, all `commit` messages (consumer records) up to their offsets are committed, using the
+    *   given `consumer`.
     */
-  def runPublishAndCommit[K, V](producer: KafkaProducer[K, V], closeWhenComplete: Boolean)(using
-      BufferCapacity
+  def runPublishAndCommit[K, V](producer: KafkaProducer[K, V], consumer: ActorRef[KafkaConsumerWrapper[K, V]], closeWhenComplete: Boolean)(
+      using BufferCapacity
   ): Flow[SendPacket[K, V]] => Unit = flow =>
     import KafkaStage.*
-    flow.mapPublishAndCommit(producer, closeWhenComplete).runDrain()
+    flow.mapPublishAndCommit(producer, consumer, closeWhenComplete).runDrain()
+
+  /** @return
+    *   A drain, which consumes all commit packets emitted by the provided [[Flow]]. For each packet, all `commit` messages (consumer
+    *   records) are committed: for each topic-partition, up to the highest observed offset, using the given `consumer`.
+    */
+  def runCommit[K, V](consumer: ActorRef[KafkaConsumerWrapper[K, V]])(using BufferCapacity): Flow[CommitPacket] => Unit = flow =>
+    import KafkaStage.*
+    flow.mapCommit(consumer).runDrain()
 end KafkaDrain

kafka/src/main/scala/ox/kafka/KafkaFlow.scala

@@ -1,9 +1,12 @@
 package ox.kafka
 
-import org.apache.kafka.clients.consumer.{ConsumerRecord, KafkaConsumer}
+import org.apache.kafka.clients.consumer.ConsumerRecord
+import org.apache.kafka.clients.consumer.KafkaConsumer
 import org.slf4j.LoggerFactory
 import ox.*
+import ox.channels.ActorRef
 import ox.flow.Flow
+import ox.flow.FlowEmit
 
 object KafkaFlow:
   private val logger = LoggerFactory.getLogger(classOf[KafkaFlow.type])
@@ -20,10 +23,24 @@ object KafkaFlow:
     Flow.usingEmit: emit =>
       supervised:
         val kafkaConsumerActor = KafkaConsumerWrapper(kafkaConsumer, closeWhenComplete)
-        kafkaConsumerActor.tell(_.subscribe(topic :: otherTopics.toList))
-        forever {
-          val records = kafkaConsumerActor.ask(_.poll())
-          records.forEach(r => emit(ReceivedMessage(kafkaConsumerActor, r)))
-        }.tapException(logger.error("Exception when polling for records", _))
+        doSubscribe(kafkaConsumerActor, topic, otherTopics*)(emit)
+
+  def subscribe[K, V](
+      kafkaConsumerActor: ActorRef[KafkaConsumerWrapper[K, V]],
+      topic: String,
+      otherTopics: String*
+  ): Flow[ReceivedMessage[K, V]] =
+    Flow.usingEmit: emit =>
+      doSubscribe(kafkaConsumerActor, topic, otherTopics*)(emit)
+
+  private def doSubscribe[K, V](kafkaConsumerActor: ActorRef[KafkaConsumerWrapper[K, V]], topic: String, otherTopics: String*)(
+      emit: FlowEmit[ReceivedMessage[K, V]]
+  ): Unit =
+    kafkaConsumerActor.tell(_.subscribe(topic :: otherTopics.toList))
+    forever {
+      val records = kafkaConsumerActor.ask(_.poll())
+      records.forEach(r => emit(ReceivedMessage(r)))
+    }.tapException(logger.error("Exception when polling for records", _))
+  end doSubscribe
 
 end KafkaFlow