fix: Cancel scheduled tasks when projection is stopped

akka-projection-r2dbc-integration/src/test/scala/akka/projection/r2dbc/EventSourcedChaosSpec.scala

@@ -67,7 +67,7 @@ object EventSourcedChaosSpec {
     private val log = LoggerFactory.getLogger(getClass)
 
     override def process(session: R2dbcSession, envelope: EventEnvelope[String]): Future[Done] = {
-      val failCount = failEvents.getOrDefault(envelope.eventOption, 0)
+      val failCount = failEvents.getOrDefault(envelope.event, 0)
       if (failCount > 0) {
         failEvents.put(envelope.event, failCount - 1)
         log.debug(
@@ -120,8 +120,8 @@ class EventSourcedChaosSpec
       val numberOfEntities = 20 + rnd.nextInt(80)
       val numberOfProjections = 1 << rnd.nextInt(4) // 1 to 8 projections
       val entityType = nextEntityType()
-      val failProbability = 0.01
-      val stopProbability = 0.01
+      val failProbability = 0.02
+      val stopProbability = 0.02
 
       var startedEntities = Map.empty[Int, ActorRef[Persister.Command]]
       def entity(i: Int): ActorRef[Persister.Command] = {

akka-projection-r2dbc-integration/src/test/scala/akka/projection/r2dbc/EventSourcedEndToEndSpec.scala

@@ -251,7 +251,7 @@ class EventSourcedEndToEndSpec
 
       val projectionName = UUID.randomUUID().toString
       val processedProbe = createTestProbe[Processed]()
-      val projections = startProjections(entityType, projectionName, nrOfProjections = 4, processedProbe.ref)
+      var projections = startProjections(entityType, projectionName, nrOfProjections = 4, processedProbe.ref)
 
       // give them some time to start before writing more events
       Thread.sleep(500)
@@ -271,7 +271,7 @@ class EventSourcedEndToEndSpec
 
         // resume projections again
         if (n == (numberOfEvents / 2) + 20)
-          startProjections(entityType, projectionName, nrOfProjections = 4, processedProbe.ref)
+          projections = startProjections(entityType, projectionName, nrOfProjections = 4, processedProbe.ref)
 
         if (n % 10 == 0)
           Thread.sleep(50)

akka-projection-r2dbc/src/main/scala/akka/projection/r2dbc/internal/R2dbcOffsetStore.scala

@@ -41,6 +41,7 @@ import scala.collection.immutable.TreeSet
 import scala.concurrent.ExecutionContext
 import scala.concurrent.Future
 
+import akka.actor.Cancellable
 import akka.stream.scaladsl.Sink
 import akka.stream.scaladsl.Source
 
@@ -207,6 +208,13 @@ private[projection] object R2dbcOffsetStore {
   val FutureDone: Future[Done] = Future.successful(Done)
 
   final case class LatestBySlice(slice: Int, pid: String, seqNr: Long)
+
+  final case class ScheduledTasks(adoptForeignOffsets: Option[Cancellable], deleteOffsets: Option[Cancellable]) {
+    def cancel(): Unit = {
+      adoptForeignOffsets.foreach(_.cancel())
+      deleteOffsets.foreach(_.cancel())
+    }
+  }
 }
 
 /**
@@ -280,21 +288,40 @@ private[projection] class R2dbcOffsetStore(
 
   private val adoptingForeignOffsets = !settings.adoptInterval.isZero && !settings.adoptInterval.isNegative
 
-  private val scheduledAdoptForeignOffsets =
-    if (adoptingForeignOffsets)
-      Some(
-        system.scheduler.scheduleWithFixedDelay(
-          settings.adoptInterval,
-          settings.adoptInterval,
-          () => adoptForeignOffsets(),
-          system.executionContext))
-    else None
-
-  private def scheduleNextDelete(): Unit = {
-    if (!settings.deleteInterval.isZero && !settings.deleteInterval.isNegative)
-      system.scheduler.scheduleOnce(settings.deleteInterval, () => deleteOldTimestampOffsets(), system.executionContext)
+  private val scheduledTasks = new AtomicReference[ScheduledTasks](ScheduledTasks(None, None))
+
+  @tailrec
+  private def scheduleTasks(): Unit = {
+    val existingTasks = scheduledTasks.get()
+    existingTasks.cancel()
+
+    val foreignOffsets =
+      if (adoptingForeignOffsets)
+        Some(
+          system.scheduler.scheduleWithFixedDelay(
+            settings.adoptInterval,
+            settings.adoptInterval,
+            () => adoptForeignOffsets(),
+            system.executionContext))
+      else None
+
+    val deleteOffsets =
+      if (!settings.deleteInterval.isZero && !settings.deleteInterval.isNegative)
+        Some(
+          system.scheduler.scheduleWithFixedDelay(
+            settings.deleteInterval,
+            settings.deleteInterval,
+            () => deleteOldTimestampOffsets(),
+            system.executionContext))
+      else
+        None
+
+    val newTasks = ScheduledTasks(foreignOffsets, deleteOffsets)
+    if (!scheduledTasks.compareAndSet(existingTasks, newTasks)) {
+      newTasks.cancel()
+      scheduleTasks() // CAS retry, concurrent update
+    }
   }
-  scheduleNextDelete()
 
   private def timestampOf(persistenceId: String, sequenceNr: Long): Future[Option[Instant]] = {
     sourceProvider match {
@@ -327,6 +354,8 @@ private[projection] class R2dbcOffsetStore(
   }
 
   def readOffset[Offset](): Future[Option[Offset]] = {
+    scheduleTasks()
+
     // look for TimestampOffset first since that is used by akka-persistence-r2dbc,
     // and then fall back to the other more primitive offset types
     sourceProvider match {
@@ -929,7 +958,7 @@ private[projection] class R2dbcOffsetStore(
         logger.debug("{} Deleted [{}] timestamp offset rows", logPrefix, rows)
       }
 
-    result.andThen(_ => scheduleNextDelete())
+    result
   }
 
   def deleteOldTimestampOffsets(slice: Int): Future[Long] = {
@@ -1008,12 +1037,12 @@ private[projection] class R2dbcOffsetStore(
 
   def adoptForeignOffsets(): Future[Long] = {
     if (!hasForeignOffsets()) {
-      scheduledAdoptForeignOffsets.foreach(_.cancel())
+      scheduledTasks.get.adoptForeignOffsets.foreach(_.cancel())
       Future.successful(0)
     } else {
       val latestTimestamp = getLatestSeen()
       val adoptableRecords = takeAdoptableForeignOffsets(latestTimestamp)
-      if (!hasForeignOffsets()) scheduledAdoptForeignOffsets.foreach(_.cancel())
+      if (!hasForeignOffsets()) scheduledTasks.get.adoptForeignOffsets.foreach(_.cancel())
       val adoptableLatestBySlice = adoptableRecords.map { adoptable =>
         LatestBySlice(adoptable.record.slice, adoptable.record.pid, adoptable.record.seqNr)
       }
@@ -1167,4 +1196,9 @@ private[projection] class R2dbcOffsetStore(
     }
   }
 
+  def stop(): Unit = {
+    scheduledTasks.get.cancel()
+    scheduledTasks.set(ScheduledTasks(None, None))
+  }
+
 }

akka-projection-r2dbc/src/main/scala/akka/projection/r2dbc/internal/R2dbcProjectionImpl.scala

@@ -728,7 +728,7 @@ private[projection] class R2dbcProjectionImpl[Offset, Envelope](
       // if the handler is retrying it will be aborted by this,
       // otherwise the stream would not be completed by the killSwitch until after all retries
       projectionState.abort.failure(AbortProjectionException)
-      streamDone
+      streamDone.andThen(_ => offsetStore.stop())(system.executionContext)
     }
 
     // RunningProjectionManagement