Kafka external sharding strategy

core/build.sbt

@@ -85,6 +85,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,14 @@ package cloudflow.akkastream.testkit
 import java.util.concurrent.atomic.AtomicReference
 
 import scala.concurrent._
-
 import akka.NotUsed
 import akka.actor._
+import akka.cluster.sharding.typed.scaladsl.{ 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._
 
@@ -172,6 +171,16 @@ private[testkit] case class TestContext(
 
   def metricTags(): Map[String, String] =
     Map()
+
+  override private[akkastream] def shardedSourceWithCommittableContext[T, M, E](inlet: CodecInlet[T],
+                                                                                shardEntity: Entity[M, E],
+                                                                                entityIdExtractor: M => String) = ???
+
+  override private[akkastream] def shardedPlainSource[T, E](inlet: CodecInlet[T],
+                                                            typeKey: EntityTypeKey[E],
+                                                            entityIdExtractor: E => String,
+                                                            resetPosition: ResetPosition) =
+    ???
 }
 
 case class TestContextException(portName: String, msg: String) extends RuntimeException(msg)

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,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],
+      entityIdExtractor: M => String
+  ): 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, E](
+      inlet: CodecInlet[T],
+      typeKey: EntityTypeKey[E],
+      entityIdExtractor: E => String,
+      resetPosition: ResetPosition
+  ): (Source[T, NotUsed], Future[KafkaClusterSharding.KafkaShardingNoEnvelopeExtractor[E]])
 
   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]
@@ -116,3 +129,6 @@ case object Earliest extends ResetPosition {
 case object Latest extends ResetPosition {
   val autoOffsetReset = "latest"
 }
+
+case class ShardedSourceEnvelope[T, E](source: SourceWithContext[T, CommittableOffset, _],
+                                       kafkaShardingExtractor: KafkaClusterSharding.KafkaShardingNoEnvelopeExtractor[E])

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.cluster.sharding.external.ExternalShardAllocationStrategy
+import akka.cluster.sharding.typed.ClusterShardingSettings
+import akka.cluster.sharding.typed.scaladsl.{ ClusterSharding, Entity, EntityTypeKey }
 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,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],
+                                                            entityIdExtractor: M => String): 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.KafkaShardingNoEnvelopeExtractor[M]] =
+      KafkaClusterSharding(system).messageExtractorNoEnvelope(
+        timeout = 10.seconds,
+        topic = topic.name,
+        entityIdExtractor = entityIdExtractor,
+        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],
+      entityIdExtractor: M => String
+  ): SourceWithContext[T, CommittableOffset, _] =
+    shardedSourceWithContext[T, M, E](inlet, shardEntity, entityIdExtractor)
+
   @deprecated("Use sourceWithCommittableContext", "1.3.4")
   override def sourceWithOffsetContext[T](inlet: CodecInlet[T]): cloudflow.akkastream.scaladsl.SourceWithOffsetContext[T] =
     sourceWithContext[T](inlet)
@@ -174,6 +240,49 @@ final class AkkaStreamletContextImpl(
       .mapMaterializedValue(_ ⇒ NotUsed)
   }
 
