WASM ABI: add datastore_table_row_count (#1636)

Centril · web-flow · commit c26c65873bb4 · 2024-08-27T23:02:18.000Z
diff --git a/crates/bindings-csharp/Runtime/Internal/FFI.cs b/crates/bindings-csharp/Runtime/Internal/FFI.cs
@@ -120,6 +120,12 @@ public static partial CheckedStatus _table_id_from_name(
         out TableId out_
     );
 
+    [LibraryImport(StdbNamespace)]
+    public static partial CheckedStatus _datastore_table_row_count(
+        TableId table_id,
+        out ulong out_
+    );
+
     [LibraryImport(StdbNamespace)]
     public static partial CheckedStatus _iter_by_col_eq(
         TableId table_id,
diff --git a/crates/bindings-csharp/Runtime/bindings.c b/crates/bindings-csharp/Runtime/bindings.c
@@ -45,6 +45,9 @@ IMPORT(void, _console_log,
 IMPORT(Status, _table_id_from_name,
        (const uint8_t* name, uint32_t name_len, TableId* id),
        (name, name_len, id));
+IMPORT(Status, _datastore_table_row_count,
+       (TableId table_id, uint64_t* count),
+       (table_id, count));
 IMPORT(Status, _iter_by_col_eq,
        (TableId table_id, ColId col_id, const uint8_t* value,
         uint32_t value_len, RowIter* iter),
diff --git a/crates/bindings-sys/src/lib.rs b/crates/bindings-sys/src/lib.rs
@@ -40,6 +40,21 @@ pub mod raw {
         /// - `NO_SUCH_TABLE`, when `name` is not the name of a table.
         pub fn _table_id_from_name(name: *const u8, name_len: usize, out: *mut TableId) -> u16;
 
+        /// Writes the number of rows currently in table identified by `table_id` to `out`.
+        ///
+        /// # Traps
+        ///
+        /// Traps if:
+        /// - `out` is NULL or `out[..size_of::<u64>()]` is not in bounds of WASM memory.
+        ///
+        /// # Errors
+        ///
+        /// Returns an error:
+        ///
+        /// - `NOT_IN_TRANSACTION`, when called outside of a transaction.
+        /// - `NO_SUCH_TABLE`, when `table_id` is not a known ID of a table.
+        pub fn _datastore_table_row_count(table_id: TableId, out: *mut u64) -> u16;
+
         /// Finds all rows in the table identified by `table_id`,
         /// where the row has a column, identified by `col_id`,
         /// with data matching the byte string, in WASM memory, pointed to at by `val`.
@@ -523,6 +538,14 @@ pub fn table_id_from_name(name: &str) -> Result<TableId, Errno> {
     unsafe { call(|out| raw::_table_id_from_name(name.as_ptr(), name.len(), out)) }
 }
 
+/// Queries and returns the number of rows in the table identified by `table_id`.
+///
+/// Returns an error if the table does not exist or if not in a transaction.
+#[inline]
+pub fn datastore_table_row_count(table_id: TableId) -> Result<u64, Errno> {
+    unsafe { call(|out| raw::_datastore_table_row_count(table_id, out)) }
+}
+
 /// Finds all rows in the table identified by `table_id`,
 /// where the row has a column, identified by `col_id`,
 /// with data matching the byte string `val`.
diff --git a/crates/core/src/db/datastore/locking_tx_datastore/committed_state.rs b/crates/core/src/db/datastore/locking_tx_datastore/committed_state.rs
@@ -57,6 +57,11 @@ impl StateView for CommittedState {
     fn get_schema(&self, table_id: TableId) -> Option<&Arc<TableSchema>> {
         self.tables.get(&table_id).map(|table| table.get_schema())
     }
+
+    fn table_row_count(&self, table_id: TableId) -> Option<u64> {
+        self.get_table(table_id).map(|table| table.row_count)
+    }
+
     fn iter<'a>(&'a self, ctx: &'a ExecutionContext, table_id: TableId) -> Result<Iter<'a>> {
         if let Some(table_name) = self.table_name(table_id) {
             return Ok(Iter::new(ctx, table_id, table_name, None, self));
diff --git a/crates/core/src/db/datastore/locking_tx_datastore/mut_tx.rs b/crates/core/src/db/datastore/locking_tx_datastore/mut_tx.rs
@@ -7,16 +7,13 @@ use super::{
     tx_state::TxState,
     SharedMutexGuard, SharedWriteGuard,
 };
-use crate::db::{
-    datastore::{
-        system_tables::{
-            table_name_is_system, StColumnFields, StColumnRow, StConstraintFields, StConstraintRow, StFields as _,
-            StIndexFields, StIndexRow, StScheduledRow, StSequenceFields, StSequenceRow, StTableFields, StTableRow,
-            SystemTable, ST_COLUMN_ID, ST_CONSTRAINT_ID, ST_INDEX_ID, ST_SCHEDULED_ID, ST_SEQUENCE_ID, ST_TABLE_ID,
-        },
-        traits::{RowTypeForTable, TxData},
+use crate::db::datastore::{
+    system_tables::{
+        table_name_is_system, StColumnFields, StColumnRow, StConstraintFields, StConstraintRow, StFields as _,
+        StIndexFields, StIndexRow, StScheduledRow, StSequenceFields, StSequenceRow, StTableFields, StTableRow,
+        SystemTable, ST_COLUMN_ID, ST_CONSTRAINT_ID, ST_INDEX_ID, ST_SCHEDULED_ID, ST_SEQUENCE_ID, ST_TABLE_ID,
     },
-    db_metrics::table_num_rows,
+    traits::{RowTypeForTable, TxData},
 };
 use crate::{
     error::{DBError, IndexError, SequenceError, TableError},
@@ -1082,6 +1079,18 @@ impl StateView for MutTxId {
             .map(|table| table.get_schema())
     }
 
+    fn table_row_count(&self, table_id: TableId) -> Option<u64> {
+        let commit_count = self.committed_state_write_lock.table_row_count(table_id);
+        let (tx_ins_count, tx_del_count) = self.tx_state.table_row_count(table_id);
+        let commit_count = commit_count.map(|cc| cc - tx_del_count);
+        // Keep track of whether `table_id` exists.
+        match (commit_count, tx_ins_count) {
+            (Some(cc), Some(ic)) => Some(cc + ic),
+            (Some(c), None) | (None, Some(c)) => Some(c),
+            (None, None) => None,
+        }
+    }
+
     fn iter<'a>(&'a self, ctx: &'a ExecutionContext, table_id: TableId) -> Result<Iter<'a>> {
         if let Some(table_name) = self.table_name(table_id) {
             return Ok(Iter::new(
@@ -1137,29 +1146,31 @@ impl StateView for MutTxId {
                 ))),
                 None => {
                     #[cfg(feature = "unindexed_iter_by_col_range_warn")]
-                    match self.schema_for_table(ctx, table_id) {
+                    match self.table_row_count(table_id) {
                         // TODO(ux): log these warnings to the module logs rather than host logs.
-                        Err(e) => log::error!(
-                            "iter_by_col_range on unindexed column, but got error from `schema_for_table` during diagnostics: {e:?}",
+                        None => log::error!(
+                            "iter_by_col_range on unindexed column, but couldn't fetch table `{table_id}`s row count",
                         ),
-                        Ok(schema) => {
+                        Some(num_rows) => {
                             const TOO_MANY_ROWS_FOR_SCAN: u64 = 1000;
-
-                            let table_name = &schema.table_name;
-                            let num_rows = table_num_rows(ctx.database(), table_id, table_name);
-
                             if num_rows >= TOO_MANY_ROWS_FOR_SCAN {
-                                let col_names = cols.iter()
-                                    .map(|col_id| schema.columns()
-                                         .get(col_id.idx())
-                                         .map(|col| &col.col_name[..])
-                                         .unwrap_or("[unknown column]"))
+                                let schema = self.schema_for_table(ctx, table_id).unwrap();
+                                let table_name = &schema.table_name;
+                                let col_names = cols
+                                    .iter()
+                                    .map(|col_id| {
+                                        schema
+                                            .columns()
+                                            .get(col_id.idx())
+                                            .map(|col| &col.col_name[..])
+                                            .unwrap_or("[unknown column]")
+                                    })
                                     .collect::<Vec<_>>();
                                 log::warn!(
                                     "iter_by_col_range without index: table {table_name} has {num_rows} rows; scanning columns {col_names:?}",
                                 );
                             }
-                        },
+                        }
                     }
 
                     Ok(IterByColRange::Scan(ScanIterByColRange::new(
diff --git a/crates/core/src/db/datastore/locking_tx_datastore/state_view.rs b/crates/core/src/db/datastore/locking_tx_datastore/state_view.rs
@@ -32,6 +32,9 @@ pub trait StateView {
         Ok(row.map(|row| row.read_col(StTableFields::TableId).unwrap()))
     }
 
+    /// Returns the number of rows in the table identified by `table_id`.
+    fn table_row_count(&self, table_id: TableId) -> Option<u64>;
+
     fn iter<'a>(&'a self, ctx: &'a ExecutionContext, table_id: TableId) -> Result<Iter<'a>>;
 
     fn table_name(&self, table_id: TableId) -> Option<&str> {
diff --git a/crates/core/src/db/datastore/locking_tx_datastore/tx.rs b/crates/core/src/db/datastore/locking_tx_datastore/tx.rs
@@ -26,6 +26,10 @@ impl StateView for TxId {
         self.committed_state_shared_lock.get_schema(table_id)
     }
 
+    fn table_row_count(&self, table_id: TableId) -> Option<u64> {
+        self.committed_state_shared_lock.table_row_count(table_id)
+    }
+
     fn iter<'a>(&'a self, ctx: &'a ExecutionContext, table_id: TableId) -> Result<Iter<'a>> {
         self.committed_state_shared_lock.iter(ctx, table_id)
     }
@@ -60,12 +64,6 @@ impl TxId {
         record_metrics(ctx, self.timer, self.lock_wait_time, true, None, None);
     }
 
-    pub(crate) fn get_row_count(&self, table_id: TableId) -> Option<u64> {
-        self.committed_state_shared_lock
-            .get_table(table_id)
-            .map(|table| table.row_count)
-    }
-
     /// The Number of Distinct Values (NDV) for a column or list of columns,
     /// if there's an index available on `cols`.
     pub(crate) fn num_distinct_values(&self, table_id: TableId, cols: &ColList) -> Option<u64> {
diff --git a/crates/core/src/db/datastore/locking_tx_datastore/tx_state.rs b/crates/core/src/db/datastore/locking_tx_datastore/tx_state.rs
@@ -64,6 +64,14 @@ pub(super) struct TxState {
 }
 
 impl TxState {
+    /// Returns the row count in insert tables
+    /// and the number of rows deleted from committed state.
+    pub(super) fn table_row_count(&self, table_id: TableId) -> (Option<u64>, u64) {
+        let del_count = self.delete_tables.get(&table_id).map(|dt| dt.len() as u64).unwrap_or(0);
+        let ins_count = self.insert_tables.get(&table_id).map(|it| it.row_count);
+        (ins_count, del_count)
+    }
+
     /// When there's an index on `cols`,
     /// returns an iterator over the [BTreeIndex] that yields all the `RowId`s
     /// that match the specified `value` in the indexed column.
diff --git a/crates/core/src/db/relational_db.rs b/crates/core/src/db/relational_db.rs
@@ -937,6 +937,12 @@ impl RelationalDB {
         self.inner.table_name_from_id_mut_tx(ctx, tx, table_id)
     }
 
+    pub fn table_row_count_mut(&self, tx: &MutTx, table_id: TableId) -> Option<u64> {
+        // TODO(Centril): Go via MutTxDatastore trait instead.
+        // Doing this for now to ship this quicker.
+        tx.table_row_count(table_id)
+    }
+
     pub fn column_constraints(
         &self,
         tx: &mut MutTx,
@@ -1821,7 +1827,7 @@ mod tests {
 
         let stdb = stdb.reopen()?;
         let tx = stdb.begin_tx();
-        assert_eq!(tx.get_row_count(table_id).unwrap(), 2);
+        assert_eq!(tx.table_row_count(table_id).unwrap(), 2);
         Ok(())
     }
 
diff --git a/crates/core/src/estimation.rs b/crates/core/src/estimation.rs
@@ -1,4 +1,4 @@
-use crate::db::relational_db::Tx;
+use crate::db::{datastore::locking_tx_datastore::state_view::StateView as _, relational_db::Tx};
 use spacetimedb_primitives::{ColList, TableId};
 use spacetimedb_vm::expr::{Query, QueryExpr, SourceExpr};
 
@@ -11,7 +11,7 @@ pub fn num_rows(tx: &Tx, expr: &QueryExpr) -> u64 {
 fn row_est(tx: &Tx, src: &SourceExpr, ops: &[Query]) -> u64 {
     match ops {
         // The base case is the table row count.
-        [] => src.table_id().and_then(|id| tx.get_row_count(id)).unwrap_or(0),
+        [] => src.table_id().and_then(|id| tx.table_row_count(id)).unwrap_or(0),
         // Walk in reverse from the end (`op`) to the beginning.
         [input @ .., op] => match op {
             // How selective is an index lookup?
@@ -61,7 +61,7 @@ fn row_est(tx: &Tx, src: &SourceExpr, ops: &[Query]) -> u64 {
 /// Note this method is not applicable to range scans.
 fn index_row_est(tx: &Tx, table_id: TableId, cols: &ColList) -> u64 {
     tx.num_distinct_values(table_id, cols)
-        .map_or(0, |ndv| tx.get_row_count(table_id).unwrap_or(0) / ndv)
+        .map_or(0, |ndv| tx.table_row_count(table_id).unwrap_or(0) / ndv)
 }
 
 #[cfg(test)]
diff --git a/crates/core/src/host/instance_env.rs b/crates/core/src/host/instance_env.rs
@@ -195,11 +195,21 @@ impl InstanceEnv {
         let tx = &mut *self.get_tx()?;
 
         // Query the table id from the name.
-        let table_id = stdb
-            .table_id_from_name_mut(tx, table_name)?
-            .ok_or(NodesError::TableNotFound)?;
+        stdb.table_id_from_name_mut(tx, table_name)?
+            .ok_or(NodesError::TableNotFound)
+    }
+
+    /// Returns the number of rows in the table identified by `table_id`.
+    ///
+    /// Errors with `GetTxError` if not in a transaction
+    /// and `TableNotFound` if the table does not exist.
+    #[tracing::instrument(skip_all)]
+    pub fn datastore_table_row_count(&self, table_id: TableId) -> Result<u64, NodesError> {
+        let stdb = &*self.dbic.relational_db;
+        let tx = &mut *self.get_tx()?;
 
-        Ok(table_id)
+        // Query the row count for id.
+        stdb.table_row_count_mut(tx, table_id).ok_or(NodesError::TableNotFound)
     }
 
     /// Finds all rows in the table identified by `table_id`
diff --git a/crates/core/src/host/mod.rs b/crates/core/src/host/mod.rs
@@ -142,6 +142,7 @@ fn from_json_seed<'de, T: serde::de::DeserializeSeed<'de>>(s: &'de str, seed: T)
 #[derive(Debug, Display, Enum, Clone, Copy, strum::AsRefStr)]
 pub enum AbiCall {
     TableIdFromName,
+    DatastoreTableRowCount,
     RowIterBsatnAdvance,
     RowIterBsatnClose,
     BytesSourceRead,
diff --git a/crates/core/src/host/wasm_common.rs b/crates/core/src/host/wasm_common.rs
@@ -347,6 +347,7 @@ macro_rules! abi_funcs {
     ($mac:ident) => {
         $mac! {
             "spacetime_10.0"::table_id_from_name,
+            "spacetime_10.0"::datastore_table_row_count,
             "spacetime_10.0"::row_iter_bsatn_advance,
             "spacetime_10.0"::row_iter_bsatn_close,
             "spacetime_10.0"::bytes_source_read,
diff --git a/crates/core/src/host/wasmtime/wasm_instance_env.rs b/crates/core/src/host/wasmtime/wasm_instance_env.rs
@@ -483,6 +483,27 @@ impl WasmInstanceEnv {
         })
     }
 
+    /// Writes the number of rows currently in table identified by `table_id` to `out`.
+    ///
+    /// # Traps
+    ///
+    /// Traps if:
+    /// - `out` is NULL or `out[..size_of::<u64>()]` is not in bounds of WASM memory.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error:
+    ///
+    /// - `NOT_IN_TRANSACTION`, when called outside of a transaction.
+    /// - `NO_SUCH_TABLE`, when `table_id` is not a known ID of a table.
+    #[tracing::instrument(skip_all)]
+    pub fn datastore_table_row_count(caller: Caller<'_, Self>, table_id: u32, out: WasmPtr<u64>) -> RtResult<u32> {
+        Self::cvt_ret::<u64>(caller, AbiCall::DatastoreTableRowCount, out, |caller| {
+            let (_, env) = Self::mem_env(caller);
+            Ok(env.instance_env.datastore_table_row_count(table_id.into())?)
+        })
+    }
+
     /// Finds all rows in the table identified by `table_id`,
     /// where the row has a column, identified by `cols`,
     /// with data matching the byte string, in WASM memory, pointed to at by `val`.

Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,9 @@ pub trait StateView {`
`32`	`32`	`Ok(row.map(\|row\| row.read_col(StTableFields::TableId).unwrap()))`
`33`	`33`	`}`
`34`	`34`
	`35`	+ /// Returns the number of rows in the table identified by `table_id`.
	`36`	`+ fn table_row_count(&self, table_id: TableId) -> Option<u64>;`
	`37`	`+`
`35`	`38`	`fn iter<'a>(&'a self, ctx: &'a ExecutionContext, table_id: TableId) -> Result<Iter<'a>>;`
`36`	`39`
`37`	`40`	`fn table_name(&self, table_id: TableId) -> Option<&str> {`