Allocate pages using a mult-tenant lock-free pool (#2587)

Author: Centril

Cargo.lock

@@ -150,6 +150,7 @@ convert_case = "0.6.0"
 crc32c = "0.6.4"
 criterion = { version = "0.5.1", features = ["async", "async_tokio", "html_reports"] }
 crossbeam-channel = "0.5"
+crossbeam-queue = "0.3.12"
 cursive = { version = "0.20", default-features = false, features = ["crossterm-backend"] }
 decorum = { version = "0.3.1", default-features = false, features = ["std"] }
 derive_more = "0.99"
@@ -197,6 +198,7 @@ mimalloc = "0.1.39"
 nohash-hasher = "0.2"
 once_cell = "1.16"
 parking_lot = { version = "0.12.1", features = ["send_guard", "arc_lock"] }
+parse-size = "1.1.0"
 paste = "1.0"
 percent-encoding = "2.3"
 petgraph = { version = "0.6.5", default-features = false }
@@ -316,4 +318,3 @@ features = [
   "broadcast",
   "ondemand",
 ]
-

Cargo.toml

@@ -8,6 +8,7 @@ use spacetimedb_bench::{
 use spacetimedb_lib::sats::{self, bsatn};
 use spacetimedb_lib::{bsatn::ToBsatn as _, ProductValue};
 use spacetimedb_schema::schema::TableSchema;
+use spacetimedb_table::page_pool::PagePool;
 use spacetimedb_testing::modules::{Csharp, ModuleLanguage, Rust};
 use std::sync::Arc;
 use std::sync::OnceLock;
@@ -140,14 +141,15 @@ fn serialize_benchmarks<
         Arc::new(table_schema),
         spacetimedb_table::indexes::SquashedOffset::COMMITTED_STATE,
     );
+    let pool = PagePool::default();
     let mut blob_store = spacetimedb_table::blob_store::HashMapBlobStore::default();
 
     let ptrs = data_pv
         .elements
         .iter()
         .map(|row| {
             table
-                .insert(&mut blob_store, row.as_product().unwrap())
+                .insert(&pool, &mut blob_store, row.as_product().unwrap())
                 .unwrap()
                 .1
                 .pointer()

crates/bench/benches/special.rs

@@ -57,6 +57,7 @@ impl<L: ModuleLanguage> BenchDatabase for SpacetimeModule<L> {
         let runtime = start_runtime();
         let config = Config {
             storage: if in_memory { Storage::Memory } else { Storage::Disk },
+            page_pool_max_size: None,
         };
 
         let module = runtime.block_on(async {

crates/bench/src/spacetime_module.rs

@@ -42,6 +42,7 @@ use spacetimedb_schema::{def::IndexAlgorithm, schema::TableSchema};
 use spacetimedb_table::{
     blob_store::{BlobStore, HashMapBlobStore},
     indexes::{RowPointer, SquashedOffset},
+    page_pool::PagePool,
     table::{IndexScanRangeIter, InsertError, RowRef, Table, TableAndIndex},
     MemoryUsage,
 };
@@ -53,13 +54,21 @@ use std::sync::Arc;
 /// logs directly. For normal usage, see the RelationalDB struct instead.
 ///
 /// NOTE: unstable API, this may change at any point in the future.
-#[derive(Default)]
 pub struct CommittedState {
     pub(crate) next_tx_offset: u64,
     pub(crate) tables: IntMap<TableId, Table>,
     pub(crate) blob_store: HashMapBlobStore,
     /// Provides fast lookup for index id -> an index.
     pub(super) index_id_map: IndexIdMap,
+    /// The page pool used to retrieve new/unused pages for tables.
+    ///
+    /// Between transactions, this is untouched.
+    /// During transactions, the [`MutTxId`] can steal pages from the committed state.
+    ///
+    /// This is a handle on a shared structure.
+    /// Pages are shared between all modules running on a particular host,
+    /// not allocated per-module.
+    pub(super) page_pool: PagePool,
 }
 
 impl MemoryUsage for CommittedState {
@@ -69,7 +78,9 @@ impl MemoryUsage for CommittedState {
             tables,
             blob_store,
             index_id_map,
+            page_pool: _,
         } = self;
+        // NOTE(centril): We do not want to include the heap usage of `page_pool` as it's a shared resource.
         next_tx_offset.heap_usage() + tables.heap_usage() + blob_store.heap_usage() + index_id_map.heap_usage()
     }
 }
@@ -143,6 +154,16 @@ fn ignore_duplicate_insert_error<T>(res: std::result::Result<T, InsertError>) ->
 }
 
 impl CommittedState {
+    pub(super) fn new(page_pool: PagePool) -> Self {
+        Self {
+            next_tx_offset: <_>::default(),
+            tables: <_>::default(),
+            blob_store: <_>::default(),
+            index_id_map: <_>::default(),
+            page_pool,
+        }
+    }
+
     /// Extremely delicate function to bootstrap the system tables.
     /// Don't update this unless you know what you're doing.
     pub(super) fn bootstrap_system_tables(&mut self, database_identity: Identity) -> Result<()> {
@@ -158,7 +179,8 @@ impl CommittedState {
         let ref_schemas = schemas.each_ref().map(|s| &**s);
 
         // Insert the table row into st_tables, creating st_tables if it's missing.
-        let (st_tables, blob_store) = self.get_table_and_blob_store_or_create(ST_TABLE_ID, &schemas[ST_TABLE_IDX]);
+        let (st_tables, blob_store, pool) =
+            self.get_table_and_blob_store_or_create(ST_TABLE_ID, &schemas[ST_TABLE_IDX]);
         // Insert the table row into `st_tables` for all system tables
         for schema in ref_schemas {
             let table_id = schema.table_id;
@@ -175,11 +197,12 @@ impl CommittedState {
             let row = ProductValue::from(row);
             // Insert the meta-row into the in-memory ST_TABLES.
             // If the row is already there, no-op.
-            ignore_duplicate_insert_error(st_tables.insert(blob_store, &row))?;
+            ignore_duplicate_insert_error(st_tables.insert(pool, blob_store, &row))?;
         }
 
         // Insert the columns into `st_columns`
-        let (st_columns, blob_store) = self.get_table_and_blob_store_or_create(ST_COLUMN_ID, &schemas[ST_COLUMN_IDX]);
+        let (st_columns, blob_store, pool) =
+            self.get_table_and_blob_store_or_create(ST_COLUMN_ID, &schemas[ST_COLUMN_IDX]);
         for col in ref_schemas.iter().flat_map(|x| x.columns()).cloned() {
             let row = StColumnRow {
                 table_id: col.table_id,
@@ -190,15 +213,15 @@ impl CommittedState {
             let row = ProductValue::from(row);
             // Insert the meta-row into the in-memory ST_COLUMNS.
             // If the row is already there, no-op.
-            ignore_duplicate_insert_error(st_columns.insert(blob_store, &row))?;
+            ignore_duplicate_insert_error(st_columns.insert(pool, blob_store, &row))?;
             // Increment row count for st_columns.
             with_label_values(ST_COLUMN_ID, ST_COLUMN_NAME).inc();
         }
 
         // Insert the FK sorted by table/column so it show together when queried.
 
         // Insert constraints into `st_constraints`
-        let (st_constraints, blob_store) =
+        let (st_constraints, blob_store, pool) =
             self.get_table_and_blob_store_or_create(ST_CONSTRAINT_ID, &schemas[ST_CONSTRAINT_IDX]);
         for (i, constraint) in ref_schemas
             .iter()
@@ -219,13 +242,14 @@ impl CommittedState {
             let row = ProductValue::from(row);
             // Insert the meta-row into the in-memory ST_CONSTRAINTS.
             // If the row is already there, no-op.
-            ignore_duplicate_insert_error(st_constraints.insert(blob_store, &row))?;
+            ignore_duplicate_insert_error(st_constraints.insert(pool, blob_store, &row))?;
             // Increment row count for st_constraints.
             with_label_values(ST_CONSTRAINT_ID, ST_CONSTRAINT_NAME).inc();
         }
 
         // Insert the indexes into `st_indexes`
-        let (st_indexes, blob_store) = self.get_table_and_blob_store_or_create(ST_INDEX_ID, &schemas[ST_INDEX_IDX]);
+        let (st_indexes, blob_store, pool) =
+            self.get_table_and_blob_store_or_create(ST_INDEX_ID, &schemas[ST_INDEX_IDX]);
         for (i, mut index) in ref_schemas
             .iter()
             .flat_map(|x| &x.indexes)
@@ -240,7 +264,7 @@ impl CommittedState {
             let row = ProductValue::from(row);
             // Insert the meta-row into the in-memory ST_INDEXES.
             // If the row is already there, no-op.
-            ignore_duplicate_insert_error(st_indexes.insert(blob_store, &row))?;
+            ignore_duplicate_insert_error(st_indexes.insert(pool, blob_store, &row))?;
             // Increment row count for st_indexes.
             with_label_values(ST_INDEX_ID, ST_INDEX_NAME).inc();
         }
@@ -260,7 +284,7 @@ impl CommittedState {
         // IMPORTANT: It is crucial that the `st_sequences` table is created last
 
         // Insert the sequences into `st_sequences`
-        let (st_sequences, blob_store) =
+        let (st_sequences, blob_store, pool) =
             self.get_table_and_blob_store_or_create(ST_SEQUENCE_ID, &schemas[ST_SEQUENCE_IDX]);
         // We create sequences last to get right the starting number
         // so, we don't sort here
@@ -286,7 +310,7 @@ impl CommittedState {
             let row = ProductValue::from(row);
             // Insert the meta-row into the in-memory ST_SEQUENCES.
             // If the row is already there, no-op.
-            ignore_duplicate_insert_error(st_sequences.insert(blob_store, &row))?;
+            ignore_duplicate_insert_error(st_sequences.insert(pool, blob_store, &row))?;
             // Increment row count for st_sequences
             with_label_values(ST_SEQUENCE_ID, ST_SEQUENCE_NAME).inc();
         }
@@ -320,7 +344,7 @@ impl CommittedState {
             .ok_or(TableError::IdNotFoundState(table_id))?;
         let blob_store = &mut self.blob_store;
         table
-            .delete_equal_row(blob_store, rel)
+            .delete_equal_row(&self.page_pool, blob_store, rel)
             .map_err(TableError::Bflatn)?
             .ok_or_else(|| anyhow!("Delete for non-existent row when replaying transaction"))?;
         Ok(())
@@ -332,8 +356,8 @@ impl CommittedState {
         schema: &Arc<TableSchema>,
         row: &ProductValue,
     ) -> Result<()> {
-        let (table, blob_store) = self.get_table_and_blob_store_or_create(table_id, schema);
-        table.insert(blob_store, row).map(drop).map_err(|e| match e {
+        let (table, blob_store, pool) = self.get_table_and_blob_store_or_create(table_id, schema);
+        table.insert(pool, blob_store, row).map(drop).map_err(|e| match e {
             InsertError::Bflatn(e) => TableError::Bflatn(e).into(),
             InsertError::Duplicate(e) => TableError::Duplicate(e).into(),
             InsertError::IndexError(e) => IndexError::UniqueConstraintViolation(e).into(),
@@ -553,9 +577,8 @@ impl CommittedState {
                     deletes.push(pv);
                 }
 
-                let table_name = &*table.get_schema().table_name;
-
                 if !deletes.is_empty() {
+                    let table_name = &*table.get_schema().table_name;
                     tx_data.set_deletes_for_table(table_id, table_name, deletes.into());
                 }
             } else if !row_ptrs.is_empty() {
@@ -579,43 +602,47 @@ impl CommittedState {
         //             and the fullness of the page.
 
         for (table_id, tx_table) in insert_tables {
-            let (commit_table, commit_blob_store) =
+            let (commit_table, commit_blob_store, page_pool) =
                 self.get_table_and_blob_store_or_create(table_id, tx_table.get_schema());
 
-            // TODO(perf): Allocate with capacity?
-            let mut inserts = vec![];
             // For each newly-inserted row, insert it into the committed state.
+            let mut inserts = Vec::with_capacity(tx_table.row_count as usize);
             for row_ref in tx_table.scan_rows(&tx_blob_store) {
                 let pv = row_ref.to_product_value();
                 commit_table
-                    .insert(commit_blob_store, &pv)
+                    .insert(page_pool, commit_blob_store, &pv)
                     .expect("Failed to insert when merging commit");
 
                 inserts.push(pv);
             }
 
-            let table_name = &*commit_table.get_schema().table_name;
-
+            // Add the table to `TxData` if there were insertions.
             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());
             }
 
+            let (schema, indexes, pages) = tx_table.consume_for_merge();
+
             // Add all newly created indexes to the committed state.
-            for (cols, mut index) in tx_table.indexes {
-                if !commit_table.indexes.contains_key(&cols) {
+            for (index_id, mut index) in indexes {
+                if !commit_table.indexes.contains_key(&index_id) {
                     index.clear();
                     // SAFETY: `tx_table` is derived from `commit_table`,
                     // so they have the same row type.
                     // This entails that all indices in `tx_table`
                     // were constructed with the same row type/layout as `commit_table`.
-                    unsafe { commit_table.insert_index(commit_blob_store, cols, index) };
+                    unsafe { commit_table.insert_index(commit_blob_store, index_id, index) };
                 }
             }
 
             // The schema may have been modified in the transaction.
             // Update this last to placate borrowck and avoid a clone.
             // None of the above operations will inspect the schema.
-            commit_table.schema = tx_table.schema;
+            commit_table.schema = schema;
+
+            // Put all the pages in the table back into the pool.
+            self.page_pool.put_many(pages);
         }
     }
 
@@ -642,8 +669,8 @@ impl CommittedState {
         self.tables.get(&table_id)
     }
 
-    pub(super) fn get_table_mut(&mut self, table_id: TableId) -> Option<&mut Table> {
-        self.tables.get_mut(&table_id)
+    pub(super) fn get_table_mut(&mut self, table_id: TableId) -> (Option<&mut Table>, &PagePool) {
+        (self.tables.get_mut(&table_id), &self.page_pool)
     }
 
     pub fn get_table_and_blob_store_immutable(&self, table_id: TableId) -> Option<(&Table, &dyn BlobStore)> {
@@ -670,13 +697,14 @@ impl CommittedState {
         &'this mut self,
         table_id: TableId,
         schema: &Arc<TableSchema>,
-    ) -> (&'this mut Table, &'this mut dyn BlobStore) {
+    ) -> (&'this mut Table, &'this mut dyn BlobStore, &'this PagePool) {
         let table = self
             .tables
             .entry(table_id)
             .or_insert_with(|| Self::make_table(schema.clone()));
         let blob_store = &mut self.blob_store;
-        (table, blob_store)
+        let pool = &mut self.page_pool;
+        (table, blob_store, pool)
     }
 
     pub fn report_data_size(&self, database_identity: Identity) {