perf(505): Track the number of rows in each table (#525)

Closes #505.

crates/core/src/db/commit_log.rs

@@ -111,33 +111,32 @@ impl CommitLog {
         let mut unwritten_commit = self.unwritten_commit.lock().unwrap();
         let mut writes = Vec::with_capacity(tx_data.records.len());
 
-        let rows_inserted = &DB_METRICS.rdb_num_rows_inserted;
-        let rows_deleted = &DB_METRICS.rdb_num_rows_deleted;
+        let txn_type = &ctx.txn_type();
+        let db = &ctx.database();
+        let reducer = &ctx.reducer_name().unwrap_or_default();
 
         for record in &tx_data.records {
             let table_id: u32 = record.table_id.into();
 
             let operation = match record.op {
                 TxOp::Insert(_) => {
                     // Increment rows inserted metric
-                    let metric = rows_inserted.with_label_values(
-                        &ctx.txn_type(),
-                        &ctx.database(),
-                        ctx.reducer_name().unwrap_or_default(),
-                        &table_id,
-                    );
-                    metric.inc();
+                    DB_METRICS
+                        .rdb_num_rows_inserted
+                        .with_label_values(txn_type, db, reducer, &table_id)
+                        .inc();
+                    // Increment table rows gauge
+                    DB_METRICS.rdb_num_table_rows.with_label_values(db, &table_id).inc();
                     Operation::Insert
                 }
                 TxOp::Delete => {
                     // Increment rows deleted metric
-                    let metric = rows_deleted.with_label_values(
-                        &ctx.txn_type(),
-                        &ctx.database(),
-                        ctx.reducer_name().unwrap_or_default(),
-                        &table_id,
-                    );
-                    metric.inc();
+                    DB_METRICS
+                        .rdb_num_rows_deleted
+                        .with_label_values(txn_type, db, reducer, &table_id)
+                        .inc();
+                    // Decrement table rows gauge
+                    DB_METRICS.rdb_num_table_rows.with_label_values(db, &table_id).dec();
                     Operation::Delete
                 }
             };

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

@@ -364,6 +364,12 @@ impl Inner {
     }
 
     fn bootstrap_system_table(&mut self, schema: TableSchema) -> Result<(), DBError> {
+        // Reset the row count metric for this system table
+        DB_METRICS
+            .rdb_num_table_rows
+            .with_label_values(&self.database_address, &schema.table_id.into())
+            .set(0);
+
         let table_id = schema.table_id;
 
         // Insert the table row into st_tables, creating st_tables if it's missing
@@ -380,6 +386,12 @@ impl Inner {
         let data_key = row.to_data_key();
         st_tables.rows.insert(RowId(data_key), row);
 
+        // Increment row count for st_tables
+        DB_METRICS
+            .rdb_num_table_rows
+            .with_label_values(&self.database_address, &ST_TABLES_ID.into())
+            .inc();
+
         // Insert the columns into st_columns
         let first_col_id = schema.columns.first().unwrap().col_id;
         for (i, col) in schema.columns.into_iter().enumerate() {
@@ -399,6 +411,11 @@ impl Inner {
                     self.committed_state
                         .get_or_create_table(ST_COLUMNS_ID, &ST_COLUMNS_ROW_TYPE, &st_columns_schema());
                 st_columns.rows.insert(RowId(data_key), row);
+                // Increment row count for st_columns
+                DB_METRICS
+                    .rdb_num_table_rows
+                    .with_label_values(&self.database_address, &ST_COLUMNS_ID.into())
+                    .inc();
             }
 
             // If any columns are auto incrementing, we need to create a sequence
@@ -435,6 +452,11 @@ impl Inner {
                 let row = ProductValue::from(row);
                 let data_key = row.to_data_key();
                 st_sequences.rows.insert(RowId(data_key), row);
+                // Increment row count for st_sequences
+                DB_METRICS
+                    .rdb_num_table_rows
+                    .with_label_values(&self.database_address, &ST_SEQUENCES_ID.into())
+                    .inc();
             }
         }
 
@@ -463,6 +485,12 @@ impl Inner {
             let data_key = row.to_data_key();
             st_constraints.rows.insert(RowId(data_key), row);
 
+            // Increment row count for st_constraints
+            DB_METRICS
+                .rdb_num_table_rows
+                .with_label_values(&self.database_address, &ST_CONSTRAINTS_ID.into())
+                .inc();
+
             //Check if add an index:
             match constraint.kind {
                 x if x.is_unique() => IndexSchema {
@@ -498,6 +526,12 @@ impl Inner {
             let row = ProductValue::from(row);
             let data_key = row.to_data_key();
             st_indexes.rows.insert(RowId(data_key), row);
+
+            // Increment row count for st_indexes
+            DB_METRICS
+                .rdb_num_table_rows
+                .with_label_values(&self.database_address, &ST_INDEXES_ID.into())
+                .inc();
         }
 
         Ok(())
@@ -1640,6 +1674,10 @@ impl Locking {
             match write.operation {
                 Operation::Delete => {
                     Self::table_rows(&mut inner, table_id, schema, row_type).remove(&RowId(write.data_key));
+                    DB_METRICS
+                        .rdb_num_table_rows
+                        .with_label_values(&inner.database_address, &table_id.into())
+                        .dec();
                 }
                 Operation::Insert => {
                     let product_value = match write.data_key {
@@ -1657,6 +1695,10 @@ impl Locking {
                     };
                     Self::table_rows(&mut inner, table_id, schema, row_type)
                         .insert(RowId(write.data_key), product_value);
+                    DB_METRICS
+                        .rdb_num_table_rows
+                        .with_label_values(&inner.database_address, &table_id.into())
+                        .inc();
                 }
             }
         }

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

@@ -1,6 +1,6 @@
 use crate::{execution_context::TransactionType, host::AbiCall, util::typed_prometheus::metrics_group};
 use once_cell::sync::Lazy;
-use prometheus::{Histogram, HistogramVec, IntCounterVec};
+use prometheus::{Histogram, HistogramVec, IntCounterVec, IntGaugeVec};
 use spacetimedb_lib::Address;
 
 metrics_group!(
@@ -56,6 +56,11 @@ metrics_group!(
         #[labels(db: Address, reducer: str)]
         pub scheduled_reducer_delay_sec: HistogramVec,
 
+        #[name = spacetime_num_table_rows]
+        #[help = "The number of rows in a table"]
+        #[labels(db: Address, table_id: u32)]
+        pub rdb_num_table_rows: IntGaugeVec,
+
         #[name = spacetime_num_rows_inserted_cumulative]
         #[help = "The cumulative number of rows inserted into a table"]
         #[labels(txn_type: TransactionType, db: Address, reducer: str, table_id: u32)]

crates/core/src/db/relational_db.rs

@@ -395,7 +395,12 @@ impl RelationalDB {
             .rdb_drop_table_time
             .with_label_values(&table_id.0)
             .start_timer();
-        self.inner.drop_table_mut_tx(tx, table_id)
+        self.inner.drop_table_mut_tx(tx, table_id).map(|_| {
+            DB_METRICS
+                .rdb_num_table_rows
+                .with_label_values(&self.address, &table_id.into())
+                .set(0)
+        })
     }
 
     /// Rename a table.