Fine-grained logging for PartitionFlow owned classes

core/src/main/scala/com/evolutiongaming/kafka/flow/AdditionalStatePersist.scala

@@ -4,6 +4,7 @@ import cats.Applicative
 import cats.effect.syntax.all.*
 import cats.effect.{Clock, MonadCancel, MonadCancelThrow, Ref}
 import cats.syntax.all.*
+import com.evolutiongaming.catshelper.Log
 import com.evolutiongaming.kafka.flow.kafka.OffsetToCommit
 import com.evolutiongaming.kafka.flow.persistence.Persistence
 import com.evolutiongaming.skafka.consumer.ConsumerRecord
@@ -56,9 +57,10 @@ object AdditionalStatePersist {
     ignorePersistErrors: Boolean = false,
   ): F[AdditionalStatePersist[F, S, ConsumerRecord[String, ByteVector]]] = {
     for {
+      log              <- keyContext.log(classOf[AdditionalStatePersist[F, S, ConsumerRecord[String, ByteVector]]])
       requestedRef     <- Ref.of(false)
       lastPersistedRef <- Ref.of(none[Instant])
-    } yield of(persistence, keyContext, cooldown, requestedRef, lastPersistedRef, ignorePersistErrors)
+    } yield of(persistence, keyContext, cooldown, requestedRef, lastPersistedRef, ignorePersistErrors, log)
   }
 
   private[flow] def of[F[_]: MonadCancelThrow: Clock, S](
@@ -68,6 +70,7 @@ object AdditionalStatePersist {
     requestedRef: Ref[F, Boolean],
     lastPersistedRef: Ref[F, Option[Instant]],
     ignorePersistErrors: Boolean,
+    log: Log[F],
   ): AdditionalStatePersist[F, S, ConsumerRecord[String, ByteVector]] =
     new AdditionalStatePersist[F, S, ConsumerRecord[String, ByteVector]] {
       private val F          = MonadCancel[F, Throwable]
@@ -76,7 +79,7 @@ object AdditionalStatePersist {
       private val charsToPrint = 1024
 
       override def request: F[Unit] =
-        requestedRef.set(true) >> keyContext.log.info("Additional persisting requested")
+        requestedRef.set(true) >> log.info("Additional persisting requested")
 
       override def persistIfNeeded(record: ConsumerRecord[String, ByteVector], state: S): F[Unit] = {
         for {
@@ -90,16 +93,14 @@ object AdditionalStatePersist {
                   _ <- persistence.flush.attempt.flatMap {
                     case Left(e) if ignorePersistErrors =>
                       val trimmedState = state.toString.take(charsToPrint)
-                      keyContext
-                        .log
+                      log
                         .warn(
                           s"Additional persisting failed, error ignored, error: $e, first $charsToPrint chars of state: $trimmedState",
                           e
                         )
                     case Left(e) =>
                       val trimmedState = state.toString.take(charsToPrint)
-                      keyContext
-                        .log
+                      log
                         .error(
                           s"Additional persisting failed, error: $e, first $charsToPrint chars of state: $trimmedState",
                           e
@@ -108,7 +109,7 @@ object AdditionalStatePersist {
                       for {
                         _ <- OffsetToCommit[F](record.offset).flatMap(keyContext.hold)
                         _ <- lastPersistedRef.set(Instant.ofEpochMilli(now).some)
-                        _ <- keyContext.log.info("Additional persisting success")
+                        _ <- log.info("Additional persisting success")
                       } yield ()
                   }
                 } yield ()

core/src/main/scala/com/evolutiongaming/kafka/flow/KeyContext.scala

@@ -4,7 +4,7 @@ import cats.effect.{Ref, Resource}
 import cats.mtl.Stateful
 import cats.syntax.all.*
 import cats.{Applicative, Monad}
-import com.evolutiongaming.catshelper.Log
+import com.evolutiongaming.catshelper.{Log, LogOf}
 import com.evolutiongaming.kafka.flow.effect.CatsEffectMtlInstances.*
 import com.evolutiongaming.skafka.Offset
 
@@ -16,40 +16,39 @@ trait KeyContext[F[_]] {
   def holding: F[Option[Offset]]
   def hold(offset: Offset): F[Unit]
   def remove: F[Unit]
-  def log: Log[F]
+  def log(source: Class[_]): F[Log[F]]
 }
 object KeyContext {
 
   def apply[F[_]](implicit F: KeyContext[F]): KeyContext[F] = F
 
   def empty[F[_]: Applicative]: KeyContext[F] = new KeyContext[F] {
-    def log                  = Log.empty
-    def holding              = none[Offset].pure[F]
-    def hold(offset: Offset) = ().pure[F]
-    def remove               = ().pure[F]
+    def holding               = none[Offset].pure[F]
+    def hold(offset: Offset)  = ().pure[F]
+    def remove                = ().pure[F]
+    def log(source: Class[_]) = Log.empty[F].pure[F]
   }
 
-  def of[F[_]: Ref.Make: Monad: Log](removeFromCache: F[Unit]): F[KeyContext[F]] =
+  def of[F[_]: Ref.Make: Monad: LogOf](removeFromCache: F[Unit], mdc: Log.Mdc): F[KeyContext[F]] =
     Ref.of[F, Option[Offset]](None) map { storage =>
-      KeyContext(storage.stateInstance, removeFromCache)
+      KeyContext(storage.stateInstance, removeFromCache, mdc)
     }
 
-  def apply[F[_]: Monad: Log](
+  def apply[F[_]: Monad: LogOf](
     storage: Stateful[F, Option[Offset]],
-    removeFromCache: F[Unit]
+    removeFromCache: F[Unit],
+    mdc: Log.Mdc
   ): KeyContext[F] = new KeyContext[F] {
-    def holding              = storage.get
-    def hold(offset: Offset) = storage.set(Some(offset))
-    def remove               = storage.set(None) *> removeFromCache
-    def log                  = Log[F]
+    def holding               = storage.get
+    def hold(offset: Offset)  = storage.set(Some(offset))
+    def remove                = storage.set(None) *> removeFromCache
+    def log(source: Class[_]) = LogOf[F].apply(source).map(_.withMdc(mdc))
   }
 
-  def resource[F[_]: Ref.Make: Monad](
+  def resource[F[_]: Ref.Make: Monad: LogOf](
     removeFromCache: F[Unit],
-    log: Log[F]
-  ): Resource[F, KeyContext[F]] = {
-    implicit val _log = log
-    Resource.eval(of(removeFromCache))
-  }
+    mdc: Log.Mdc
+  ): Resource[F, KeyContext[F]] =
+    Resource.eval(of(removeFromCache, mdc))
 
 }

core/src/main/scala/com/evolutiongaming/kafka/flow/KeyFlow.scala

@@ -74,7 +74,8 @@ object KeyFlow {
     registry: EntityRegistry[F, KafkaKey, S],
   ): Resource[F, KeyFlow[F, A]] =
     for {
-      state <- persistence.read(KeyContext[F].log).toResource
+      log   <- KeyContext[F].log(classOf[KeyFlow[F, A]]).toResource
+      state <- persistence.read(log).toResource
       _     <- storage.set(state).toResource
       // we should not run any timers if there was decision
       // by fold or tick to run the state, because in this

core/src/main/scala/com/evolutiongaming/kafka/flow/KeyFlowOf.scala

@@ -50,11 +50,19 @@ object KeyFlowOf {
     timerFlowOf: TimerFlowOf[F],
     fold: EnhancedFold[F, S, A],
     tick: TickOption[F, S],
-  ): KeyFlowOf[F, S, A] = { (key, context, persistence, timers, additionalPersist, registry) =>
-    implicit val _context = context
-    timerFlowOf(context, persistence, timers) flatMap { timerFlow =>
-      KeyFlow.of(key, fold, tick, persistence, additionalPersist, timerFlow, registry)
+  ): KeyFlowOf[F, S, A] = new KeyFlowOf[F, S, A] {
+    override def apply(
+      key: KafkaKey,
+      context: KeyContext[F],
+      persistence: Persistence[F, S, A],
+      timers: TimerContext[F],
+      additionalPersist: AdditionalStatePersist[F, S, A],
+      registry: EntityRegistry[F, KafkaKey, S]
+    ): Resource[F, KeyFlow[F, A]] = {
+      implicit val _context = context
+      timerFlowOf(context, persistence, timers) flatMap { timerFlow =>
+        KeyFlow.of(key, fold, tick, persistence, additionalPersist, timerFlow, registry)
+      }
     }
   }
-
 }

core/src/main/scala/com/evolutiongaming/kafka/flow/PartitionFlow.scala

@@ -66,7 +66,7 @@ object PartitionFlow {
       }
     }
 
-  def of[F[_]: Async](
+  def of[F[_]: Async: LogOf](
     topicPartition: TopicPartition,
     assignedAt: Offset,
     keyStateOf: KeyStateOf[F],
@@ -97,7 +97,7 @@ object PartitionFlow {
   } yield flow
 
   // TODO: put most `Ref` variables into one state class?
-  def of[F[_]: Async](
+  def of[F[_]: Async: LogOf](
     topicPartition: TopicPartition,
     keyStateOf: KeyStateOf[F],
     committedOffset: Ref[F, Offset],
@@ -116,7 +116,7 @@ object PartitionFlow {
         for {
           context <- KeyContext.resource[F](
             removeFromCache = cache.remove(key).flatten.void,
-            log             = log.prefixed(key)
+            mdc             = Log.Mdc.Eager("key" -> key, "topicPartition" -> topicPartition.toString)
           )
           keyState <- keyStateOf(topicPartition, key, createdAt, context)
         } yield PartitionKey(keyState, context)

core/src/main/scala/com/evolutiongaming/kafka/flow/timer/TimerFlowOf.scala

@@ -2,8 +2,10 @@ package com.evolutiongaming.kafka.flow.timer
 
 import cats.{Applicative, Monad, MonadThrow}
 import cats.effect.Resource
+import cats.effect.syntax.all.*
 import cats.effect.kernel.Resource.ExitCase
 import cats.syntax.all.*
+import com.evolutiongaming.catshelper.Log
 import com.evolutiongaming.kafka.flow.KeyContext
 import com.evolutiongaming.kafka.flow.persistence.FlushBuffers
 import com.evolutiongaming.skafka.Offset
@@ -17,7 +19,6 @@ trait TimerFlowOf[F[_]] {
     persistence: FlushBuffers[F],
     timers: TimerContext[F]
   ): Resource[F, TimerFlow[F]]
-
 }
 object TimerFlowOf {
 
@@ -39,18 +40,19 @@ object TimerFlowOf {
     maxIdle: FiniteDuration   = 10.minutes,
     flushOnRevoke: Boolean    = false,
   ): TimerFlowOf[F] = { (context, persistence, timers) =>
-    def register(touchedAt: Timestamp) =
+    def register(touchedAt: Timestamp): F[Unit] =
       timers.registerProcessing(touchedAt.clock plusMillis fireEvery.toMillis)
 
     val acquire = Resource.eval {
       for {
+        log         <- context.log(classOf[TimerFlowOf[F]])
         current     <- timers.current
         persistedAt <- timers.persistedAt
         committedAt  = persistedAt getOrElse current
         _           <- context.hold(committedAt.offset)
         _           <- register(committedAt)
       } yield new TimerFlow[F] {
-        def onTimer = for {
+        def onTimer: F[Unit] = for {
           current         <- timers.current
           processedAt     <- timers.processedAt
           touchedAt        = processedAt getOrElse committedAt
@@ -60,7 +62,7 @@ object TimerFlowOf {
           canUnload        = expired || offsetDifference > maxOffsetDifference
           _ <-
             if (canUnload) {
-              context.log.info(s"flush, offset difference: $offsetDifference") *>
+              log.info(s"flush, offset difference: $offsetDifference") *>
                 persistence.flush *>
                 context.remove
             } else {
@@ -73,7 +75,6 @@ object TimerFlowOf {
     val cancel = flushOnCancel.apply(context, persistence, timers)
 
     if (flushOnRevoke) acquire <* cancel else acquire
-
   }
 
   /** Performs flush periodically.
@@ -108,6 +109,7 @@ object TimerFlowOf {
 
     val acquire = Resource.eval {
       for {
+        log         <- context.log(classOf[TimerFlowOf[F]])
         current     <- timers.current
         persistedAt <- timers.persistedAt
         committedAt  = persistedAt getOrElse current
@@ -120,13 +122,14 @@ object TimerFlowOf {
           flushedAt      = persistedAt getOrElse committedAt
           triggerFlushAt = flushedAt.clock plusMillis persistEvery.toMillis
           canPersist     = (current.clock compareTo triggerFlushAt) >= 0
-          _ <- MonadThrow[F].whenA(canPersist)(
-            persistence.attemptToPersist(
-              ignorePersistErrors = ignorePersistErrors,
-              context             = context,
-              currentOffset       = current.offset
+          _ <- MonadThrow[F]
+            .whenA(canPersist)(
+              persistence.attemptToPersist(
+                ignorePersistErrors = ignorePersistErrors,
+                context             = context,
+                currentOffset       = current.offset
+              )(log)
             )
-          )
           _ <- register(current)
         } yield ()
       }
@@ -135,7 +138,6 @@ object TimerFlowOf {
     val cancel = flushOnCancel.apply(context, persistence, timers)
 
     if (flushOnRevoke) acquire <* cancel else acquire
-
   }
 
   /** Combines [[unloadOrphaned]] with [[persistPeriodically]] in a single TimerFlow
@@ -162,12 +164,13 @@ object TimerFlowOf {
     maxIdle: FiniteDuration      = 10.minutes,
     flushOnRevoke: Boolean       = false,
     ignorePersistErrors: Boolean = false,
-  ): TimerFlowOf[F] = (context, persistence, timers) => {
+  ): TimerFlowOf[F] = { (context, persistence, timers) =>
     def register(touchedAt: Timestamp): F[Unit] =
       timers.registerProcessing(touchedAt.clock plusMillis fireEvery.toMillis)
 
     val acquire: Resource[F, TimerFlow[F]] = Resource.eval {
       for {
+        log         <- context.log(classOf[TimerFlowOf[F]])
         current     <- timers.current
         persistedAt <- timers.persistedAt
         committedAt  = persistedAt getOrElse current
@@ -190,10 +193,10 @@ object TimerFlowOf {
               ignorePersistErrors = ignorePersistErrors,
               context             = context,
               currentOffset       = current.offset
-            )
+            )(log)
           )
           _ <- Applicative[F].whenA(canUnload)(
-            context.log.info(s"flush, offset difference: $offsetDifference") *> context.remove
+            log.info(s"flush, offset difference: $offsetDifference") *> context.remove
           )
           _ <- register(current)
         } yield ()
@@ -206,39 +209,41 @@ object TimerFlowOf {
   }
 
   /** Performs flush when `Resource` is cancelled only */
-  def flushOnCancel[F[_]: Monad]: TimerFlowOf[F] = { (context, persistence, _) =>
-    val cancel = context.holding flatMap { holding =>
-      Applicative[F].whenA(holding.isDefined) {
-        context.log.info(s"flush on revoke, holding offset: $holding") *>
-          persistence.flush *>
-          context.remove
+  def flushOnCancel[F[_]: Monad]: TimerFlowOf[F] =
+    (context: KeyContext[F], persistence: FlushBuffers[F], _: TimerContext[F]) =>
+      context.log(classOf[TimerFlowOf[F]]).toResource.flatMap { log =>
+        val cancel = context.holding flatMap { holding =>
+          Applicative[F].whenA(holding.isDefined) {
+            log.info(s"flush on revoke, holding offset: $holding") *>
+              persistence.flush *>
+              context.remove
+          }
+        }
+
+        Resource.makeCase(TimerFlow.empty.pure) {
+          case (_, ExitCase.Succeeded) =>
+            cancel
+          case (_, ExitCase.Canceled) =>
+            cancel
+          // there is no point to try flushing if it failed with an error
+          // the state might not be consistend and storage not accessible
+          // plus this is a concurrent operation, and we do not want anything
+          // to happen concurrently for a specific key
+          case (_, _) => ().pure[F]
+        }
       }
-    }
-
-    Resource.makeCase(TimerFlow.empty.pure) {
-      case (_, ExitCase.Succeeded) =>
-        cancel
-      case (_, ExitCase.Canceled) =>
-        cancel
-      // there is no point to try flushing if it failed with an error
-      // the state might not be consistend and storage not accessible
-      // plus this is a concurrent operation, and we do not want anything
-      // to happen concurrently for a specific key
-      case (_, _) => ().pure[F]
-    }
-  }
 
   private implicit class AttemptToPersist[F[_]: MonadThrow](persistence: FlushBuffers[F]) {
-    def attemptToPersist(ignorePersistErrors: Boolean, context: KeyContext[F], currentOffset: Offset): F[Unit] =
+    def attemptToPersist(ignorePersistErrors: Boolean, context: KeyContext[F], currentOffset: Offset)(
+      log: Log[F]
+    ): F[Unit] =
       persistence.flush.attempt.flatMap {
         case Left(err) if ignorePersistErrors =>
           // 'context' will continue holding the previous offset from the last time the state was persisted
           // and offsets committed (or just the last committed offset if no state has ever been persisted before).
           // Thus, when calculating the next offset to commit in `PartitionFlow#offsetToCommit` it will take
           // the minimal one (previous) and won't commit any offsets
-          context
-            .log
-            .info(s"Failed to persist state, the error is ignored and offsets won't be committed, error: $err")
+          log.info(s"Failed to persist state, the error is ignored and offsets won't be committed, error: $err")
         case Left(err) => err.raiseError[F, Unit]
         case Right(_)  => context.hold(currentOffset)
       }