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,15 +19,16 @@ 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.ShardingMessageExtractor
+import akka.cluster.sharding.typed.scaladsl.{ ClusterSharding, Entity, EntityTypeKey }
 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._
 
@@ -73,6 +74,11 @@ 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]) = {
+    ClusterSharding(system.toTyped).init(shardEntity)
+    sourceWithContext(inlet)
+  }
+
   private def flowWithCommittableContext[T](outlet: CodecOutlet[T]): cloudflow.akkastream.scaladsl.FlowWithCommittableContext[T, T] = {
     val flow = Flow[T]
 
@@ -109,6 +115,12 @@ 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): Source[T, _] = {
+    ClusterSharding(system.toTyped).init(shardEntity)
+    plainSource(inlet, resetPosition)
+  }
+
   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,11 +17,12 @@
 package cloudflow.akkastream
 
 import scala.concurrent.Future
-
 import akka.NotUsed
 import akka.actor.ActorSystem
+import akka.cluster.sharding.typed.scaladsl.{ Entity, EntityTypeKey }
 import akka.kafka.ConsumerMessage.{ Committable, CommittableOffset }
 import akka.kafka.CommitterSettings
+import akka.kafka.cluster.sharding.KafkaClusterSharding
 import akka.stream.scaladsl._
 import cloudflow.streamlets._
 
@@ -38,11 +39,21 @@ 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]
+  ): SourceWithContext[T, CommittableOffset, _]
+
   @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
+  ): Source[T, _]
 
   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,21 +21,21 @@ import java.nio.file.{ Files, Paths }
 
 import scala.concurrent._
 import scala.util._
-
 import akka._
 import akka.actor.ActorSystem
+import akka.actor.typed.scaladsl.adapter._
+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._
 
 /**
@@ -45,7 +45,8 @@ final class AkkaStreamletContextImpl(
     private[cloudflow] override val streamletDefinition: StreamletDefinition,
     sys: ActorSystem
 ) extends AkkaStreamletContext {
-  implicit val system: ActorSystem = sys
+  implicit val system: ActorSystem  = sys
+  implicit val ec: ExecutionContext = sys.dispatcher
 
   override def config: Config = streamletDefinition.config
 
@@ -91,6 +92,68 @@ 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]): SourceWithContext[T, CommittableOffset, _] = {
+    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 committable source for group: $gId topic: ${topic.name}")
+
+    import scala.concurrent.duration._
+    val messageExtractor: Future[KafkaClusterSharding.KafkaShardingMessageExtractor[M]] =
+      KafkaClusterSharding(system).messageExtractor(
+        topic = topic.name,
+        timeout = 10.seconds,
+        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
+            }
+        }
+      }
+      .asSourceWithContext { case (_, committableOffset) ⇒ committableOffset }
+      .map { case (record, _) ⇒ record }
+  }
+
+  override def shardedSourceWithCommittableContext[T, M, E](
+      inlet: CodecInlet[T],
+      shardEntity: Entity[M, E]
+  ): SourceWithContext[T, CommittableOffset, _] =
+    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 +219,56 @@ final class AkkaStreamletContextImpl(
       }
   }
 
+  def shardedPlainSource[T, M, E](inlet: CodecInlet[T], shardEntity: Entity[M, E], resetPosition: ResetPosition = Latest): Source[T, _] = {
+    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 committable source for group: $gId topic: ${topic.name}")
+
+    import scala.concurrent.duration._
+    val messageExtractor: Future[KafkaClusterSharding.KafkaShardingMessageExtractor[M]] =
+      KafkaClusterSharding(system).messageExtractor(
+        topic = topic.name,
+        timeout = 10.seconds,
+        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)
+            }
+        }
+      }
+  }
+
   def plainSink[T](outlet: CodecOutlet[T]): Sink[T, NotUsed] = {
     val topic = findTopicForPort(outlet)
     val producerSettings = ProducerSettings(system, new ByteArraySerializer, new ByteArraySerializer)

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

@@ -20,13 +20,13 @@ import java.nio.file.Path
 
 import akka.NotUsed
 import akka.actor.ActorSystem
+import akka.actor.typed.scaladsl.adapter._
+import akka.annotation.ApiMayChange
+import akka.cluster.sharding.typed.scaladsl.{ ClusterSharding, Entity }
 import akka.stream.scaladsl._
-
 import akka.kafka._
 import akka.kafka.ConsumerMessage._
-
 import com.typesafe.config.Config
-
 import cloudflow.streamlets._
 import cloudflow.akkastream.scaladsl._
 
@@ -89,6 +89,12 @@ abstract class AkkaStreamletLogic(implicit val context: AkkaStreamletContext) ex
    */
   def getExecutionContext() = executionContext
 
