Count bytes_written for updates (#2616)

Author: jsdt

crates/core/src/db/relational_db.rs

@@ -884,6 +884,7 @@ impl RelationalDB {
         name: &str,
         schema: &[(&str, AlgebraicType)],
         indexes: &[ColList],
+        unique_constraints: &[ColList],
         access: StAccess,
     ) -> Result<TableId, DBError> {
         let mut module_def_builder = RawModuleDefV9Builder::new();
@@ -895,6 +896,9 @@ impl RelationalDB {
         for columns in indexes {
             table_builder = table_builder.with_index(btree(columns.clone()), "accessor_name_doesnt_matter");
         }
+        for columns in unique_constraints {
+            table_builder = table_builder.with_unique_constraint(columns.clone());
+        }
         table_builder.finish();
         let module_def: ModuleDef = module_def_builder.finish().try_into()?;
 
@@ -915,7 +919,7 @@ impl RelationalDB {
         access: StAccess,
     ) -> Result<TableId, DBError> {
         let indexes: Vec<ColList> = indexes.iter().map(|col_id| (*col_id).into()).collect();
-        self.create_table_for_test_with_the_works(name, schema, &indexes[..], access)
+        self.create_table_for_test_with_the_works(name, schema, &indexes[..], &[], access)
     }
 
     pub fn create_table_for_test(
@@ -933,7 +937,7 @@ impl RelationalDB {
         schema: &[(&str, AlgebraicType)],
         idx_cols: ColList,
     ) -> Result<TableId, DBError> {
-        self.create_table_for_test_with_the_works(name, schema, &[idx_cols], StAccess::Public)
+        self.create_table_for_test_with_the_works(name, schema, &[idx_cols], &[], StAccess::Public)
     }
 
     pub fn create_table_for_test_mix_indexes(
@@ -949,7 +953,7 @@ impl RelationalDB {
             .chain(std::iter::once(idx_cols_multi))
             .collect();
 
-        self.create_table_for_test_with_the_works(name, schema, &indexes[..], StAccess::Public)
+        self.create_table_for_test_with_the_works(name, schema, &indexes[..], &[], StAccess::Public)
     }
 
     pub fn drop_table(&self, tx: &mut MutTx, table_id: TableId) -> Result<(), DBError> {

crates/core/src/host/instance_env.rs

@@ -295,6 +295,8 @@ impl InstanceEnv {
         if update_flags.is_scheduler_table {
             self.schedule_row(stdb, tx, table_id, row_ptr)?;
         }
+        tx.metrics.bytes_written += buffer.len();
+        tx.metrics.rows_updated += 1;
 
         Ok(row_len)
     }
@@ -490,14 +492,6 @@ impl From<GetTxError> for NodesError {
 mod test {
     use std::{ops::Bound, sync::Arc};
 
-    use anyhow::{anyhow, Result};
-    use parking_lot::RwLock;
-    use spacetimedb_lib::{bsatn::to_vec, AlgebraicType, AlgebraicValue, Hash, Identity, ProductValue};
-    use spacetimedb_paths::{server::ModuleLogsDir, FromPathUnchecked};
-    use spacetimedb_primitives::{IndexId, TableId};
-    use spacetimedb_sats::product;
-    use tempfile::TempDir;
-
     use crate::{
         database_logger::DatabaseLogger,
         db::{
@@ -512,6 +506,14 @@ mod test {
             module_subscription_actor::ModuleSubscriptions, module_subscription_manager::SubscriptionManager,
         },
     };
+    use anyhow::{anyhow, Result};
+    use parking_lot::RwLock;
+    use spacetimedb_lib::db::auth::StAccess;
+    use spacetimedb_lib::{bsatn::to_vec, AlgebraicType, AlgebraicValue, Hash, Identity, ProductValue};
+    use spacetimedb_paths::{server::ModuleLogsDir, FromPathUnchecked};
+    use spacetimedb_primitives::{IndexId, TableId};
+    use spacetimedb_sats::product;
+    use tempfile::TempDir;
 
     use super::{ChunkPool, InstanceEnv, TxSlot};
 
@@ -622,6 +624,37 @@ mod test {
         Ok((table_id, index_id))
     }
 
+    fn create_table_with_unique_index(db: &RelationalDB) -> Result<(TableId, IndexId)> {
+        let table_id = db.create_table_for_test_with_the_works(
+            "t",
+            &[("id", AlgebraicType::U64), ("str", AlgebraicType::String)],
+            &[0.into()],
+            &[0.into()],
+            StAccess::Public,
+        )?;
+        let index_id = db.with_read_only(Workload::ForTests, |tx| {
+            db.schema_for_table(tx, table_id)?
+                .indexes
+                .iter()
+                .find(|schema| {
+                    schema
+                        .index_algorithm
+                        .columns()
+                        .as_singleton()
+                        .is_some_and(|col_id| col_id.idx() == 0)
+                })
+                .map(|schema| schema.index_id)
+                .ok_or_else(|| anyhow!("Index not found for ColId `{}`", 0))
+        })?;
+        db.with_auto_commit(Workload::ForTests, |tx| -> Result<_> {
+            for i in 1..=5 {
+                db.insert(tx, table_id, &bsatn_row(i)?)?;
+            }
+            Ok(())
+        })?;
+        Ok((table_id, index_id))
+    }
+
     #[test]
     fn table_scan_metrics() -> Result<()> {
         let db = relational_db()?;
@@ -743,6 +776,33 @@ mod test {
         Ok(())
     }
 
+    #[test]
+    fn update_metrics() -> Result<()> {
+        let db = relational_db()?;
+        let env = instance_env(db.clone())?;
+
+        let (table_id, index_id) = create_table_with_unique_index(&db)?;
+
+        let mut tx_slot = env.tx.clone();
+
+        let row_id: u64 = 1;
+        let row_val: String = "string".to_string();
+        let mut new_row_bytes = to_vec(&product!(row_id, row_val))?;
+        let new_row_len = new_row_bytes.len();
+        // Delete a single row via the index
+        let f = || -> Result<_> {
+            env.update(table_id, index_id, new_row_bytes.as_mut_slice())?;
+            Ok(())
+        };
+        let tx = db.begin_mut_tx(IsolationLevel::Serializable, Workload::ForTests);
+        let (tx, res) = tx_slot.set(tx, f);
+
+        res?;
+
+        assert_eq!(new_row_len, tx.metrics.bytes_written);
+        Ok(())
+    }
+
     #[test]
     fn delete_by_index_metrics() -> Result<()> {
         let db = relational_db()?;