perf(461): Record metrics for rows [inserted|deleted] on commit (#462)

Closes #456.
Closes #461.

Define an execution context for database operations.
Store the database id in RelationalDB.
Record metrics for rows [inserted|deleted] on commit.

crates/bench/src/spacetime_raw.rs

@@ -3,8 +3,11 @@ use crate::{
     schemas::{table_name, BenchTable, IndexStrategy},
     ResultBench,
 };
-use spacetimedb::db::datastore::traits::{IndexDef, TableDef};
 use spacetimedb::db::relational_db::{open_db, RelationalDB};
+use spacetimedb::{
+    db::datastore::traits::{IndexDef, TableDef},
+    execution_context::ExecutionContext,
+};
 use spacetimedb_lib::sats::AlgebraicValue;
 use spacetimedb_primitives::{ColId, TableId};
 use std::hint::black_box;
@@ -38,7 +41,7 @@ impl BenchDatabase for SpacetimeRaw {
 
     fn create_table<T: BenchTable>(&mut self, index_strategy: IndexStrategy) -> ResultBench<Self::TableId> {
         let name = table_name::<T>(index_strategy);
-        self.db.with_auto_commit(|tx| {
+        self.db.with_auto_commit(&ExecutionContext::default(), |tx| {
             let table_def = TableDef::from(T::product_type());
             let table_id = self.db.create_table(tx, table_def)?;
             self.db.rename_table(tx, table_id, &name)?;
@@ -63,30 +66,31 @@ impl BenchDatabase for SpacetimeRaw {
     }
 
     fn clear_table(&mut self, table_id: &Self::TableId) -> ResultBench<()> {
-        self.db.with_auto_commit(|tx| {
+        self.db.with_auto_commit(&ExecutionContext::default(), |tx| {
             self.db.clear_table(tx, *table_id)?;
             Ok(())
         })
     }
 
     fn count_table(&mut self, table_id: &Self::TableId) -> ResultBench<u32> {
-        self.db
-            .with_auto_commit(|tx| Ok(self.db.iter(tx, *table_id)?.map(|_| 1u32).sum()))
+        self.db.with_auto_commit(&ExecutionContext::default(), |tx| {
+            Ok(self.db.iter(tx, *table_id)?.map(|_| 1u32).sum())
+        })
     }
 
     fn empty_transaction(&mut self) -> ResultBench<()> {
-        self.db.with_auto_commit(|_tx| Ok(()))
+        self.db.with_auto_commit(&ExecutionContext::default(), |_tx| Ok(()))
     }
 
     fn insert<T: BenchTable>(&mut self, table_id: &Self::TableId, row: T) -> ResultBench<()> {
-        self.db.with_auto_commit(|tx| {
+        self.db.with_auto_commit(&ExecutionContext::default(), |tx| {
             self.db.insert(tx, *table_id, row.into_product_value())?;
             Ok(())
         })
     }
 
     fn insert_bulk<T: BenchTable>(&mut self, table_id: &Self::TableId, rows: Vec<T>) -> ResultBench<()> {
-        self.db.with_auto_commit(|tx| {
+        self.db.with_auto_commit(&ExecutionContext::default(), |tx| {
             for row in rows {
                 self.db.insert(tx, *table_id, row.into_product_value())?;
             }
@@ -95,7 +99,7 @@ impl BenchDatabase for SpacetimeRaw {
     }
 
     fn iterate(&mut self, table_id: &Self::TableId) -> ResultBench<()> {
-        self.db.with_auto_commit(|tx| {
+        self.db.with_auto_commit(&ExecutionContext::default(), |tx| {
             for row in self.db.iter(tx, *table_id)? {
                 black_box(row);
             }
@@ -110,7 +114,7 @@ impl BenchDatabase for SpacetimeRaw {
         value: AlgebraicValue,
     ) -> ResultBench<()> {
         let col: ColId = column_index.into();
-        self.db.with_auto_commit(|tx| {
+        self.db.with_auto_commit(&ExecutionContext::default(), |tx| {
             for row in self.db.iter_by_col_eq(tx, *table_id, col, value)? {
                 black_box(row);
             }

crates/core/src/database_instance_context.rs

@@ -84,7 +84,15 @@ impl DatabaseInstanceContext {
             address,
             logger: Arc::new(Mutex::new(DatabaseLogger::open(log_path))),
             relational_db: Arc::new(
-                RelationalDB::open(db_path, message_log, odb, address, config.fsync != FsyncPolicy::Never).unwrap(),
+                RelationalDB::open(
+                    db_path,
+                    message_log,
+                    odb,
+                    database_id,
+                    address,
+                    config.fsync != FsyncPolicy::Never,
+                )
+                .unwrap(),
             ),
             publisher_address,
         })

crates/core/src/db/commit_log.rs

@@ -7,12 +7,14 @@ use super::{
 use crate::{
     db::{
         datastore::{locking_tx_datastore::RowId, traits::TxOp},
+        db_metrics::DB_METRICS,
         messages::{
             transaction::Transaction,
             write::{Operation, Write},
         },
     },
     error::DBError,
+    execution_context::ExecutionContext,
 };
 
 use anyhow::Context;
@@ -52,13 +54,18 @@ impl CommitLog {
     ///
     /// Returns `Some(n_bytes_written)` if `commit_result` was persisted, `None` if it doesn't have bytes to write.
     #[tracing::instrument(skip_all)]
-    pub fn append_tx<D>(&self, tx_data: &TxData, datastore: &D) -> Result<Option<usize>, DBError>
+    pub fn append_tx<D>(
+        &self,
+        ctx: &ExecutionContext,
+        tx_data: &TxData,
+        datastore: &D,
+    ) -> Result<Option<usize>, DBError>
     where
         D: MutTxDatastore<RowId = RowId>,
     {
         if let Some(mlog) = &self.mlog {
             let mut mlog = mlog.lock().unwrap();
-            self.generate_commit(tx_data, datastore)
+            self.generate_commit(ctx, tx_data, datastore)
                 .as_deref()
                 .map(|bytes| self.append_commit_bytes(&mut mlog, bytes))
                 .transpose()
@@ -88,7 +95,12 @@ impl CommitLog {
         Ok(commit.len())
     }
 
-    fn generate_commit<D: MutTxDatastore<RowId = RowId>>(&self, tx_data: &TxData, _datastore: &D) -> Option<Vec<u8>> {
+    fn generate_commit<D: MutTxDatastore<RowId = RowId>>(
+        &self,
+        ctx: &ExecutionContext,
+        tx_data: &TxData,
+        _datastore: &D,
+    ) -> Option<Vec<u8>> {
         // We are not creating a commit for empty transactions.
         // The reason for this is that empty transactions get encoded as 0 bytes,
         // so a commit containing an empty transaction contains no useful information.
@@ -97,18 +109,46 @@ impl CommitLog {
         }
 
         let mut unwritten_commit = self.unwritten_commit.lock().unwrap();
-        let writes = tx_data
-            .records
-            .iter()
-            .map(|record| Write {
-                operation: match record.op {
-                    TxOp::Insert(_) => Operation::Insert,
-                    TxOp::Delete => Operation::Delete,
-                },
-                set_id: record.table_id.0,
+        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;
+
+        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_id(),
+                        &ctx.reducer_id().unwrap_or(0),
+                        &table_id,
+                    );
+                    metric.inc();
+                    Operation::Insert
+                }
+                TxOp::Delete => {
+                    // Increment rows deleted metric
+                    let metric = rows_deleted.with_label_values(
+                        &ctx.txn_type(),
+                        &ctx.database_id(),
+                        &ctx.reducer_id().unwrap_or(0),
+                        &table_id,
+                    );
+                    metric.inc();
+                    Operation::Delete
+                }
+            };
+
+            writes.push(Write {
+                operation,
+                set_id: table_id,
                 data_key: record.key,
             })
-            .collect();
+        }
+
         let transaction = Transaction { writes };
         unwritten_commit.transactions.push(Arc::new(transaction));
 

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

@@ -1,6 +1,6 @@
-use crate::util::typed_prometheus::metrics_group;
+use crate::{execution_context::TransactionType, util::typed_prometheus::metrics_group};
 use once_cell::sync::Lazy;
-use prometheus::{Histogram, HistogramVec};
+use prometheus::{Histogram, HistogramVec, IntCounterVec};
 
 metrics_group!(
     #[non_exhaustive]
@@ -49,6 +49,16 @@ metrics_group!(
         #[help = "The time spent deleting values in a set from a table"]
         #[labels(table_id: u32)]
         pub rdb_delete_by_rel_time: HistogramVec,
+
+        #[name = spacetime_num_rows_inserted]
+        #[help = "The number of rows inserted into a table"]
+        #[labels(txn_type: TransactionType, database_id: u64, reducer_id: u64, table_id: u32)]
+        pub rdb_num_rows_inserted: IntCounterVec,
+
+        #[name = spacetime_num_rows_deleted]
+        #[help = "The number of rows deleted from a table"]
+        #[labels(txn_type: TransactionType, database_id: u64, reducer_id: u64, table_id: u32)]
+        pub rdb_num_rows_deleted: IntCounterVec,
     }
 );
 

crates/core/src/db/relational_db.rs

@@ -13,6 +13,7 @@ use crate::db::messages::commit::Commit;
 use crate::db::ostorage::hashmap_object_db::HashMapObjectDB;
 use crate::db::ostorage::ObjectDB;
 use crate::error::{DBError, DatabaseError, IndexError, TableError};
+use crate::execution_context::ExecutionContext;
 use crate::hash::Hash;
 use fs2::FileExt;
 use nonempty::NonEmpty;
@@ -44,6 +45,7 @@ pub struct RelationalDB {
     pub(crate) inner: Locking,
     commit_log: CommitLog,
     _lock: Arc<File>,
+    id: u64,
 }
 
 impl DataRow for RelationalDB {
@@ -66,6 +68,7 @@ impl RelationalDB {
         root: impl AsRef<Path>,
         message_log: Option<Arc<Mutex<MessageLog>>>,
         odb: Arc<Mutex<Box<dyn ObjectDB + Send>>>,
+        id: u64,
         address: Address,
         fsync: bool,
     ) -> Result<Self, DBError> {
@@ -159,12 +162,18 @@ impl RelationalDB {
             inner: datastore,
             commit_log,
             _lock: Arc::new(lock),
+            id,
         };
 
         log::trace!("[{}] DATABASE: OPENED", address);
         Ok(db)
     }
 
+    /// Returns the id for this database
+    pub fn id(&self) -> u64 {
+        self.id
+    }
+
     /// Obtain a read-only view of this database's [`CommitLog`].
     pub fn commit_log(&self) -> CommitLogView {
         CommitLogView::from(&self.commit_log)
@@ -259,10 +268,10 @@ impl RelationalDB {
     }
 
     #[tracing::instrument(skip_all)]
-    pub fn commit_tx(&self, tx: MutTxId) -> Result<Option<(TxData, Option<usize>)>, DBError> {
+    pub fn commit_tx(&self, ctx: &ExecutionContext, tx: MutTxId) -> Result<Option<(TxData, Option<usize>)>, DBError> {
         log::trace!("COMMIT TX");
         if let Some(tx_data) = self.inner.commit_mut_tx(tx)? {
-            let bytes_written = self.commit_log.append_tx(&tx_data, &self.inner)?;
+            let bytes_written = self.commit_log.append_tx(ctx, &tx_data, &self.inner)?;
             return Ok(Some((tx_data, bytes_written)));
         }
         Ok(None)
@@ -296,14 +305,14 @@ impl RelationalDB {
     ///     Ok(())
     /// })?;
     /// ```
-    pub fn with_auto_commit<F, A, E>(&self, f: F) -> Result<A, E>
+    pub fn with_auto_commit<F, A, E>(&self, ctx: &ExecutionContext, f: F) -> Result<A, E>
     where
         F: FnOnce(&mut MutTxId) -> Result<A, E>,
         E: From<DBError>,
     {
         let mut tx = self.begin_tx();
         let res = f(&mut tx);
-        self.finish_tx(tx, res)
+        self.finish_tx(ctx, tx, res)
     }
 
     /// Run a fallible function in a transaction, rolling it back if the
@@ -343,14 +352,14 @@ impl RelationalDB {
 
     /// Perform the transactional logic for the `tx` according to the `res`
     #[tracing::instrument(skip_all)]
-    pub fn finish_tx<A, E>(&self, tx: MutTxId, res: Result<A, E>) -> Result<A, E>
+    pub fn finish_tx<A, E>(&self, ctx: &ExecutionContext, tx: MutTxId, res: Result<A, E>) -> Result<A, E>
     where
         E: From<DBError>,
     {
         if res.is_err() {
             self.rollback_tx(tx);
         } else {
-            match self.commit_tx(tx).map_err(E::from)? {
+            match self.commit_tx(ctx, tx).map_err(E::from)? {
                 Some(_) => (),
                 None => panic!("TODO: retry?"),
             }
@@ -613,7 +622,7 @@ pub fn open_db(path: impl AsRef<Path>, in_memory: bool, fsync: bool) -> Result<R
         Some(Arc::new(Mutex::new(MessageLog::open(path.join("mlog"))?)))
     };
     let odb = Arc::new(Mutex::new(make_default_ostorage(in_memory, path.join("odb"))?));
-    let stdb = RelationalDB::open(path, mlog, odb, Address::zero(), fsync)?;
+    let stdb = RelationalDB::open(path, mlog, odb, 0, Address::zero(), fsync)?;
 
     Ok(stdb)
 }
@@ -658,6 +667,7 @@ mod tests {
     use crate::db::datastore::traits::TableDef;
     use crate::db::message_log::MessageLog;
     use crate::db::relational_db::{open_db, ST_TABLES_ID};
+    use crate::execution_context::ExecutionContext;
 
     use super::RelationalDB;
     use crate::db::relational_db::make_default_ostorage;
@@ -704,7 +714,7 @@ mod tests {
         let mut schema = TableDef::from(ProductType::from([("my_col", AlgebraicType::I32)]));
         schema.table_name = "MyTable".to_string();
         stdb.create_table(&mut tx, schema)?;
-        stdb.commit_tx(tx)?;
+        stdb.commit_tx(&ExecutionContext::default(), tx)?;
 
         Ok(())
     }
@@ -719,7 +729,7 @@ mod tests {
         schema.table_name = "MyTable".to_string();
         stdb.create_table(&mut tx, schema)?;
 
-        stdb.commit_tx(tx)?;
+        stdb.commit_tx(&ExecutionContext::default(), tx)?;
 
         let mlog = Some(Arc::new(Mutex::new(MessageLog::open(tmp_dir.path().join("mlog"))?)));
         let in_memory = false;
@@ -728,7 +738,7 @@ mod tests {
             tmp_dir.path().join("odb"),
         )?));
 
-        match RelationalDB::open(tmp_dir.path(), mlog, odb, Address::zero(), true) {
+        match RelationalDB::open(tmp_dir.path(), mlog, odb, 0, Address::zero(), true) {
             Ok(_) => {
                 panic!("Allowed to open database twice")
             }
@@ -821,7 +831,7 @@ mod tests {
         stdb.insert(&mut tx, table_id, product![AlgebraicValue::I32(-1)])?;
         stdb.insert(&mut tx, table_id, product![AlgebraicValue::I32(0)])?;
         stdb.insert(&mut tx, table_id, product![AlgebraicValue::I32(1)])?;
-        stdb.commit_tx(tx)?;
+        stdb.commit_tx(&ExecutionContext::default(), tx)?;
 
         let tx = stdb.begin_tx();
         let mut rows = stdb
@@ -871,7 +881,7 @@ mod tests {
         stdb.insert(&mut tx, table_id, product![AlgebraicValue::I32(-1)])?;
         stdb.insert(&mut tx, table_id, product![AlgebraicValue::I32(0)])?;
         stdb.insert(&mut tx, table_id, product![AlgebraicValue::I32(1)])?;
-        stdb.commit_tx(tx)?;
+        stdb.commit_tx(&ExecutionContext::default(), tx)?;
 
         let tx = stdb.begin_tx();
         let mut rows = stdb
@@ -910,7 +920,7 @@ mod tests {
         let mut schema = TableDef::from(ProductType::from([("my_col", AlgebraicType::I32)]));
         schema.table_name = "MyTable".to_string();
         let table_id = stdb.create_table(&mut tx, schema)?;
-        stdb.commit_tx(tx)?;
+        stdb.commit_tx(&ExecutionContext::default(), tx)?;
 
         let mut tx = stdb.begin_tx();
         stdb.insert(&mut tx, table_id, product![AlgebraicValue::I32(-1)])?;
@@ -1031,7 +1041,7 @@ mod tests {
 
         assert_eq!(rows, vec![1]);
 
-        stdb.commit_tx(tx)?;
+        stdb.commit_tx(&ExecutionContext::default(), tx)?;
         drop(stdb);
 
         dbg!("reopen...");