Kafka external sharding strategy

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

@@ -21,7 +21,7 @@ import java.util.concurrent.atomic.AtomicReference
 import scala.concurrent._
 import akka.NotUsed
 import akka.actor._
-import akka.cluster.sharding.typed.scaladsl.EntityTypeKey
+import akka.cluster.sharding.typed.scaladsl.{ Entity, EntityTypeKey }
 import akka.kafka.CommitterSettings
 import akka.kafka.ConsumerMessage._
 import akka.stream._
@@ -172,9 +172,9 @@ private[testkit] case class TestContext(
   def metricTags(): Map[String, String] =
     Map()
 
-  override private[akkastream] def shardedSourceWithCommittableContext[T, E](inlet: CodecInlet[T],
-                                                                             typeKey: EntityTypeKey[E],
-                                                                             entityIdExtractor: E => String) = ???
+  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],

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

@@ -19,7 +19,7 @@ package cloudflow.akkastream
 import scala.concurrent.Future
 import akka.NotUsed
 import akka.actor.ActorSystem
-import akka.cluster.sharding.typed.scaladsl.EntityTypeKey
+import akka.cluster.sharding.typed.scaladsl.{ Entity, EntityTypeKey }
 import akka.kafka.ConsumerMessage.{ Committable, CommittableOffset }
 import akka.kafka.CommitterSettings
 import akka.kafka.cluster.sharding.KafkaClusterSharding
@@ -39,11 +39,11 @@ trait AkkaStreamletContext extends StreamletContext {
       inlet: CodecInlet[T]
   ): cloudflow.akkastream.scaladsl.SourceWithCommittableContext[T]
 
-  private[akkastream] def shardedSourceWithCommittableContext[T, E](
+  private[akkastream] def shardedSourceWithCommittableContext[T, M, E](
       inlet: CodecInlet[T],
-      typeKey: EntityTypeKey[E],
-      entityIdExtractor: E => String
-  ): Source[ShardedSourceEnvelope[T, E], _]
+      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]

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

@@ -23,7 +23,9 @@ import scala.concurrent._
 import scala.util._
 import akka._
 import akka.actor.ActorSystem
-import akka.cluster.sharding.typed.scaladsl.EntityTypeKey
+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
@@ -90,9 +92,9 @@ final class AkkaStreamletContextImpl(
   override def sourceWithCommittableContext[T](inlet: CodecInlet[T]): cloudflow.akkastream.scaladsl.SourceWithCommittableContext[T] =
     sourceWithContext[T](inlet)
 
-  private[akkastream] def shardedSourceWithContext[T, E](inlet: CodecInlet[T],
-                                                         typeKey: EntityTypeKey[E],
-                                                         entityIdExtractor: E => String): Source[ShardedSourceEnvelope[T, E], _] = {
+  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)
 
@@ -103,7 +105,7 @@ final class AkkaStreamletContextImpl(
       .withProperties(topic.kafkaConsumerProperties)
 
     val rebalanceListener: akka.actor.typed.ActorRef[ConsumerRebalanceEvent] =
-      KafkaClusterSharding(system).rebalanceListener(typeKey)
+      KafkaClusterSharding(system).rebalanceListener(shardEntity.typeKey)
 
     import akka.actor.typed.scaladsl.adapter._
     val subscription = Subscriptions
@@ -112,40 +114,48 @@ final class AkkaStreamletContextImpl(
 
     system.log.info(s"Creating committable source for group: $gId topic: ${topic.name}")
 
-    val consumer = 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 }
-
     import scala.concurrent.duration._
-    val messageExtractor: Future[KafkaClusterSharding.KafkaShardingNoEnvelopeExtractor[E]] =
+    val messageExtractor: Future[KafkaClusterSharding.KafkaShardingNoEnvelopeExtractor[M]] =
       KafkaClusterSharding(system).messageExtractorNoEnvelope(
         timeout = 10.seconds,
         topic = topic.name,
         entityIdExtractor = entityIdExtractor,
         settings = consumerSettings
       )
 
-    Source.future {
-      messageExtractor.map { messageExtractor =>
-        ShardedSourceEnvelope(consumer, messageExtractor)
+    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, E](
+  override def shardedSourceWithCommittableContext[T, M, E](
       inlet: CodecInlet[T],
-      typeKey: EntityTypeKey[E],
-      entityIdExtractor: E => String
-  ): Source[ShardedSourceEnvelope[T, E], _] =
-    shardedSourceWithContext[T, E](inlet, typeKey, entityIdExtractor)
+      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] =

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

@@ -20,7 +20,7 @@ import java.nio.file.Path
 
 import akka.NotUsed
 import akka.actor.ActorSystem
-import akka.cluster.sharding.typed.scaladsl.EntityTypeKey
+import akka.cluster.sharding.typed.scaladsl.{ Entity, EntityTypeKey }
 import akka.stream.scaladsl._
 import akka.kafka._
 import akka.kafka.ConsumerMessage._
@@ -119,12 +119,12 @@ abstract class AkkaStreamletLogic(implicit val context: AkkaStreamletContext) ex
   def sourceWithCommittableContext[T](inlet: CodecInlet[T]): SourceWithCommittableContext[T] =
     context.sourceWithCommittableContext(inlet)
 
-  def shardedSourceWithCommittableContext[T, E](
+  def shardedSourceWithCommittableContext[T, M, E](
       inlet: CodecInlet[T],
-      typeKey: EntityTypeKey[E],
-      entityIdExtractor: E => String
-  ): Source[ShardedSourceEnvelope[T, E], _] =
-    context.shardedSourceWithCommittableContext(inlet, typeKey, entityIdExtractor)
+      shardEntity: Entity[M, E],
+      entityIdExtractor: M => String
+  ): SourceWithContext[T, CommittableOffset, _] =
+    context.shardedSourceWithCommittableContext(inlet, shardEntity, entityIdExtractor)
 
   /**
    * Java API

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/scala/connectedcar/streamlets/ConnectedCarCluster.scala

@@ -25,26 +25,14 @@ object ConnectedCarCluster extends AkkaStreamlet with Clustering {
   override def createLogic = new RunnableGraphStreamletLogic() {
     val typeKey = EntityTypeKey[ConnectedCarERecordWrapper]("Car")
 
-    val source = shardedSourceWithCommittableContext(in,
-                                      typeKey: EntityTypeKey[ConnectedCarERecordWrapper],
-                                      (msg: ConnectedCarERecordWrapper) => msg.record.carId+""
-                                    )
+    val messageExtractor = (msg: ConnectedCarERecordWrapper) => msg.record.carId+""
+    val entity = Entity(typeKey)(createBehavior = entityContext => ConnectedCarActor())
 
-    val clusterSharding = ClusterSharding(system.toTyped)
-
-    def runnableGraph = source.map(kafkaEnvelope => {
+    val source = shardedSourceWithCommittableContext(in, entity, messageExtractor)
 
-      clusterSharding.init(
-        Entity(typeKey)(createBehavior = _ => ConnectedCarActor())
-          .withAllocationStrategy(new ExternalShardAllocationStrategy(system, typeKey.name))
-          .withMessageExtractor(kafkaEnvelope.kafkaShardingExtractor)
-          .withSettings(ClusterShardingSettings(system.toTyped)))
+    val clusterSharding = ClusterSharding(system.toTyped)
 
-      kafkaEnvelope.source
-        .via(flow)
-        .to(committableSink(out))
-        .run()
-    }).to(Sink.ignore)
+    def runnableGraph = source.via(flow).to(Sink.ignore)
 
     implicit val timeout: Timeout = 3.seconds
     def flow =