datastore: add support for truncation (#3215)

# Description of Changes

Adds support for truncation in the datastore.
The commitlog format already supported it.
A table is considered truncated in a transaction when there are deletes
and when the row count post-deletes is 0.

# API and ABI breaking changes

None

# Expected complexity level and risk

3?

# Testing

Some existing datastore tests have been amended to check whether
`TxData` contain truncation or not.

---------

Signed-off-by: Mazdak Farrokhzad <twingoow@gmail.com>
Co-authored-by: Shubham Mishra <shubham@clockworklabs.io>
Co-authored-by: Shubham Mishra <shivam828787@gmail.com>

Author: Centril

crates/core/src/db/relational_db.rs

@@ -10,6 +10,7 @@ use futures::channel::mpsc;
 use futures::StreamExt;
 use parking_lot::RwLock;
 use spacetimedb_commitlog as commitlog;
+use spacetimedb_data_structures::map::IntSet;
 use spacetimedb_datastore::db_metrics::DB_METRICS;
 use spacetimedb_datastore::error::{DatastoreError, TableError};
 use spacetimedb_datastore::execution_context::{ReducerContext, Workload, WorkloadType};
@@ -889,12 +890,17 @@ impl RelationalDB {
                     rowdata: rowdata.clone(),
                 })
                 .collect();
+
+            let truncates: IntSet<TableId> = tx_data.truncates().collect();
+
             let deletes: Box<_> = tx_data
                 .deletes()
                 .map(|(table_id, rowdata)| Ops {
                     table_id: *table_id,
                     rowdata: rowdata.clone(),
                 })
+                // filter out deletes for tables that are truncated in the same transaction.
+                .filter(|ops| !truncates.contains(&ops.table_id))
                 .collect();
 
             let inputs = reducer_context.map(|rcx| rcx.into());
@@ -905,7 +911,7 @@ impl RelationalDB {
                 mutations: Some(Mutations {
                     inserts,
                     deletes,
-                    truncates: [].into(),
+                    truncates: truncates.into_iter().collect(),
                 }),
             };
 

crates/core/src/subscription/tx.rs

@@ -133,8 +133,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 +143,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())
         })
     }
 

crates/datastore/src/locking_tx_datastore/committed_state.rs

@@ -315,19 +315,35 @@ impl CommittedState {
         Ok(())
     }
 
+    pub(super) fn replay_truncate(&mut self, table_id: TableId) -> Result<()> {
+        // (1) Table dropped? Avoid an error and just ignore the row instead.
+        if self.table_dropped.contains(&table_id) {
+            return Ok(());
+        }
+
+        // Get the table for mutation.
+        let (table, blob_store, ..) = self.get_table_and_blob_store_mut(table_id)?;
+
+        // We do not need to consider a truncation of `st_table` itself,
+        // as if that happens, the database is bricked.
+
+        table.clear(blob_store);
+
+        Ok(())
+    }
+
     pub(super) fn replay_delete_by_rel(&mut self, table_id: TableId, row: &ProductValue) -> Result<()> {
+        // (1) Table dropped? Avoid an error and just ignore the row instead.
+        if self.table_dropped.contains(&table_id) {
+            return Ok(());
+        }
+
         // Get the table for mutation.
-        let table = match self.tables.get_mut(&table_id) {
-            Some(t) => t,
-            // (1) If it was dropped, avoid an error and just ignore the row instead.
-            None if self.table_dropped.contains(&table_id) => return Ok(()),
-            None => return Err(TableError::IdNotFoundState(table_id).into()),
-        };
+        let (table, blob_store, _, page_pool) = self.get_table_and_blob_store_mut(table_id)?;
 
         // Delete the row.
-        let blob_store = &mut self.blob_store;
         table
-            .delete_equal_row(&self.page_pool, blob_store, row)
+            .delete_equal_row(page_pool, blob_store, row)
             .map_err(TableError::Bflatn)?
             .ok_or_else(|| anyhow!("Delete for non-existent row when replaying transaction"))?;
 
@@ -461,9 +477,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));
@@ -564,8 +580,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()
@@ -595,13 +610,27 @@ impl CommittedState {
 
     pub(super) fn merge(&mut self, tx_state: TxState, ctx: &ExecutionContext) -> TxData {
         let mut tx_data = TxData::default();
+        let mut truncates = IntSet::default();
 
         // First, apply deletes. This will free up space in the committed tables.
-        self.merge_apply_deletes(&mut tx_data, tx_state.delete_tables, tx_state.pending_schema_changes);
+        self.merge_apply_deletes(
+            &mut tx_data,
+            tx_state.delete_tables,
+            tx_state.pending_schema_changes,
+            &mut truncates,
+        );
 
         // Then, apply inserts. This will re-fill the holes freed by deletions
         // before allocating new pages.
-        self.merge_apply_inserts(&mut tx_data, tx_state.insert_tables, tx_state.blob_store);
+        self.merge_apply_inserts(
+            &mut tx_data,
+            tx_state.insert_tables,
+            tx_state.blob_store,
+            &mut truncates,
+        );
+
+        // Record any truncated tables in the `TxData`.
+        tx_data.add_truncates(truncates);
 
         // If the TX will be logged, record its projected tx offset,
         // then increment the counter.
@@ -618,6 +647,7 @@ impl CommittedState {
         tx_data: &mut TxData,
         delete_tables: BTreeMap<TableId, DeleteTable>,
         pending_schema_changes: ThinVec<PendingSchemaChange>,
+        truncates: &mut IntSet<TableId>,
     ) {
         fn delete_rows(
             tx_data: &mut TxData,
@@ -626,6 +656,7 @@ impl CommittedState {
             blob_store: &mut dyn BlobStore,
             row_ptrs_len: usize,
             row_ptrs: impl Iterator<Item = RowPointer>,
+            truncates: &mut IntSet<TableId>,
         ) {
             let mut deletes = Vec::with_capacity(row_ptrs_len);
 
@@ -645,17 +676,27 @@ 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;
+                if truncated {
+                    truncates.insert(table_id);
+                }
             }
         }
 
         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(),
+                    truncates,
+                ),
+                Err(_) if !row_ptrs.is_empty() => panic!("Deletion for non-existent table {table_id:?}... huh?"),
+                Err(_) => {}
             }
         }
 
@@ -672,6 +713,7 @@ impl CommittedState {
                     &mut self.blob_store,
                     row_ptrs.len(),
                     row_ptrs.into_iter(),
+                    truncates,
                 );
             }
         }
@@ -682,6 +724,7 @@ impl CommittedState {
         tx_data: &mut TxData,
         insert_tables: BTreeMap<TableId, Table>,
         tx_blob_store: impl BlobStore,
+        truncates: &mut IntSet<TableId>,
     ) {
         // TODO(perf): Consider moving whole pages from the `insert_tables` into the committed state,
         //             rather than copying individual rows out of them.
@@ -710,6 +753,11 @@ impl CommittedState {
             if !inserts.is_empty() {
                 let table_name = &*commit_table.get_schema().table_name;
                 tx_data.set_inserts_for_table(table_id, table_name, inserts.into());
+
+                // if table has inserted rows, it cannot be truncated
+                if truncates.contains(&table_id) {
+                    truncates.remove(&table_id);
+                }
             }
 
             let (schema, _indexes, pages) = tx_table.consume_for_merge();
@@ -835,12 +883,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 {