Build indexes over TxData for subscription eval (#2768)

Author: joshua-spacetime

crates/core/src/db/datastore/locking_tx_datastore/mut_tx.rs

@@ -25,7 +25,7 @@ use core::ops::RangeBounds;
 use core::{cell::RefCell, mem};
 use core::{iter, ops::Bound};
 use smallvec::SmallVec;
-use spacetimedb_execution::{dml::MutDatastore, Datastore, DeltaStore};
+use spacetimedb_execution::{dml::MutDatastore, Datastore, DeltaStore, Row};
 use spacetimedb_lib::{db::raw_def::v9::RawSql, metrics::ExecutionMetrics};
 use spacetimedb_lib::{
     db::{auth::StAccess, raw_def::SEQUENCE_ALLOCATION_STEP},
@@ -103,6 +103,30 @@ impl DeltaStore for MutTxId {
     fn deletes_for_table(&self, _: TableId) -> Option<std::slice::Iter<'_, ProductValue>> {
         None
     }
+
+    /// Subscriptions are currently evaluated using read-only transcations.
+    /// Hence this will never be called on a mutable transaction.
+    fn index_scan_range_for_delta(
+        &self,
+        _: TableId,
+        _: IndexId,
+        _: spacetimedb_lib::query::Delta,
+        _: impl RangeBounds<AlgebraicValue>,
+    ) -> impl Iterator<Item = Row> {
+        std::iter::empty()
+    }
+
+    /// Subscriptions are currently evaluated using read-only transcations.
+    /// Hence this will never be called on a mutable transaction.
+    fn index_scan_point_for_delta(
+        &self,
+        _: TableId,
+        _: IndexId,
+        _: spacetimedb_lib::query::Delta,
+        _: &AlgebraicValue,
+    ) -> impl Iterator<Item = Row> {
+        std::iter::empty()
+    }
 }
 
 impl MutDatastore for MutTxId {

crates/core/src/db/datastore/traits.rs

@@ -218,6 +218,11 @@ impl TxData {
         self.inserts.iter()
     }
 
+    /// Get the `i`th inserted row for `table_id` if it exists
+    pub fn get_ith_insert(&self, table_id: TableId, i: usize) -> Option<&ProductValue> {
+        self.inserts.get(&table_id).and_then(|rows| rows.get(i))
+    }
+
     /// Obtain an iterator over the inserted rows per table.
     ///
     /// If you don't need access to the table name, [`Self::inserts`] is
@@ -237,6 +242,11 @@ impl TxData {
         self.deletes.iter()
     }
 
+    /// 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))
+    }
+
     /// Obtain an iterator over the inserted rows per table.
     ///
     /// If you don't need access to the table name, [`Self::deletes`] is

crates/core/src/subscription/module_subscription_actor.rs

@@ -699,7 +699,7 @@ impl ModuleSubscriptions {
         // Create the delta transaction we'll use to eval updates against.
         let delta_read_tx = tx_data
             .as_ref()
-            .map(|tx_data| DeltaTx::new(&read_tx, tx_data))
+            .map(|tx_data| DeltaTx::new(&read_tx, tx_data, subscriptions.index_ids_for_subscriptions()))
             .unwrap_or_else(|| DeltaTx::from(&*read_tx));
 
         let event = Arc::new(event);

crates/core/src/subscription/module_subscription_manager.rs

@@ -23,7 +23,7 @@ use spacetimedb_client_api_messages::websocket::{
 use spacetimedb_data_structures::map::{Entry, IntMap};
 use spacetimedb_lib::metrics::ExecutionMetrics;
 use spacetimedb_lib::{AlgebraicValue, ConnectionId, Identity, ProductValue};
-use spacetimedb_primitives::{ColId, TableId};
+use spacetimedb_primitives::{ColId, IndexId, TableId};
 use spacetimedb_subscription::{SubscriptionPlan, TableName};
 use std::collections::BTreeSet;
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -78,6 +78,15 @@ impl Plan {
         self.plans[0].subscribed_table_name()
     }
 
+    /// Returns the index ids from which this subscription reads
+    pub fn index_ids(&self) -> impl Iterator<Item = (TableId, IndexId)> {
+        self.plans
+            .iter()
+            .flat_map(|plan| plan.index_ids())
+            .collect::<HashSet<_>>()
+            .into_iter()
+    }
+
     /// Returns the table ids from which this subscription reads
     pub fn table_ids(&self) -> impl Iterator<Item = TableId> + '_ {
         self.plans
@@ -174,6 +183,11 @@ impl QueryState {
         itertools::chain(&self.legacy_subscribers, &self.subscriptions)
     }
 
+    /// Return the [`Query`] for this [`QueryState`]
+    pub fn query(&self) -> &Query {
+        &self.query
+    }
+
     /// Return the search arguments for this query
     fn search_args(&self) -> impl Iterator<Item = (TableId, ColId, AlgebraicValue)> {
         let mut args = HashSet::new();
@@ -305,6 +319,79 @@ impl SearchArguments {
 
 type ClientsMap = HashMap<ClientId, ClientInfo>;
 
+/// Keeps track of the indexes that are used in subscriptions.
+#[derive(Debug, Default)]
+pub struct QueriedTableIndexIds {
+    ids: HashMap<TableId, HashMap<IndexId, usize>>,
+}
+
+impl FromIterator<(TableId, IndexId)> for QueriedTableIndexIds {
+    fn from_iter<T: IntoIterator<Item = (TableId, IndexId)>>(iter: T) -> Self {
+        let mut index_ids = Self::default();
+        for (table_id, index_id) in iter {
+            index_ids.insert_index_id(table_id, index_id);
+        }
+        index_ids
+    }
+}
+
+impl QueriedTableIndexIds {
+    /// Returns the index ids that are used in subscriptions for this table.
+    /// Note, it does not return all of the index ids that are defined on this table.
+    /// Only those that are used by at least one subscription query.
+    pub fn index_ids_for_table(&self, table_id: TableId) -> impl Iterator<Item = IndexId> + '_ {
+        self.ids
+            .get(&table_id)
+            .into_iter()
+            .flat_map(|index_ids| index_ids.keys())
+            .copied()
+    }
+
+    /// Insert a new `table_id` `index_id` pair into this container.
+    /// Note, different queries may read from the same index.
+    /// Hence we may already be tracking this index, in which case we bump its ref count.
+    pub fn insert_index_id(&mut self, table_id: TableId, index_id: IndexId) {
+        *self.ids.entry(table_id).or_default().entry(index_id).or_default() += 1;
+    }
+
+    /// Remove a `table_id` `index_id` pair from this container.
+    /// Note, different queries may read from the same index.
+    /// Hence we only remove this key from the map if its ref count goes to zero.
+    pub fn delete_index_id(&mut self, table_id: TableId, index_id: IndexId) {
+        if let Some(ids) = self.ids.get_mut(&table_id) {
+            if let Some(n) = ids.get_mut(&index_id) {
+                *n -= 1;
+
+                if *n == 0 {
+                    ids.remove(&index_id);
+
+                    if ids.is_empty() {
+                        self.ids.remove(&table_id);
+                    }
+                }
+            }
+        }
+    }
+
+    /// Insert the index ids from which a query reads into this mapping.
+    /// Note, an index may already be tracked if another query is already using it.
+    /// In this case we just bump its ref count.
+    pub fn insert_index_ids_for_query(&mut self, query: &Query) {
+        for (table_id, index_id) in query.index_ids() {
+            self.insert_index_id(table_id, index_id);
+        }
+    }
+
+    /// Delete the index ids from which a query reads from this mapping
+    /// Note, we will not remove an index id from this mapping if another query is using it.
+    /// Instead we decrement its ref count.
+    pub fn delete_index_ids_for_query(&mut self, query: &Query) {
+        for (table_id, index_id) in query.index_ids() {
+            self.delete_index_id(table_id, index_id);
+        }
+    }
+}
+
 /// Responsible for the efficient evaluation of subscriptions.
 /// It performs basic multi-query optimization,
 /// in that if a query has N subscribers,
@@ -318,17 +405,21 @@ pub struct SubscriptionManager {
     /// in order to dispatch messages to clients.
     clients: Arc<RwLock<ClientsMap>>,
 
-    // Queries for which there is at least one subscriber.
+    /// Queries for which there is at least one subscriber.
     queries: HashMap<QueryHash, QueryState>,
 
-    // If a query reads from a table,
-    // but does not have a simple equality filter on that table,
-    // we map the table to the query in this inverted index.
+    /// If a query reads from a table,
+    /// but does not have a simple equality filter on that table,
+    /// we map the table to the query in this inverted index.
     tables: IntMap<TableId, HashSet<QueryHash>>,
 
-    // If a query reads from a table,
-    // and has a simple equality filter on that table,
-    // we map the filter values to the query in this lookup table.
+    /// Tracks the indices used across all subscriptions
+    /// to enable building the appropriate indexes for row updates.
+    indexes: QueriedTableIndexIds,
+
+    /// If a query reads from a table,
+    /// and has a simple equality filter on that table,
+    /// we map the filter values to the query in this lookup table.
     search_args: SearchArguments,
 
     /// Transmit side of a channel to the manager's [`Self::send_worker`] task.
@@ -418,6 +509,7 @@ impl SubscriptionManager {
         Self {
             clients: Default::default(),
             queries: Default::default(),
+            indexes: Default::default(),
             tables: Default::default(),
             search_args: Default::default(),
             send_worker_tx,
@@ -506,6 +598,7 @@ impl SubscriptionManager {
                 if !query_state.has_subscribers() {
                     SubscriptionManager::remove_query_from_tables(
                         &mut self.tables,
+                        &mut self.indexes,
                         &mut self.search_args,
                         &query_state.query,
                     );
@@ -572,6 +665,7 @@ impl SubscriptionManager {
             if !query_state.has_subscribers() {
                 SubscriptionManager::remove_query_from_tables(
                     &mut self.tables,
+                    &mut self.indexes,
                     &mut self.search_args,
                     &query_state.query,
                 );
@@ -639,7 +733,7 @@ impl SubscriptionManager {
                 .entry(hash)
                 .or_insert_with(|| QueryState::new(query.clone()));
 
-            Self::insert_query(&mut self.tables, &mut self.search_args, query_state);
+            Self::insert_query(&mut self.tables, &mut self.indexes, &mut self.search_args, query_state);
 
             let entry = ci.subscription_ref_count.entry(hash).or_insert(0);
             *entry += 1;
@@ -687,7 +781,7 @@ impl SubscriptionManager {
                 .queries
                 .entry(hash)
                 .or_insert_with(|| QueryState::new(unit.clone()));
-            Self::insert_query(&mut self.tables, &mut self.search_args, query_state);
+            Self::insert_query(&mut self.tables, &mut self.indexes, &mut self.search_args, query_state);
             query_state.legacy_subscribers.insert(client_id);
         }
     }
@@ -697,11 +791,13 @@ impl SubscriptionManager {
     // This takes a ref to the table map instead of `self` to avoid borrowing issues.
     fn remove_query_from_tables(
         tables: &mut IntMap<TableId, HashSet<QueryHash>>,
+        index_ids: &mut QueriedTableIndexIds,
         search_args: &mut SearchArguments,
         query: &Query,
     ) {
         let hash = query.hash();
         search_args.remove_query(&hash);
+        index_ids.delete_index_ids_for_query(query);
         for table_id in query.table_ids() {
             if let Entry::Occupied(mut entry) = tables.entry(table_id) {
                 let hashes = entry.get_mut();
@@ -717,11 +813,13 @@ impl SubscriptionManager {
     // This takes a ref to the table map instead of `self` to avoid borrowing issues.
     fn insert_query(
         tables: &mut IntMap<TableId, HashSet<QueryHash>>,
+        index_ids: &mut QueriedTableIndexIds,
         search_args: &mut SearchArguments,
         query_state: &QueryState,
     ) {
         // If this is new, we need to update the table to query mapping.
         if !query_state.has_subscribers() {
+            index_ids.insert_index_ids_for_query(query_state.query());
             let hash = query_state.query.hash();
             let mut table_ids = query_state.query.table_ids().collect::<HashSet<_>>();
             for (table_id, col_id, arg) in query_state.search_args() {
@@ -760,6 +858,7 @@ impl SubscriptionManager {
                 queries_to_remove.push(*query_hash);
                 SubscriptionManager::remove_query_from_tables(
                     &mut self.tables,
+                    &mut self.indexes,
                     &mut self.search_args,
                     &query_state.query,
                 );
@@ -825,6 +924,11 @@ impl SubscriptionManager {
         queries.into_iter()
     }
 
+    /// Returns the index ids that are used in subscription queries
+    pub fn index_ids_for_subscriptions(&self) -> &QueriedTableIndexIds {
+        &self.indexes
+    }
+
     /// This method takes a set of delta tables,
     /// evaluates only the necessary queries for those delta tables,
     /// and then sends the results to each client.