Introduce partition mapping to kafka-backed persistence module (#718)

Author: nviliunov-evolution-throwaway

persistence-kafka/src/main/scala/com/evolutiongaming/kafka/flow/kafkapersistence/KafkaPersistenceModule.scala

@@ -82,21 +82,33 @@ object KafkaPersistenceModule {
     producer: Producer[F],
     consumerConfig: ConsumerConfig,
     snapshotTopicPartition: TopicPartition,
-    metrics: FlowMetrics[F]
+    metrics: FlowMetrics[F],
+    partitionMapper: KafkaPersistencePartitionMapper = KafkaPersistencePartitionMapper.identity,
   )(
     implicit fromBytesKey: FromBytes[F, String],
     fromBytesState: FromBytes[F, S],
     toBytesState: ToBytes[F, S]
   ): Resource[F, KafkaPersistenceModule[F, S]] = {
     implicit val fromTry: FromTry[F] = FromTry.lift
 
-    def readPartitionData(implicit log: Log[F]): F[BytesByKey] =
-      KafkaPartitionPersistence.readSnapshots[F](
-        consumerOf     = consumerOf,
-        consumerConfig = consumerConfig,
-        snapshotTopic  = snapshotTopicPartition.topic,
-        partition      = snapshotTopicPartition.partition
-      )
+    def readPartitionData(implicit log: Log[F]): F[BytesByKey] = {
+      val targetPartition = partitionMapper.getStatePartition(snapshotTopicPartition.partition)
+      KafkaPartitionPersistence
+        .readSnapshots[F](
+          consumerOf     = consumerOf,
+          consumerConfig = consumerConfig,
+          snapshotTopic  = snapshotTopicPartition.topic,
+          partition      = targetPartition
+        )
+        .map { snapshots =>
+          snapshots
+            .view
+            .filterKeys { key =>
+              partitionMapper.isStateKeyOwned(key, snapshotTopicPartition.partition)
+            }
+            .toMap
+        }
+    }
 
     def makeKeysOf(cache: Cache[F, String, ByteVector]): F[KeysOf[F, KafkaKey]] = {
       LogOf[F].apply(classOf[KeysOf[F, KafkaKey]]).map { implicit log =>
@@ -127,7 +139,7 @@ object KafkaPersistenceModule {
 
         val snapshotDatabase = SnapshotDatabase(
           read  = read,
-          write = KafkaSnapshotWriteDatabase.of[F, S](snapshotTopicPartition, producer)
+          write = KafkaSnapshotWriteDatabase.of[F, S](snapshotTopicPartition, producer, partitionMapper)
         ).withMetricsK(metrics.snapshotDatabaseMetrics)
 
         PersistenceOf.snapshotsOnly[F, KafkaKey, S, ConsumerRecord[String, ByteVector]](

persistence-kafka/src/main/scala/com/evolutiongaming/kafka/flow/kafkapersistence/KafkaPersistenceModuleOf.scala

@@ -36,7 +36,8 @@ object KafkaPersistenceModuleOf {
     producer: Producer[F],
     consumerConfig: ConsumerConfig,
     snapshotTopic: Topic,
-    metrics: FlowMetrics[F]
+    metrics: FlowMetrics[F],
+    partitionMapper: KafkaPersistencePartitionMapper = KafkaPersistencePartitionMapper.identity,
   )(
     implicit fromBytesKey: FromBytes[F, String],
     fromBytesState: FromBytes[F, S],
@@ -47,7 +48,8 @@ object KafkaPersistenceModuleOf {
       producer               = producer,
       consumerConfig         = consumerConfig,
       snapshotTopicPartition = TopicPartition(snapshotTopic, partition),
-      metrics                = metrics
+      metrics                = metrics,
+      partitionMapper        = partitionMapper,
     )
   }
 
@@ -60,6 +62,7 @@ object KafkaPersistenceModuleOf {
     implicit fromBytesKey: FromBytes[F, String],
     fromBytesState: FromBytes[F, S],
     toBytesState: ToBytes[F, S]
-  ): KafkaPersistenceModuleOf[F, S] = caching(consumerOf, producer, consumerConfig, snapshotTopic, FlowMetrics.empty[F])
+  ): KafkaPersistenceModuleOf[F, S] =
+    caching(consumerOf, producer, consumerConfig, snapshotTopic, FlowMetrics.empty[F])
 
 }

persistence-kafka/src/main/scala/com/evolutiongaming/kafka/flow/kafkapersistence/KafkaPersistencePartitionMapper.scala

@@ -0,0 +1,57 @@
+package com.evolutiongaming.kafka.flow.kafkapersistence
+
+import com.evolutiongaming.skafka.Partition
+import org.apache.kafka.clients.producer.internals.BuiltInPartitioner
+
+/** Maps partitions of source Kafka topics into persistence topics.
+  *
+  * Please be careful when using this with [[com.evolutiongaming.kafka.flow.RemapKey]]. Only the identity mapper is
+  * guaranteed to work properly with an arbitrary `RemapKey`, for other combinations you have to manually ensure that
+  * the [[isStateKeyOwned]] implementation is correct and will not allow duplicate KeyFlows.
+  *
+  * If the aggregate key depends on the record's contents, then only the identity mapper can be used.
+  */
+trait KafkaPersistencePartitionMapper {
+
+  /** Called after rebalance or initial partition assignment.
+    * @param sourcePartition
+    *   partition of the input stream, i.e. the kafka-journal topic.
+    * @return
+    *   partition of the persistence topic that has snapshots for aggregates built by events from the `sourcePartition`.
+    */
+  def getStatePartition(sourcePartition: Partition): Partition
+
+  /** Checks if the aggregate in the state partition should be initialized as a
+    * [[com.evolutiongaming.kafka.flow.KeyFlow]].
+    *
+    * If the aggregate is initialized, it will have timers and ticks started. This is not desirable if the aggregate is
+    * actually sourced from a different partition, which will also be started concurrently.
+    * @param stateKey
+    *   the aggregate's key.
+    * @param sourcePartition
+    *   partition of the input stream.
+    * @return
+    *   `true` if the aggregate is built from events in `sourcePartition`.
+    */
+  def isStateKeyOwned(stateKey: String, sourcePartition: Partition): Boolean
+}
+
+object KafkaPersistencePartitionMapper {
+  def identity: KafkaPersistencePartitionMapper = Identity
+  def modulo(sourcePartitions: Int, statePartitions: Int): KafkaPersistencePartitionMapper =
+    new Modulo(sourcePartitions, statePartitions)
+
+  private object Identity extends KafkaPersistencePartitionMapper {
+    override def getStatePartition(sourcePartition: Partition): Partition = sourcePartition
+
+    override def isStateKeyOwned(stateKey: String, sourcePartition: Partition): Boolean = true
+  }
+
+  private class Modulo(sourcePartitions: Int, statePartitions: Int) extends KafkaPersistencePartitionMapper {
+    override def getStatePartition(sourcePartition: Partition): Partition =
+      Partition.unsafe(sourcePartition.value % statePartitions)
+
+    override def isStateKeyOwned(stateKey: String, sourcePartition: Partition): Boolean =
+      BuiltInPartitioner.partitionForKey(stateKey.getBytes, sourcePartitions) == sourcePartition.value
+  }
+}

persistence-kafka/src/main/scala/com/evolutiongaming/kafka/flow/kafkapersistence/KafkaSnapshotWriteDatabase.scala

@@ -12,15 +12,17 @@ object KafkaSnapshotWriteDatabase {
   def of[F[_]: FromTry: Monad, S: ToBytes[F, *]](
     snapshotTopicPartition: TopicPartition,
     producer: Producer[F],
+    partitionMapper: KafkaPersistencePartitionMapper = KafkaPersistencePartitionMapper.identity,
   ): SnapshotWriteDatabase[F, KafkaKey, S] = new SnapshotWriteDatabase[F, KafkaKey, S] {
     override def persist(key: KafkaKey, snapshot: S): F[Unit] = produce(key, snapshot.some)
 
     override def delete(key: KafkaKey): F[Unit] = produce(key, none)
 
     private def produce(key: KafkaKey, snapshot: Option[S]): F[Unit] = {
+      val targetPartition = partitionMapper.getStatePartition(key.topicPartition.partition)
       val record = new ProducerRecord(
         topic     = snapshotTopicPartition.topic,
-        partition = snapshotTopicPartition.partition.some,
+        partition = targetPartition.some,
         key       = key.key.some,
         value     = snapshot
       )