Avoid creating the test snapshot repeatedly

Author: kim

crates/snapshot/tests/remote.rs

@@ -1,6 +1,7 @@
-use std::sync::Arc;
+use std::{sync::Arc, time::Instant};
 
 use env_logger::Env;
+use log::info;
 use pretty_assertions::assert_matches;
 use spacetimedb::{
     db::{
@@ -15,6 +16,7 @@ use spacetimedb::{
     Identity,
 };
 use spacetimedb_durability::{EmptyHistory, TxOffset};
+use spacetimedb_fs_utils::dir_trie::DirTrie;
 use spacetimedb_lib::{
     bsatn,
     db::raw_def::v9::{RawModuleDefV9Builder, RawTableDefBuilder},
@@ -32,38 +34,33 @@ use spacetimedb_snapshot::{
     Snapshot, SnapshotError, SnapshotRepository,
 };
 use spacetimedb_table::page_pool::PagePool;
-use tempfile::tempdir;
-use tokio::task::spawn_blocking;
+use tempfile::{tempdir, TempDir};
+use tokio::{sync::OnceCell, task::spawn_blocking};
 
 // TODO: Happy path for compressed snapshot, pending #2034
 #[tokio::test]
 async fn can_sync_a_snapshot() -> anyhow::Result<()> {
     enable_logging();
     let tmp = tempdir()?;
+    let src = SourceSnapshot::get_or_create().await?;
 
-    let src_path = SnapshotsPath::from_path_unchecked(tmp.path().join("src"));
-    let dst_path = SnapshotsPath::from_path_unchecked(tmp.path().join("dst"));
-
-    src_path.create()?;
+    let dst_path = SnapshotsPath::from_path_unchecked(tmp.path());
     dst_path.create()?;
 
-    let src_repo = SnapshotRepository::open(src_path, Identity::ZERO, 0).map(Arc::new)?;
     let dst_repo = SnapshotRepository::open(dst_path.clone(), Identity::ZERO, 0).map(Arc::new)?;
 
-    let snapshot_offset = create_snapshot(src_repo.clone()).await?;
-    let src_snapshot_path = src_repo.snapshot_dir_path(snapshot_offset);
-    let (mut src_snapshot, _) = Snapshot::read_from_file(&src_snapshot_path.snapshot_file(snapshot_offset))?;
+    let mut src_snapshot = src.meta.clone();
     let total_objects = src_snapshot.total_objects() as u64;
 
-    let blob_provider = SnapshotRepository::object_repo(&src_snapshot_path).map(Arc::new)?;
+    let blob_provider = src.objects.clone();
 
     // This is the first snapshot in `dst_repo`, so all objects should be written.
     let stats = synchronize_snapshot(blob_provider.clone(), dst_path.clone(), src_snapshot.clone()).await?;
     assert_eq!(stats.objects_written, total_objects);
 
     // Assert that the copied snapshot is valid.
     let pool = PagePool::default();
-    let dst_snapshot_full = dst_repo.read_snapshot(snapshot_offset, &pool)?;
+    let dst_snapshot_full = dst_repo.read_snapshot(src.offset, &pool)?;
     Locking::restore_from_snapshot(dst_snapshot_full, pool)?;
 
     // Check that `verify_snapshot` agrees.
@@ -88,38 +85,88 @@ async fn can_sync_a_snapshot() -> anyhow::Result<()> {
 
 #[tokio::test]
 async fn rejects_overwrite() -> anyhow::Result<()> {
+    enable_logging();
     let tmp = tempdir()?;
+    let src = SourceSnapshot::get_or_create().await?;
 
-    let src_path = SnapshotsPath::from_path_unchecked(tmp.path().join("src"));
-    let dst_path = SnapshotsPath::from_path_unchecked(tmp.path().join("dst"));
-
-    src_path.create()?;
+    let dst_path = SnapshotsPath::from_path_unchecked(tmp.path());
     dst_path.create()?;
 
-    let src_repo = SnapshotRepository::open(src_path, Identity::ZERO, 0).map(Arc::new)?;
-
-    let snapshot_offset = create_snapshot(src_repo.clone()).await?;
-    let src_snapshot_path = src_repo.snapshot_dir_path(snapshot_offset);
-    let (src_snapshot, _) = Snapshot::read_from_file(&src_snapshot_path.snapshot_file(snapshot_offset))?;
-
-    let blob_provider = SnapshotRepository::object_repo(&src_snapshot_path).map(Arc::new)?;
+    let src_snapshot = src.meta.clone();
+    let blob_provider = src.objects.clone();
 
     synchronize_snapshot(blob_provider.clone(), dst_path.clone(), src_snapshot.clone()).await?;
 
     // Try to overwrite with the previous snapshot.
-    let prev_offset = src_repo.latest_snapshot_older_than(snapshot_offset - 1)?.unwrap();
-    let src_snapshot_path = src_repo.snapshot_dir_path(prev_offset);
+    // A previous snapshot exists because one is created immediately after
+    // database initialization.
+    let prev_offset = src.repo.latest_snapshot_older_than(src.offset - 1)?.unwrap();
+    let src_snapshot_path = src.repo.snapshot_dir_path(prev_offset);
     let (mut src_snapshot, _) = Snapshot::read_from_file(&src_snapshot_path.snapshot_file(prev_offset))?;
     // Pretend it's the current snapshot, thereby altering the preimage.
-    src_snapshot.tx_offset = snapshot_offset;
+    src_snapshot.tx_offset = src.offset;
 
     let res = synchronize_snapshot(blob_provider, dst_path, src_snapshot).await;
     assert_matches!(res, Err(SnapshotError::HashMismatch { .. }));
 
     Ok(())
 }
 
+/// Creating a snapshot takes a long time, because we need to commit
+/// `SNAPSHOT_FREQUENCY` transactions to trigger one.
+///
+/// Until the snapshot frequency becomes configurable,
+/// avoid creating the source snapshot repeatedly.
+struct SourceSnapshot {
+    offset: TxOffset,
+    meta: Snapshot,
+    objects: Arc<DirTrie>,
+    repo: Arc<SnapshotRepository>,
+
+    #[allow(unused)]
+    tmp: TempDir,
+}
+
+impl SourceSnapshot {
+    async fn get_or_create() -> anyhow::Result<&'static Self> {
+        static SOURCE_SNAPSHOT: OnceCell<SourceSnapshot> = OnceCell::const_new();
+        SOURCE_SNAPSHOT.get_or_try_init(Self::try_init).await
+    }
+
+    async fn try_init() -> anyhow::Result<Self> {
+        let tmp = tempdir()?;
+
+        let repo_path = SnapshotsPath::from_path_unchecked(tmp.path());
+        let repo = spawn_blocking(move || {
+            repo_path.create()?;
+            SnapshotRepository::open(repo_path, Identity::ZERO, 0).map(Arc::new)
+        })
+        .await
+        .unwrap()?;
+        let offset = create_snapshot(repo.clone()).await?;
+
+        let dir_path = repo.snapshot_dir_path(offset);
+        let (meta, objects) = spawn_blocking(move || {
+            let meta = Snapshot::read_from_file(&dir_path.snapshot_file(offset)).map(|(file, _)| file)?;
+            let objects = SnapshotRepository::object_repo(&dir_path).map(Arc::new)?;
+
+            Ok::<_, SnapshotError>((meta, objects))
+        })
+        .await
+        .unwrap()?;
+
+        Ok(SourceSnapshot {
+            offset,
+            meta,
+            objects,
+            repo,
+            tmp,
+        })
+    }
+}
+
 async fn create_snapshot(repo: Arc<SnapshotRepository>) -> anyhow::Result<TxOffset> {
+    let start = Instant::now();
     let mut watch = spawn_blocking(|| {
         let tmp = TempReplicaDir::new()?;
         let db = TestDB::open_db(&tmp, EmptyHistory::new(), None, Some(repo), 0)?;
@@ -148,6 +195,10 @@ async fn create_snapshot(repo: Arc<SnapshotRepository>) -> anyhow::Result<TxOffs
         snapshot_offset = *watch.borrow_and_update();
     }
     assert!(snapshot_offset >= SNAPSHOT_FREQUENCY);
+    info!(
+        "snapshot creation took {}s",
+        Instant::now().duration_since(start).as_secs_f32()
+    );
 
     Ok(snapshot_offset)
 }