datastore: add truncation support

Centril · Centril · commit cc9ef8748c68 · 2025-09-02T09:48:34.000+02:00
diff --git a/crates/core/src/db/relational_db.rs b/crates/core/src/db/relational_db.rs
@@ -23,7 +23,7 @@ use spacetimedb_datastore::system_tables::StModuleRow;
 use spacetimedb_datastore::system_tables::{StFields, StVarFields, StVarName, StVarRow, ST_MODULE_ID, ST_VAR_ID};
 use spacetimedb_datastore::traits::{
     InsertFlags, IsolationLevel, Metadata, MutTx as _, MutTxDatastore, Program, RowTypeForTable, Tx as _, TxDatastore,
-    UpdateFlags,
+    TxTableTruncated, UpdateFlags,
 };
 use spacetimedb_datastore::{
     locking_tx_datastore::{
@@ -869,11 +869,17 @@ impl RelationalDB {
                 .collect();
             let deletes: Box<_> = tx_data
                 .deletes()
-                .map(|(table_id, rowdata)| Ops {
+                .filter(|(_, truncated, _)| *truncated == TxTableTruncated::No)
+                .map(|(table_id, _, rowdata)| Ops {
                     table_id: *table_id,
                     rowdata: rowdata.clone(),
                 })
                 .collect();
+            let truncates: Box<_> = tx_data
+                .deletes()
+                .filter(|(_, truncated, _)| *truncated == TxTableTruncated::Yes)
+                .map(|(table_id, ..)| *table_id)
+                .collect();
 
             let inputs = reducer_context.map(|rcx| rcx.into());
 
@@ -883,7 +889,7 @@ impl RelationalDB {
                 mutations: Some(Mutations {
                     inserts,
                     deletes,
-                    truncates: [].into(),
+                    truncates,
                 }),
             };
 
diff --git a/crates/core/src/subscription/tx.rs b/crates/core/src/subscription/tx.rs
@@ -76,9 +76,10 @@ impl DeltaTableIndexes {
             indexes
         }
 
+        let deletes = data.deletes().map(|(table_id, _, rows)| (table_id, rows));
         Self {
             inserts: build_indexes_for_rows(tx, meta, data.inserts()),
-            deletes: build_indexes_for_rows(tx, meta, data.deletes()),
+            deletes: build_indexes_for_rows(tx, meta, deletes),
         }
     }
 }
@@ -133,8 +134,8 @@ impl DeltaStore for DeltaTx<'_> {
         self.data
             .and_then(|data| {
                 data.inserts()
-                    .find(|(id, _)| **id == table_id)
-                    .map(|(_, rows)| rows.len())
+                    .find(|(id, ..)| **id == table_id)
+                    .map(|(.., rows)| rows.len())
             })
             .unwrap_or_default()
     }
@@ -143,25 +144,25 @@ impl DeltaStore for DeltaTx<'_> {
         self.data
             .and_then(|data| {
                 data.deletes()
-                    .find(|(id, _)| **id == table_id)
-                    .map(|(_, rows)| rows.len())
+                    .find(|(id, ..)| **id == table_id)
+                    .map(|(.., rows)| rows.len())
             })
             .unwrap_or_default()
     }
 
     fn inserts_for_table(&self, table_id: TableId) -> Option<std::slice::Iter<'_, ProductValue>> {
         self.data.and_then(|data| {
             data.inserts()
-                .find(|(id, _)| **id == table_id)
-                .map(|(_, rows)| rows.iter())
+                .find(|(id, ..)| **id == table_id)
+                .map(|(.., rows)| rows.iter())
         })
     }
 
     fn deletes_for_table(&self, table_id: TableId) -> Option<std::slice::Iter<'_, ProductValue>> {
         self.data.and_then(|data| {
             data.deletes()
-                .find(|(id, _)| **id == table_id)
-                .map(|(_, rows)| rows.iter())
+                .find(|(id, ..)| **id == table_id)
+                .map(|(.., rows)| rows.iter())
         })
     }
 
diff --git a/crates/datastore/src/locking_tx_datastore/committed_state.rs b/crates/datastore/src/locking_tx_datastore/committed_state.rs
@@ -334,14 +334,16 @@ impl CommittedState {
         Ok(())
     }
 