+  /**
+   * Helper method to make it easier to start typed cluster sharding
+   * with an classic actor system
+   */
+  def clusterSharding() = ClusterSharding(system.toTyped)
+
   /**
    * This source emits `T` records together with the offset position as context, thus makes it possible
    * to commit offset positions to Kafka (as received through the `inlet`).
@@ -132,6 +138,36 @@ abstract class AkkaStreamletLogic(implicit val context: AkkaStreamletContext) ex
   def getSourceWithCommittableContext[T](inlet: CodecInlet[T]): akka.stream.javadsl.SourceWithContext[T, Committable, _] =
     context.sourceWithCommittableContext(inlet).asJava
 
+  /**
+   * This source is designed to function the same as [[sourceWithCommittableContext]]
+   * while also leveraging Akka Kafka Cluster Sharding for stateful streaming.
+   *
+   * This source emits `T` records together with the committable context, thus makes it possible
+   * to commit offset positions to Kafka using `committableSink(outlet: CodecOutlet[T])`.
+   *
+   * It is required to use this source with Akka Cluster.  This source will start up
+   * Akka Cluster Sharding using the supplied `shardEntity` and configure the kafka external
+   * shard strategy to co-locate Kafka partition consumption with Akka Cluster shards.
+   *
+   * @param inlet the inlet to consume messages from. The inlet specifies a [[cloudflow.streamlets.Codec]] that is used to deserialize the records read from the underlying transport.
+   * @param shardEntity is used to specific the settings for the started shard region
+   **/
+  @ApiMayChange
+  def shardedSourceWithCommittableContext[T, M, E](inlet: CodecInlet[T],
+                                                   shardEntity: Entity[M, E]): SourceWithContext[T, CommittableOffset, _] =
+    context.shardedSourceWithCommittableContext(inlet, shardEntity)
+
+  /**
+   * Java API
+   * @see [[shardedSourceWithCommittableContext]]
+   */
+  @ApiMayChange
+  def getShardedSourceWithCommittableContext[T, M, E](
+      inlet: CodecInlet[T],
+      shardEntity: Entity[M, E]
+  ): akka.stream.javadsl.SourceWithContext[T, Committable, _] =
+    context.shardedSourceWithCommittableContext(inlet, shardEntity).asJava
+
   /**
    * The `plainSource` emits `T` records (as received through the `inlet`).
    *
@@ -152,6 +188,41 @@ abstract class AkkaStreamletLogic(implicit val context: AkkaStreamletContext) ex
   def getPlainSource[T](inlet: CodecInlet[T], resetPosition: ResetPosition): akka.stream.javadsl.Source[T, NotUsed] =
     plainSource(inlet, resetPosition).asJava
 
+  /**
+   * This source is designed to function the same as [[plainSource]]
+   * while also leveraging Akka Kafka Cluster Sharding for stateful streaming.
+   *
+   * The `plainSource` emits `T` records (as received through the `inlet`).
+   *
+   * It has no support for committing offsets to Kafka.
+   *
+   * It is required to use this source with Akka Cluster.  This source will start up
+   * Akka Cluster Sharding using the supplied `shardEntity` and configure the kafka external
+   * shard strategy to co-locate Kafka partition consumption with Akka Cluster shards.
+   *
+   * @param the inlet to consume messages from. The inlet specifies a [[cloudflow.streamlets.Codec]] that is used to deserialize the records read from the underlying transport.
+   * @param shardEntity is used to specific the settings for the started shard region
+   **/
+  @ApiMayChange
+  def shardedPlainSource[T, M, E](inlet: CodecInlet[T], shardEntity: Entity[M, E], resetPosition: ResetPosition = Latest): Source[T, _] =
+    context.shardedPlainSource(inlet, shardEntity, resetPosition)
+
+  /**
+   * Java API
+   */
+  @ApiMayChange
+  def getShardedPlainSource[T, M, E](inlet: CodecInlet[T], shardEntity: Entity[M, E]): akka.stream.javadsl.Source[T, _] =
+    shardedPlainSource(inlet, shardEntity, Latest).asJava
+
+  /**
+   * Java API
+   */
+  @ApiMayChange
+  def getShardedPlainSource[T, M, E](inlet: CodecInlet[T],
+                                     shardEntity: Entity[M, E],
+                                     resetPosition: ResetPosition = Latest): akka.stream.javadsl.Source[T, _] =
+    shardedPlainSource(inlet, shardEntity, resetPosition).asJava
+
   /**
    * Creates a sink for publishing `T` records to the outlet. The records are partitioned according to the `partitioner` of the `outlet`.
    * The `outlet` specifies a [[cloudflow.streamlets.Codec]] that will be used to serialize the records that are written to Kafka.

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

@@ -16,8 +16,10 @@
 
 package cloudflow.akkastream
 
+import akka.annotation.ApiMayChange
 import cloudflow.streamlets._
 
+@ApiMayChange
 trait Clustering { this: AkkaStreamlet ⇒
   protected[cloudflow] override def attributes = Set(AkkaClusterAttribute) ++ customAttributes
 }

core/project/Dependencies.scala

@@ -6,7 +6,7 @@ object Version {
   val Akka          = "2.6.6"
   val AkkaHttp      = "10.1.12"
   val AkkaMgmt      = "1.0.8"
-  val AlpakkaKafka  = "2.0.3"
+  val AlpakkaKafka  = "2.0.4"
   val Scala         = "2.12.11"
   val Spark         = "2.4.5"
   val Flink         = "1.10.0"
@@ -22,10 +22,12 @@ object Library {
   val AkkaHttpJackson        = "com.typesafe.akka"  %% "akka-http-jackson"         % Version.AkkaHttp
   val AkkaHttpSprayJson      = "com.typesafe.akka"  %% "akka-http-spray-json"      % Version.AkkaHttp
   val AkkaActor              = "com.typesafe.akka"  %% "akka-actor"                % Version.Akka
+  val AkkaShardingTyped      = "com.typesafe.akka"  %% "akka-cluster-sharding-typed"% Version.Akka
   val AkkaStream             = "com.typesafe.akka"  %% "akka-stream"               % Version.Akka
   val AkkaSlf4j              = "com.typesafe.akka"  %% "akka-slf4j"                % Version.Akka
   val AkkaStreamContrib      = "com.typesafe.akka"  %% "akka-stream-contrib"       % "0.10"
   val AkkaStreamKafka        = ("com.typesafe.akka" %% "akka-stream-kafka"         % Version.AlpakkaKafka) .exclude("com.fasterxml.jackson.core","jackson-databind") .exclude("com.fasterxml.jackson.module", "jackson-module-scala")
+  val AkkaStreamKafaSharding = "com.typesafe.akka"  %% "akka-stream-kafka-cluster-sharding" % Version.AlpakkaKafka
   val AkkaStreamKafkaTestkit = ("com.typesafe.akka" %% "akka-stream-kafka-testkit" % Version.AlpakkaKafka) .exclude("com.typesafe.akka", "akka-stream-testkit")
   val AkkaStreamTestkit      = "com.typesafe.akka"  %% "akka-stream-testkit"       % Version.Akka