Started fixing the --in-memory flag

Author: cloutiertyler

crates/cli/tests/server.rs

@@ -3,6 +3,15 @@ mod util;
 use crate::util::SpacetimeDbGuard;
 use assert_cmd::cargo::cargo_bin_cmd;
 
+#[test]
+fn cli_can_ping_spacetimedb_in_memory() {
+    let spacetime = SpacetimeDbGuard::spawn_in_memory();
+    let mut cmd = cargo_bin_cmd!("spacetimedb-cli");
+    cmd.args(["server", "ping", &spacetime.host_url.to_string()])
+        .assert()
+        .success();
+}
+
 #[test]
 fn cli_can_ping_spacetimedb_on_disk() {
     let spacetime = SpacetimeDbGuard::spawn_in_temp_data_dir();

crates/cli/tests/util.rs

@@ -30,6 +30,12 @@ impl SpacetimeDbGuard {
         Self::spawn_spacetime_start(&["start", "--data-dir", &data_dir])
     }
 
+    /// Start `spacetimedb` in-memory via:
+    /// cargo run -p spacetimedb-cli -- start --in-memory --listen-addr 127.0.0.1:<port>
+    pub fn spawn_in_memory() -> Self {
+        Self::spawn_spacetime_start(&["start", "--in-memory"])
+    }
+
     fn spawn_spacetime_start(extra_args: &[&str]) -> Self {
         let port = find_free_port();
         let addr: SocketAddr = format!("127.0.0.1:{port}").parse().unwrap();

crates/core/src/db/mod.rs

@@ -14,7 +14,7 @@ pub mod update;
 
 /// Whether SpacetimeDB is run in memory, or persists objects and
 /// a message log to disk.
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum Storage {
     /// The object store is in memory, and no message log is kept.
     Memory,

crates/core/src/host/disk_storage.rs

@@ -79,4 +79,8 @@ impl ExternalStorage for DiskStorage {
     async fn lookup(&self, program_hash: Hash) -> anyhow::Result<Option<Box<[u8]>>> {
         self.get(&program_hash).await.map_err(Into::into)
     }
+
+    async fn put(&self, program_bytes: &[u8]) -> anyhow::Result<Hash> {
+        self.put(program_bytes).await.map_err(Into::into)
+    }
 }

crates/core/src/host/host_controller.rs

@@ -38,7 +38,6 @@ use spacetimedb_sats::hash::Hash;
 use spacetimedb_schema::auto_migrate::{ponder_migrate, AutoMigrateError, MigrationPolicy, PrettyPrintStyle};
 use spacetimedb_schema::def::ModuleDef;
 use spacetimedb_table::page_pool::PagePool;
-use std::future::Future;
 use std::ops::Deref;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
@@ -61,16 +60,7 @@ pub type ExternalDurability = (Arc<dyn Durability<TxData = Txdata>>, DiskSizeFn)
 #[async_trait]
 pub trait ExternalStorage: Send + Sync + 'static {
     async fn lookup(&self, program_hash: Hash) -> anyhow::Result<Option<Box<[u8]>>>;
-}
-#[async_trait]
-impl<F, Fut> ExternalStorage for F
-where
-    F: Fn(Hash) -> Fut + Send + Sync + 'static,
-    Fut: Future<Output = anyhow::Result<Option<Box<[u8]>>>> + Send,
-{
-    async fn lookup(&self, program_hash: Hash) -> anyhow::Result<Option<Box<[u8]>>> {
-        self(program_hash).await
-    }
+    async fn put(&self, program_bytes: &[u8]) -> anyhow::Result<Hash>;
 }
 
 pub type ProgramStorage = Arc<dyn ExternalStorage>;

crates/core/src/host/memory_storage.rs

@@ -0,0 +1,55 @@
+use async_trait::async_trait;
+use spacetimedb_lib::{hash_bytes, Hash};
+use std::collections::HashMap;
+use std::io;
+use std::sync::Arc;
+use tokio::sync::RwLock;
+
+use super::ExternalStorage;
+
+/// A simple [`ExternalStorage`] that stores programs in memory.
+#[derive(Clone, Debug, Default)]
+pub struct MemoryStorage {
+    inner: Arc<RwLock<HashMap<Hash, Box<[u8]>>>>,
+}
+
+impl MemoryStorage {
+    /// Create a new empty `MemoryStorage`.
+    pub async fn new() -> io::Result<Self> {
+        Ok(Self {
+            inner: Arc::new(RwLock::new(HashMap::new())),
+        })
+    }
+
+    #[tracing::instrument(level = "trace", skip(self))]
+    pub async fn get(&self, key: &Hash) -> io::Result<Option<Box<[u8]>>> {
+        let guard = self.inner.read().await;
+        Ok(guard.get(key).cloned())
+    }
+
+    #[tracing::instrument(level = "trace", skip(self, value))]
+    pub async fn put(&self, value: &[u8]) -> io::Result<Hash> {
+        let h = hash_bytes(value);
+        let mut guard = self.inner.write().await;
+        guard.insert(h, Box::from(value));
+        Ok(h)
+    }
+
+    #[tracing::instrument(level = "trace", skip(self))]
+    pub async fn prune(&self, key: &Hash) -> anyhow::Result<()> {
+        let mut guard = self.inner.write().await;
+        guard.remove(key);
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl ExternalStorage for MemoryStorage {
+    async fn lookup(&self, program_hash: Hash) -> anyhow::Result<Option<Box<[u8]>>> {
+        self.get(&program_hash).await.map_err(Into::into)
+    }
+
+    async fn put(&self, program_bytes: &[u8]) -> anyhow::Result<Hash> {
+        self.put(program_bytes).await.map_err(Into::into)
+    }
+}

crates/core/src/host/mod.rs

@@ -11,6 +11,7 @@ use spacetimedb_schema::def::deserialize::{ArgsSeed, FunctionDef};
 
 mod disk_storage;
 mod host_controller;
+mod memory_storage;
 mod module_common;
 #[allow(clippy::too_many_arguments)]
 pub mod module_host;
@@ -27,6 +28,7 @@ pub use host_controller::{
     extract_schema, CallProcedureReturn, ExternalDurability, ExternalStorage, HostController, MigratePlanResult,
     ProcedureCallResult, ProgramStorage, ReducerCallResult, ReducerOutcome,
 };
+pub use memory_storage::MemoryStorage;
 pub use module_host::{ModuleHost, NoSuchModule, ProcedureCallError, ReducerCallError, UpdateDatabaseResult};
 pub use scheduler::Scheduler;
 

crates/standalone/src/control_db.rs

@@ -75,6 +75,15 @@ impl ControlDb {
         Ok(Self { db })
     }
 
+    pub fn new_in_memory() -> Result<Self> {
+        let config = sled::Config::default()
+            .temporary(true)
+            .flush_every_ms(Some(50))
+            .mode(sled::Mode::HighThroughput);
+        let db = config.open()?;
+        Ok(Self { db })
+    }
+
     #[cfg(test)]
     pub fn at(path: impl AsRef<std::path::Path>) -> Result<Self> {
         let config = sled::Config::default()

crates/standalone/src/lib.rs

@@ -13,7 +13,9 @@ use spacetimedb::config::{CertificateAuthority, MetadataFile};
 use spacetimedb::db;
 use spacetimedb::db::persistence::LocalPersistenceProvider;
 use spacetimedb::energy::{EnergyBalance, EnergyQuanta, NullEnergyMonitor};
-use spacetimedb::host::{DiskStorage, HostController, MigratePlanResult, UpdateDatabaseResult};
+use spacetimedb::host::{
+    DiskStorage, HostController, MemoryStorage, MigratePlanResult, ProgramStorage, UpdateDatabaseResult,
+};
 use spacetimedb::identity::{AuthCtx, Identity};
 use spacetimedb::messages::control_db::{Database, Node, Replica};
 use spacetimedb::util::jobs::JobCores;
@@ -41,11 +43,11 @@ pub struct StandaloneOptions {
 
 pub struct StandaloneEnv {
     control_db: ControlDb,
-    program_store: Arc<DiskStorage>,
+    program_store: ProgramStorage,
     host_controller: HostController,
     client_actor_index: ClientActorIndex,
     metrics_registry: prometheus::Registry,
-    _pid_file: PidFile,
+    _pid_file: Option<PidFile>,
     auth_provider: auth::DefaultJwtAuthProvider,
     websocket_options: WebSocketOptions,
 }
@@ -57,17 +59,24 @@ impl StandaloneEnv {
         data_dir: Arc<ServerDataDir>,
         db_cores: JobCores,
     ) -> anyhow::Result<Arc<Self>> {
-        let _pid_file = data_dir.pid_file()?;
-        let meta_path = data_dir.metadata_toml();
-        let mut meta = MetadataFile::new("standalone");
-        if let Some(existing_meta) = MetadataFile::read(&meta_path).context("failed reading metadata.toml")? {
-            meta = existing_meta.check_compatibility_and_update(meta)?;
-        }
-        meta.write(&meta_path).context("failed writing metadata.toml")?;
+        let (pid_file, control_db, program_store): (Option<PidFile>, ControlDb, ProgramStorage) =
+            if config.db_config.storage == db::Storage::Disk {
+                let meta_path = data_dir.metadata_toml();
+                let mut meta = MetadataFile::new("standalone");
+                if let Some(existing_meta) = MetadataFile::read(&meta_path).context("failed reading metadata.toml")? {
+                    meta = existing_meta.check_compatibility_and_update(meta)?;
+                }
+                meta.write(&meta_path).context("failed writing metadata.toml")?;
+                let control_db = ControlDb::new(&data_dir.control_db()).context("failed to initialize control db")?;
+                let program_store = Arc::new(DiskStorage::new(data_dir.program_bytes().0).await?);
+                (Some(data_dir.pid_file()?), control_db, program_store)
+            } else {
+                let control_db = ControlDb::new_in_memory().context("failed to initialize in-memory control db")?;
+                let program_store = Arc::new(MemoryStorage::new().await?);
+                (None, control_db, program_store)
+            };
 
-        let control_db = ControlDb::new(&data_dir.control_db()).context("failed to initialize control db")?;
         let energy_monitor = Arc::new(NullEnergyMonitor);
-        let program_store = Arc::new(DiskStorage::new(data_dir.program_bytes().0).await?);
 
         let persistence_provider = Arc::new(LocalPersistenceProvider::new(data_dir.clone()));
         let host_controller = HostController::new(
@@ -94,7 +103,7 @@ impl StandaloneEnv {
             host_controller,
             client_actor_index,
             metrics_registry,
-            _pid_file,
+            _pid_file: pid_file,
             auth_provider: auth_env,
             websocket_options: config.websocket,
         }))