diff --git a/crates/client-api-messages/src/energy.rs b/crates/client-api-messages/src/energy.rs index 9b82ef69b1..f4f47491f3 100644 --- a/crates/client-api-messages/src/energy.rs +++ b/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 { diff --git a/crates/core/src/db/datastore/locking_tx_datastore/committed_state.rs b/crates/core/src/db/datastore/locking_tx_datastore/committed_state.rs index c909d95dd3..d877d2426c 100644 --- a/crates/core/src/db/datastore/locking_tx_datastore/committed_state.rs +++ b/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 memory_usage(&self) -> usize { + let Self { + next_tx_offset, + tables, + blob_store, + index_id_map, + } = self; + next_tx_offset.memory_usage() + tables.memory_usage() + blob_store.memory_usage() + index_id_map.memory_usage() + } +} + impl StateView for CommittedState { fn get_schema(&self, table_id: TableId) -> Option<&Arc> { self.tables.get(&table_id).map(|table| table.get_schema()) diff --git a/crates/core/src/db/datastore/locking_tx_datastore/datastore.rs b/crates/core/src/db/datastore/locking_tx_datastore/datastore.rs index f9b74fc82e..12d931f6a5 100644 --- a/crates/core/src/db/datastore/locking_tx_datastore/datastore.rs +++ b/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 { database_address: Address, } +impl MemoryUsage for Locking { + fn memory_usage(&self) -> usize { + let Self { + committed_state, + sequence_state, + database_address, + } = self; + std::mem::size_of_val(&**committed_state) + + committed_state.read().memory_usage() + + std::mem::size_of_val(&**sequence_state) + + sequence_state.lock().memory_usage() + + database_address.memory_usage() + } +} + impl Locking { pub fn new(database_address: Address) -> Self { Self { diff --git a/crates/core/src/db/datastore/locking_tx_datastore/sequence.rs b/crates/core/src/db/datastore/locking_tx_datastore/sequence.rs index 4331354263..0ffa952e52 100644 --- a/crates/core/src/db/datastore/locking_tx_datastore/sequence.rs +++ b/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 memory_usage(&self) -> usize { + // MEMUSE: intentionally ignoring schema + let Self { schema: _, value } = self; + value.memory_usage() + } +} + impl Sequence { pub(super) fn new(schema: SequenceSchema) -> Self { Self { @@ -102,6 +111,13 @@ pub(super) struct SequencesState { sequences: IntMap, } +impl MemoryUsage for SequencesState { + fn memory_usage(&self) -> usize { + let Self { sequences } = self; + sequences.memory_usage() + } +} + impl SequencesState { pub(super) fn get_sequence_mut(&mut self, seq_id: SequenceId) -> Option<&mut Sequence> { self.sequences.get_mut(&seq_id) diff --git a/crates/core/src/db/relational_db.rs b/crates/core/src/db/relational_db.rs index 874f8af231..5a49221062 100644 --- a/crates/core/src/db/relational_db.rs +++ b/crates/core/src/db/relational_db.rs @@ -36,6 +36,7 @@ use spacetimedb_schema::schema::{IndexSchema, Schema, SequenceSchema, TableSchem 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.memory_usage() + } + pub fn encode_row(row: &ProductValue, bytes: &mut Vec) { // TODO: large file storage of the row elements row.encode(bytes); diff --git a/crates/core/src/energy.rs b/crates/core/src/energy.rs index 9fe3fd82f5..0e796e1bb2 100644 --- a/crates/core/src/energy.rs +++ b/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) {} } diff --git a/crates/core/src/host/host_controller.rs b/crates/core/src/host/host_controller.rs index 692beda21a..316e3afbc1 100644 --- a/crates/core/src/host/host_controller.rs +++ b/crates/core/src/host/host_controller.rs @@ -755,7 +755,7 @@ impl Host { } scheduler_starter.start(&module_host)?; - let metrics_task = tokio::spawn(disk_monitor(replica_ctx.clone(), energy_monitor.clone())).abort_handle(); + let metrics_task = tokio::spawn(storage_monitor(replica_ctx.clone(), energy_monitor.clone())).abort_handle(); Ok(Host { module: watch::Sender::new(module_host), @@ -826,22 +826,23 @@ impl Drop for Host { } } -const DISK_METERING_INTERVAL: Duration = Duration::from_secs(5); +const STORAGE_METERING_INTERVAL: Duration = Duration::from_secs(5); /// Periodically collect the disk usage of `replica_ctx` and update metrics as well as /// the `energy_monitor` accordingly. -async fn disk_monitor(replica_ctx: Arc, energy_monitor: Arc) { - let mut interval = tokio::time::interval(DISK_METERING_INTERVAL); +async fn storage_monitor(replica_ctx: Arc, energy_monitor: Arc) { + let mut interval = tokio::time::interval(STORAGE_METERING_INTERVAL); // We don't care about happening precisely every 5 seconds - it just matters // that the time between ticks is accurate. interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); - let mut prev_disk_usage = replica_ctx.total_disk_usage(); + let mut prev_disk_usage = tokio::task::block_in_place(|| replica_ctx.total_disk_usage()); let mut prev_tick = interval.tick().await; loop { let tick = interval.tick().await; let dt = tick - prev_tick; - let disk_usage = tokio::task::block_in_place(|| replica_ctx.total_disk_usage()); + let (disk_usage, mem_usage) = + tokio::task::block_in_place(|| (replica_ctx.total_disk_usage(), replica_ctx.mem_usage())); if let Some(num_bytes) = disk_usage.durability { DB_METRICS .message_log_size @@ -856,6 +857,7 @@ async fn disk_monitor(replica_ctx: Arc, energy_monitor: Arc usize { + self.relational_db.size_in_memory() + } } impl Deref for ReplicaContext { diff --git a/crates/primitives/src/col_list.rs b/crates/primitives/src/col_list.rs index cc1ecd7cb7..42d0933155 100644 --- a/crates/primitives/src/col_list.rs +++ b/crates/primitives/src/col_list.rs @@ -276,6 +276,14 @@ impl ColList { let addr = unsafe { self.check }; addr & 1 != 0 } + + #[doc(hidden)] + pub fn heap_size(&self) -> usize { + match self.as_inline() { + Ok(_) => 0, + Err(heap) => heap.capacity() as usize, + } + } } impl Drop for ColList { diff --git a/crates/standalone/src/energy_monitor.rs b/crates/standalone/src/energy_monitor.rs index 74ed2e4d6a..dd3be4a027 100644 --- a/crates/standalone/src/energy_monitor.rs +++ b/crates/standalone/src/energy_monitor.rs @@ -40,6 +40,11 @@ impl EnergyMonitor for StandaloneEnergyMonitor { let amount = EnergyQuanta::from_disk_usage(disk_usage, period); self.withdraw_energy(database.owner_identity, amount) } + + fn record_memory_usage(&self, database: &Database, _instance_id: u64, mem_usage: u64, period: Duration) { + let amount = EnergyQuanta::from_memory_usage(mem_usage, period); + self.withdraw_energy(database.owner_identity, amount) + } } impl StandaloneEnergyMonitor { diff --git a/crates/table/src/bflatn_to_bsatn_fast_path.rs b/crates/table/src/bflatn_to_bsatn_fast_path.rs index c953ba2654..0b9dacaef9 100644 --- a/crates/table/src/bflatn_to_bsatn_fast_path.rs +++ b/crates/table/src/bflatn_to_bsatn_fast_path.rs @@ -20,6 +20,8 @@ //! one of 20 bytes to copy the leading `(u64, u64, u32)`, which contains no padding, //! and then one of 8 bytes to copy the trailing `u64`, skipping over 4 bytes of padding in between. +use crate::MemoryUsage; + use super::{ indexes::{Byte, Bytes}, layout::{ @@ -47,6 +49,13 @@ pub(crate) struct StaticBsatnLayout { fields: Box<[MemcpyField]>, } +impl MemoryUsage for StaticBsatnLayout { + fn memory_usage(&self) -> usize { + let Self { bsatn_length, fields } = self; + bsatn_length.memory_usage() + fields.memory_usage() + } +} + impl StaticBsatnLayout { /// Serialize `row` from BFLATN to BSATN into `buf`. /// @@ -156,6 +165,8 @@ struct MemcpyField { length: u16, } +impl MemoryUsage for MemcpyField {} + impl MemcpyField { /// Copies the bytes at `row[self.bflatn_offset .. self.bflatn_offset + self.length]` /// into `buf[self.bsatn_offset + self.length]`. diff --git a/crates/table/src/blob_store.rs b/crates/table/src/blob_store.rs index b0883a8a30..303f37a4b2 100644 --- a/crates/table/src/blob_store.rs +++ b/crates/table/src/blob_store.rs @@ -15,6 +15,8 @@ use blake3::hash; use spacetimedb_data_structures::map::{Entry, HashMap}; use spacetimedb_lib::{de::Deserialize, ser::Serialize}; +use crate::MemoryUsage; + /// The content address of a blob-stored object. #[derive(Eq, PartialEq, PartialOrd, Ord, Clone, Copy, Hash, Debug, Serialize, Deserialize)] pub struct BlobHash { @@ -24,6 +26,8 @@ pub struct BlobHash { pub data: [u8; BlobHash::SIZE], } +impl MemoryUsage for BlobHash {} + impl BlobHash { /// The size of the hash function's output in bytes. pub const SIZE: usize = 32; @@ -142,6 +146,13 @@ pub struct HashMapBlobStore { map: HashMap, } +impl MemoryUsage for HashMapBlobStore { + fn memory_usage(&self) -> usize { + let Self { map } = self; + map.memory_usage() + } +} + /// A blob object including a reference count and the data. struct BlobObject { /// Reference count of the blob. @@ -150,6 +161,13 @@ struct BlobObject { blob: Box<[u8]>, } +impl MemoryUsage for BlobObject { + fn memory_usage(&self) -> usize { + let Self { uses, blob } = self; + uses.memory_usage() + blob.memory_usage() + } +} + impl BlobStore for HashMapBlobStore { fn clone_blob(&mut self, hash: &BlobHash) -> Result<(), NoSuchBlobError> { self.map.get_mut(hash).ok_or(NoSuchBlobError)?.uses += 1; diff --git a/crates/table/src/btree_index.rs b/crates/table/src/btree_index.rs index 6ab256d07d..939f129775 100644 --- a/crates/table/src/btree_index.rs +++ b/crates/table/src/btree_index.rs @@ -23,7 +23,7 @@ use super::indexes::RowPointer; use super::table::RowRef; -use crate::{read_column::ReadColumn, static_assert_size}; +use crate::{read_column::ReadColumn, static_assert_size, MemoryUsage}; use core::ops::RangeBounds; use spacetimedb_primitives::{ColList, IndexId}; use spacetimedb_sats::{ @@ -127,6 +127,28 @@ enum TypedIndex { AlgebraicValue(Index), } +impl MemoryUsage for TypedIndex { + fn memory_usage(&self) -> usize { + match self { + TypedIndex::Bool(this) => this.memory_usage(), + TypedIndex::U8(this) => this.memory_usage(), + TypedIndex::I8(this) => this.memory_usage(), + TypedIndex::U16(this) => this.memory_usage(), + TypedIndex::I16(this) => this.memory_usage(), + TypedIndex::U32(this) => this.memory_usage(), + TypedIndex::I32(this) => this.memory_usage(), + TypedIndex::U64(this) => this.memory_usage(), + TypedIndex::I64(this) => this.memory_usage(), + TypedIndex::U128(this) => this.memory_usage(), + TypedIndex::I128(this) => this.memory_usage(), + TypedIndex::U256(this) => this.memory_usage(), + TypedIndex::I256(this) => this.memory_usage(), + TypedIndex::String(this) => this.memory_usage(), + TypedIndex::AlgebraicValue(this) => this.memory_usage(), + } + } +} + impl TypedIndex { /// Add the row referred to by `row_ref` to the index `self`, /// which must be keyed at `cols`. @@ -329,6 +351,18 @@ pub struct BTreeIndex { pub key_type: AlgebraicType, } +impl MemoryUsage for BTreeIndex { + fn memory_usage(&self) -> usize { + let Self { + index_id, + is_unique, + idx, + key_type, + } = self; + index_id.memory_usage() + is_unique.memory_usage() + idx.memory_usage() + key_type.memory_usage() + } +} + static_assert_size!(BTreeIndex, 64); impl BTreeIndex { diff --git a/crates/table/src/btree_index/multimap.rs b/crates/table/src/btree_index/multimap.rs index c232f8b44a..ef2c3f6b4f 100644 --- a/crates/table/src/btree_index/multimap.rs +++ b/crates/table/src/btree_index/multimap.rs @@ -3,6 +3,8 @@ use core::slice; use smallvec::SmallVec; use std::collections::btree_map::{BTreeMap, Range}; +use crate::MemoryUsage; + /// A multi map that relates a `K` to a *set* of `V`s. #[derive(Default)] pub struct MultiMap { @@ -15,6 +17,13 @@ pub struct MultiMap { map: BTreeMap>, } +impl MemoryUsage for MultiMap { + fn memory_usage(&self) -> usize { + let Self { map } = self; + map.memory_usage() + } +} + impl MultiMap { /// Returns an empty multi map. pub fn new() -> Self { diff --git a/crates/table/src/fixed_bit_set.rs b/crates/table/src/fixed_bit_set.rs index 431735dfaa..cd6896c735 100644 --- a/crates/table/src/fixed_bit_set.rs +++ b/crates/table/src/fixed_bit_set.rs @@ -5,6 +5,8 @@ use core::{ pub use internal_unsafe::FixedBitSet; use internal_unsafe::Len; +use crate::MemoryUsage; + /// A type used to represent blocks in a bit set. /// A smaller type, compared to usize, /// means taking less advantage of native operations. @@ -243,6 +245,12 @@ impl FixedBitSet { } } +impl MemoryUsage for FixedBitSet { + fn memory_usage(&self) -> usize { + std::mem::size_of_val(self.storage()) + } +} + /// An iterator that yields the set indices of a [`FixedBitSet`]. pub struct IterSet<'a, B = DefaultBitBlock> { /// The block iterator. diff --git a/crates/table/src/indexes.rs b/crates/table/src/indexes.rs index 2dd4ec8106..27e97a455d 100644 --- a/crates/table/src/indexes.rs +++ b/crates/table/src/indexes.rs @@ -2,7 +2,7 @@ //! bytes, row hashes, (page) sizes, offsets, and indices. use super::util::range_move; -use crate::static_assert_size; +use crate::{static_assert_size, MemoryUsage}; use ahash::RandomState; use core::fmt; use core::ops::{AddAssign, Div, Mul, Range, SubAssign}; @@ -53,6 +53,8 @@ pub const PAGE_DATA_SIZE: usize = PAGE_SIZE - PAGE_HEADER_SIZE; #[cfg_attr(any(test, feature = "proptest"), derive(proptest_derive::Arbitrary))] pub struct RowHash(pub u64); +impl MemoryUsage for RowHash {} + static_assert_size!(RowHash, 8); /// `RowHash` is already a hash, so no need to hash again. @@ -70,6 +72,8 @@ impl RowHash { #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Add, Sub)] pub struct Size(pub u16); +impl MemoryUsage for Size {} + // We need to be able to serialize and deserialize `Size` because they appear in the `PageHeader`. impl_serialize!([] Size, (self, ser) => self.0.serialize(ser)); impl_deserialize!([] Size, de => u16::deserialize(de).map(Size)); @@ -100,6 +104,8 @@ pub struct PageOffset( #[cfg_attr(any(test, feature = "proptest"), proptest(strategy = "0..PageOffset::PAGE_END.0"))] pub u16, ); +impl MemoryUsage for PageOffset {} + static_assert_size!(PageOffset, 2); // We need to ser/de `PageOffset`s because they appear within the `PageHeader`. @@ -201,6 +207,8 @@ impl fmt::LowerHex for PageOffset { #[derive(Clone, Copy, PartialEq, Eq, Debug)] pub struct PageIndex(#[cfg_attr(any(test, feature = "proptest"), proptest(strategy = "0..MASK_PI"))] pub u64); +impl MemoryUsage for PageIndex {} + static_assert_size!(PageIndex, 8); impl PageIndex { @@ -230,6 +238,8 @@ impl PageIndex { #[cfg_attr(any(test, feature = "proptest"), derive(proptest_derive::Arbitrary))] pub struct SquashedOffset(pub u8); +impl MemoryUsage for SquashedOffset {} + static_assert_size!(SquashedOffset, 1); impl SquashedOffset { @@ -258,6 +268,8 @@ impl SquashedOffset { #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct RowPointer(pub u64); +impl MemoryUsage for RowPointer {} + static_assert_size!(RowPointer, 8); // Offsets and bits for the various components of `RowPointer`. diff --git a/crates/table/src/layout.rs b/crates/table/src/layout.rs index 2628cd4688..f4386a25bf 100644 --- a/crates/table/src/layout.rs +++ b/crates/table/src/layout.rs @@ -6,6 +6,8 @@ //! These, and others, determine what the layout of objects typed at those types are. //! They also implement [`HasLayout`] which generalizes over layout annotated types. +use crate::MemoryUsage; + use super::{ indexes::Size, var_len::{VarLenGranule, VarLenRef}, @@ -55,6 +57,8 @@ pub struct Layout { pub align: u16, } +impl MemoryUsage for Layout {} + /// A type which knows what its layout is. /// /// This does not refer to layout in Rust. @@ -107,6 +111,17 @@ pub enum AlgebraicTypeLayout { VarLen(VarLenType), } +impl MemoryUsage for AlgebraicTypeLayout { + fn memory_usage(&self) -> usize { + match self { + AlgebraicTypeLayout::Sum(x) => x.memory_usage(), + AlgebraicTypeLayout::Product(x) => x.memory_usage(), + AlgebraicTypeLayout::Primitive(x) => x.memory_usage(), + AlgebraicTypeLayout::VarLen(x) => x.memory_usage(), + } + } +} + impl HasLayout for AlgebraicTypeLayout { fn layout(&self) -> &Layout { match self { @@ -170,6 +185,13 @@ pub const fn row_size_for_type() -> Size { #[derive(Debug, PartialEq, Eq, Clone)] pub struct RowTypeLayout(ProductTypeLayout); +impl MemoryUsage for RowTypeLayout { + fn memory_usage(&self) -> usize { + let Self(layout) = self; + layout.memory_usage() + } +} + impl RowTypeLayout { /// Returns a view of this row type as a product type. pub fn product(&self) -> &ProductTypeLayout { @@ -217,6 +239,13 @@ pub struct ProductTypeLayout { pub elements: Collection, } +impl MemoryUsage for ProductTypeLayout { + fn memory_usage(&self) -> usize { + let Self { layout, elements } = self; + layout.memory_usage() + elements.memory_usage() + } +} + impl HasLayout for ProductTypeLayout { fn layout(&self) -> &Layout { &self.layout @@ -239,6 +268,13 @@ pub struct ProductTypeElementLayout { pub name: Option>, } +impl MemoryUsage for ProductTypeElementLayout { + fn memory_usage(&self) -> usize { + let Self { offset, ty, name } = self; + offset.memory_usage() + ty.memory_usage() + name.memory_usage() + } +} + /// A mirrior of [`SumType`] annotated with a [`Layout`]. #[derive(Debug, PartialEq, Eq, Clone)] pub struct SumTypeLayout { @@ -251,6 +287,17 @@ pub struct SumTypeLayout { pub payload_offset: u16, } +impl MemoryUsage for SumTypeLayout { + fn memory_usage(&self) -> usize { + let Self { + layout, + variants, + payload_offset, + } = self; + layout.memory_usage() + variants.memory_usage() + payload_offset.memory_usage() + } +} + impl HasLayout for SumTypeLayout { fn layout(&self) -> &Layout { &self.layout @@ -270,6 +317,15 @@ pub struct SumTypeVariantLayout { pub name: Option>, } +impl MemoryUsage for SumTypeVariantLayout { + fn memory_usage(&self) -> usize { + let Self { ty, name } = self; + ty.memory_usage() + name.memory_usage() + } +} + +impl MemoryUsage for PrimitiveType {} + impl HasLayout for PrimitiveType { fn layout(&self) -> &'static Layout { match self { @@ -301,6 +357,16 @@ pub enum VarLenType { Map(Box), } +impl MemoryUsage for VarLenType { + fn memory_usage(&self) -> usize { + match self { + VarLenType::String => 0, + VarLenType::Array(x) => x.memory_usage(), + VarLenType::Map(x) => x.memory_usage(), + } + } +} + /// The layout of var-len objects. Aligned at a `u16` which it has 2 of. const VAR_LEN_REF_LAYOUT: Layout = Layout { size: 4, align: 2 }; const _: () = assert!(VAR_LEN_REF_LAYOUT.size as usize == mem::size_of::()); diff --git a/crates/table/src/lib.rs b/crates/table/src/lib.rs index 3f25cc6530..b702af487e 100644 --- a/crates/table/src/lib.rs +++ b/crates/table/src/lib.rs @@ -25,5 +25,8 @@ pub mod row_type_visitor; pub mod table; pub mod var_len; +mod memory_usage; +pub use memory_usage::MemoryUsage; + #[doc(hidden)] // Used in tests and benchmarks. pub mod util; diff --git a/crates/table/src/memory_usage.rs b/crates/table/src/memory_usage.rs new file mode 100644 index 0000000000..dcd29dc3dc --- /dev/null +++ b/crates/table/src/memory_usage.rs @@ -0,0 +1,252 @@ +use std::hash::{BuildHasher, Hash}; +use std::mem; + +use spacetimedb_sats::{ + algebraic_value::Packed, i256, u256, AlgebraicType, AlgebraicValue, ArrayType, ArrayValue, MapType, ProductType, + ProductTypeElement, ProductValue, SumType, SumTypeVariant, SumValue, +}; + +/// For inspecting how much memory a value is using. +/// +/// This trait specifically measures heap memory. If you want to measure stack memory too, add +/// `mem::size_of_val()` to it. (This only really matters for the outermost type in a hierarchy.) +pub trait MemoryUsage { + /// The **heap** memory usage of this type. The default implementation returns 0. + #[inline(always)] + fn memory_usage(&self) -> usize { + 0 + } +} + +impl MemoryUsage for bool {} +impl MemoryUsage for u8 {} +impl MemoryUsage for u16 {} +impl MemoryUsage for u32 {} +impl MemoryUsage for u64 {} +impl MemoryUsage for u128 {} +impl MemoryUsage for u256 {} +impl MemoryUsage for usize {} +impl MemoryUsage for i8 {} +impl MemoryUsage for i16 {} +impl MemoryUsage for i32 {} +impl MemoryUsage for i64 {} +impl MemoryUsage for i128 {} +impl MemoryUsage for i256 {} +impl MemoryUsage for isize {} +impl MemoryUsage for f32 {} +impl MemoryUsage for f64 {} + +impl MemoryUsage for spacetimedb_sats::F32 {} +impl MemoryUsage for spacetimedb_sats::F64 {} + +impl MemoryUsage for Box { + fn memory_usage(&self) -> usize { + mem::size_of_val::(self) + T::memory_usage(self) + } +} + +impl MemoryUsage for std::sync::Arc { + fn memory_usage(&self) -> usize { + let refcounts = mem::size_of::() * 2; + refcounts + mem::size_of_val::(self) + T::memory_usage(self) + } +} + +impl MemoryUsage for std::rc::Rc { + fn memory_usage(&self) -> usize { + let refcounts = mem::size_of::() * 2; + refcounts + mem::size_of_val::(self) + T::memory_usage(self) + } +} + +impl MemoryUsage for [T] { + fn memory_usage(&self) -> usize { + self.iter().map(T::memory_usage).sum() + } +} + +impl MemoryUsage for str {} + +impl MemoryUsage for Option { + fn memory_usage(&self) -> usize { + self.as_ref().map_or(0, T::memory_usage) + } +} + +impl MemoryUsage for (A, B) { + fn memory_usage(&self) -> usize { + self.0.memory_usage() + self.1.memory_usage() + } +} + +impl MemoryUsage for String { + fn memory_usage(&self) -> usize { + self.capacity() + } +} + +impl MemoryUsage for Vec { + fn memory_usage(&self) -> usize { + self.capacity() * mem::size_of::() + self.iter().map(T::memory_usage).sum::() + } +} + +impl MemoryUsage + for spacetimedb_data_structures::map::HashMap +{ + fn memory_usage(&self) -> usize { + self.allocation_size() + + self + .iter() + .map(|(k, v)| k.memory_usage() + v.memory_usage()) + .sum::() + } +} + +impl MemoryUsage for std::collections::BTreeMap { + fn memory_usage(&self) -> usize { + self.iter() + .map(|(k, v)| k.memory_usage() + v.memory_usage()) + .sum::() + } +} + +impl MemoryUsage for smallvec::SmallVec +where + A::Item: MemoryUsage, +{ + fn memory_usage(&self) -> usize { + self.as_slice().memory_usage() + + if self.spilled() { + self.capacity() * mem::size_of::() + } else { + 0 + } + } +} + +impl MemoryUsage for spacetimedb_primitives::TableId {} +impl MemoryUsage for spacetimedb_primitives::SequenceId {} +impl MemoryUsage for spacetimedb_primitives::ConstraintId {} +impl MemoryUsage for spacetimedb_primitives::IndexId {} +impl MemoryUsage for spacetimedb_primitives::ColId {} +impl MemoryUsage for spacetimedb_primitives::ColList { + fn memory_usage(&self) -> usize { + self.heap_size() + } +} + +impl MemoryUsage for AlgebraicValue { + fn memory_usage(&self) -> usize { + match self { + AlgebraicValue::Sum(x) => x.memory_usage(), + AlgebraicValue::Product(x) => x.memory_usage(), + AlgebraicValue::Array(x) => x.memory_usage(), + AlgebraicValue::Map(x) => x.memory_usage(), + AlgebraicValue::String(x) => x.memory_usage(), + _ => 0, + } + } +} +impl MemoryUsage for SumValue { + fn memory_usage(&self) -> usize { + self.value.memory_usage() + } +} +impl MemoryUsage for ProductValue { + fn memory_usage(&self) -> usize { + self.elements.memory_usage() + } +} +impl MemoryUsage for ArrayValue { + fn memory_usage(&self) -> usize { + match self { + ArrayValue::Sum(v) => v.memory_usage(), + ArrayValue::Product(v) => v.memory_usage(), + ArrayValue::Bool(v) => v.memory_usage(), + ArrayValue::I8(v) => v.memory_usage(), + ArrayValue::U8(v) => v.memory_usage(), + ArrayValue::I16(v) => v.memory_usage(), + ArrayValue::U16(v) => v.memory_usage(), + ArrayValue::I32(v) => v.memory_usage(), + ArrayValue::U32(v) => v.memory_usage(), + ArrayValue::I64(v) => v.memory_usage(), + ArrayValue::U64(v) => v.memory_usage(), + ArrayValue::I128(v) => v.memory_usage(), + ArrayValue::U128(v) => v.memory_usage(), + ArrayValue::I256(v) => v.memory_usage(), + ArrayValue::U256(v) => v.memory_usage(), + ArrayValue::F32(v) => v.memory_usage(), + ArrayValue::F64(v) => v.memory_usage(), + ArrayValue::String(v) => v.memory_usage(), + ArrayValue::Array(v) => v.memory_usage(), + ArrayValue::Map(v) => v.memory_usage(), + } + } +} +impl MemoryUsage for AlgebraicType { + fn memory_usage(&self) -> usize { + match self { + AlgebraicType::Ref(_) => 0, + AlgebraicType::Sum(x) => x.memory_usage(), + AlgebraicType::Product(x) => x.memory_usage(), + AlgebraicType::Array(x) => x.memory_usage(), + AlgebraicType::Map(x) => x.memory_usage(), + AlgebraicType::String + | AlgebraicType::Bool + | AlgebraicType::I8 + | AlgebraicType::U8 + | AlgebraicType::I16 + | AlgebraicType::U16 + | AlgebraicType::I32 + | AlgebraicType::U32 + | AlgebraicType::I64 + | AlgebraicType::U64 + | AlgebraicType::I128 + | AlgebraicType::U128 + | AlgebraicType::I256 + | AlgebraicType::U256 + | AlgebraicType::F32 + | AlgebraicType::F64 => todo!(), + } + } +} +impl MemoryUsage for SumType { + fn memory_usage(&self) -> usize { + self.variants.memory_usage() + } +} +impl MemoryUsage for SumTypeVariant { + fn memory_usage(&self) -> usize { + self.name.memory_usage() + self.algebraic_type.memory_usage() + } +} +impl MemoryUsage for ProductType { + fn memory_usage(&self) -> usize { + self.elements.memory_usage() + } +} +impl MemoryUsage for ProductTypeElement { + fn memory_usage(&self) -> usize { + self.name.memory_usage() + self.algebraic_type.memory_usage() + } +} +impl MemoryUsage for ArrayType { + fn memory_usage(&self) -> usize { + self.elem_ty.memory_usage() + } +} +impl MemoryUsage for MapType { + fn memory_usage(&self) -> usize { + self.key_ty.memory_usage() + self.ty.memory_usage() + } +} + +impl MemoryUsage for Packed { + fn memory_usage(&self) -> usize { + { self.0 }.memory_usage() + } +} + +impl MemoryUsage for spacetimedb_lib::Address {} +impl MemoryUsage for spacetimedb_lib::Identity {} diff --git a/crates/table/src/page.rs b/crates/table/src/page.rs index 1b915a69ba..60cbdd9e74 100644 --- a/crates/table/src/page.rs +++ b/crates/table/src/page.rs @@ -39,7 +39,7 @@ use super::{ layout::MIN_ROW_SIZE, var_len::{is_granule_offset_aligned, VarLenGranule, VarLenGranuleHeader, VarLenMembers, VarLenRef}, }; -use crate::{fixed_bit_set::IterSet, static_assert_size, table::BlobNumBytes}; +use crate::{fixed_bit_set::IterSet, static_assert_size, table::BlobNumBytes, MemoryUsage}; use core::{mem, ops::ControlFlow, ptr}; use spacetimedb_lib::{de::Deserialize, ser::Serialize}; use thiserror::Error; @@ -63,6 +63,13 @@ struct FreeCellRef { next: PageOffset, } +impl MemoryUsage for FreeCellRef { + fn memory_usage(&self) -> usize { + let Self { next } = self; + next.memory_usage() + } +} + impl FreeCellRef { /// The sentinel for NULL cell references. const NIL: Self = Self { @@ -157,6 +164,21 @@ struct FixedHeader { fixed_row_size: Size, } +impl MemoryUsage for FixedHeader { + fn memory_usage(&self) -> usize { + let Self { + next_free, + last, + num_rows, + present_rows, + // MEMUSE: it's just a u16, ok to ignore + #[cfg(debug_assertions)] + fixed_row_size: _, + } = self; + next_free.memory_usage() + last.memory_usage() + num_rows.memory_usage() + present_rows.memory_usage() + } +} + #[cfg(debug_assertions)] static_assert_size!(FixedHeader, 18); @@ -250,6 +272,17 @@ struct VarHeader { first: PageOffset, } +impl MemoryUsage for VarHeader { + fn memory_usage(&self) -> usize { + let Self { + next_free, + freelist_len, + first, + } = self; + next_free.memory_usage() + freelist_len.memory_usage() + first.memory_usage() + } +} + static_assert_size!(VarHeader, 6); impl Default for VarHeader { @@ -293,6 +326,18 @@ struct PageHeader { unmodified_hash: Option, } +impl MemoryUsage for PageHeader { + fn memory_usage(&self) -> usize { + let Self { + fixed, + var, + // MEMUSE: no allocation, ok to ignore + unmodified_hash: _, + } = self; + fixed.memory_usage() + var.memory_usage() + } +} + static_assert_size!(PageHeader, PAGE_HEADER_SIZE); impl PageHeader { @@ -372,6 +417,12 @@ pub struct Page { row_data: [Byte; PageOffset::PAGE_END.idx()], } +impl MemoryUsage for Page { + fn memory_usage(&self) -> usize { + self.header.memory_usage() + } +} + static_assert_size!(Page, PAGE_SIZE); /// A mutable view of the fixed-len section of a [`Page`]. diff --git a/crates/table/src/pages.rs b/crates/table/src/pages.rs index 87333d9b49..4c982e1b1c 100644 --- a/crates/table/src/pages.rs +++ b/crates/table/src/pages.rs @@ -1,5 +1,7 @@ //! Provides [`Pages`], a page manager dealing with [`Page`]s as a collection. +use crate::MemoryUsage; + use super::blob_store::BlobStore; use super::indexes::{Bytes, PageIndex, PageOffset, RowPointer, Size}; use super::page::Page; @@ -40,6 +42,13 @@ pub struct Pages { non_full_pages: Vec, } +impl MemoryUsage for Pages { + fn memory_usage(&self) -> usize { + let Self { pages, non_full_pages } = self; + pages.memory_usage() + non_full_pages.memory_usage() + } +} + impl Pages { /// Is there space to allocate another page? pub fn can_allocate_new_page(&self) -> Result { diff --git a/crates/table/src/pointer_map.rs b/crates/table/src/pointer_map.rs index 15972555e5..01bcbb1eab 100644 --- a/crates/table/src/pointer_map.rs +++ b/crates/table/src/pointer_map.rs @@ -14,7 +14,7 @@ //! retrieval is probably no more than 100% slower. use super::indexes::{PageIndex, PageOffset, RowHash, RowPointer, SquashedOffset}; -use crate::static_assert_size; +use crate::{static_assert_size, MemoryUsage}; use core::{hint, slice}; use spacetimedb_data_structures::map::{ Entry, @@ -25,6 +25,8 @@ use spacetimedb_data_structures::map::{ #[derive(Clone, Copy, PartialEq, Eq, Debug)] struct ColliderSlotIndex(u32); +impl MemoryUsage for ColliderSlotIndex {} + impl ColliderSlotIndex { /// Returns a new slot index based on `idx`. fn new(idx: usize) -> Self { @@ -43,6 +45,8 @@ impl ColliderSlotIndex { #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)] struct PtrOrCollider(RowPointer); +impl MemoryUsage for PtrOrCollider {} + /// An unpacked representation of [`&mut PtrOrCollider`](PtrOrCollider). enum MapSlotRef<'map> { /// The hash has no collision and is associated to a single row pointer. @@ -149,6 +153,17 @@ pub struct PointerMap { emptied_collider_slots: Vec, } +impl MemoryUsage for PointerMap { + fn memory_usage(&self) -> usize { + let Self { + map, + colliders, + emptied_collider_slots, + } = self; + map.memory_usage() + colliders.memory_usage() + emptied_collider_slots.memory_usage() + } +} + static_assert_size!(PointerMap, 80); // Provides some type invariant checks. diff --git a/crates/table/src/row_type_visitor.rs b/crates/table/src/row_type_visitor.rs index 55b9da5bf9..99fa49c3f6 100644 --- a/crates/table/src/row_type_visitor.rs +++ b/crates/table/src/row_type_visitor.rs @@ -27,6 +27,8 @@ //! The `VarLenMembers` impl for `VarLenVisitorProgram` //! implements a simple interpreter loop for the var-len visitor bytecode. +use crate::MemoryUsage; + use super::{ indexes::{Byte, Bytes, PageOffset}, layout::{align_to, AlgebraicTypeLayout, HasLayout, ProductTypeLayout, RowTypeLayout, SumTypeLayout}, @@ -315,6 +317,8 @@ impl Insn { const FIXUP: Self = Self::Goto(u16::MAX); } +impl MemoryUsage for Insn {} + #[allow(clippy::disallowed_macros)] // This is for test code. pub fn dump_visitor_program(program: &VarLenVisitorProgram) { for (idx, insn) in program.insns.iter().enumerate() { @@ -354,6 +358,13 @@ pub struct VarLenVisitorProgram { insns: Arc<[Insn]>, } +impl MemoryUsage for VarLenVisitorProgram { + fn memory_usage(&self) -> usize { + let Self { insns } = self; + insns.memory_usage() + } +} + /// Evalutes the `program`, /// provided the `instr_ptr` as its program counter / intruction pointer, /// and a callback `read_tag` to extract a tag at the given offset, diff --git a/crates/table/src/table.rs b/crates/table/src/table.rs index c9ed1a0c99..7830c2b2a5 100644 --- a/crates/table/src/table.rs +++ b/crates/table/src/table.rs @@ -16,7 +16,7 @@ use super::{ read_column::{ReadColumn, TypeError}, row_hash::hash_row_in_page, row_type_visitor::{row_type_visitor, VarLenVisitorProgram}, - static_assert_size, + static_assert_size, MemoryUsage, }; use core::hash::{Hash, Hasher}; use core::ops::RangeBounds; @@ -41,6 +41,8 @@ use thiserror::Error; #[derive(Copy, Clone, PartialEq, Eq, Debug, Default, From, Add, Sub, AddAssign)] pub struct BlobNumBytes(usize); +impl MemoryUsage for BlobNumBytes {} + /// A database table containing the row schema, the rows, and indices. /// /// The table stores the rows into a page manager @@ -125,6 +127,42 @@ impl TableInner { static_assert_size!(Table, 256); +impl MemoryUsage for Table { + fn memory_usage(&self) -> usize { + let Self { + inner, + pointer_map, + indexes, + // MEMUSE: intentionally ignoring schema + schema: _, + squashed_offset, + row_count, + blob_store_bytes, + } = self; + inner.memory_usage() + + pointer_map.memory_usage() + + indexes.memory_usage() + + squashed_offset.memory_usage() + + row_count.memory_usage() + + blob_store_bytes.memory_usage() + } +} + +impl MemoryUsage for TableInner { + fn memory_usage(&self) -> usize { + let Self { + row_layout, + static_bsatn_layout, + visitor_prog, + pages, + } = self; + row_layout.memory_usage() + + static_bsatn_layout.memory_usage() + + visitor_prog.memory_usage() + + pages.memory_usage() + } +} + /// Various error that can happen on table insertion. #[derive(Error, Debug)] pub enum InsertError {