Kafka external sharding strategy

core/build.sbt

@@ -84,6 +84,8 @@ lazy val akkastream =
       libraryDependencies ++= Vector(
             AkkaStream,
             AkkaStreamKafka,
+            AkkaStreamKafaSharding,
+            AkkaShardingTyped,
             AkkaCluster,
             AkkaManagement,
             AkkaHttp,

core/cloudflow-akka-testkit/src/main/scala/cloudflow/akkastream/testkit/TestContext.scala

@@ -19,18 +19,20 @@ package cloudflow.akkastream.testkit
 import java.util.concurrent.atomic.AtomicReference
 
 import scala.concurrent._
-
 import akka.NotUsed
 import akka.actor._
+import akka.actor.typed.scaladsl.adapter._
+import akka.cluster.sharding.typed.scaladsl.{ ClusterSharding, Entity }
 import akka.kafka.CommitterSettings
 import akka.kafka.ConsumerMessage._
 import akka.stream._
 import akka.stream.scaladsl._
 import com.typesafe.config._
-
 import cloudflow.akkastream._
 import cloudflow.streamlets._
 
+import scala.concurrent.duration.{ DurationInt, FiniteDuration }
+
 private[testkit] abstract class Completed
 
 private[testkit] case class TestContext(
@@ -73,6 +75,25 @@ private[testkit] case class TestContext(
       )
       .getOrElse(throw TestContextException(inlet.name, s"Bad test context, could not find source for inlet ${inlet.name}"))
 
+  def shardedSourceWithCommittableContext[T, M, E](
+      inlet: CodecInlet[T],
+      shardEntity: Entity[M, E],
+      kafkaTimeout: FiniteDuration = 10.seconds
+  ): SourceWithContext[T, CommittableOffset, Future[NotUsed]] = {
+    ClusterSharding(system.toTyped).init(shardEntity)
+
+    Source
+      .futureSource(
+        Future {
+          sourceWithContext(inlet).asSource
+            .asInstanceOf[Source[(T, CommittableOffset), NotUsed]]
+        }(system.dispatcher)
+      )
+      .asSourceWithContext { case (_, committableOffset) ⇒ committableOffset }
+      .map { case (record, _) ⇒ record }
+
+  }
+
   private def flowWithCommittableContext[T](outlet: CodecOutlet[T]): cloudflow.akkastream.scaladsl.FlowWithCommittableContext[T, T] = {
     val flow = Flow[T]
 
@@ -109,6 +130,19 @@ private[testkit] case class TestContext(
 
   def plainSource[T](inlet: CodecInlet[T], resetPosition: ResetPosition): Source[T, NotUsed] =
     sourceWithCommittableContext[T](inlet).asSource.map(_._1).mapMaterializedValue(_ ⇒ NotUsed)
+
+  def shardedPlainSource[T, M, E](inlet: CodecInlet[T],
+                                  shardEntity: Entity[M, E],
+                                  resetPosition: ResetPosition = Latest,
+                                  kafkaTimeout: FiniteDuration = 10.seconds): Source[T, Future[NotUsed]] = {
+    ClusterSharding(system.toTyped).init(shardEntity)
+    Source.futureSource(
+      Future {
+        plainSource(inlet, resetPosition)
+      }(system.dispatcher)
+    )
+  }
+
   def plainSink[T](outlet: CodecOutlet[T]): Sink[T, NotUsed] = sinkRef[T](outlet).sink.contramap { el ⇒
     (el, TestCommittableOffset())
   }

core/cloudflow-akka/src/main/scala/cloudflow/akkastream/AkkaStreamletContext.scala

@@ -17,14 +17,16 @@
 package cloudflow.akkastream
 
 import scala.concurrent.Future
-
 import akka.NotUsed
 import akka.actor.ActorSystem
+import akka.cluster.sharding.typed.scaladsl.Entity
 import akka.kafka.ConsumerMessage.{ Committable, CommittableOffset }
 import akka.kafka.CommitterSettings
 import akka.stream.scaladsl._
 import cloudflow.streamlets._
 
+import scala.concurrent.duration.{ DurationInt, FiniteDuration }
+
 /**
  * Runtime context for [[AkkaStreamlet]]s, which provides means to create [[akka.stream.scaladsl.Source Source]]s and [[akka.stream.scaladsl.Sink Sink]]s respectively
  * for [[cloudflow.streamlets.CodecInlet CodeInlet]]s and [[cloudflow.streamlets.CodecOutlet CodeOutlet]]s.
@@ -38,11 +40,23 @@ trait AkkaStreamletContext extends StreamletContext {
       inlet: CodecInlet[T]
   ): cloudflow.akkastream.scaladsl.SourceWithCommittableContext[T]
 
+  private[akkastream] def shardedSourceWithCommittableContext[T, M, E](
+      inlet: CodecInlet[T],
+      shardEntity: Entity[M, E],
+      kafkaTimeout: FiniteDuration = 10.seconds
+  ): SourceWithContext[T, CommittableOffset, Future[NotUsed]]
+
   @deprecated("Use `sourceWithCommittableContext` instead.", "1.3.4")
   private[akkastream] def sourceWithOffsetContext[T](inlet: CodecInlet[T]): cloudflow.akkastream.scaladsl.SourceWithOffsetContext[T]
 
   private[akkastream] def plainSource[T](inlet: CodecInlet[T], resetPosition: ResetPosition): Source[T, NotUsed]
   private[akkastream] def plainSink[T](outlet: CodecOutlet[T]): Sink[T, NotUsed]
+  private[akkastream] def shardedPlainSource[T, M, E](
+      inlet: CodecInlet[T],
+      shardEntity: Entity[M, E],
+      resetPosition: ResetPosition = Latest,
+      kafkaTimeout: FiniteDuration = 10.seconds
+  ): Source[T, Future[NotUsed]]
 
   private[akkastream] def committableSink[T](outlet: CodecOutlet[T], committerSettings: CommitterSettings): Sink[(T, Committable), NotUsed]
   private[akkastream] def committableSink[T](committerSettings: CommitterSettings): Sink[(T, Committable), NotUsed]

core/cloudflow-akka/src/main/scala/cloudflow/akkastream/AkkaStreamletContextImpl.scala

@@ -21,23 +21,24 @@ import java.nio.file.{ Files, Paths }
 
 import scala.concurrent._
 import scala.util._
-
 import akka._
 import akka.actor.ActorSystem
+import akka.cluster.sharding.external.ExternalShardAllocationStrategy
+import akka.cluster.sharding.typed.scaladsl.{ ClusterSharding, Entity }
 import akka.kafka._
 import akka.kafka.ConsumerMessage._
+import akka.kafka.cluster.sharding.KafkaClusterSharding
 import akka.kafka.scaladsl._
 import akka.stream._
 import akka.stream.scaladsl._
-
 import com.typesafe.config._
-
 import org.apache.kafka.clients.consumer.ConsumerConfig
 import org.apache.kafka.clients.producer.ProducerRecord
 import org.apache.kafka.common.serialization._
-
 import cloudflow.streamlets._
 
+import scala.concurrent.duration.{ DurationInt, FiniteDuration }
+
 /**
  * Implementation of the StreamletContext trait.
  */
@@ -91,6 +92,71 @@ final class AkkaStreamletContextImpl(
   override def sourceWithCommittableContext[T](inlet: CodecInlet[T]): cloudflow.akkastream.scaladsl.SourceWithCommittableContext[T] =
     sourceWithContext[T](inlet)
 
+  private[akkastream] def shardedSourceWithContext[T, M, E](
+      inlet: CodecInlet[T],
+      shardEntity: Entity[M, E],
+      kafkaTimeout: FiniteDuration = 10.seconds
+  ): SourceWithContext[T, CommittableOffset, Future[NotUsed]] = {
+    val topic = findTopicForPort(inlet)
+    val gId   = topic.groupId(streamletDefinition.appId, streamletRef, inlet)
+
+    val consumerSettings = ConsumerSettings(system, new ByteArrayDeserializer, new ByteArrayDeserializer)
+      .withBootstrapServers(topic.bootstrapServers.getOrElse(internalKafkaBootstrapServers))
+      .withGroupId(gId)
+      .withProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
+      .withProperties(topic.kafkaConsumerProperties)
+
+    val rebalanceListener: akka.actor.typed.ActorRef[ConsumerRebalanceEvent] =
+      KafkaClusterSharding(system).rebalanceListener(shardEntity.typeKey)
+
+    import akka.actor.typed.scaladsl.adapter._
+    val subscription = Subscriptions
+      .topics(topic.name)
+      .withRebalanceListener(rebalanceListener.toClassic)
+
+    system.log.info(s"Creating sharded committable source for group: $gId topic: ${topic.name}")
+
+    val messageExtractor: Future[KafkaClusterSharding.KafkaShardingMessageExtractor[M]] =
+      KafkaClusterSharding(system).messageExtractor(
+        topic = topic.name,
+        timeout = kafkaTimeout,
+        settings = consumerSettings
+      )
+
+    Source
+      .futureSource {
+        messageExtractor.map { m =>
+          ClusterSharding(system.toTyped).init(
+            shardEntity
+              .withAllocationStrategy(
+                shardEntity.allocationStrategy
+                  .getOrElse(new ExternalShardAllocationStrategy(system, shardEntity.typeKey.name))
+              )
+              .withMessageExtractor(m)
+          )
+
+          Consumer
+            .sourceWithOffsetContext(consumerSettings, subscription)
+            // TODO clean this up, once SourceWithContext has mapError and mapMaterializedValue
+            .asSource
+            .mapMaterializedValue(_ ⇒ NotUsed) // TODO we should likely use control to gracefully stop.
+            .via(handleTermination)
+            .map {
+              case (record, committableOffset) ⇒ inlet.codec.decode(record.value) -> committableOffset
+            }
+        }(system.dispatcher)
+      }
+      .asSourceWithContext { case (_, committableOffset) ⇒ committableOffset }
+      .map { case (record, _) ⇒ record }
+  }
+
+  override def shardedSourceWithCommittableContext[T, M, E](
+      inlet: CodecInlet[T],
+      shardEntity: Entity[M, E],
+      kafkaTimeout: FiniteDuration = 10.seconds
+  ): SourceWithContext[T, CommittableOffset, Future[NotUsed]] =
+    shardedSourceWithContext(inlet, shardEntity)
+
   @deprecated("Use sourceWithCommittableContext", "1.3.4")
   override def sourceWithOffsetContext[T](inlet: CodecInlet[T]): cloudflow.akkastream.scaladsl.SourceWithOffsetContext[T] =
     sourceWithContext[T](inlet)
@@ -156,6 +222,58 @@ final class AkkaStreamletContextImpl(
       }
   }
 
+  def shardedPlainSource[T, M, E](inlet: CodecInlet[T],
+                                  shardEntity: Entity[M, E],
+                                  resetPosition: ResetPosition = Latest,
+                                  kafkaTimeout: FiniteDuration = 10.seconds): Source[T, Future[NotUsed]] = {
+    val topic = findTopicForPort(inlet)
+    val gId   = topic.groupId(streamletDefinition.appId, streamletRef, inlet)
+    val consumerSettings = ConsumerSettings(system, new ByteArrayDeserializer, new ByteArrayDeserializer)
+      .withBootstrapServers(topic.bootstrapServers.getOrElse(internalKafkaBootstrapServers))
+      .withGroupId(gId)
+      .withProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, resetPosition.autoOffsetReset)
+      .withProperties(topic.kafkaConsumerProperties)
+
+    val rebalanceListener: akka.actor.typed.ActorRef[ConsumerRebalanceEvent] =
+      KafkaClusterSharding(system).rebalanceListener(shardEntity.typeKey)
+
+    import akka.actor.typed.scaladsl.adapter._
+    val subscription = Subscriptions
+      .topics(topic.name)
+      .withRebalanceListener(rebalanceListener.toClassic)
+
+    system.log.info(s"Creating sharded plain source for group: $gId topic: ${topic.name}")
+
+    val messageExtractor: Future[KafkaClusterSharding.KafkaShardingMessageExtractor[M]] =
+      KafkaClusterSharding(system).messageExtractor(
+        topic = topic.name,
+        timeout = kafkaTimeout,
+        settings = consumerSettings
+      )
+
+    Source
+      .futureSource {
+        messageExtractor.map { m =>
+          ClusterSharding(system.toTyped).init(
+            shardEntity
+              .withAllocationStrategy(
+                shardEntity.allocationStrategy
+                  .getOrElse(new ExternalShardAllocationStrategy(system, shardEntity.typeKey.name))
+              )
+              .withMessageExtractor(m)
+          )
+
+          Consumer
+            .plainSource(consumerSettings, subscription)
+            .mapMaterializedValue(_ ⇒ NotUsed) // TODO we should likely use control to gracefully stop.
+            .via(handleTermination)
+            .map { record ⇒
+              inlet.codec.decode(record.value)
+            }
+        }(system.dispatcher)
+      }
+  }
+
   def plainSink[T](outlet: CodecOutlet[T]): Sink[T, NotUsed] = {
     val topic = findTopicForPort(outlet)
     val producerSettings = ProducerSettings(system, new ByteArraySerializer, new ByteArraySerializer)