Data size SKU for billing

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

@@ -659,6 +659,39 @@ impl CommittedState {
         let blob_store = &mut self.blob_store;
         (table, blob_store)
     }
+
+    pub(super) fn report_data_size(&self, database_identity: Identity) {
+        use crate::db::db_metrics::data_size::DATA_SIZE_METRICS;
+
+        for (table_id, table) in &self.tables {
+            let table_name = &table.schema.table_name;
+            DATA_SIZE_METRICS
+                .data_size_table_num_rows
+                .with_label_values(&database_identity, &table_id.0, table_name)
+                .set(table.num_rows() as _);
+            DATA_SIZE_METRICS
+                .data_size_table_bytes_used_by_rows
+                .with_label_values(&database_identity, &table_id.0, table_name)
+                .set(table.bytes_used_by_rows() as _);
+            DATA_SIZE_METRICS
+                .data_size_table_num_rows_in_indexes
+                .with_label_values(&database_identity, &table_id.0, table_name)
+                .set(table.num_rows_in_indexes() as _);
+            DATA_SIZE_METRICS
+                .data_size_table_bytes_used_by_index_keys
+                .with_label_values(&database_identity, &table_id.0, table_name)
+                .set(table.bytes_used_by_index_keys() as _);
+        }
+
+        DATA_SIZE_METRICS
+            .data_size_blob_store_num_blobs
+            .with_label_values(&database_identity)
+            .set(self.blob_store.num_blobs() as _);
+        DATA_SIZE_METRICS
+            .data_size_blob_store_bytes_used_by_blobs
+            .with_label_values(&database_identity)
+            .set(self.blob_store.bytes_used_by_blobs() as _);
+    }
 }
 
 pub struct CommittedIndexIterWithDeletedMutTx<'a> {

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

@@ -684,6 +684,12 @@ pub(super) fn record_metrics(
                 .inc_by(deletes.len() as u64);
         }
     }
+
+    if let Some(committed_state) = committed_state {
+        // TODO(cleanliness,bikeshedding): Consider inlining `report_data_size` here,
+        // or moving the above metric writes into it, for consistency of organization.
+        committed_state.report_data_size(*db);
+    }
 }
 
 impl MutTx for Locking {

crates/core/src/db/db_metrics/data_size.rs

@@ -0,0 +1,42 @@
+use once_cell::sync::Lazy;
+use prometheus::IntGaugeVec;
+use spacetimedb_lib::Identity;
+use spacetimedb_metrics::metrics_group;
+
+metrics_group!(
+    #[non_exhaustive]
+    pub struct DbDataSize {
+        #[name = spacetime_data_size_table_num_rows]
+        #[help = "The number of rows in a table"]
+        #[labels(db: Identity, table_id: u32, table_name: str)]
+        pub data_size_table_num_rows: IntGaugeVec,
+
+        #[name = spacetime_data_size_bytes_used_by_rows]
+        #[help = "The number of bytes used by rows in pages in a table"]
+        #[labels(db: Identity, table_id: u32, table_name: str)]
+        pub data_size_table_bytes_used_by_rows: IntGaugeVec,
+
+        #[name = spacetime_data_size_table_num_rows_in_indexes]
+        #[help = "The number of rows stored in indexes in a table"]
+        // TODO: Consider partitioning by index ID or index name.
+        #[labels(db: Identity, table_id: u32, table_name: str)]
+        pub data_size_table_num_rows_in_indexes: IntGaugeVec,
+
+        #[name = spacetime_data_size_table_bytes_used_by_index_keys]
+        #[help = "The number of bytes used by keys stored in indexes in a table"]
+        #[labels(db: Identity, table_id: u32, table_name: str)]
+        pub data_size_table_bytes_used_by_index_keys: IntGaugeVec,
+
+        #[name = spacetime_data_size_blob_store_num_blobs]
+        #[help = "The number of large blobs stored in a database's blob store"]
+        #[labels(db: Identity)]
+        pub data_size_blob_store_num_blobs: IntGaugeVec,
+
+        #[name = spacetime_data_size_blob_store_bytes_used_by_blobs]
+        #[help = "The number of bytes used by large blobs stored in a database's blob store"]
+        #[labels(db: Identity)]
+        pub data_size_blob_store_bytes_used_by_blobs: IntGaugeVec,
+    }
+);
+
+pub static DATA_SIZE_METRICS: Lazy<DbDataSize> = Lazy::new(DbDataSize::new);

crates/core/src/db/db_metrics/mod.rs

@@ -5,6 +5,8 @@ use spacetimedb_lib::Identity;
 use spacetimedb_metrics::metrics_group;
 use spacetimedb_primitives::TableId;
 
+pub mod data_size;
+
 metrics_group!(
     #[non_exhaustive]
     pub struct DbMetrics {

crates/sats/src/proptest.rs

@@ -207,6 +207,18 @@ pub fn generate_typed_row() -> impl Strategy<Value = (ProductType, ProductValue)
     generate_row_type(0..=SIZE).prop_flat_map(|ty| (Just(ty.clone()), generate_product_value(ty)))
 }
 
+pub fn generate_typed_row_vec(
+    num_rows_min: usize,
+    num_rows_max: usize,
+) -> impl Strategy<Value = (ProductType, Vec<ProductValue>)> {
+    generate_row_type(0..=SIZE).prop_flat_map(move |ty| {
+        (
+            Just(ty.clone()),
+            vec(generate_product_value(ty), num_rows_min..num_rows_max),
+        )
+    })
+}
+
 /// Generates a type `ty` and a value typed at `ty`.
 pub fn generate_typed_value() -> impl Strategy<Value = (AlgebraicType, AlgebraicValue)> {
     generate_algebraic_type().prop_flat_map(|ty| (Just(ty.clone()), generate_algebraic_value(ty)))

crates/table/src/blob_store.rs

@@ -104,6 +104,27 @@ pub trait BlobStore: Sync {
     ///
     /// Used when capturing a snapshot.
     fn iter_blobs(&self) -> BlobsIter<'_>;
+
+    /// Returns the amount of memory in bytes used by blobs in this `BlobStore`.
+    ///
+    /// Duplicate blobs are counted a number of times equal to their refcount.
+    /// This is in order to preserve the property that inserting a large blob
+    /// causes this quantity to increase by that blob's size,
+    /// and deleting a large blob causes it to decrease the same amount.
+    fn bytes_used_by_blobs(&self) -> u64 {
+        self.iter_blobs()
+            .map(|(_, uses, data)| data.len() as u64 * uses as u64)
+            .sum()
+    }
+
+    /// Returns the number of blobs, or more precisely, blob-usages, recorded in this `BlobStore`.
+    ///
+    /// Duplicate blobs are counted a number of times equal to their refcount.
+    /// This is in order to preserve the property that inserting a large blob
+    /// causes this quantity to increase by 1, and deleting a large blob causes it to decrease by 1.
+    fn num_blobs(&self) -> u64 {
+        self.iter_blobs().map(|(_, uses, _)| uses as u64).sum()
+    }
 }
 
 /// A blob store that panics on all operations.