+    pub(super) fn replay_truncate(&mut self, table_id: TableId) -> Result<()> {
+        let (table, blob_store, ..) = self.get_table_and_blob_store_mut(table_id)?;
+        table.clear(blob_store);
+        Ok(())
+    }
+
     pub(super) fn replay_delete_by_rel(&mut self, table_id: TableId, rel: &ProductValue) -> Result<()> {
-        let table = self
-            .tables
-            .get_mut(&table_id)
-            .ok_or(TableError::IdNotFoundState(table_id))?;
-        let blob_store = &mut self.blob_store;
+        let (table, blob_store, _, page_pool) = self.get_table_and_blob_store_mut(table_id)?;
         table
-            .delete_equal_row(&self.page_pool, blob_store, rel)
+            .delete_equal_row(page_pool, blob_store, rel)
             .map_err(TableError::Bflatn)?
             .ok_or_else(|| anyhow!("Delete for non-existent row when replaying transaction"))?;
 
@@ -459,9 +461,9 @@ impl CommittedState {
         for index_row in rows {
             let index_id = index_row.index_id;
             let table_id = index_row.table_id;
-            let (Some(table), blob_store, index_id_map) = self.get_table_and_blob_store_mut(table_id) else {
-                panic!("Cannot create index for table which doesn't exist in committed state");
-            };
+            let (table, blob_store, index_id_map, _) = self
+                .get_table_and_blob_store_mut(table_id)
+                .expect("index should exist in committed state; cannot create it");
             let algo: IndexAlgorithm = index_row.index_algorithm.into();
             let columns: ColSet = algo.columns().into();
             let is_unique = unique_constraints.contains(&(table_id, columns));
@@ -562,8 +564,7 @@ impl CommittedState {
             "Cannot get TX_STATE RowPointer from CommittedState.",
         );
         let table = self
-            .tables
-            .get(&table_id)
+            .get_table(table_id)
             .expect("Attempt to get COMMITTED_STATE row from table not present in tables.");
         // TODO(perf, deep-integration): Use `get_row_ref_unchecked`.
         table.get_row_ref(&self.blob_store, row_ptr).unwrap()
@@ -643,17 +644,18 @@ impl CommittedState {
 
             if !deletes.is_empty() {
                 let table_name = &*table.get_schema().table_name;
-                // TODO(centril): Pass this along to record truncated tables.
-                let _truncated = table.row_count == 0;
-                tx_data.set_deletes_for_table(table_id, table_name, deletes.into());
+                let truncated = table.row_count == 0;
+                tx_data.set_deletes_for_table(table_id, table_name, deletes.into(), truncated);
             }
         }
 
         for (table_id, row_ptrs) in delete_tables {
-            if let (Some(table), blob_store, _) = self.get_table_and_blob_store_mut(table_id) {
-                delete_rows(tx_data, table_id, table, blob_store, row_ptrs.len(), row_ptrs.iter());
-            } else if !row_ptrs.is_empty() {
-                panic!("Deletion for non-existent table {table_id:?}... huh?");
+            match self.get_table_and_blob_store_mut(table_id) {
+                Ok((table, blob_store, ..)) => {
+                    delete_rows(tx_data, table_id, table, blob_store, row_ptrs.len(), row_ptrs.iter())
+                }
+                Err(_) if !row_ptrs.is_empty() => panic!("Deletion for non-existent table {table_id:?}... huh?"),
+                Err(_) => {}
             }
         }
 
@@ -832,12 +834,21 @@ impl CommittedState {
     pub(super) fn get_table_and_blob_store_mut(
         &mut self,
         table_id: TableId,
-    ) -> (Option<&mut Table>, &mut dyn BlobStore, &mut IndexIdMap) {
-        (
-            self.tables.get_mut(&table_id),
+    ) -> Result<(&mut Table, &mut dyn BlobStore, &mut IndexIdMap, &PagePool)> {
+        // NOTE(centril): `TableError` is a fairly large type.
+        // Not making this lazy made `TableError::drop` show up in perf.
+        // TODO(centril): Box all the errors.
+        #[allow(clippy::unnecessary_lazy_evaluations)]
+        let table = self
+            .tables
+            .get_mut(&table_id)
+            .ok_or_else(|| TableError::IdNotFoundState(table_id))?;
+        Ok((
+            table,
             &mut self.blob_store as &mut dyn BlobStore,
             &mut self.index_id_map,
-        )
+            &self.page_pool,
+        ))
     }
 
     fn make_table(schema: Arc<TableSchema>) -> Table {
diff --git a/crates/datastore/src/locking_tx_datastore/datastore.rs b/crates/datastore/src/locking_tx_datastore/datastore.rs
@@ -1098,8 +1098,26 @@ impl<F: FnMut(u64)> spacetimedb_commitlog::payload::txdata::Visitor for ReplayVi
         Ok(row)
     }
 
-    fn visit_truncate(&mut self, _table_id: TableId) -> std::result::Result<(), Self::Error> {
-        Err(anyhow!("visit: truncate not yet supported").into())
+    fn visit_truncate(&mut self, table_id: TableId) -> std::result::Result<(), Self::Error> {
+        let schema = self.committed_state.schema_for_table(table_id)?;
+        // TODO: avoid clone
+        let table_name = schema.table_name.clone();
+
+        self.committed_state.replay_truncate(table_id).with_context(|| {
+            format!(
+                "Error truncating table {:?} during transaction {:?} playback",
+                table_id, self.committed_state.next_tx_offset
+            )
+        })?;
+
+        // NOTE: the `rdb_num_table_rows` metric is used by the query optimizer,
+        // and therefore has performance implications and must not be disabled.
+        DB_METRICS
+            .rdb_num_table_rows
+            .with_label_values(self.database_identity, &table_id.into(), &table_name)
+            .set(0);
+
+        Ok(())
     }
 
     fn visit_tx_start(&mut self, offset: u64) -> std::result::Result<(), Self::Error> {
@@ -1153,7 +1171,7 @@ mod tests {
         ST_ROW_LEVEL_SECURITY_NAME, ST_SCHEDULED_ID, ST_SCHEDULED_NAME, ST_SEQUENCE_ID, ST_SEQUENCE_NAME,
         ST_TABLE_NAME, ST_VAR_ID, ST_VAR_NAME,
     };
-    use crate::traits::{IsolationLevel, MutTx};
+    use crate::traits::{IsolationLevel, MutTx, TxTableTruncated};
     use crate::Result;
     use bsatn::to_vec;
     use core::{fmt, mem};
@@ -2775,7 +2793,7 @@ mod tests {
         let tx_data_2 = commit(&datastore, tx)?;
         // Ensure that none of the commits deleted rows in our table.
         for tx_data in [&tx_data_1, &tx_data_2] {
-            assert_eq!(tx_data.deletes().find(|(tid, _)| **tid == table_id), None);
+            assert_eq!(tx_data.deletes().find(|(tid, ..)| **tid == table_id), None);
         }
         // Ensure that the first commit added the row but that the second didn't.
         for (tx_data, expected_rows) in [(&tx_data_1, vec![row.clone()]), (&tx_data_2, vec![])] {
@@ -3099,11 +3117,12 @@ mod tests {
         // Now drop the table again and commit.
         assert!(datastore.drop_table_mut_tx(&mut tx, table_id).is_ok());
         let tx_data = commit(&datastore, tx)?;
-        let (_, deleted) = tx_data
+        let (_, truncated, deleted_rows) = tx_data
             .deletes()
-            .find(|(id, _)| **id == table_id)
+            .find(|(id, ..)| **id == table_id)
             .expect("should have deleted rows for `table_id`");
-        assert_eq!(&**deleted, [row]);
+        assert_eq!(&**deleted_rows, [row]);
+        assert_eq!(truncated, TxTableTruncated::Yes);
 
         // In the next transaction, the table doesn't exist.
         assert!(
@@ -3307,8 +3326,9 @@ mod tests {
         let to_product = |col: &ColumnSchema| value_serialize(&StColumnRow::from(col.clone())).into_product().unwrap();
         let (_, inserts) = tx_data.inserts().find(|(id, _)| **id == ST_COLUMN_ID).unwrap();
         assert_eq!(&**inserts, [to_product(&columns[1])].as_slice());
-        let (_, deletes) = tx_data.deletes().find(|(id, _)| **id == ST_COLUMN_ID).unwrap();
+        let (_, truncated, deletes) = tx_data.deletes().find(|(id, ..)| **id == ST_COLUMN_ID).unwrap();
         assert_eq!(&**deletes, [to_product(&columns_original[1])].as_slice());
+        assert_eq!(truncated, TxTableTruncated::No);
 
         // Check that we can successfully scan using the new schema type post commit.
         let tx = begin_tx(&datastore);
diff --git a/crates/datastore/src/locking_tx_datastore/mut_tx.rs b/crates/datastore/src/locking_tx_datastore/mut_tx.rs
@@ -414,13 +414,8 @@ impl MutTxId {
         TxTableForInsertion<'_>,
         (&mut Table, &mut dyn BlobStore, &mut IndexIdMap),
     )> {
-        let (commit_table, commit_bs, idx_map) = self.committed_state_write_lock.get_table_and_blob_store_mut(table_id);
-        // NOTE(centril): `TableError` is a fairly large type.
-        // Not making this lazy made `TableError::drop` show up in perf.
-        // TODO(centril): Box all the errors.
-        #[allow(clippy::unnecessary_lazy_evaluations)]
-        let commit_table = commit_table.ok_or_else(|| TableError::IdNotFoundState(table_id))?;
-
+        let (commit_table, commit_bs, idx_map, _) =
+            self.committed_state_write_lock.get_table_and_blob_store_mut(table_id)?;
         // Get the insert table, so we can write the row into it.
         let tx = self
             .tx_state
diff --git a/crates/datastore/src/traits.rs b/crates/datastore/src/traits.rs
@@ -174,7 +174,9 @@ pub struct TxData {
     /// The inserted rows per table.
     inserts: BTreeMap<TableId, Arc<[ProductValue]>>,
     /// The deleted rows per table.
-    deletes: BTreeMap<TableId, Arc<[ProductValue]>>,
+    ///
+    /// Also stores per table whether it has been truncated.
+    deletes: BTreeMap<TableId, TxDeleteEntry>,
     /// Map of all `TableId`s in both `inserts` and `deletes` to their
     /// corresponding table name.
     tables: IntMap<TableId, String>,
@@ -186,6 +188,24 @@ pub struct TxData {
     // TODO: Store an `Arc<String>` or equivalent instead.
 }
 
+/// A record of a list of deletes for and potential truncation of a table,
+/// within a transaction.
+pub struct TxDeleteEntry {
+    /// Were all rows previously in the table deleted within this transaction?
+    truncated: TxTableTruncated,
+    /// The deleted rows of the table.
+    rows: Arc<[ProductValue]>,
+}
+
+/// Whether a table was truncated in a transaction.
+#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
+pub enum TxTableTruncated {
+    /// It was truncated.
+    Yes,
+    /// It was not truncated.
+    No,
+}
+
 impl TxData {
     /// Set `tx_offset` as the expected on-disk transaction offset of this transaction.
     pub fn set_tx_offset(&mut self, tx_offset: u64) {
@@ -208,8 +228,22 @@ impl TxData {
     }
 
     /// Set `rows` as the deleted rows for `(table_id, table_name)`.
-    pub fn set_deletes_for_table(&mut self, table_id: TableId, table_name: &str, rows: Arc<[ProductValue]>) {
-        self.deletes.insert(table_id, rows);
+    ///
+    /// When `truncated` is set, the table has been emptied in this transaction.
+    pub fn set_deletes_for_table(
+        &mut self,
+        table_id: TableId,
+        table_name: &str,
+        rows: Arc<[ProductValue]>,
+        truncated: bool,
+    ) {
+        let truncated = if truncated {
+            TxTableTruncated::Yes
+        } else {
+            TxTableTruncated::No
+        };
+        let entry = TxDeleteEntry { truncated, rows };
+        self.deletes.insert(table_id, entry);
         self.tables.entry(table_id).or_insert_with(|| table_name.to_owned());
     }
 
@@ -238,26 +272,28 @@ impl TxData {
     }
 
     /// Obtain an iterator over the deleted rows per table.
-    pub fn deletes(&self) -> impl Iterator<Item = (&TableId, &Arc<[ProductValue]>)> + '_ {
-        self.deletes.iter()
+    pub fn deletes(&self) -> impl Iterator<Item = (&TableId, TxTableTruncated, &Arc<[ProductValue]>)> + '_ {
+        self.deletes
+            .iter()
+            .map(|(table_id, entry)| (table_id, entry.truncated, &entry.rows))
     }
 
     /// Get the `i`th deleted row for `table_id` if it exists
     pub fn get_ith_delete(&self, table_id: TableId, i: usize) -> Option<&ProductValue> {
-        self.deletes.get(&table_id).and_then(|rows| rows.get(i))
+        self.deletes.get(&table_id).and_then(|entry| entry.rows.get(i))
     }
 
     /// Obtain an iterator over the inserted rows per table.
     ///
     /// If you don't need access to the table name, [`Self::deletes`] is
     /// slightly more efficient.
     pub fn deletes_with_table_name(&self) -> impl Iterator<Item = (&TableId, &str, &Arc<[ProductValue]>)> + '_ {
-        self.deletes.iter().map(|(table_id, rows)| {
+        self.deletes.iter().map(|(table_id, entry)| {
             let table_name = self
                 .tables
                 .get(table_id)
                 .expect("invalid `TxData`: partial table name mapping");
-            (table_id, table_name.as_str(), rows)
+            (table_id, table_name.as_str(), &entry.rows)
         })
     }
 
@@ -266,7 +302,7 @@ impl TxData {
     /// This is used to determine if a transaction should be written to disk.
     pub fn has_rows_or_connect_disconnect(&self, reducer_context: Option<&ReducerContext>) -> bool {
         self.inserts().any(|(_, inserted_rows)| !inserted_rows.is_empty())
-            || self.deletes().any(|(_, deleted_rows)| !deleted_rows.is_empty())
+            || self.deletes().any(|(.., deleted_rows)| !deleted_rows.is_empty())
             || matches!(
                 reducer_context.map(|rcx| rcx.name.strip_prefix("__identity_")),
                 Some(Some("connected__" | "disconnected__"))
