Energy metering for persistent memory usage

Cargo.toml

@@ -152,7 +152,7 @@ futures = "0.3"
 futures-channel = "0.3"
 getrandom = "0.2.7"
 glob = "0.3.1"
-hashbrown = "0.14"
+hashbrown = { version = "0.15", default-features = false, features = ["equivalent", "inline-more"] }
 headers = "0.4"
 heck = "0.4"
 hex = "0.4.3"

crates/bindings/tests/snapshots/deps__spacetimedb_bindings_dependencies.snap

@@ -64,14 +64,15 @@ spacetimedb
 │   ├── itertools (*)
 │   ├── spacetimedb_bindings_macro (*)
 │   ├── spacetimedb_data_structures
+│   │   ├── ahash
+│   │   │   ├── cfg_if
+│   │   │   ├── getrandom (*)
+│   │   │   ├── once_cell
+│   │   │   └── zerocopy (*)
+│   │   │   [build-dependencies]
+│   │   │   └── version_check
 │   │   ├── hashbrown
-│   │   │   ├── ahash
-│   │   │   │   ├── cfg_if
-│   │   │   │   ├── once_cell
-│   │   │   │   └── zerocopy (*)
-│   │   │   │   [build-dependencies]
-│   │   │   │   └── version_check
-│   │   │   └── allocator_api2
+│   │   │   └── equivalent
 │   │   ├── nohash_hasher
 │   │   ├── smallvec
 │   │   └── thiserror

crates/client-api-messages/src/energy.rs

@@ -36,6 +36,13 @@ impl EnergyQuanta {
         let energy = bytes_stored * sec + (bytes_stored * nsec) / 1_000_000_000;
         Self::new(energy)
     }
