DST: Test harness

Cargo.toml

@@ -12,6 +12,7 @@ members = [
   "crates/commitlog",
   "crates/core",
   "crates/data-structures",
+  "crates/dst",
   "crates/datastore",
   "crates/durability",
   "crates/execution",
@@ -25,6 +26,7 @@ members = [
   "crates/physical-plan",
   "crates/primitives",
   "crates/query",
+  "crates/runtime",
   "crates/sats",
   "crates/schema",
   "crates/smoketests",
@@ -139,6 +141,7 @@ spacetimedb-pg = { path = "crates/pg", version = "=2.2.0" }
 spacetimedb-physical-plan = { path = "crates/physical-plan", version = "=2.2.0" }
 spacetimedb-primitives = { path = "crates/primitives", version = "=2.2.0" }
 spacetimedb-query = { path = "crates/query", version = "=2.2.0" }
+spacetimedb-runtime = { path = "crates/runtime", version = "=2.2.0" }
 spacetimedb-sats = { path = "crates/sats", version = "=2.2.0" }
 spacetimedb-schema = { path = "crates/schema", version = "=2.2.0" }
 spacetimedb-standalone = { path = "crates/standalone", version = "=2.2.0" }

crates/commitlog/src/lib.rs

@@ -156,7 +156,7 @@ impl Options {
 /// The canonical commitlog API over a repository backend `R`.
 ///
 /// The default backend is the on-disk filesystem repository
-/// [`repo::Fs`], but tests may supply another [`Repo`]
+/// [`repo::Fs`], but tests and simulators may supply another [`Repo`]
 /// implementation.
 ///
 /// Records in the log are of type `T`, which canonically is instantiated to
@@ -203,7 +203,7 @@ where
 {
     /// Open the log in `repo` with [`Options`].
     ///
-    /// This is useful for tests which provide a repository
+    /// This is useful for tests and simulators which provide a repository
     /// implementation other than [`repo::Fs`].
     pub fn open_with_repo(repo: R, opts: Options) -> io::Result<Self> {
         let inner = commitlog::Generic::open(repo, opts)?;

crates/commitlog/src/repo/mod.rs

@@ -161,6 +161,8 @@ pub trait RepoWithoutLockFile: Repo {}
 
 impl<T: RepoWithoutLockFile> RepoWithoutLockFile for &T {}
 
+impl RepoWithoutLockFile for Fs {}
+
 #[cfg(any(test, feature = "test"))]
 impl RepoWithoutLockFile for Memory {}
 

crates/core/Cargo.toml

@@ -28,6 +28,7 @@ spacetimedb-primitives.workspace = true
 spacetimedb-paths.workspace = true
 spacetimedb-physical-plan.workspace = true
 spacetimedb-query.workspace = true
+spacetimedb-runtime = { workspace = true, features = ["tokio"] }
 spacetimedb-sats = { workspace = true, features = ["serde"] }
 spacetimedb-schema.workspace = true
 spacetimedb-table.workspace = true

crates/core/src/database_logger.rs

@@ -11,7 +11,7 @@ use std::path::Path;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
-use tokio::io::{AsyncRead, BufReader};
+use tokio::io::{AsyncRead, BufReader, ReadBuf};
 use tokio::sync::{broadcast, mpsc, oneshot};
 use tokio_stream::wrappers::errors::BroadcastStreamRecvError;
 use tokio_stream::wrappers::BroadcastStream;
@@ -592,7 +592,7 @@ fn seek_to(file: &mut File, buf: &mut [u8], num_lines: u32) -> io::Result<()> {
     Ok(())
 }
 
-fn read_exact_at(file: &std::fs::File, buf: &mut [u8], offset: u64) -> io::Result<()> {
+fn read_exact_at(file: &File, buf: &mut [u8], offset: u64) -> io::Result<()> {
     #[cfg(unix)]
     {
         use std::os::unix::fs::FileExt;
@@ -641,7 +641,7 @@ impl MaybeFile {
 }
 
 impl AsyncRead for MaybeFile {
-    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut tokio::io::ReadBuf<'_>) -> Poll<io::Result<()>> {
+    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
         match self.project() {
             MaybeFileProj::File { inner } => inner.poll_read(cx, buf),
             MaybeFileProj::Empty => Poll::Ready(Ok(())),

crates/core/src/db/durability.rs

@@ -9,9 +9,9 @@ use spacetimedb_datastore::{execution_context::ReducerContext, traits::TxData};
 use spacetimedb_durability::Transaction;
 use spacetimedb_lib::Identity;
 use spacetimedb_sats::ProductValue;
-use tokio::{runtime, time::timeout};
 
 use crate::db::persistence::Durability;
+use spacetimedb_runtime::Handle;
 
 pub(super) fn request_durability(
     durability: &Durability,
@@ -32,18 +32,19 @@ pub(super) fn request_durability(
     }));
 }
 
-pub(super) fn spawn_close(durability: Arc<Durability>, runtime: &runtime::Handle, database_identity: Identity) {
-    let rt = runtime.clone();
-    rt.spawn(async move {
-        let label = format!("[{database_identity}]");
-        match timeout(Duration::from_secs(10), durability.close()).await {
+pub(super) fn spawn_close(durability: Arc<Durability>, runtime: &Handle, database_identity: Identity) {
+    let label = format!("[{database_identity}]");
+    let runtime = runtime.clone();
+    runtime.clone().spawn(async move {
+        match runtime.timeout(Duration::from_secs(10), durability.close()).await {
             Err(_elapsed) => {
                 error!("{label} timeout waiting for durability shutdown");
             }
             Ok(offset) => {
                 info!("{label} durability shut down at tx offset: {offset:?}");
             }
         }
+        log::info!("closing spawn close");
     });
 }
 

crates/core/src/db/persistence.rs

@@ -4,7 +4,8 @@ use async_trait::async_trait;
 use spacetimedb_commitlog::SizeOnDisk;
 use spacetimedb_durability::{DurabilityExited, TxOffset};
 use spacetimedb_paths::server::ServerDataDir;
-use spacetimedb_snapshot::DynSnapshotRepo;
+use spacetimedb_runtime::Handle;
+use spacetimedb_snapshot::{DynSnapshotRepo, SnapshotStore};
 
 use crate::{messages::control_db::Database, util::asyncify};
 
@@ -35,14 +36,16 @@ pub struct Persistence {
     /// Currently the expectation is that the reported size is the commitlog
     /// size only.
     pub disk_size: DiskSizeFn,
+    /// Optional snapshot store used during database restore.
+    pub snapshot_store: Option<Arc<dyn SnapshotStore>>,
     /// An optional [SnapshotWorker].
     ///
     /// The current expectation is that snapshots are only enabled for
     /// persistent (as opposed to in-memory) databases. This is enforced by
     /// this type.
     pub snapshots: Option<SnapshotWorker>,
-    /// The tokio runtime onto which durability-related tasks shall be spawned.
-    pub runtime: tokio::runtime::Handle,
+    /// Runtime onto which durability-related tasks shall be spawned.
+    pub runtime: Handle,
 }
 
 impl Persistence {
@@ -53,9 +56,20 @@ impl Persistence {
         snapshots: Option<SnapshotWorker>,
         runtime: tokio::runtime::Handle,
     ) -> Self {
+        Self::new_with_runtime(durability, disk_size, snapshots, Handle::tokio(runtime))
+    }
+
+    pub fn new_with_runtime(
+        durability: impl spacetimedb_durability::Durability<TxData = Txdata> + 'static,
+        disk_size: impl Fn() -> io::Result<SizeOnDisk> + Send + Sync + 'static,
+        snapshots: Option<SnapshotWorker>,
+        runtime: Handle,
+    ) -> Self {
+        let snapshot_store = snapshots.as_ref().map(SnapshotWorker::snapshot_store);
         Self {
             durability: Arc::new(durability),
             disk_size: Arc::new(disk_size),
+            snapshot_store,
             snapshots,
             runtime,
         }
@@ -66,6 +80,13 @@ impl Persistence {
         self.snapshots.as_ref().map(|worker| worker.snapshot_repo())
     }
 
+    /// If snapshot restore is enabled, get the [SnapshotStore] to read from.
+    pub fn snapshot_store(&self) -> Option<Arc<dyn SnapshotStore>> {
+        self.snapshot_store
+            .clone()
+            .or_else(|| self.snapshots.as_ref().map(SnapshotWorker::snapshot_store))
+    }
+
     /// Get the [TxOffset] reported as durable by the [Durability] impl.
     ///
     /// Returns `Ok(None)` if no offset is durable yet, and `Err(DurabilityExited)`
@@ -91,12 +112,13 @@ impl Persistence {
         Option<Arc<Durability>>,
         Option<DiskSizeFn>,
         Option<SnapshotWorker>,
-        Option<tokio::runtime::Handle>,
+        Option<Handle>,
     ) {
         this.map(
             |Self {
                  durability,
                  disk_size,
+                 snapshot_store: _,
                  snapshots,
                  runtime,
              }| (Some(durability), Some(disk_size), snapshots, Some(runtime)),
@@ -143,13 +165,15 @@ impl PersistenceProvider for LocalPersistenceProvider {
     async fn persistence(&self, database: &Database, replica_id: u64) -> anyhow::Result<Persistence> {
         let replica_dir = self.data_dir.replica(replica_id);
         let snapshot_dir = replica_dir.snapshots();
+        let runtime = Handle::tokio_current();
 
         let database_identity = database.database_identity;
         let snapshot_worker =
             asyncify(move || relational_db::open_snapshot_repo(snapshot_dir, database_identity, replica_id))
                 .await
-                .map(|repo| SnapshotWorker::new(repo, snapshot::Compression::Enabled))?;
-        let (durability, disk_size) = relational_db::local_durability(replica_dir, Some(&snapshot_worker)).await?;
+                .map(|repo| SnapshotWorker::new(repo, snapshot::Compression::Enabled, runtime.clone()))?;
+        let (durability, disk_size) =
+            relational_db::local_durability(replica_dir, runtime.clone(), Some(&snapshot_worker)).await?;
 
         tokio::spawn(relational_db::snapshot_watching_commitlog_compressor(
             snapshot_worker.subscribe(),
@@ -161,8 +185,9 @@ impl PersistenceProvider for LocalPersistenceProvider {
         Ok(Persistence {
             durability,
             disk_size,
+            snapshot_store: Some(snapshot_worker.snapshot_store()),
             snapshots: Some(snapshot_worker),
-            runtime: tokio::runtime::Handle::current(),
+            runtime,
         })
     }
 }