Don't drop events from pubsub; queries no longer heartbeat

leviramsey · leviramsey · commit 47e9f98c98d5 · 2026-03-10T20:34:37.000-04:00
diff --git a/core/src/main/scala/akka/persistence/dynamodb/internal/BySliceQuery.scala b/core/src/main/scala/akka/persistence/dynamodb/internal/BySliceQuery.scala
@@ -109,7 +109,6 @@ import org.slf4j.Logger
     dao: BySliceQuery.Dao[Item],
     createEnvelope: (TimestampOffset, Item) => Envelope,
     extractOffset: Envelope => TimestampOffset,
-    createHeartbeat: Instant => Option[Envelope],
     clock: Clock,
     settings: DynamoDBSettings,
     log: Logger)(implicit val ec: ExecutionContext) {
@@ -401,22 +400,6 @@ import org.slf4j.Logger
           .via(deserializeAndAddOffset(newState.currentOffset)))
     }
 
-    def heartbeat(state: QueryState): Option[Envelope] = {
-      if (state.idleCountBeforeHeartbeat >= 2 && state.previousQueryWallClock != Instant.EPOCH) {
-        // use wall clock to measure duration since start, up to idle backtracking limit
-        val timestamp = state.startTimestamp.plus(
-          JDuration.between(state.startWallClock, state.previousQueryWallClock.minus(backtrackingBehindCurrentTime)))
-
-        val h = createHeartbeat(timestamp)
-        if (h.isDefined)
-          log.debug("{} heartbeat timestamp [{}]", logPrefix, timestamp)
-        h
-      } else None
-    }
-
-    val nextHeartbeat: QueryState => Option[Envelope] =
-      if (settings.journalPublishEvents) heartbeat else _ => None
-
     val currentTimestamp = InstantFactory.now() // Can we use DDB as a timestamp source?
     val currentWallClock = clock.instant()
 
@@ -426,8 +409,7 @@ import org.slf4j.Logger
       updateState = nextOffset,
       delayNextQuery = delayNextQuery,
       nextQuery = nextQuery,
-      beforeQuery = _ => None,
-      heartbeat = nextHeartbeat)
+      beforeQuery = _ => None)
   }
 
   private def deserializeAndAddOffset(timestampOffset: TimestampOffset): Flow[Item, Envelope, NotUsed] = {
diff --git a/core/src/main/scala/akka/persistence/dynamodb/internal/ContinuousQuery.scala b/core/src/main/scala/akka/persistence/dynamodb/internal/ContinuousQuery.scala
@@ -30,10 +30,8 @@ private[dynamodb] object ContinuousQuery {
       updateState: (S, T) => S,
       delayNextQuery: S => Option[FiniteDuration],
       nextQuery: S => (S, Option[Source[T, NotUsed]]),
-      beforeQuery: S => Option[Future[S]] = (_: S) => None,
-      heartbeat: S => Option[T] = (_: S) => None): Source[T, NotUsed] =
-    Source.fromGraph(
-      new ContinuousQuery[S, T](initialState, updateState, delayNextQuery, nextQuery, beforeQuery, heartbeat))
+      beforeQuery: S => Option[Future[S]] = (_: S) => None): Source[T, NotUsed] =
+    Source.fromGraph(new ContinuousQuery[S, T](initialState, updateState, delayNextQuery, nextQuery, beforeQuery))
 
   private case object NextQuery
 
@@ -71,8 +69,7 @@ final private[dynamodb] class ContinuousQuery[S, T](
     updateState: (S, T) => S,
     delayNextQuery: S => Option[FiniteDuration],
     nextQuery: S => (S, Option[Source[T, NotUsed]]),
-    beforeQuery: S => Option[Future[S]],
-    heartbeat: S => Option[T])
+    beforeQuery: S => Option[Future[S]])
     extends GraphStage[SourceShape[T]] {
   import ContinuousQuery._
 
@@ -155,13 +152,8 @@ final private[dynamodb] class ContinuousQuery[S, T](
                 }
             })
 
