Clear DB size gauges when shutting a module down (#2259)

Author: jsdt

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

@@ -682,23 +682,23 @@ impl CommittedState {
     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 {
+        for (_, 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)
+                .with_label_values(&database_identity, 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)
+                .with_label_values(&database_identity, 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)
+                .with_label_values(&database_identity, 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)
+                .with_label_values(&database_identity, table_name)
                 .set(table.bytes_used_by_index_keys() as _);
         }
 

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

@@ -3,28 +3,30 @@ use prometheus::IntGaugeVec;
 use spacetimedb_lib::Identity;
 use spacetimedb_metrics::metrics_group;
 
+use crate::worker_metrics::WORKER_METRICS;
+
 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)]
+        #[labels(db: Identity, 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)]
+        #[labels(db: Identity, 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)]
+        #[labels(db: Identity, 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)]
+        #[labels(db: Identity, table_name: str)]
         pub data_size_table_bytes_used_by_index_keys: IntGaugeVec,
 
         #[name = spacetime_data_size_blob_store_num_blobs]
@@ -40,3 +42,32 @@ metrics_group!(
 );
 
 pub static DATA_SIZE_METRICS: Lazy<DbDataSize> = Lazy::new(DbDataSize::new);
+
+// Remove all gauges associated with a database.
+// This is useful if a database is being deleted.
+pub fn remove_database_gauges<'a, I>(db: &Identity, table_names: I)
+where
+    I: IntoIterator<Item = &'a str>,
+{
+    // Remove the per-table gauges.
+    for table_name in table_names {
+        let _ = DATA_SIZE_METRICS
+            .data_size_table_num_rows
+            .remove_label_values(db, table_name);
+        let _ = DATA_SIZE_METRICS
+            .data_size_table_bytes_used_by_rows
+            .remove_label_values(db, table_name);
+        let _ = DATA_SIZE_METRICS
+            .data_size_table_num_rows_in_indexes
+            .remove_label_values(db, table_name);
+        let _ = DATA_SIZE_METRICS
+            .data_size_table_bytes_used_by_index_keys
+            .remove_label_values(db, table_name);
+    }
+    // Remove the per-db gauges.
+    let _ = DATA_SIZE_METRICS.data_size_blob_store_num_blobs.remove_label_values(db);
+    let _ = DATA_SIZE_METRICS
+        .data_size_blob_store_bytes_used_by_blobs
+        .remove_label_values(db);
+    let _ = WORKER_METRICS.wasm_memory_bytes.remove_label_values(db);
+}

crates/core/src/host/host_controller.rs

@@ -3,10 +3,10 @@ use super::scheduler::SchedulerStarter;
 use super::wasmtime::WasmtimeRuntime;
 use super::{Scheduler, UpdateDatabaseResult};
 use crate::database_logger::DatabaseLogger;
-use crate::db;
 use crate::db::datastore::traits::Program;
 use crate::db::db_metrics::DB_METRICS;
 use crate::db::relational_db::{self, DiskSizeFn, RelationalDB, Txdata};
+use crate::db::{self, db_metrics};
 use crate::energy::{EnergyMonitor, EnergyQuanta};
 use crate::messages::control_db::{Database, HostType};
 use crate::module_host_context::ModuleCreationContext;
@@ -25,6 +25,7 @@ use spacetimedb_lib::hash_bytes;
 use spacetimedb_paths::server::{ReplicaDir, ServerDataDir};
 use spacetimedb_sats::hash::Hash;
 use std::future::Future;
+use std::ops::Deref;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tokio::sync::{watch, OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock as AsyncRwLock};
@@ -419,6 +420,8 @@ impl HostController {
             if let Some(host) = lock.write_owned().await.take() {
                 let module = host.module.borrow().clone();
                 module.exit().await;
+                let table_names = module.info().module_def.tables().map(|t| t.name.deref());
+                db_metrics::data_size::remove_database_gauges(&module.info().database_identity, table_names);
             }
         }
 

crates/metrics/src/typed_prometheus.rs

@@ -82,6 +82,11 @@ macro_rules! metrics_histogram_vec {
                 use $crate::typed_prometheus::AsPrometheusLabel as _;
                 self.0.with_label_values(&[ $($labels.as_prometheus_str().as_ref()),+ ])
             }
+
+            pub fn remove_label_values(&self, $($labels: &$labelty),+) -> prometheus::Result<()> {
+                use $crate::typed_prometheus::AsPrometheusLabel as _;
+                self.0.remove_label_values(&[ $($labels.as_prometheus_str().as_ref()),+ ])
+            }
         }
 
         impl prometheus::core::Collector for $name {
@@ -111,6 +116,10 @@ macro_rules! metrics_vec {
                 use $crate::typed_prometheus::AsPrometheusLabel as _;
                 self.0.with_label_values(&[ $($labels.as_prometheus_str().as_ref()),+ ])
             }
+            pub fn remove_label_values(&self, $($labels: &$labelty),+) -> prometheus::Result<()> {
+                use $crate::typed_prometheus::AsPrometheusLabel as _;
+                self.0.remove_label_values(&[ $($labels.as_prometheus_str().as_ref()),+ ])
+            }
         }
 
         impl prometheus::core::Collector for $name {