fix: stuck queries when too many events with same timestamp (#586)

* fix: stuck queries when too many events with same timestamp
* ci: add mima filters
* ci: pin sqlserver to the previous version
* test: same timestamp, overlapping seq numbers
* suggestions from code review
* track previous timestamp to only filter by seq nr when same timestamp
* ci: update docker compose command for yugabyte tests

Author: pvlugter

.github/workflows/build-test.yml

@@ -203,7 +203,7 @@ jobs:
 
       - name: Start DB
         run: |-
-          docker-compose -f docker/docker-compose-yugabyte.yml up -d
+          docker compose -f docker/docker-compose-yugabyte.yml up -d
           # TODO: could we poll the port instead of sleep?
           sleep 10
           docker exec -i yb-tserver-n1 /home/yugabyte/bin/ysqlsh -h yb-tserver-n1 -t < ddl-scripts/create_tables_yugabyte.sql

core/src/main/mima-filters/1.2.4.backwards.excludes/rows-by-slices.excludes

@@ -0,0 +1,3 @@
+# internals: fromSeqNr added to rowsBySlices
+ProblemFilters.exclude[DirectMissingMethodProblem]("akka.persistence.r2dbc.internal.BySliceQuery#Dao.rowsBySlices")
+ProblemFilters.exclude[ReversedMissingMethodProblem]("akka.persistence.r2dbc.internal.BySliceQuery#Dao.rowsBySlices")

core/src/main/scala/akka/persistence/r2dbc/internal/BySliceQuery.scala

@@ -31,7 +31,19 @@ import org.slf4j.Logger
 
   object QueryState {
     val empty: QueryState =
-      QueryState(TimestampOffset.Zero, 0, 0, 0, 0, backtrackingCount = 0, TimestampOffset.Zero, 0, 0, Buckets.empty)
+      QueryState(
+        latest = TimestampOffset.Zero,
+        rowCount = 0,
+        rowCountSinceBacktracking = 0,
+        queryCount = 0,
+        idleCount = 0,
+        backtrackingCount = 0,
+        latestBacktracking = TimestampOffset.Zero,
+        latestBacktrackingSeenCount = 0,
+        backtrackingExpectFiltered = 0,
+        buckets = Buckets.empty,
+        previous = TimestampOffset.Zero,
+        previousBacktracking = TimestampOffset.Zero)
   }
 
   final case class QueryState(
@@ -44,7 +56,9 @@ import org.slf4j.Logger
       latestBacktracking: TimestampOffset,
       latestBacktrackingSeenCount: Int,
       backtrackingExpectFiltered: Int,
-      buckets: Buckets) {
+      buckets: Buckets,
+      previous: TimestampOffset,
+      previousBacktracking: TimestampOffset) {
 
     def backtracking: Boolean = backtrackingCount > 0
 
@@ -56,6 +70,10 @@ import org.slf4j.Logger
       if (backtracking) latestBacktracking.timestamp
       else latest.timestamp
 
+    def nextQueryFromSeqNr: Option[Long] =
+      if (backtracking) highestSeenSeqNr(previousBacktracking, latestBacktracking)
+      else highestSeenSeqNr(previous, latest)
+
     def nextQueryToTimestamp(atLeastNumberOfEvents: Int): Option[Instant] = {
       buckets.findTimeForLimit(nextQueryFromTimestamp, atLeastNumberOfEvents) match {
         case Some(t) =>
@@ -70,6 +88,12 @@ import org.slf4j.Logger
     }
   }
 
+  // only filter by highest seen seq nr when the next query is the same timestamp (or when unknown for initial queries)
+  private def highestSeenSeqNr(previous: TimestampOffset, latest: TimestampOffset): Option[Long] =
+    Option.when((previous == TimestampOffset.Zero || previous.timestamp == latest.timestamp) && latest.seen.nonEmpty) {
+      latest.seen.values.max
+    }
+
   object Buckets {
     type EpochSeconds = Long
     type Count = Long
@@ -157,6 +181,7 @@ import org.slf4j.Logger
         minSlice: Int,
         maxSlice: Int,
         fromTimestamp: Instant,
+        fromSeqNr: Option[Long], // for events with same timestamp as `fromTimestamp`
         toTimestamp: Option[Instant],
         behindCurrentTime: FiniteDuration,
         backtracking: Boolean): Source[SerializedRow, NotUsed]
@@ -212,7 +237,10 @@ import org.slf4j.Logger
       // so continue until rowCount is 0. That means an extra query at the end to make sure there are no
       // more to fetch.
       if (state.queryCount == 0L || state.rowCount > 0) {
-        val newState = state.copy(rowCount = 0, queryCount = state.queryCount + 1)
+        val newState = state.copy(rowCount = 0, queryCount = state.queryCount + 1, previous = state.latest)
+
+        val fromTimestamp = state.latest.timestamp
+        val fromSeqNr = highestSeenSeqNr(state.previous, state.latest)
 
         val toTimestamp = newState.nextQueryToTimestamp(settings.querySettings.bufferSize) match {
           case Some(t) =>
@@ -228,7 +256,7 @@ import org.slf4j.Logger
             state.queryCount,
             minSlice,
             maxSlice,
-            state.latest.timestamp,
+            fromTimestamp,
             toTimestamp,
             state.rowCount)
 
@@ -238,7 +266,8 @@ import org.slf4j.Logger
               entityType,
               minSlice,
               maxSlice,
-              state.latest.timestamp,
+              fromTimestamp,
+              fromSeqNr,
               toTimestamp = Some(toTimestamp),
               behindCurrentTime = Duration.Zero,
               backtracking = false)
@@ -312,7 +341,11 @@ import org.slf4j.Logger
             s"Unexpected offset [$offset] before latestBacktracking [${state.latestBacktracking}].")
 
         val newSeenCount =
-          if (offset.timestamp == state.latestBacktracking.timestamp) state.latestBacktrackingSeenCount + 1 else 1
+          if (offset.timestamp == state.latestBacktracking.timestamp &&
+            highestSeenSeqNr(state.previousBacktracking, offset) ==
+              highestSeenSeqNr(state.previousBacktracking, state.latestBacktracking))
+            state.latestBacktrackingSeenCount + 1
+          else 1
 
         state.copy(
           latestBacktracking = offset,
@@ -420,6 +453,7 @@ import org.slf4j.Logger
         else settings.querySettings.behindCurrentTime
 
       val fromTimestamp = newState.nextQueryFromTimestamp
+      val fromSeqNr = newState.nextQueryFromSeqNr
       val toTimestamp = newState.nextQueryToTimestamp(settings.querySettings.bufferSize)
 
       if (log.isDebugEnabled()) {
@@ -446,14 +480,19 @@ import org.slf4j.Logger
           else s"Found [${state.rowCount}] rows in previous query.")
       }
 
-      newState ->
+      val newStateWithPrevious =
+        if (newState.backtracking) newState.copy(previousBacktracking = newState.latestBacktracking)
+        else newState.copy(previous = newState.latest)
+
+      newStateWithPrevious ->
       Some(
         dao
           .rowsBySlices(
             entityType,
             minSlice,
             maxSlice,
             fromTimestamp,
+            fromSeqNr,
             toTimestamp,
             behindCurrentTime,
             backtracking = newState.backtracking)

core/src/main/scala/akka/persistence/r2dbc/internal/h2/H2QueryDao.scala

@@ -24,13 +24,17 @@ private[r2dbc] class H2QueryDao(executorProvider: R2dbcExecutorProvider) extends
   override protected lazy val log: Logger = LoggerFactory.getLogger(classOf[H2QueryDao])
 
   override protected def eventsBySlicesRangeSql(
+      fromSeqNrParam: Boolean,
       toDbTimestampParam: Boolean,
       behindCurrentTime: FiniteDuration,
       backtracking: Boolean,
       minSlice: Int,
       maxSlice: Int): String = {
     // not caching, too many combinations
 
+    def fromSeqNrParamCondition =
+      if (fromSeqNrParam) "AND (db_timestamp != ? OR seq_nr >= ?)" else ""
+
     def toDbTimestampParamCondition =
       if (toDbTimestampParam) "AND db_timestamp <= ?" else ""
 
@@ -51,7 +55,7 @@ private[r2dbc] class H2QueryDao(executorProvider: R2dbcExecutorProvider) extends
       FROM ${journalTable(minSlice)}
       WHERE entity_type = ?
       AND ${sliceCondition(minSlice, maxSlice)}
-      AND db_timestamp >= ? $toDbTimestampParamCondition $behindCurrentTimeIntervalCondition
+      AND db_timestamp >= ? $fromSeqNrParamCondition $toDbTimestampParamCondition $behindCurrentTimeIntervalCondition
       AND deleted = false
       ORDER BY db_timestamp, seq_nr
       LIMIT ?"""

core/src/main/scala/akka/persistence/r2dbc/internal/postgres/PostgresDurableStateDao.scala

@@ -653,6 +653,7 @@ private[r2dbc] class PostgresDurableStateDao(executorProvider: R2dbcExecutorProv
       minSlice: Int,
       maxSlice: Int,
       fromTimestamp: Instant,
+      fromSeqNr: Option[Long],
       toTimestamp: Option[Instant],
       behindCurrentTime: FiniteDuration,
       backtracking: Boolean): Source[SerializedStateRow, NotUsed] = {

core/src/main/scala/akka/persistence/r2dbc/internal/postgres/PostgresQueryDao.scala

@@ -69,13 +69,22 @@ private[r2dbc] class PostgresQueryDao(executorProvider: R2dbcExecutorProvider) e
     "SELECT CURRENT_TIMESTAMP AS db_timestamp"
 
   protected def eventsBySlicesRangeSql(
+      fromSeqNrParam: Boolean,
       toDbTimestampParam: Boolean,
       behindCurrentTime: FiniteDuration,
       backtracking: Boolean,
       minSlice: Int,
       maxSlice: Int): String = {
     // not caching, too many combinations
 
+    // If more events than the buffer size are all on the same timestamp, then the query will get stuck on that same
+    // timestamp. Avoid this by also starting from the highest seen sequence number for that timestamp, using the fact
+    // that events are ordered by db_timestamp, seq_nr. Note that sequence numbers are per persistence id, so a later
+    // timestamp can have an earlier sequence number. Add a logical conditional only when db_timestamp = fromTimestamp
+    // to also filter for seq_nr >= fromSeqNr. Expressed in a logically equivalent form, where A -> B === ~A v B.
+    def fromSeqNrParamCondition =
+      if (fromSeqNrParam) "AND (db_timestamp != ? OR seq_nr >= ?)" else ""
+
     def toDbTimestampParamCondition =
       if (toDbTimestampParam) "AND db_timestamp <= ?" else ""
 
@@ -96,7 +105,7 @@ private[r2dbc] class PostgresQueryDao(executorProvider: R2dbcExecutorProvider) e
       FROM ${journalTable(minSlice)}
       WHERE entity_type = ?
       AND ${sliceCondition(minSlice, maxSlice)}
-      AND db_timestamp >= ? $toDbTimestampParamCondition $behindCurrentTimeIntervalCondition
+      AND db_timestamp >= ? $fromSeqNrParamCondition $toDbTimestampParamCondition $behindCurrentTimeIntervalCondition
       AND deleted = false
       ORDER BY db_timestamp, seq_nr
       LIMIT ?"""
@@ -209,16 +218,25 @@ private[r2dbc] class PostgresQueryDao(executorProvider: R2dbcExecutorProvider) e
       stmt: Statement,
       entityType: String,
       fromTimestamp: Instant,
+      fromSeqNr: Option[Long],
       toTimestamp: Option[Instant]): Statement = {
     stmt
       .bind(0, entityType)
       .bindTimestamp(1, fromTimestamp)
+    val index1 = 2
+    val index2 = fromSeqNr match {
+      case Some(seqNr) =>
+        stmt.bindTimestamp(index1, fromTimestamp)
+        stmt.bind(index1 + 1, seqNr)
+        index1 + 2
+      case None => index1
+    }
     toTimestamp match {
       case Some(until) =>
-        stmt.bindTimestamp(2, until)
-        stmt.bind(3, settings.querySettings.bufferSize)
+        stmt.bindTimestamp(index2, until)
+        stmt.bind(index2 + 1, settings.querySettings.bufferSize)
       case None =>
-        stmt.bind(2, settings.querySettings.bufferSize)
+        stmt.bind(index2, settings.querySettings.bufferSize)
     }
   }
 
@@ -227,6 +245,7 @@ private[r2dbc] class PostgresQueryDao(executorProvider: R2dbcExecutorProvider) e
       minSlice: Int,
       maxSlice: Int,
       fromTimestamp: Instant,
+      fromSeqNr: Option[Long],
       toTimestamp: Option[Instant],
       behindCurrentTime: FiniteDuration,
       backtracking: Boolean): Source[SerializedJournalRow, NotUsed] = {
@@ -241,12 +260,13 @@ private[r2dbc] class PostgresQueryDao(executorProvider: R2dbcExecutorProvider) e
         val stmt = connection
           .createStatement(
             eventsBySlicesRangeSql(
+              fromSeqNrParam = fromSeqNr.isDefined,
               toDbTimestampParam = toTimestamp.isDefined,
               behindCurrentTime,
               backtracking,
               minSlice,
               maxSlice))
-        bindEventsBySlicesRangeSql(stmt, entityType, fromTimestamp, toTimestamp)
+        bindEventsBySlicesRangeSql(stmt, entityType, fromTimestamp, fromSeqNr, toTimestamp)
       },
       row =>
         if (backtracking)

core/src/main/scala/akka/persistence/r2dbc/internal/postgres/PostgresSnapshotDao.scala

@@ -388,6 +388,7 @@ private[r2dbc] class PostgresSnapshotDao(executorProvider: R2dbcExecutorProvider
       minSlice: Int,
       maxSlice: Int,
       fromTimestamp: Instant,
+      fromSeqNr: Option[Long],
       toTimestamp: Option[Instant],
       behindCurrentTime: FiniteDuration,
       backtracking: Boolean): Source[SerializedSnapshotRow, NotUsed] = {

core/src/main/scala/akka/persistence/r2dbc/internal/sqlserver/SqlServerQueryDao.scala

@@ -112,13 +112,17 @@ private[r2dbc] class SqlServerQueryDao(executorProvider: R2dbcExecutorProvider)
   }
 
   override protected def eventsBySlicesRangeSql(
+      fromSeqNrParam: Boolean,
       toDbTimestampParam: Boolean,
       behindCurrentTime: FiniteDuration,
       backtracking: Boolean,
       minSlice: Int,
       maxSlice: Int): String = {
     // not caching, too many combinations
 
+    def fromSeqNrParamCondition =
+      if (fromSeqNrParam) "AND (db_timestamp != @from OR seq_nr >= @fromSeqNr)" else ""
+
     def toDbTimestampParamCondition =
       if (toDbTimestampParam) "AND db_timestamp <= @until" else ""
 
@@ -139,7 +143,7 @@ private[r2dbc] class SqlServerQueryDao(executorProvider: R2dbcExecutorProvider)
         FROM ${journalTable(minSlice)}
         WHERE entity_type = @entityType
         AND ${sliceCondition(minSlice, maxSlice)}
-        AND db_timestamp >= @from $toDbTimestampParamCondition $behindCurrentTimeIntervalCondition
+        AND db_timestamp >= @from $fromSeqNrParamCondition $toDbTimestampParamCondition $behindCurrentTimeIntervalCondition
         AND deleted = $sqlFalse
         ORDER BY db_timestamp, seq_nr"""
   }
@@ -148,11 +152,13 @@ private[r2dbc] class SqlServerQueryDao(executorProvider: R2dbcExecutorProvider)
       stmt: Statement,
       entityType: String,
       fromTimestamp: Instant,
+      fromSeqNr: Option[Long],
       toTimestamp: Option[Instant]): Statement = {
     stmt
       .bind("@limit", settings.querySettings.bufferSize)
       .bind("@entityType", entityType)
       .bindTimestamp("@from", fromTimestamp)
+    fromSeqNr.foreach(seqNr => stmt.bind("@fromSeqNr", seqNr))
     toTimestamp.foreach(timestamp => stmt.bindTimestamp("@until", timestamp))
     stmt
   }