-            val sourceWithHeartbeat = heartbeat(newState) match {
-              case None    => source
-              case Some(h) => Source.single(h).concat(source)
-            }
-
             val graph = Source
-              .fromGraph(sourceWithHeartbeat)
+              .fromGraph(source)
               .to(sinkIn.sink)
             interpreter.subFusingMaterializer.materialize(graph)
             sinkIn.pull()
diff --git a/core/src/main/scala/akka/persistence/dynamodb/query/scaladsl/DynamoDBReadJournal.scala b/core/src/main/scala/akka/persistence/dynamodb/query/scaladsl/DynamoDBReadJournal.scala
@@ -130,10 +130,8 @@ final class DynamoDBReadJournal(system: ExtendedActorSystem, config: Config, cfg
       correlationId: Option[String]): BySliceQuery[SerializedJournalItem, EventEnvelope[Event]] = {
     val createEnvelope: (TimestampOffset, SerializedJournalItem) => EventEnvelope[Event] = createEventEnvelope
     val extractOffset = (env: EventEnvelope[Event]) => env.offset.asInstanceOf[TimestampOffset]
-    val createHeartbeat = (timestamp: Instant) =>
-      Some(createEventEnvelopeHeartbeat[Event](entityType, slice, timestamp, correlationId))
 
-    new BySliceQuery(queryDao, createEnvelope, extractOffset, createHeartbeat, clock, settings, log)
+    new BySliceQuery(queryDao, createEnvelope, extractOffset, clock, settings, log)
   }
 
   private def snapshotsBySlice[Snapshot, Event](
@@ -145,10 +143,8 @@ final class DynamoDBReadJournal(system: ExtendedActorSystem, config: Config, cfg
       (offset, row) => createEnvelopeFromSnapshot(row, offset, transformSnapshot)
 
     val extractOffset: EventEnvelope[Event] => TimestampOffset = env => env.offset.asInstanceOf[TimestampOffset]
-    val createHeartbeat = (timestamp: Instant) =>
-      Some(createEventEnvelopeHeartbeat[Event](entityType, slice, timestamp, correlationId))
 
-    new BySliceQuery(snapshotDao, createEnvelope, extractOffset, createHeartbeat, clock, settings, log)
+    new BySliceQuery(snapshotDao, createEnvelope, extractOffset, clock, settings, log)
   }
 
   private def createEnvelopeFromSnapshot[Snapshot, Event](
@@ -324,8 +320,20 @@ final class DynamoDBReadJournal(system: ExtendedActorSystem, config: Config, cfg
 
     val dbSource = bySliceQueries.head.mergeAll(bySliceQueries.tail, eagerComplete = false)
     if (settings.journalPublishEvents) {
+      val initialHeartbeats = (minSlice to maxSlice).flatMap { slice =>
+        sliceStartOffset(slice, offset) match {
+          case t: TimestampOffset =>
+            Some(
+              createEventEnvelopeHeartbeat[Event](
+                entityType,
+                slice,
+                t.timestamp.minus(JDuration.ofMillis(settings.querySettings.backtrackingWindow.toMillis)),
+                correlationId))
+          case _ => None
+        }
+      }
       val pubSubSource = eventsBySlicesPubSubSource[Event](entityType, minSlice, maxSlice)
-      mergeDbAndPubSubSources(dbSource, pubSubSource, correlationId)
+      mergeDbAndPubSubSources(dbSource, pubSubSource, initialHeartbeats, correlationId)
     } else
       dbSource
   }
@@ -480,8 +488,20 @@ final class DynamoDBReadJournal(system: ExtendedActorSystem, config: Config, cfg
 
     val dbSource = bySliceQueries.head.mergeAll(bySliceQueries.tail, eagerComplete = false)
     if (settings.journalPublishEvents) {
+      val initialHeartbeats = (minSlice to maxSlice).flatMap { slice =>
+        sliceStartOffset(slice, offset) match {
+          case t: TimestampOffset =>
+            Some(
+              createEventEnvelopeHeartbeat[Event](
+                entityType,
+                slice,
+                t.timestamp.minus(JDuration.ofMillis(settings.querySettings.backtrackingWindow.toMillis)),
+                correlationId))
+          case _ => None
+        }
+      }
       val pubSubSource = eventsBySlicesPubSubSource[Event](entityType, minSlice, maxSlice)
-      mergeDbAndPubSubSources(dbSource, pubSubSource, correlationId)
+      mergeDbAndPubSubSources(dbSource, pubSubSource, initialHeartbeats, correlationId)
     } else
       dbSource
   }
@@ -554,15 +574,24 @@ final class DynamoDBReadJournal(system: ExtendedActorSystem, config: Config, cfg
   private def mergeDbAndPubSubSources[Event, Snapshot](
       dbSource: Source[EventEnvelope[Event], NotUsed],
       pubSubSource: Source[EventEnvelope[Event], NotUsed],
+      initialHeartbeats: Iterable[EventEnvelope[Event]],
       correlationId: Option[String]) = {
-    dbSource
-      .mergePrioritized(pubSubSource, leftPriority = 1, rightPriority = 10)
-      .via(
-        skipPubSubTooFarAhead(
-          settings.querySettings.backtrackingEnabled,
-          JDuration.ofMillis(settings.querySettings.backtrackingWindow.toMillis),
-          correlationId))
-      .via(deduplicate(settings.querySettings.deduplicateCapacity))
+    val dbMergedWithPubsub =
+      dbSource
+        .mergePrioritized(pubSubSource, leftPriority = 1, rightPriority = 10)
+
+    val base =
+      if (settings.querySettings.backtrackingEnabled) {
+        Source
+          .fromIterator(() => initialHeartbeats.iterator)
+          .concat(dbMergedWithPubsub)
+          .via(
+            handlePubSubTooFarAhead(
+              JDuration.ofMillis(settings.querySettings.backtrackingWindow.toMillis),
+              correlationId))
+      } else dbMergedWithPubsub
+
+    base.via(deduplicate(settings.querySettings.deduplicateCapacity))
   }
 
   /**
@@ -605,77 +634,85 @@ final class DynamoDBReadJournal(system: ExtendedActorSystem, config: Config, cfg
     }
   }
 
-  /**
-   * INTERNAL API
-   */
-  @InternalApi private[akka] def skipPubSubTooFarAhead[Event](
-      enabled: Boolean,
-      maxAheadOfBacktracking: JDuration,
+  /** INTERNAL API */
+  @nowarn("msg=eventMetadata in class EventEnvelope is deprecated")
+  @InternalApi private[akka] def handlePubSubTooFarAhead[Event](
+      backtrackingWindow: JDuration,
       correlationId: Option[String]): Flow[EventEnvelope[Event], EventEnvelope[Event], NotUsed] = {
     def correlationIdLogText = CorrelationId.toLogText(correlationId)
-    if (!enabled)
-      Flow[EventEnvelope[Event]]
-    else
-      Flow[EventEnvelope[Event]]
-        .statefulMapConcat(() => {
-          // track backtracking offset per slice
-          var latestBacktrackingPerSlice = Map.empty[Int, Instant]
-          def latestBacktracking(slice: Int): Instant = latestBacktrackingPerSlice.get(slice) match {
-            case Some(instant) => instant
-            case None          => Instant.EPOCH
-          }
-          env => {
-            val slice = persistenceExt.sliceForPersistenceId(env.persistenceId)
-            env.offset match {
-              case t: TimestampOffset =>
-                if (EnvelopeOrigin.fromQuery(env)) {
-                  if (log.isDebugEnabled()) {
-                    val l = latestBacktracking(slice)
-                    if (l.isAfter(t.timestamp))
-                      log.debug(
-                        "event from query for persistenceId [{}] seqNr [{}] " +
-                        s"timestamp [{}]{} was before last event from backtracking or heartbeat [{}].",
-                        env.persistenceId,
-                        env.sequenceNr,
-                        t.timestamp,
-                        correlationIdLogText,
-                        l)
-                  }
 
-                  env :: Nil
-                } else {
-                  if (EnvelopeOrigin.fromBacktracking(env)) {
+    Flow[EventEnvelope[Event]]
+      .statefulMapConcat(() => {
+        // track backtracking offsets per slice
+        var latestBacktrackingPerSlice = Map.empty[Int, Instant]
+        def latestBacktracking(slice: Int): Instant = latestBacktrackingPerSlice.get(slice) match {
+          case Some(instant) => instant
+          case None          => Instant.EPOCH
+        }
+
+        env => {
+          val slice = persistenceExt.sliceForPersistenceId(env.persistenceId)
+          val l = latestBacktracking(slice)
+          env.offset match {
+            case t: TimestampOffset =>
+              if (EnvelopeOrigin.fromQuery(env)) {
+                if (log.isDebugEnabled() && l.isAfter(t.timestamp)) {
+                  log.debug(
+                    "event from query for persistenceId [{}] seqNr [{}] " +
+                    s"timestamp [{}] was before last event from backtracking or heartbeat [{}].",
+                    env.persistenceId,
+                    env.sequenceNr,
+                    t.timestamp,
+                    l)
+                }
+                env :: Nil
+              } else {
+                if (EnvelopeOrigin.fromBacktracking(env)) {
+                  if (l.isBefore(t.timestamp)) {
                     latestBacktrackingPerSlice = latestBacktrackingPerSlice.updated(slice, t.timestamp)
-                    env :: Nil
-                  } else if (EnvelopeOrigin.fromHeartbeat(env)) {
+                  }
+                  env :: Nil
+                } else if (EnvelopeOrigin.fromHeartbeat(env)) {
+                  if (l.isBefore(t.timestamp)) {
                     latestBacktrackingPerSlice = latestBacktrackingPerSlice.updated(slice, t.timestamp)
-                    Nil // always drop heartbeats
-                  } else if (EnvelopeOrigin.fromPubSub(env) && latestBacktracking(slice) == Instant.EPOCH) {
+                  }
+                  Nil // heartbeats are an internal implementation and must never leak out
+                } else if (EnvelopeOrigin.fromPubSub(env)) {
+                  if (l == Instant.EPOCH) {
                     log.trace(
                       "Dropping pubsub event for persistenceId [{}] seqNr [{}]{} because no event from backtracking yet.",
                       env.persistenceId,
                       env.sequenceNr,
                       correlationIdLogText)
                     Nil
-                  } else if (EnvelopeOrigin.fromPubSub(env) && JDuration
-                      .between(latestBacktracking(slice), t.timestamp)
-                      .compareTo(maxAheadOfBacktracking) > 0) {
-                    // drop from pubsub when too far ahead from backtracking
-                    log.debug(
-                      "Dropping pubsub event for persistenceId [{}] seqNr [{}]{} because too far ahead of backtracking.",
+                  } else if (JDuration.between(l, t.timestamp).compareTo(backtrackingWindow) > 0) {
+                    // far ahead of backtracking, so adjust the offset (the timestamp in the envelope is unaffected)
+                    // so as to prevent the offset from being moved too far ahead
+                    val offset = TimestampOffset(l.plus(backtrackingWindow), t.readTimestamp, t.seen)
+                    val emittedEnv = new EventEnvelope(
+                      offset,
                       env.persistenceId,
                       env.sequenceNr,
-                      correlationIdLogText)
-                    Nil
+                      env.eventOption,
+                      env.timestamp,
+                      env.eventMetadata,
+                      env.entityType,
+                      env.slice,
+                      env.filtered,
+                      env.source,
+                      env.tags)
+                    emittedEnv :: Nil
                   } else {
                     env :: Nil
                   }
+                } else {
+                  env :: Nil
                 }
-              case _ =>
-                env :: Nil
-            }
+              }
+            case _ => env :: Nil
           }
-        })
+        }
+      })
   }
 
   // EventTimestampQuery
diff --git a/core/src/test/scala/akka/persistence/dynamodb/query/EventsBySlicePubSubSpec.scala b/core/src/test/scala/akka/persistence/dynamodb/query/EventsBySlicePubSubSpec.scala
@@ -48,6 +48,7 @@ import akka.stream.typed.scaladsl.ActorFlow
 import com.typesafe.config.Config
 import com.typesafe.config.ConfigFactory
 import org.scalatest.wordspec.AnyWordSpecLike
+import akka.persistence.query.Offset
 
 object EventsBySlicePubSubSpec {
   def config: Config = ConfigFactory
@@ -102,6 +103,18 @@ class EventsBySlicePubSubSpec
       source = EnvelopeOrigin.SourcePubSub)
   }
 
+  private def changeOffset(envelope: EventEnvelope[String], newOffset: Offset): EventEnvelope[String] =
+    EventEnvelope(
+      newOffset,
+      envelope.persistenceId,
+      envelope.sequenceNr,
+      envelope.event,
+      envelope.timestamp,
+      envelope.entityType,
+      envelope.slice,
+      envelope.filtered,
+      envelope.source)
+
   def backtrackingEnvelope(env: EventEnvelope[String]): EventEnvelope[String] =
     new EventEnvelope[String](
       env.offset,
@@ -215,16 +228,17 @@ class EventsBySlicePubSubSpec
       envelope.eventOption should be(empty)
     }
 
-    "skipPubSubTooFarAhead" in {
+    "handlePubSubTooFarAhead" in {
       persistenceExt.sliceForPersistenceId(envA1.persistenceId) should not be persistenceExt.sliceForPersistenceId(
         envB1.persistenceId)
 
+      val backtrackWindowMillis = settings.querySettings.backtrackingWindow.toMillis
+
       val (in, out) =
         TestSource[EventEnvelope[String]]()
           .via(
-            query.skipPubSubTooFarAhead(
-              enabled = true,
-              maxAheadOfBacktracking = JDuration.ofMillis(settings.querySettings.backtrackingWindow.toMillis),
+            query.handlePubSubTooFarAhead(
+              backtrackingWindow = JDuration.ofMillis(backtrackWindowMillis),
               correlationId = None))
           .toMat(TestSink[EventEnvelope[String]]())(Keep.both)
           .run()
@@ -255,21 +269,21 @@ class EventsBySlicePubSubSpec
       val time1 = envA1.offset
         .asInstanceOf[TimestampOffset]
         .timestamp
-      val time2 = time1
-        .plusMillis(settings.querySettings.backtrackingWindow.toMillis)
+      val time2 = time1.plusMillis(backtrackWindowMillis)
       val envA3 = createEnvelope(pidA, 3L, "a3", time2.plusMillis(1))
       val envA4 = createEnvelope(pidA, 4L, "a4", time2.plusMillis(2))
       in.sendNext(envA3)
       in.sendNext(envA4)
-      // dropped because > backtrackingWindow
-      out.expectNoMessage()
+      // emitted with adjusted offset based on latest backtracking
+      out.expectNext(changeOffset(envA3, TimestampOffset(time2, time2.plusMillis(1), Map(pidA.id -> 3L))))
+      out.expectNext(changeOffset(envA4, TimestampOffset(time2, time2.plusMillis(2), Map(pidA.id -> 4L))))
 
       val pidCSameSlice =
         randomPersistenceIdForSlice(entityType, persistenceExt.sliceForPersistenceId(pidA.id))
       val envC1 = createEnvelope(pidCSameSlice, 1L, "c1", time2.plusMillis(1))
       in.sendNext(envC1)
-      // dropped because > backtrackingWindow
-      out.expectNoMessage()
+      // emitted with adjusted offset based on latest backtracking
+      out.expectNext(changeOffset(envC1, TimestampOffset(time2, time2.plusMillis(1), Map(pidCSameSlice.id -> 1L))))
 
       val pidDSameSlice =
         randomPersistenceIdForSlice(entityType, persistenceExt.sliceForPersistenceId(pidA.id))
