[ci] works, with canton restartSource

Signed-off-by: Itai Segall <itai.segall@digitalasset.com>

Author: isegall-da

apps/scan/src/main/scala/org/lfdecentralizedtrust/splice/scan/store/ScanKeyValueProvider.scala

@@ -21,26 +21,27 @@ class ScanKeyValueProvider(val store: KeyValueStore, val loggerFactory: NamedLog
   private val latestAcsSnapshotInBulkStorageKey = "latest_acs_snapshot_in_bulk_storage"
 
   final def setLatestAcsSnapshotsInBulkStorage(
-      timestamp: CantonTimestamp,
       migrationId: Long,
+      timestamp: CantonTimestamp,
   )(implicit tc: TraceContext): Future[Unit] = store.setValue(
     latestAcsSnapshotInBulkStorageKey,
-    AcsSnapshotTimestampMigration(timestamp, migrationId),
+    AcsSnapshotTimestampMigration(migrationId, timestamp),
   )
 
   final def getLatestAcsSnapshotInBulkStorage()(implicit
       tc: TraceContext,
       ec: ExecutionContext,
-  ): OptionT[Future, (CantonTimestamp, Long)] = {
-    val result: OptionT[Future, AcsSnapshotTimestampMigration] = store.readValueAndLogOnDecodingFailure(latestAcsSnapshotInBulkStorageKey)
-    result.map(result => (result.timestamp, result.migrationId))
+  ): OptionT[Future, (Long, CantonTimestamp)] = {
+    val result: OptionT[Future, AcsSnapshotTimestampMigration] =
+      store.readValueAndLogOnDecodingFailure(latestAcsSnapshotInBulkStorageKey)
+    result.map(result => (result.migrationId, result.timestamp))
   }
 }
 
 object ScanKeyValueProvider {
   final case class AcsSnapshotTimestampMigration(
-      timestamp: CantonTimestamp,
       migrationId: Long,
+      timestamp: CantonTimestamp,
   )
   implicit val timestampCodec: Codec[CantonTimestamp] =
     Codec

apps/scan/src/main/scala/org/lfdecentralizedtrust/splice/scan/store/bulk/AcsSnapshotBulkStorage.scala

@@ -13,7 +13,7 @@ import org.apache.pekko.actor.ActorSystem
 import org.apache.pekko.stream.scaladsl.{Keep, Source}
 import org.apache.pekko.pattern.after
 import org.apache.pekko.stream.{KillSwitch, KillSwitches}
-import org.lfdecentralizedtrust.splice.scan.store.AcsSnapshotStore
+import org.lfdecentralizedtrust.splice.scan.store.{AcsSnapshotStore, ScanKeyValueProvider}
 
 import scala.concurrent.{ExecutionContext, Future}
 import scala.concurrent.duration.*
@@ -22,18 +22,19 @@ class AcsSnapshotBulkStorage(
     val config: BulkStorageConfig,
     val acsSnapshotStore: AcsSnapshotStore,
     val s3Connection: S3BucketConnection,
+    val kvProvider: ScanKeyValueProvider,
     override val loggerFactory: NamedLoggerFactory,
 )(implicit actorSystem: ActorSystem, tc: TraceContext, ec: ExecutionContext)
     extends NamedLogging {
 
   // TODO(#3429): persist progress (or conclude it from the S3 storage), and start from latest successfully dumped snapshot upon restart
-  private def getStartTimestamp: (Long, CantonTimestamp) = (0, CantonTimestamp.MinValue)
+  private def getStartTimestamp: Future[Option[(Long, CantonTimestamp)]] = kvProvider.getLatestAcsSnapshotInBulkStorage().value
 
   // When new snapshot is not yet available, how long to wait for a new one.
   // TODO(#3429): make it longer for prod (so consider making it configurable/overridable for tests)
   private val snapshotPollingInterval = 5.seconds
 
-  private def getAcsSnapshotTimestampAfter(
+  private def getAcsSnapshotTimestampsAfter(
       startMigrationId: Long,
       startTimestamp: CantonTimestamp,
   ): Source[(Long, CantonTimestamp), NotUsed] = {
@@ -66,15 +67,15 @@ class AcsSnapshotBulkStorage(
   /** This is the main implementation of the pipeline. It is a Pekko Source that gets a `start` timestamp, and starts dumping to S3
     *  all snapshots (strictly) after `start`. It is an infinite source that should never complete.
     */
-  private def mksrc(
-      start: (Long, CantonTimestamp)
-  ): Source[(Long, CantonTimestamp), (KillSwitch, Future[Done])] = start match {
-    case (startMigrationId, startAfterTimestamp) =>
-      logger.debug(
-        s"Starting ACS snapshot dump source, from migration $startMigrationId, timestamp $startAfterTimestamp"
-      )
+  private def mksrc(): Source[(Long, CantonTimestamp), (KillSwitch, Future[Done])] =  {
+      logger.debug("Starting ACS snapshot dump source")
       val base =
-        getAcsSnapshotTimestampAfter(startMigrationId, startAfterTimestamp)
+        Source.single[Unit](())
+          .mapAsync(1){_ => getStartTimestamp}
+          .flatMapConcat {
+            case Some((startMigrationId, startAfterTimestamp)) => getAcsSnapshotTimestampsAfter(startMigrationId, startAfterTimestamp)
+            case None => getAcsSnapshotTimestampsAfter(0, CantonTimestamp.MinValue)
+          }
           .via(
             SingleAcsSnapshotBulkStorage.asFlow(
               config,
@@ -83,8 +84,16 @@ class AcsSnapshotBulkStorage(
               loggerFactory,
             )
           )
+          .mapAsync(1){case (migrationId, timestamp) =>
+            for {
+              _ <- kvProvider.setLatestAcsSnapshotsInBulkStorage(migrationId, timestamp)
+            } yield {
+              (migrationId, timestamp)
+            }
+          }
+
 
-      val withKs = base.viaMat(KillSwitches.single)(Keep.right)
+    val withKs = base.viaMat(KillSwitches.single)(Keep.right)
       withKs.watchTermination() { case (ks, done) => (ks: KillSwitch, done) }
   }
 
@@ -96,28 +105,28 @@ class AcsSnapshotBulkStorage(
     //  to Pekko's built-in RestartSource (for now, it's convenient to use Canton's ability to track state via lastEmittedElement)
     // TODO(#3429): tweak the retry parameters here
     val delay = FiniteDuration(5, "seconds")
-    val policy = new RetrySourcePolicy[(Long, CantonTimestamp), (Long, CantonTimestamp)] {
+    val policy = new RetrySourcePolicy[Unit, (Long, CantonTimestamp)] {
       override def shouldRetry(
-          lastState: (Long, CantonTimestamp),
+          lastState: Unit,
           lastEmittedElement: Option[(Long, CantonTimestamp)],
           lastFailure: Option[Throwable],
-      ): Option[(scala.concurrent.duration.FiniteDuration, (Long, CantonTimestamp))] = {
+      ): Option[(scala.concurrent.duration.FiniteDuration, Unit)] = {
         lastFailure.map { t =>
           logger.warn(
             s"Writing ACS snapshot to bulk storage failed with : ${ErrorUtil
                 .messageWithStacktrace(t)}, will retry after delay of $delay from last successful timestamp $lastEmittedElement"
           )
           // Always retry (TODO(#3429): consider a max number of retries?)
-          delay -> lastEmittedElement.fold(lastState)(identity)
+          delay -> ()
         }
       }
     }
 
     PekkoUtil
       .restartSource(
         name = "acs-snapshot-dump",
-        initial = getStartTimestamp,
-        mkSource = (start: (Long, CantonTimestamp)) => mksrc(start),
+        initial = (),
+        mkSource = (_: Unit) => mksrc(),
         policy = policy,
       )
 

apps/scan/src/test/scala/org/lfdecentralizedtrust/splice/scan/store/ScanKeyValueProviderTest.scala

@@ -13,6 +13,7 @@ import org.lfdecentralizedtrust.splice.store.db.SplicePostgresTest
 import org.scalatest.matchers.should.Matchers
 
 import scala.concurrent.Future
+import scala.concurrent.duration._
 
 class ScanKeyValueProviderTest
     extends StoreTest
@@ -22,13 +23,15 @@ class ScanKeyValueProviderTest
   "ScanKeyValueProvider" should {
     "set and get acs snapshots timestamps" in {
       val ts = CantonTimestamp.now()
+      val ts2 = ts.add(1.minute)
       val migrationId = 7L
       for {
         provider <- mkProvider
-        _ <- provider.setLatestAcsSnapshotsInBulkStorage(ts, migrationId)
+        _ <- provider.setLatestAcsSnapshotsInBulkStorage(migrationId, ts)
+        _ <- provider.setLatestAcsSnapshotsInBulkStorage(migrationId, ts2)
         readBack <- provider.getLatestAcsSnapshotInBulkStorage().value
       } yield {
-        readBack.value shouldBe (ts, migrationId)
+        readBack.value shouldBe (migrationId, ts2)
       }
     }
   }

apps/scan/src/test/scala/org/lfdecentralizedtrust/splice/scan/store/bulk/AcsSnapshotBulkStorageTest.scala

@@ -4,17 +4,24 @@
 package org.lfdecentralizedtrust.splice.scan.store.bulk
 
 import com.digitalasset.canton.data.CantonTimestamp
+import com.digitalasset.canton.lifecycle.FutureUnlessShutdown
 import com.digitalasset.canton.logging.NamedLoggerFactory
 import com.digitalasset.canton.protocol.LfContractId
+import com.digitalasset.canton.resource.DbStorage
 import com.digitalasset.canton.topology.PartyId
 import com.digitalasset.canton.tracing.TraceContext
 import com.digitalasset.canton.util.PekkoUtil.WithKillSwitch
 import com.digitalasset.canton.{HasActorSystem, HasExecutionContext}
 import org.apache.pekko.stream.scaladsl.Sink
 import org.apache.pekko.stream.testkit.scaladsl.TestSink
 import org.lfdecentralizedtrust.splice.http.v0.definitions as httpApi
-import org.lfdecentralizedtrust.splice.scan.store.AcsSnapshotStore
+import org.lfdecentralizedtrust.splice.scan.store.{
+  AcsSnapshotStore,
+  ScanKeyValueProvider,
+  ScanKeyValueStore,
+}
 import org.lfdecentralizedtrust.splice.scan.store.AcsSnapshotStore.QueryAcsSnapshotResult
+import org.lfdecentralizedtrust.splice.store.db.SplicePostgresTest
 import org.lfdecentralizedtrust.splice.store.events.SpliceCreatedEvent
 import org.lfdecentralizedtrust.splice.store.{HardLimit, Limit, StoreTest}
 import org.lfdecentralizedtrust.splice.util.PackageQualifiedName
@@ -34,7 +41,8 @@ class AcsSnapshotBulkStorageTest
     extends StoreTest
     with HasExecutionContext
     with HasActorSystem
-    with HasS3Mock {
+    with HasS3Mock
+    with SplicePostgresTest {
 
   val acsSnapshotSize = 48500
   val bulkStorageTestConfig = BulkStorageConfig(
@@ -99,33 +107,38 @@ class AcsSnapshotBulkStorageTest
       withS3Mock {
         val store = new MockAcsSnapshotStore()
         val s3BucketConnection = getS3BucketConnectionWithInjectedErrors(loggerFactory)
-        val probe = new AcsSnapshotBulkStorage(
-          bulkStorageTestConfig,
-          store.store,
-          s3BucketConnection,
-          loggerFactory,
-        ).getSource
-          .runWith(TestSink.probe[WithKillSwitch[(Long, CantonTimestamp)]])
+        val ts1 = CantonTimestamp.tryFromInstant(Instant.ofEpochSecond(10))
+        val ts2 = CantonTimestamp.tryFromInstant(Instant.ofEpochSecond(20))
+        for {
+          kvProvider <- mkProvider
+          probe = new AcsSnapshotBulkStorage(
+            bulkStorageTestConfig,
+            store.store,
+            s3BucketConnection,
+            kvProvider,
+            loggerFactory,
+          ).getSource
+            .runWith(TestSink.probe[WithKillSwitch[(Long, CantonTimestamp)]])
 
-        clue("Initially, a single snapshot is dumped") {
-          probe.request(2)
-          probe.expectNext(2.minutes).value shouldBe (0, CantonTimestamp.tryFromInstant(
-            Instant.ofEpochSecond(10)
-          ))
-          probe.expectNoMessage(10.seconds)
-        }
+          _ = clue("Initially, a single snapshot is dumped") {
+            probe.request(2)
+            probe.expectNext(2.minutes).value shouldBe(0, ts1)
+            probe.expectNoMessage(10.seconds)
+          }
+          persistedTs1 <- kvProvider.getLatestAcsSnapshotInBulkStorage().value
+          _ = persistedTs1.value shouldBe (0, ts1)
 
-        clue("Add another snapshot to the store, it is also dumped") {
-          store.addSnapshot(CantonTimestamp.tryFromInstant(Instant.ofEpochSecond(20)))
-          val next = probe.expectNext(2.minutes)
-          next.value shouldBe (0, CantonTimestamp.tryFromInstant(
-            Instant.ofEpochSecond(20)
-          ))
-          probe.expectNoMessage(10.seconds)
-          next.killSwitch.shutdown()
+          _ = clue("Add another snapshot to the store, it is also dumped") {
+            store.addSnapshot(CantonTimestamp.tryFromInstant(Instant.ofEpochSecond(20)))
+            val next = probe.expectNext(2.minutes)
+            next.value shouldBe(0, ts2)
+            probe.expectNoMessage(10.seconds)
+            next.killSwitch.shutdown()
+          }
+          persistedTs2 <- kvProvider.getLatestAcsSnapshotInBulkStorage().value
+        } yield {
+          persistedTs2.value shouldBe (0, ts2)
         }
-
-        succeed
       }
     }
   }
@@ -254,4 +267,16 @@ class AcsSnapshotBulkStorageTest
     s3BucketConnectionWithErrors
   }
 
+  def mkProvider: Future[ScanKeyValueProvider] = {
+    ScanKeyValueStore(
+      dsoParty = dsoParty,
+      participantId = mkParticipantId("participant"),
+      storage,
+      loggerFactory,
+    ).map(new ScanKeyValueProvider(_, loggerFactory))
+  }
+
+  override protected def cleanDb(
+      storage: DbStorage
+  )(implicit traceContext: TraceContext): FutureUnlessShutdown[?] = resetAllAppTables(storage)
 }