Emit all messages across sequence number gaps in messagesWithBatch

core/src/main/scala/org/apache/pekko/persistence/jdbc/journal/dao/BaseJournalDaoWithReadMessages.scala

@@ -19,7 +19,6 @@ import pekko.NotUsed
 import pekko.actor.Scheduler
 import pekko.annotation.InternalApi
 import pekko.persistence.PersistentRepr
-import pekko.persistence.jdbc.journal.dao.FlowControl.{ Continue, ContinueDelayed, Stop }
 import pekko.stream.Materializer
 import pekko.stream.scaladsl.{ Sink, Source }
 
@@ -28,6 +27,7 @@ import scala.concurrent.{ ExecutionContext, Future }
 import scala.util.{ Failure, Success, Try }
 
 trait BaseJournalDaoWithReadMessages extends JournalDaoWithReadMessages {
+  import BaseJournalDaoWithReadMessages._
 
   implicit val ec: ExecutionContext
   implicit val mat: Materializer
@@ -42,62 +42,102 @@ trait BaseJournalDaoWithReadMessages extends JournalDaoWithReadMessages {
   }
 
   /**
-   * separate this method for unit tests.
+   * Separate this method for unit tests.
    */
   @InternalApi
   private[dao] def internalBatchStream(
       persistenceId: String,
       fromSequenceNr: Long,
       toSequenceNr: Long,
       batchSize: Int,
-      refreshInterval: Option[(FiniteDuration, Scheduler)]) = {
-    val firstSequenceNr: Long = Math.max(1, fromSequenceNr)
+      refreshInterval: Option[(FiniteDuration, Scheduler)]
+  ): Source[Seq[Try[(PersistentRepr, Long)]], NotUsed] = {
+
+    val normalizedBatchSize = Math.max(1, batchSize)
+    val normalizedFromSequenceNr = Math.max(1, fromSequenceNr)
+
+    def pollOrComplete(from: Long): QueryPlan = refreshInterval match {
+      case Some((delay, scheduler)) => PollRemaining(from, delay, scheduler)
+      case None                     => Complete
+    }
+
+    // A windowed query spans [from, min(from + batchSize, toSequenceNr)]; at full width that is
+    // batchSize + 1 sequence numbers, which lets a window tolerate one missing sequence number
+    // and still return a full batch.
+    def windowedQuery(from: Long): QueryWindow = {
+      val endInclusive =
+        if (from <= (Long.MaxValue - normalizedBatchSize)) Math.min(from + normalizedBatchSize, toSequenceNr)
+        else toSequenceNr // from + batchSize overflows; fall back to the full remaining range
+      QueryWindow(from, endInclusive)
+    }
+
+    def retrieveBatch(from: Long, endInclusive: Long): Future[Option[(QueryPlan, Seq[Try[(PersistentRepr, Long)]])]] =
+      messages(persistenceId, from, endInclusive, normalizedBatchSize).runWith(Sink.seq).map { batch =>
+        // Messages are ordered by sequence number, therefore the last one is the largest
+        val lastSeqNrInBatch: Option[Long] = batch.lastOption match {
+          case Some(Success((repr, _))) => Some(repr.sequenceNr)
+          case Some(Failure(cause))     => throw cause // fail the returned Future
+          case None                     => None
+        }
+        val hasReachedToSequenceNr = lastSeqNrInBatch.exists(_ >= toSequenceNr)
+        val isFullBatch = batch.size == normalizedBatchSize
+        val wasWindowedQuery = endInclusive < toSequenceNr
+        // fromSequenceNr beyond toSequenceNr requests an empty range, which must complete
+        // rather than poll a range that can never produce messages
+        val isEmptyRange = from > toSequenceNr
+        val nextFrom: Long = lastSeqNrInBatch.map(_ + 1).getOrElse(from)
+
+        // Deleted messages leave gaps, so a window may hold fewer live messages than batchSize
+        // even when more messages exist beyond it. A short batch is conclusive only when the
+        // queried range reached toSequenceNr, otherwise requery the full remaining range.
+        val nextQueryPlan: QueryPlan =
+          if (hasReachedToSequenceNr || isEmptyRange) Complete
+          else if (isFullBatch) windowedQuery(nextFrom)
+          else if (wasWindowedQuery) QueryRemaining(nextFrom)
+          else pollOrComplete(nextFrom)
+
+        Some((nextQueryPlan, batch))
+      }
+
+    // A full-range first query crosses a leading gap (e.g. journal purged up to a snapshot)
+    // in a single round trip; the LIMITed query costs the same as a windowed one when dense.
+    // Once a full batch shows the journal to be dense, reads switch to bounded windows, the
+    // perf safeguard introduced by PR #180, and fall back to a full-range query on a gap.
     Source
-      .unfoldAsync[(Long, FlowControl), Seq[Try[(PersistentRepr, Long)]]]((firstSequenceNr, Continue)) {
-        case (from, control) =>
-          def limitWindow(from: Long): Long = {
-            if (from == firstSequenceNr || batchSize <= 0 || (Long.MaxValue - batchSize) < from) {
-              toSequenceNr
-            } else {
-              Math.min(from + batchSize, toSequenceNr)
-            }
-          }
-
-          def retrieveNextBatch(): Future[Option[((Long, FlowControl), Seq[Try[(PersistentRepr, Long)]])]] = {
-            for {
-              xs <- messages(persistenceId, from, limitWindow(from), batchSize).runWith(Sink.seq)
-            } yield {
-              val hasMoreEvents = xs.size == batchSize
-              // Events are ordered by sequence number, therefore the last one is the largest)
-              val lastSeqNrInBatch: Option[Long] = xs.lastOption match {
-                case Some(Success((repr, _))) => Some(repr.sequenceNr)
-                case Some(Failure(e))         => throw e // fail the returned Future
-                case None                     => None
-              }
-              val hasLastEvent = lastSeqNrInBatch.exists(_ >= toSequenceNr)
-              val nextControl: FlowControl =
-                if (hasLastEvent || from > toSequenceNr) Stop
-                else if (hasMoreEvents) Continue
-                else if (refreshInterval.isEmpty) Stop
-                else ContinueDelayed
-
-              val nextFrom: Long = lastSeqNrInBatch match {
-                // Continue querying from the last sequence number (the events are ordered)
-                case Some(lastSeqNr) => lastSeqNr + 1
-                case None            => from
-              }
-              Some(((nextFrom, nextControl), xs))
-            }
-          }
-
-          control match {
-            case Stop            => Future.successful(None)
-            case Continue        => retrieveNextBatch()
-            case ContinueDelayed =>
-              val (delay, scheduler) = refreshInterval.get
-              pekko.pattern.after(delay, scheduler)(retrieveNextBatch())
-          }
+      .unfoldAsync[QueryPlan, Seq[Try[(PersistentRepr, Long)]]](QueryRemaining(normalizedFromSequenceNr)) {
+        case Complete =>
+          Future.successful(None)
+        case QueryWindow(from, endInclusive) =>
+          retrieveBatch(from, endInclusive)
+        case QueryRemaining(from) =>
+          retrieveBatch(from, toSequenceNr)
+        case PollRemaining(from, delay, scheduler) =>
+          pekko.pattern.after(delay, scheduler)(retrieveBatch(from, toSequenceNr))
       }
   }
+}
+
+private[dao] object BaseJournalDaoWithReadMessages {
+
+  /** The query the batch stream will run. */
+  private sealed trait QueryPlan
+
+  /**
+   * Query the window [from, endInclusive]: the dense fast path, planned only after a full batch
+   * showed the messages to be dense. Windowing bounds the range a query scans; it is
+   * not load-bearing for correctness, since a short window falls back to [[QueryRemaining]].
+   */
+  private final case class QueryWindow(from: Long, endInclusive: Long) extends QueryPlan
+
+  /**
+   * Query the whole remaining range [from, toSequenceNr]: planned first, and whenever a short
+   * windowed batch leaves the remainder undetermined, since gaps cannot hide messages from an
+   * unwindowed query.
+   */
+  private final case class QueryRemaining(from: Long) extends QueryPlan
+
+  /** Poll the remaining range after `delay`: the live tail, when the range is exhausted but may grow. */
+  private final case class PollRemaining(from: Long, delay: FiniteDuration, scheduler: Scheduler) extends QueryPlan
 
+  private case object Complete extends QueryPlan
 }

core/src/test/scala/org/apache/pekko/persistence/jdbc/journal/dao/LimitWindowingStreamTest.scala

@@ -73,6 +73,7 @@ abstract class LimitWindowingStreamTest(configFile: String)
 
       lastInsert.futureValue(Timeout(totalMessages.seconds))
       val readMessagesDao = dao.asInstanceOf[BaseJournalDaoWithReadMessages]
+      // Asserts batch/message counts only. Windows are intentionally batchSize+1 wide (see BaseJournalDaoWithReadMessages.windowedQuery).
       val messagesSrc =
         readMessagesDao.internalBatchStream(persistenceId, 0, totalMessages, batchSize = fetchSize, None)
 

core/src/test/scala/org/apache/pekko/persistence/jdbc/journal/dao/MessagesWithBatchDatabaseContractTest.scala

@@ -0,0 +1,195 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.pekko.persistence.jdbc.journal.dao
+
+import org.apache.pekko
+import pekko.NotUsed
+import pekko.actor.ActorSystem
+import pekko.persistence.jdbc.query.{ H2Cleaner, QueryTestSpec }
+import pekko.persistence.{ AtomicWrite, PersistentRepr }
+import pekko.stream.scaladsl.{ Sink, Source }
+import pekko.stream.{ Materializer, SystemMaterializer }
+
+import java.io.NotSerializableException
+import java.util.UUID
+import scala.collection.immutable
+import scala.concurrent.duration._
+import scala.concurrent.{ ExecutionContext, Future }
+
+/**
+ * Covers the database-coupled side of the batched message stream: hard deletes as a
+ * representative contract check of the `messages` query (ascending order, limit, inclusive
+ * range), soft deletes (the `deleted` flag filter), the journal's own prefix purge
+ * (`delete`), and serialization failures. The full gap state-machine specification lives
+ * in MessagesWithBatchTest.
+ */
+abstract class MessagesWithBatchDatabaseContractTest(configFile: String) extends QueryTestSpec(configFile) {
+
+  private lazy val journalQueries =
+    new JournalQueries(profile, journalConfig.eventJournalTableConfiguration, journalConfig.eventTagTableConfiguration)
+
+  it should "emit all messages when a hard-deleted gap is wider than the batch size" in
+  withActorSystem { implicit system =>
+    withSetup { (dao, persistenceId) =>
+      persistMessages(dao, persistenceId, count = 4)
+      hardDelete(persistenceId, 2, 3)
+
+      val emitted = emittedSequenceNrs(dao, persistenceId, toSequenceNr = 4, batchSize = 1)
+
+      emitted shouldBe Seq(1L, 4L)
+    }
+  }
+
+  it should "emit messages beyond a hard-deleted gap wider than the batch size when polling with a refresh interval" in
+  withActorSystem { implicit system =>
+    withSetup { (dao, persistenceId) =>
+      persistMessages(dao, persistenceId, count = 4)
+      hardDelete(persistenceId, 2, 3)
+
+      val emitted =
+        emittedSequenceNrs(dao, persistenceId, toSequenceNr = 4, batchSize = 1, count = Some(2),
+          refreshInterval = Some(50.millis))
+
+      emitted shouldBe Seq(1L, 4L)
+    }
+  }
+
+  it should "emit all messages when a gap mixes hard-deleted and soft-deleted messages" in
+  withActorSystem { implicit system =>
+    withSetup { (dao, persistenceId) =>
+      persistMessages(dao, persistenceId, count = 5)
+      hardDelete(persistenceId, 2)
+      softDelete(persistenceId, 3)
+      hardDelete(persistenceId, 4)
+
+      val emitted = emittedSequenceNrs(dao, persistenceId, toSequenceNr = 5, batchSize = 1)
+
+      emitted shouldBe Seq(1L, 5L)
+    }
+  }
+
+  it should "complete with no messages when every message is soft-deleted" in
+  withActorSystem { implicit system =>
+    withSetup { (dao, persistenceId) =>
+      persistMessages(dao, persistenceId, count = 3)
+      softDelete(persistenceId, 1, 2, 3)
+
+      val emitted = emittedSequenceNrs(dao, persistenceId, toSequenceNr = 3, batchSize = 1)
+
+      emitted shouldBe empty
+    }
+  }
+
+  // The journal's prefix purge hard-deletes the messages below the highest affected sequence
+  // number and soft-deletes that one, leaving a leading gap: the scenario behind issue #516
+  it should "emit the remaining messages after a prefix purge through the journal's delete" in
+  withActorSystem { implicit system =>
+    withSetup { (dao, persistenceId) =>
+      persistMessages(dao, persistenceId, count = 4)
+      dao.delete(persistenceId, 3).futureValue
+
+      val emitted = emittedSequenceNrs(dao, persistenceId, toSequenceNr = 4, batchSize = 1)
+
+      emitted shouldBe Seq(4L)
+    }
+  }
+
+  it should "fail the stream when the last message in a batch cannot be deserialized" in
+  withActorSystem { implicit system =>
+    withSetup { (dao, persistenceId) =>
+      persistMessages(dao, persistenceId, count = 1)
+      persistCorruptMessage(persistenceId, sequenceNr = 2)
+
+      val result = dao
+        .messagesWithBatch(persistenceId, fromSequenceNr = 1, toSequenceNr = 2, batchSize = 2,
+          refreshInterval = None)
+        .runWith(Sink.seq)
+
+      result.failed.futureValue shouldBe a[NotSerializableException]
+    }
+  }
+
+  private implicit def materializer(implicit system: ActorSystem): Materializer =
+    SystemMaterializer(system).materializer
+
+  private def withSetup(f: (JournalDao, String) => Unit)(implicit system: ActorSystem): Unit = {
+    implicit val ec: ExecutionContext = system.dispatcher
+    withDao(dao => f(dao, UUID.randomUUID().toString))
+  }
+
+  private def persistMessages(dao: JournalDao, persistenceId: String, count: Int): Unit = {
+    val writerUuid = UUID.randomUUID().toString
+    val payload = Array.fill(8)('a'.toByte)
+    val writes = (1 to count).map { sequenceNr =>
+      AtomicWrite(immutable.Seq(PersistentRepr(payload, sequenceNr, persistenceId, writerUuid = writerUuid)))
+    }
+    dao.asyncWriteMessages(writes).futureValue
+  }
+
+  private def hardDelete(persistenceId: String, sequenceNrs: Long*): Unit = {
+    import profile.api._
+    val deleteRows = journalQueries.JournalTable
+      .filter(row => row.persistenceId === persistenceId && row.sequenceNumber.inSet(sequenceNrs))
+      .delete
+    db.run(deleteRows).futureValue
+  }
+
+  private def softDelete(persistenceId: String, sequenceNrs: Long*): Unit =
+    sequenceNrs.foreach { sequenceNr =>
+      db.run(journalQueries.markSeqNrJournalMessagesAsDeleted(persistenceId, sequenceNr)).futureValue
+    }
+
+  /** Inserts a journal row whose serializer id is unknown, so reading it yields a Failure. */
+  private def persistCorruptMessage(persistenceId: String, sequenceNr: Long): Unit = {
+    import profile.api._
+    val unknownSerializerId = 999999
+    val corruptRow = JournalTables.JournalPekkoSerializationRow(Long.MinValue, deleted = false, persistenceId,
+      sequenceNr, writer = UUID.randomUUID().toString, writeTimestamp = 0L, adapterManifest = "",
+      eventPayload = Array.fill(8)('x'.toByte),
+      eventSerId = unknownSerializerId, eventSerManifest = "", metaPayload = None, metaSerId = None,
+      metaSerManifest = None)
+    db.run(journalQueries.JournalTable += corruptRow).futureValue
+  }
+
+  /**
+   * Collects the sequence numbers the batched message stream emits. A polling stream never
+   * completes on its own, so pass `count` to bound it.
+   */
+  private def emittedSequenceNrs(dao: JournalDao, persistenceId: String, toSequenceNr: Long, batchSize: Int,
+      refreshInterval: Option[FiniteDuration] = None, count: Option[Int] = None)(
+      implicit system: ActorSystem): Seq[Long] = {
+    val sequenceNrs = sequenceNrSource(dao, persistenceId, toSequenceNr, batchSize, refreshInterval)
+    count
+      .fold(sequenceNrs)(sequenceNrs.take(_))
+      .completionTimeout(10.seconds)
+      .runWith(Sink.seq)
+      .futureValue
+  }
+
+  private def sequenceNrSource(dao: JournalDao, persistenceId: String, toSequenceNr: Long,
+      batchSize: Int, refreshInterval: Option[FiniteDuration])(implicit system: ActorSystem): Source[Long, NotUsed] =
+    dao
+      .messagesWithBatch(persistenceId, fromSequenceNr = 1, toSequenceNr, batchSize,
+        refreshInterval.map(_ -> system.scheduler))
+      .mapAsync(1)(Future.fromTry)
+      .map { case (repr, _) => repr.sequenceNr }
+}
+
+final class H2MessagesWithBatchDatabaseContractTest
+    extends MessagesWithBatchDatabaseContractTest("h2-application.conf")
+    with H2Cleaner