RAII things

crates/core/src/storage/backend.rs

@@ -0,0 +1,108 @@
+use crate::storage::{Storage, StorageHashMap, StorageResult};
+
+/// Per-surfnet storage backend. Owns the connection resources for one surfnet
+/// (the SQLite pool, or a lease on the shared PostgreSQL pool) and mints the
+/// per-table stores that share them. Constructed once per surfnet from the
+/// database URL; dropping the surfnet drops the backend and its connections.
+#[derive(Clone)]
+pub enum StorageBackend {
+    /// No database URL: stores are independent in-process maps.
+    Memory,
+    #[cfg(feature = "sqlite")]
+    Sqlite(super::sqlite::SqliteBackend),
+    #[cfg(feature = "postgres")]
+    Postgres(super::postgres::PostgresBackend),
+}
+
+impl StorageBackend {
+    /// Selects and connects the backend for `database_url`: `None` is
+    /// in-memory, a `postgres://`/`postgresql://` URL is PostgreSQL, and
+    /// anything else is treated as an SQLite path (including `:memory:`).
+    pub fn open(database_url: &Option<&str>, surfnet_id: &str) -> StorageResult<Self> {
+        let Some(url) = database_url else {
+            return Ok(StorageBackend::Memory);
+        };
+        if url.starts_with("postgres://") || url.starts_with("postgresql://") {
+            #[cfg(feature = "postgres")]
+            {
+                let backend = super::postgres::PostgresBackend::open(url, surfnet_id)?;
+                return Ok(StorageBackend::Postgres(backend));
+            }
+            #[cfg(not(feature = "postgres"))]
+            return Err(super::StorageError::PostgresNotEnabled);
+        }
+        #[cfg(feature = "sqlite")]
+        {
+            let backend = super::sqlite::SqliteBackend::open(url, surfnet_id)?;
+            Ok(StorageBackend::Sqlite(backend))
+        }
+        #[cfg(not(feature = "sqlite"))]
+        {
+            let _ = surfnet_id;
+            Err(super::StorageError::SqliteNotEnabled)
+        }
+    }
+
+    /// Opens the kv store for `table_name`, backed by a hash map when the
+    /// backend is [`StorageBackend::Memory`].
+    pub fn open_store<K, V>(&self, table_name: &str) -> StorageResult<Box<dyn Storage<K, V>>>
+    where
+        K: serde::Serialize
+            + serde::de::DeserializeOwned
+            + Send
+            + Sync
+            + 'static
+            + Clone
+            + Eq
+            + std::hash::Hash,
+        V: serde::Serialize + serde::de::DeserializeOwned + Send + Sync + 'static + Clone,
+    {
+        self.open_store_with_default(table_name, || Box::new(StorageHashMap::new()))
+    }
+
+    /// Opens the kv store for `table_name`, using `default_storage_constructor`
+    /// when the backend is [`StorageBackend::Memory`].
+    pub fn open_store_with_default<K, V, F>(
+        &self,
+        table_name: &str,
+        default_storage_constructor: F,
+    ) -> StorageResult<Box<dyn Storage<K, V>>>
+    where
+        K: serde::Serialize
+            + serde::de::DeserializeOwned
+            + Send
+            + Sync
+            + 'static
+            + Clone
+            + Eq
+            + std::hash::Hash,
+        V: serde::Serialize + serde::de::DeserializeOwned + Send + Sync + 'static + Clone,
+        F: FnOnce() -> Box<dyn Storage<K, V>>,
+    {
+        #[cfg(not(any(feature = "sqlite", feature = "postgres")))]
+        let _ = table_name;
+        match self {
+            StorageBackend::Memory => Ok(default_storage_constructor()),
+            #[cfg(feature = "sqlite")]
+            StorageBackend::Sqlite(backend) => Ok(Box::new(backend.open_store(table_name)?)),
+            #[cfg(feature = "postgres")]
+            StorageBackend::Postgres(backend) => Ok(Box::new(backend.open_store(table_name)?)),
+        }
+    }
+
+    /// Whether stores opened on this backend survive process restart when
+    /// pointed at the same database.
+    pub fn is_persistent(&self) -> bool {
+        !matches!(self, StorageBackend::Memory)
+    }
+
+    /// Releases backend resources that need explicit cleanup before exit.
+    /// For SQLite this checkpoints the WAL and removes the `-wal`/`-shm`
+    /// files; the other backends have nothing to flush.
+    pub fn shutdown(&self) {
+        #[cfg(feature = "sqlite")]
+        if let StorageBackend::Sqlite(backend) = self {
+            backend.checkpoint();
+        }
+    }
+}

