flatten db

Author: Shubham8287

crates/core/src/database_logger.rs

@@ -674,7 +674,7 @@ impl SystemLogger {
     }
 }
 
-impl spacetimedb_engine::db::update::UpdateLogger for SystemLogger {
+impl spacetimedb_engine::update::UpdateLogger for SystemLogger {
     fn info(&self, msg: &str) {
         self.info(msg);
     }

crates/core/src/db.rs

@@ -0,0 +1,17 @@
+pub use spacetimedb_engine::relational_db::{spawn_tx_metrics_recorder, Config, MetricsRecorderQueue, Storage};
+
+pub mod persistence {
+    pub use spacetimedb_engine::persistence::*;
+}
+
+pub mod relational_db {
+    pub use spacetimedb_engine::relational_db::*;
+}
+
+pub mod snapshot {
+    pub use spacetimedb_engine::snapshot::*;
+}
+
+pub mod update {
+    pub use spacetimedb_engine::update::*;
+}

crates/core/src/host/module_host.rs

@@ -601,7 +601,7 @@ fn init_database_inner(
             table_defs.sort_by_key(|x| &x.name);
             for def in table_defs {
                 logger.info(&format!("Creating table `{}`", &def.name));
-                spacetimedb_engine::db::update::create_table_from_def(stdb, tx, module_def, def)?;
+                spacetimedb_engine::update::create_table_from_def(stdb, tx, module_def, def)?;
             }
 
             // Create all in-memory views defined by the module.

crates/core/src/lib.rs

@@ -4,7 +4,7 @@ pub mod energy;
 pub mod sql;
 
 pub mod auth;
-pub use spacetimedb_engine::db;
+pub mod db;
 pub use spacetimedb_engine::metrics;
 pub mod messages;
 pub use spacetimedb_lib::Identity;

crates/engine/src/sql/ast.rs crates/engine/src/ast.rscrates/engine/src/sql/ast.rs renamed to crates/engine/src/ast.rs

@@ -10,7 +10,7 @@ use spacetimedb_durability::Transaction;
 use spacetimedb_lib::Identity;
 use spacetimedb_sats::ProductValue;
 
-use crate::db::persistence::Durability;
+use crate::persistence::Durability;
 use spacetimedb_runtime::Handle;
 
 pub(super) fn request_durability(

crates/engine/src/db/mod.rs

@@ -1,8 +1,12 @@
-pub mod db;
 pub mod error;
 pub mod metrics;
 pub mod rls;
-mod sql;
+mod ast;
+pub(crate) mod durability;
+pub mod persistence;
+pub mod relational_db;
+pub mod snapshot;
+pub mod update;
 pub mod util;
 
 pub use spacetimedb_lib::identity;

crates/engine/src/db/durability.rs crates/engine/src/durability.rscrates/engine/src/db/durability.rs renamed to crates/engine/src/durability.rs

@@ -1,8 +1,119 @@
-use crate::db::durability::{request_durability, spawn_close as spawn_durability_close};
-use crate::db::MetricsRecorderQueue;
+use crate::durability::{request_durability, spawn_close as spawn_durability_close};
 use crate::error::{DBError, RestoreSnapshotError};
 use crate::metrics::ExecutionCounters;
 use crate::metrics::ENGINE_METRICS;
+
+/// Whether SpacetimeDB is run in memory, or persists objects and
+/// a message log to disk.
+#[derive(Clone, Copy)]
+pub enum Storage {
+    /// The object store is in memory, and no message log is kept.
+    Memory,
+
+    /// The object store is persisted to disk, and a message log is kept.
+    Disk,
+}
+
+/// Internal database config parameters
+#[derive(Clone, Copy)]
+pub struct Config {
+    /// Specifies the object storage model.
+    pub storage: Storage,
+    /// Specifies the page pool max size in bytes.
+    pub page_pool_max_size: Option<usize>,
+}
+
+/// A message that is processed by the [`spawn_metrics_recorder`] actor.
+/// We use a separate task to record metrics to avoid blocking transactions.
+pub struct MetricsMessage {
+    /// The reducer the produced these metrics.
+    reducer: Option<ReducerName>,
+    /// Metrics from a mutable transaction.
+    metrics_for_writer: Option<TxMetrics>,
+    /// Metrics from a read-only transaction.
+    /// A message may have metrics for both types of transactions,
+    /// because metrics for a reducer and its subscription updates are recorded together.
+    metrics_for_reader: Option<TxMetrics>,
+    /// The row updates for an immutable transaction.
+    /// Needed for insert and delete counters.
+    tx_data: Option<Arc<TxData>>,
+    /// Cached metrics counters for each workload type.
+    counters: Arc<EnumMap<WorkloadType, ExecutionCounters>>,
+}
+
+/// The handle used to send work to the tx metrics recorder.
+#[derive(Clone)]
+pub struct MetricsRecorderQueue {
+    tx: spacetimedb_runtime::sync::mpsc::UnboundedSender<MetricsMessage>,
+}
+
+impl MetricsRecorderQueue {
+    pub fn send_metrics(
+        &self,
+        reducer: Option<ReducerName>,
+        metrics_for_writer: Option<TxMetrics>,
+        metrics_for_reader: Option<TxMetrics>,
+        tx_data: Option<Arc<TxData>>,
+        counters: Arc<EnumMap<WorkloadType, ExecutionCounters>>,
+    ) {
+        if let Err(err) = self.tx.send(MetricsMessage {
+            reducer,
+            metrics_for_writer,
+            metrics_for_reader,
+            tx_data,
+            counters,
+        }) {
+            log::warn!("failed to send metrics: {err}");
+        }
+    }
+}
+
+fn record_metrics(
+    MetricsMessage {
+        reducer,
+        metrics_for_writer,
+        metrics_for_reader,
+        tx_data,
+        counters,
+    }: MetricsMessage,
+) {
+    if let Some(tx_metrics) = metrics_for_writer {
+        tx_metrics.report(
+            tx_data.as_deref(),
+            reducer.as_ref(),
+            |wl| &counters[wl],
+        );
+    }
+    if let Some(tx_metrics) = metrics_for_reader {
+        tx_metrics.report(
+            None,
+            reducer.as_ref(),
+            |wl| &counters[wl],
+        );
+    }
+}
+
+const TX_METRICS_RECORDING_INTERVAL: std::time::Duration = std::time::Duration::from_millis(5);
+
+/// Spawns a task for recording transaction metrics.
+/// Returns the handle for pushing metrics to the recorder.
+pub fn spawn_tx_metrics_recorder(
+    handle: &spacetimedb_runtime::Handle,
+) -> (MetricsRecorderQueue, spacetimedb_runtime::AbortHandle) {
+    let handle_clone = handle.clone();
+    let (tx, mut rx) = spacetimedb_runtime::sync::mpsc::unbounded_channel();
+    let abort_handle = handle
+        .spawn(async move {
+            loop {
+                handle_clone.sleep(TX_METRICS_RECORDING_INTERVAL).await;
+                while let Ok(metrics) = rx.try_recv() {
+                    record_metrics(metrics);
+                }
+            }
+        })
+        .abort_handle();
+    (MetricsRecorderQueue { tx }, abort_handle)
+}
 use crate::util::asyncify;
 use anyhow::{anyhow, Context};
 use enum_map::EnumMap;
@@ -1830,8 +1941,8 @@ fn default_row_count_fn(db: Identity) -> RowCountFn {
 
 #[cfg(any(test, feature = "test"))]
 pub mod tests_utils {
-    use crate::db::snapshot;
-    use crate::db::snapshot::SnapshotWorker;
+    use crate::snapshot;
+    use crate::snapshot::SnapshotWorker;
 
     use super::*;
     use core::ops::Deref;
@@ -2331,7 +2442,7 @@ mod tests {
 
     use super::tests_utils::begin_mut_tx;
     use super::*;
-    use crate::db::relational_db::tests_utils::{begin_tx, create_view_for_test, insert, make_snapshot, TestDB};
+    use crate::relational_db::tests_utils::{begin_tx, create_view_for_test, insert, make_snapshot, TestDB};
     use anyhow::bail;
     use bytes::Bytes;
     use commitlog::payload::txdata;