+  def shardedPlainSource[T, E](
+      inlet: CodecInlet[T],
+      typeKey: EntityTypeKey[E],
+      entityIdExtractor: E => String,
+      resetPosition: ResetPosition = Latest
+  ): (Source[T, NotUsed], Future[KafkaClusterSharding.KafkaShardingNoEnvelopeExtractor[E]]) = {
+    // TODO clean this up, lot of copying code, refactor.
+    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(typeKey)
+
+    import akka.actor.typed.scaladsl.adapter._
+    val subscription = Subscriptions
+      .topics(topic.name)
+      .withRebalanceListener(rebalanceListener.toClassic)
+
+    val consumer = Consumer
+      .plainSource(consumerSettings, subscription)
+      .mapMaterializedValue(_ ⇒ NotUsed) // TODO we should likely use control to gracefully stop.
+      .via(handleTermination)
+      .map { record ⇒
+        inlet.codec.decode(record.value)
+      }
+
+    import scala.concurrent.duration._
+    val messageExtractor: Future[KafkaClusterSharding.KafkaShardingNoEnvelopeExtractor[E]] =
+      KafkaClusterSharding(system).messageExtractorNoEnvelope(
+        timeout = 10.seconds,
+        topic = "user-topic",
+        entityIdExtractor = entityIdExtractor,
+        settings = consumerSettings
+      )
+
+    (consumer, messageExtractor)
+  }
+
   def sinkRef[T](outlet: CodecOutlet[T]): WritableSinkRef[T] = {
     val topic = findTopicForPort(outlet)
 

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

@@ -20,16 +20,17 @@ import java.nio.file.Path
 
 import akka.NotUsed
 import akka.actor.ActorSystem
+import akka.cluster.sharding.typed.scaladsl.{ Entity, EntityTypeKey }
 import akka.stream.scaladsl._
-
 import akka.kafka._
 import akka.kafka.ConsumerMessage._
-
+import akka.kafka.cluster.sharding.KafkaClusterSharding
 import com.typesafe.config.Config
-
 import cloudflow.streamlets._
 import cloudflow.akkastream.scaladsl._
 
+import scala.concurrent.Future
+
 /**
  * Provides an entry-point for defining the behavior of an AkkaStreamlet.
  * Override the `run` method to implement the specific logic / code that should be executed once the streamlet deployed
@@ -118,6 +119,13 @@ abstract class AkkaStreamletLogic(implicit val context: AkkaStreamletContext) ex
   def sourceWithCommittableContext[T](inlet: CodecInlet[T]): SourceWithCommittableContext[T] =
     context.sourceWithCommittableContext(inlet)
 
+  def shardedSourceWithCommittableContext[T, M, E](
+      inlet: CodecInlet[T],
+      shardEntity: Entity[M, E],
+      entityIdExtractor: M => String
+  ): SourceWithContext[T, CommittableOffset, _] =
+    context.shardedSourceWithCommittableContext(inlet, shardEntity, entityIdExtractor)
+
   /**
    * Java API
    */
@@ -141,6 +149,14 @@ abstract class AkkaStreamletLogic(implicit val context: AkkaStreamletContext) ex
   def plainSource[T](inlet: CodecInlet[T], resetPosition: ResetPosition = Latest): akka.stream.scaladsl.Source[T, NotUsed] =
     context.plainSource(inlet, resetPosition)
 
+  def shardedPlainSource[T, E](
+      inlet: CodecInlet[T],
+      typeKey: EntityTypeKey[E],
+      entityIdExtractor: E => String,
+      resetPosition: ResetPosition = Latest
+  ): (Source[T, NotUsed], Future[KafkaClusterSharding.KafkaShardingNoEnvelopeExtractor[E]]) =
+    context.shardedPlainSource(inlet, typeKey, entityIdExtractor, resetPosition)
+
   /**
    * Java API
    */

core/project/Dependencies.scala

@@ -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
 

examples/call-record-aggregator/.sbtopts

@@ -0,0 +1 @@
+-mem 2048

examples/call-record-aggregator/project/plugins.sbt

@@ -1 +1 @@
-addSbtPlugin("org.scalameta" % "sbt-scalafmt" % "2.2.1")
+addSbtPlugin("com.julianpeeters" % "sbt-avrohugger" % "2.0.0-RC18")

examples/connected-car-cluster-sharding/akka-connected-car-streamlet/src/main/resources/application.conf

@@ -0,0 +1,4 @@
+akka {
+  stdout-loglevel = "DEBUG"
+  loglevel = "DEBUG"
+}

examples/connected-car-cluster-sharding/akka-connected-car-streamlet/src/main/scala/connectedcar/actors/ConnectedCarActor.scala

@@ -1,49 +1,30 @@
 package connectedcar.actors
 
-import akka.actor.{ Actor, ActorLogging, ActorRef }
-import akka.cluster.sharding.ShardRegion
-import connectedcar.data.{ ConnectedCarAgg, ConnectedCarERecord }
+import akka.actor.typed.{ActorRef, Behavior}
+import akka.actor.typed.scaladsl.Behaviors
 
-import scala.concurrent.ExecutionContext
+import connectedcar.data.{ConnectedCarAgg, ConnectedCarERecord}
 
-object ConnectedCarActor {
-  val extractEntityId: ShardRegion.ExtractEntityId = {
-    case msg: ConnectedCarERecord ⇒ (msg.carId.toString, msg)
-  }
+case class ConnectedCarERecordWrapper(record: ConnectedCarERecord, sender:ActorRef[ConnectedCarAgg])
 
-  private val numberOfShards = 100
+object ConnectedCarActor {
+  def apply(): Behavior[ConnectedCarERecordWrapper] = {
+    def updated(numberOfRecords: Int, driverName: String, carId:Int, averageSpeed: Double, currentSpeed: Double): Behavior[ConnectedCarERecordWrapper] = {
+      Behaviors.receive { (context, message) => {
+        context.log.info("Updated CarId: " + carId +
+          " Driver Name: " + driverName + " CarSpeed: " + currentSpeed + " From Actor:" + message.sender.path)
 
-  val extractShardId: ShardRegion.ExtractShardId = {
-    case msg: ConnectedCarERecord ⇒ (msg.carId % numberOfShards).toString
-  }
-}
-
-class ConnectedCarActor extends Actor with ActorLogging {
-
-  val carId: String      = "Car-" + self.path.name
-  var driverName: String = null
-  var currentSpeed       = 0.0
-  var averageSpeed       = 0.0
-  var numberOfRecords    = 0
-
-  var treeActor: ActorRef           = null
-  implicit val ec: ExecutionContext = context.dispatcher
-
-  override def receive: Receive = {
-    case record: ConnectedCarERecord ⇒ {
-      if (numberOfRecords == 0) {
-        driverName = record.driver
-        averageSpeed = record.speed
-      } else {
-        averageSpeed = ((averageSpeed * numberOfRecords) + record.speed) / (numberOfRecords + 1)
-      }
 
-      numberOfRecords += 1
-      currentSpeed = record.speed
+        val newAverage = ((averageSpeed * numberOfRecords) + message.record.speed) / (numberOfRecords + 1)
+        val newNumberOfRecords = numberOfRecords+1
 
-      log.info("Updated CarId: " + carId + " Driver Name: " + driverName + " CarSpeed: " + currentSpeed + " From Actor:" + sender().path)
+        message.sender ! ConnectedCarAgg(message.record.carId, message.record.driver, averageSpeed, newNumberOfRecords)
 
-      sender() ! ConnectedCarAgg(record.carId, record.driver, averageSpeed, numberOfRecords)
+        updated(newNumberOfRecords, message.record.driver, message.record.carId, newAverage, message.record.speed)
+      }
+      }
     }
+
+    updated(0, "", 0, 0, 0.0)
   }
-}
+}

examples/connected-car-cluster-sharding/akka-connected-car-streamlet/src/main/scala/connectedcar/streamlets/CarDataPrinter.scala

@@ -13,7 +13,7 @@ object CarDataPrinter extends AkkaStreamlet {
   override def createLogic() = new RunnableGraphStreamletLogic() {
     val flow = FlowWithCommittableContext[ConnectedCarAgg]
       .map { record ⇒
-        log.info("CarId: " + record.carId)
+        println("CarId: " + record.carId)
       }
 
     def runnableGraph =