crates/core/src/storage/census.rs

@@ -0,0 +1,198 @@
+//! Process-wide counters for storage resource accounting.
+//!
+//! `lsof` and `ps` sample the process from outside; these counters record the
+//! storage layer's own resource events, so a workload's before/after delta is
+//! exact and assertable in tests. SQLite only: the PostgreSQL pool is
+//! process-shared by design and its sessions live on the server.
+
+use std::sync::atomic::{AtomicU64, Ordering};
+
+static POOLS_CREATED: AtomicU64 = AtomicU64::new(0);
+static POOLS_REUSED: AtomicU64 = AtomicU64::new(0);
+static CONNECTIONS_OPENED: AtomicU64 = AtomicU64::new(0);
+static CONNECTIONS_CLOSED: AtomicU64 = AtomicU64::new(0);
+static CONNECTIONS_PEAK: AtomicU64 = AtomicU64::new(0);
+
+pub fn pool_created() {
+    POOLS_CREATED.fetch_add(1, Ordering::Relaxed);
+}
+
+/// Reuse exists only while pools live in a shared cache; a backend that owns
+/// its pool never reuses, and reports zero here.
+pub fn pool_reused() {
+    POOLS_REUSED.fetch_add(1, Ordering::Relaxed);
+}
+
+pub fn connection_opened() {
+    CONNECTIONS_OPENED.fetch_add(1, Ordering::Relaxed);
+    CONNECTIONS_PEAK.fetch_max(live_connections(), Ordering::Relaxed);
+}
+
+pub fn connection_closed() {
+    CONNECTIONS_CLOSED.fetch_add(1, Ordering::Relaxed);
+}
+
+fn live_connections() -> u64 {
+    CONNECTIONS_OPENED
+        .load(Ordering::Relaxed)
+        .saturating_sub(CONNECTIONS_CLOSED.load(Ordering::Relaxed))
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct CensusSnapshot {
+    pub pools_created: u64,
+    pub pools_reused: u64,
+    pub connections_opened: u64,
+    pub connections_closed: u64,
+    /// Connections opened and not yet closed at snapshot time.
+    pub connections_live: u64,
+    /// High-water mark of live connections over the process lifetime.
+    pub connections_peak: u64,
+}
+
+pub fn snapshot() -> CensusSnapshot {
+    CensusSnapshot {
+        pools_created: POOLS_CREATED.load(Ordering::Relaxed),
+        pools_reused: POOLS_REUSED.load(Ordering::Relaxed),
+        connections_opened: CONNECTIONS_OPENED.load(Ordering::Relaxed),
+        connections_closed: CONNECTIONS_CLOSED.load(Ordering::Relaxed),
+        connections_live: live_connections(),
+        connections_peak: CONNECTIONS_PEAK.load(Ordering::Relaxed),
+    }
+}
+
+impl CensusSnapshot {
+    /// The counter movement between `earlier` and `self`. Monotonic fields
+    /// subtract; `connections_live` is the live count at `self`, and
+    /// `connections_peak` is the process-lifetime high-water mark.
+    pub fn since(&self, earlier: &CensusSnapshot) -> CensusSnapshot {
+        CensusSnapshot {
+            pools_created: self.pools_created - earlier.pools_created,
+            pools_reused: self.pools_reused - earlier.pools_reused,
+            connections_opened: self.connections_opened - earlier.connections_opened,
+            connections_closed: self.connections_closed - earlier.connections_closed,
+            connections_live: self.connections_live,
+            connections_peak: self.connections_peak,
+        }
+    }
+}
+
+/// A SQLite connection that reports its own close: r2d2's
+/// `CustomizeConnection::on_release` fires when the pool discards a broken or
+/// reaped connection, not at pool drop, so `Drop` on the connection itself is
+/// the only signal that counts every close.
+#[cfg(feature = "sqlite")]
+pub struct CountedConnection(surfpool_db::diesel::SqliteConnection);
+
+#[cfg(feature = "sqlite")]
+impl Drop for CountedConnection {
+    fn drop(&mut self) {
+        connection_closed();
+    }
+}
+
+#[cfg(feature = "sqlite")]
+impl std::ops::Deref for CountedConnection {
+    type Target = surfpool_db::diesel::SqliteConnection;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+#[cfg(feature = "sqlite")]
+impl std::ops::DerefMut for CountedConnection {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+#[cfg(feature = "sqlite")]
+pub struct CountingSqliteManager(
+    surfpool_db::diesel::r2d2::ConnectionManager<surfpool_db::diesel::SqliteConnection>,
+);
+
+#[cfg(feature = "sqlite")]
+impl CountingSqliteManager {
+    pub fn new(connection_string: &str) -> Self {
+        Self(surfpool_db::diesel::r2d2::ConnectionManager::new(
+            connection_string,
+        ))
+    }
+}
+
+#[cfg(feature = "sqlite")]
+impl surfpool_db::diesel::r2d2::ManageConnection for CountingSqliteManager {
+    type Connection = CountedConnection;
+    type Error = surfpool_db::diesel::r2d2::Error;
+
+    fn connect(&self) -> Result<Self::Connection, Self::Error> {
+        let conn = self.0.connect()?;
+        connection_opened();
+        Ok(CountedConnection(conn))
+    }
+
+    fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> {
+        self.0.is_valid(&mut conn.0)
+    }
+
+    fn has_broken(&self, conn: &mut Self::Connection) -> bool {
+        self.0.has_broken(&mut conn.0)
+    }
+}
+
+#[cfg(all(test, feature = "sqlite"))]
+mod tests {
+    use super::*;
+
+    /// Dropping the backend (and its stores) must close every connection it
+    /// opened.
+    #[test]
+    #[ignore = "process-global counters; run alone with --ignored --nocapture"]
+    fn dropping_backend_closes_connections() {
+        let temp_file = tempfile::NamedTempFile::new().unwrap();
+        let db_path = temp_file.path().to_str().unwrap();
+
+        let before = snapshot();
+        let backend = crate::storage::sqlite::SqliteBackend::open(db_path, "census").unwrap();
+        let store: crate::storage::SqliteStorage<String, String> =
+            backend.open_store("census_table").unwrap();
+        let opened = snapshot().since(&before);
+        assert!(
+            opened.connections_opened > 0,
+            "pool should open connections"
+        );
+
+        drop(store);
+        drop(backend);
+        let after = snapshot().since(&before);
+        assert_eq!(
+            after.connections_closed, after.connections_opened,
+            "every connection opened by the backend should close at drop"
+        );
+    }
+
+    /// Builds and drops N surfnets against on-disk and in-memory SQLite,
+    /// printing the exact resource movement per phase.
+    #[test]
+    #[ignore = "process-global counters; run alone with --ignored --nocapture"]
+    fn census_workload() {
+        const N: usize = 10;
+
+        let t0 = snapshot();
+        for _ in 0..N {
+            let tt = crate::storage::tests::TestType::sqlite();
+            let (svm, _simnet_rx, _geyser_rx) = tt.initialize_svm();
+            drop(svm);
+        }
+        let t1 = snapshot();
+        println!("on-disk  x{}: {:?}", N, t1.since(&t0));
+
+        for _ in 0..N {
+            let tt = crate::storage::tests::TestType::in_memory();
+            let (svm, _simnet_rx, _geyser_rx) = tt.initialize_svm();
+            drop(svm);
+        }
+        let t2 = snapshot();
+        println!("in-memory x{}: {:?}", N, t2.since(&t1));
+    }
+}