+
+    const ENERGY_PER_MEM_BYTE_SEC: u128 = 100;
+
+    pub fn from_memory_usage(bytes_stored: u64, storage_period: Duration) -> Self {
+        let byte_seconds = Self::from_disk_usage(bytes_stored, storage_period).get();
+        Self::new(byte_seconds * Self::ENERGY_PER_MEM_BYTE_SEC)
+    }
 }
 
 impl fmt::Display for EnergyQuanta {

crates/client-api/src/routes/prometheus.rs

@@ -2,7 +2,7 @@ use crate::{log_and_500, ControlStateReadAccess};
 use axum::extract::State;
 use axum::response::IntoResponse;
 use serde::{Deserialize, Serialize};
-use spacetimedb_data_structures::map::HashMap;
+use spacetimedb_data_structures::map::{HashCollectionExt, HashMap};
 
 #[derive(Serialize, Deserialize)]
 struct SDConfig {

crates/core/src/db/datastore/locking_tx_datastore/committed_state.rs

@@ -37,6 +37,7 @@ use spacetimedb_table::{
     blob_store::{BlobStore, HashMapBlobStore},
     indexes::{RowPointer, SquashedOffset},
     table::{IndexScanIter, InsertError, RowRef, Table},
+    MemoryUsage,
 };
 use std::collections::BTreeMap;
 use std::sync::Arc;
@@ -55,6 +56,18 @@ pub struct CommittedState {
     pub(super) index_id_map: IndexIdMap,
 }
 
+impl MemoryUsage for CommittedState {
+    fn heap_usage(&self) -> usize {
+        let Self {
+            next_tx_offset,
+            tables,
+            blob_store,
+            index_id_map,
+        } = self;
+        next_tx_offset.heap_usage() + tables.heap_usage() + blob_store.heap_usage() + index_id_map.heap_usage()
+    }
+}
+
 impl StateView for CommittedState {
     fn get_schema(&self, table_id: TableId) -> Option<&Arc<TableSchema>> {
         self.tables.get(&table_id).map(|table| table.get_schema())

crates/core/src/db/datastore/locking_tx_datastore/datastore.rs

@@ -37,6 +37,7 @@ use spacetimedb_snapshot::ReconstructedSnapshot;
 use spacetimedb_table::{
     indexes::RowPointer,
     table::{RowRef, Table},
+    MemoryUsage,
 };
 use std::time::{Duration, Instant};
 use std::{borrow::Cow, sync::Arc};
@@ -64,6 +65,21 @@ pub struct Locking {
     pub(crate) database_address: Address,
 }
 
+impl MemoryUsage for Locking {
+    fn heap_usage(&self) -> usize {
+        let Self {
+            committed_state,
+            sequence_state,
+            database_address,
+        } = self;
+        std::mem::size_of_val(&**committed_state)
+            + committed_state.read().heap_usage()
+            + std::mem::size_of_val(&**sequence_state)
+            + sequence_state.lock().heap_usage()
+            + database_address.heap_usage()
+    }
+}
+
 impl Locking {
     pub fn new(database_address: Address) -> Self {
         Self {

crates/core/src/db/datastore/locking_tx_datastore/sequence.rs

@@ -1,12 +1,21 @@
 use spacetimedb_data_structures::map::IntMap;
 use spacetimedb_primitives::SequenceId;
 use spacetimedb_schema::schema::SequenceSchema;
+use spacetimedb_table::MemoryUsage;
 
 pub(super) struct Sequence {
     schema: SequenceSchema,
     pub(super) value: i128,
 }
 
+impl MemoryUsage for Sequence {
+    fn heap_usage(&self) -> usize {
+        // MEMUSE: intentionally ignoring schema
+        let Self { schema: _, value } = self;
+        value.heap_usage()
+    }
+}
+
 impl Sequence {
     pub(super) fn new(schema: SequenceSchema) -> Self {
         Self {
@@ -102,6 +111,13 @@ pub(super) struct SequencesState {
     sequences: IntMap<SequenceId, Sequence>,
 }
 
+impl MemoryUsage for SequencesState {
+    fn heap_usage(&self) -> usize {
+        let Self { sequences } = self;
+        sequences.heap_usage()
+    }
+}
+
 impl SequencesState {
     pub(super) fn get_sequence_mut(&mut self, seq_id: SequenceId) -> Option<&mut Sequence> {
         self.sequences.get_mut(&seq_id)

crates/core/src/db/relational_db.rs

@@ -36,6 +36,7 @@ use spacetimedb_schema::schema::{IndexSchema, RowLevelSecuritySchema, Schema, Se
 use spacetimedb_snapshot::{SnapshotError, SnapshotRepository};
 use spacetimedb_table::indexes::RowPointer;
 use spacetimedb_table::table::RowRef;
+use spacetimedb_table::MemoryUsage;
 use std::borrow::Cow;
 use std::collections::HashSet;
 use std::fmt;
@@ -490,6 +491,11 @@ impl RelationalDB {
         self.disk_size_fn.as_ref().map_or(Ok(0), |f| f())
     }
 
+    /// The size in bytes of all of the in-memory data in this database.
+    pub fn size_in_memory(&self) -> usize {
+        self.inner.heap_usage()
+    }
+
     pub fn encode_row(row: &ProductValue, bytes: &mut Vec<u8>) {
         // TODO: large file storage of the row elements
         row.encode(bytes);

crates/core/src/db/relational_operators.rs

@@ -1,5 +1,5 @@
 use core::marker::PhantomData;
-use spacetimedb_data_structures::map::HashSet;
+use spacetimedb_data_structures::map::{HashCollectionExt, HashSet};
 use spacetimedb_sats::ProductValue;
 
 // NOTE

crates/core/src/energy.rs

@@ -21,6 +21,7 @@ pub trait EnergyMonitor: Send + Sync + 'static {
         execution_duration: Duration,
     );
     fn record_disk_usage(&self, database: &Database, replica_id: u64, disk_usage: u64, period: Duration);
+    fn record_memory_usage(&self, database: &Database, replica_id: u64, mem_usage: u64, period: Duration);
 }
 
 #[derive(Default)]
@@ -40,4 +41,6 @@ impl EnergyMonitor for NullEnergyMonitor {
     }
 
     fn record_disk_usage(&self, _database: &Database, _replica_id: u64, _disk_usage: u64, _period: Duration) {}
+
+    fn record_memory_usage(&self, _database: &Database, _replica_id: u64, _mem_usage: u64, _period: Duration) {}
 }