Materialize views for subscriptions

Author: joshua-spacetime

crates/core/src/host/module_host.rs

@@ -1118,7 +1118,7 @@ impl ModuleHost {
         // Decrement the number of subscribers for each view this caller is subscribed to
         let dec_view_subscribers = |tx: &mut MutTxId| {
             if drop_view_subscribers {
-                if let Err(err) = tx.dec_st_view_subscribers(caller_identity) {
+                if let Err(err) = tx.st_view_dec_subscriber_count_for_sender(caller_identity) {
                     log::error!("`call_identity_disconnected`: failed to delete client view data: {err}");
                 }
             }

crates/core/src/subscription/module_subscription_actor.rs

@@ -33,9 +33,10 @@ use spacetimedb_datastore::db_metrics::DB_METRICS;
 use spacetimedb_datastore::execution_context::{Workload, WorkloadType};
 use spacetimedb_datastore::locking_tx_datastore::datastore::TxMetrics;
 use spacetimedb_datastore::locking_tx_datastore::{MutTxId, TxId};
-use spacetimedb_datastore::traits::TxData;
+use spacetimedb_datastore::traits::{IsolationLevel, TxData};
 use spacetimedb_durability::TxOffset;
 use spacetimedb_execution::pipelined::{PipelinedProject, ViewProject};
+use spacetimedb_expr::expr::CollectViews;
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_lib::metrics::ExecutionMetrics;
 use spacetimedb_lib::Identity;
@@ -400,7 +401,7 @@ impl ModuleSubscriptions {
         let hash = QueryHash::from_string(&sql, auth.caller, false);
         let hash_with_param = QueryHash::from_string(&sql, auth.caller, true);
 
-        let (tx, tx_offset) = self.begin_tx(Workload::Subscribe);
+        let (mut_tx, _) = self.begin_mut_tx(Workload::Subscribe);
 
         let existing_query = {
             let guard = self.subscriptions.read();
@@ -410,7 +411,7 @@ impl ModuleSubscriptions {
         let query = return_on_err_with_sql!(
             existing_query.map(Ok).unwrap_or_else(|| compile_query_with_hashes(
                 &auth,
-                &tx,
+                &*mut_tx,
                 &sql,
                 hash,
                 hash_with_param
@@ -420,6 +421,9 @@ impl ModuleSubscriptions {
             send_err_msg
         );
 
+        let mut_tx = ScopeGuard::<MutTxId, _>::into_inner(mut_tx);
+        let (tx, tx_offset) = self.materialize_views_and_downgrade_tx(mut_tx, &query, auth.caller)?;
+
         let (table_rows, metrics) = return_on_err_with_sql!(
             self.evaluate_initial_subscription(sender.clone(), query.clone(), &tx, &auth, TableUpdateType::Subscribe),
             query.sql(),
@@ -501,8 +505,10 @@ impl ModuleSubscriptions {
             return Ok(None);
         };
 
-        let (tx, tx_offset) = self.begin_tx(Workload::Unsubscribe);
         let auth = AuthCtx::new(self.owner_identity, sender.id.identity);
+
+        let (tx, tx_offset) = self.drop_views(query, auth.caller)?;
+
         let (table_rows, metrics) = return_on_err_with_sql!(
             self.evaluate_initial_subscription(sender.clone(), query.clone(), &tx, &auth, TableUpdateType::Unsubscribe),
             query.sql(),
@@ -558,11 +564,12 @@ impl ModuleSubscriptions {
             )
         };
 
+        let auth = AuthCtx::new(self.owner_identity, sender.id.identity);
         let database_identity = self.relational_db.database_identity();
         let subscription_metrics = SubscriptionMetrics::new(&database_identity, &WorkloadType::Unsubscribe);
 
         // Always lock the db before the subscription lock to avoid deadlocks.
-        let (tx, tx_offset) = self.begin_tx(Workload::Unsubscribe);
+        let (mut_tx, _) = self.begin_mut_tx(Workload::Unsubscribe);
 
         let removed_queries = {
             let _compile_timer = subscription_metrics.compilation_time.start_timer();
@@ -580,12 +587,15 @@ impl ModuleSubscriptions {
             )
         };
 
+        let mut_tx = ScopeGuard::<MutTxId, _>::into_inner(mut_tx);
+        let (tx, tx_offset) = self.drop_views_and_downgrade_tx(mut_tx, &removed_queries, auth.caller)?;
+
         let (update, metrics) = return_on_err!(
             self.evaluate_queries(
                 sender.clone(),
                 &removed_queries,
                 &tx,
-                &AuthCtx::new(self.owner_identity, sender.id.identity),
+                &auth,
                 TableUpdateType::Unsubscribe,
             ),
             send_err_msg,
@@ -637,7 +647,7 @@ impl ModuleSubscriptions {
         queries: &[Box<str>],
         num_queries: usize,
         metrics: &SubscriptionMetrics,
-    ) -> Result<(Vec<Arc<Plan>>, AuthCtx, TxId, HistogramTimer), DBError> {
+    ) -> Result<(Vec<Arc<Plan>>, AuthCtx, MutTxId, HistogramTimer), DBError> {
         let mut subscribe_to_all_tables = false;
         let mut plans = Vec::with_capacity(num_queries);
         let mut query_hashes = Vec::with_capacity(num_queries);
@@ -656,7 +666,7 @@ impl ModuleSubscriptions {
         let auth = AuthCtx::new(self.owner_identity, sender);
 
         // We always get the db lock before the subscription lock to avoid deadlocks.
-        let (tx, _tx_offset) = self.begin_tx(Workload::Subscribe);
+        let (mut_tx, _tx_offset) = self.begin_mut_tx(Workload::Subscribe);
 
         let compile_timer = metrics.compilation_time.start_timer();
 
@@ -669,9 +679,14 @@ impl ModuleSubscriptions {
 
         if subscribe_to_all_tables {
             plans.extend(
-                super::subscription::get_all(&self.relational_db, &tx, &auth)?
-                    .into_iter()
-                    .map(Arc::new),
+                super::subscription::get_all(
+                    |relational_db, tx| relational_db.get_all_tables_mut(tx).map(|schemas| schemas.into_iter()),
+                    &self.relational_db,
+                    &*mut_tx,
+                    &auth,
+                )?
+                .into_iter()
+                .map(Arc::new),
             );
         }
 
@@ -684,7 +699,7 @@ impl ModuleSubscriptions {
                 plans.push(unit);
             } else {
                 plans.push(Arc::new(
-                    compile_query_with_hashes(&auth, &tx, sql, hash, hash_with_param).map_err(|err| {
+                    compile_query_with_hashes(&auth, &*mut_tx, sql, hash, hash_with_param).map_err(|err| {
                         DBError::WithSql {
                             error: Box::new(DBError::Other(err.into())),
                             sql: sql.into(),
@@ -698,7 +713,7 @@ impl ModuleSubscriptions {
         // How many queries in this subscription are not cached?
         metrics.num_new_queries_subscribed.inc_by(new_queries);
 
-        Ok((plans, auth, scopeguard::ScopeGuard::into_inner(tx), compile_timer))
+        Ok((plans, auth, ScopeGuard::<MutTxId, _>::into_inner(mut_tx), compile_timer))
     }
 
     /// Send a message to a client connection.
@@ -763,7 +778,7 @@ impl ModuleSubscriptions {
         // How many queries make up this subscription?
         subscription_metrics.num_queries_subscribed.inc_by(num_queries as _);
 
-        let (queries, auth, tx, compile_timer) = return_on_err!(
+        let (queries, auth, mut_tx, compile_timer) = return_on_err!(
             self.compile_queries(
                 sender.id.identity,
                 &request.query_strings,
@@ -773,7 +788,7 @@ impl ModuleSubscriptions {
             send_err_msg,
             None
         );
-        let (tx, tx_offset) = self.guard_tx(tx, <_>::default());
+        let (mut_tx, _) = self.guard_mut_tx(mut_tx, <_>::default());
 
         // We minimize locking so that other clients can add subscriptions concurrently.
         // We are protected from race conditions with broadcasts, because we have the db lock,
@@ -793,6 +808,9 @@ impl ModuleSubscriptions {
         // Record how long it took to compile the subscription
         drop(compile_timer);
 
+        let mut_tx = ScopeGuard::<MutTxId, _>::into_inner(mut_tx);
+        let (tx, tx_offset) = self.materialize_views_and_downgrade_tx(mut_tx, &queries, auth.caller)?;
+
         let Ok((update, metrics)) =
             self.evaluate_queries(sender.clone(), &queries, &tx, &auth, TableUpdateType::Subscribe)
         else {
@@ -859,13 +877,14 @@ impl ModuleSubscriptions {
         // How many queries make up this subscription?
         subscription_metrics.num_queries_subscribed.inc_by(num_queries as _);
 
-        let (queries, auth, tx, compile_timer) = self.compile_queries(
+        let (queries, auth, mut_tx, compile_timer) = self.compile_queries(
             sender.id.identity,
             &subscription.query_strings,
             num_queries,
             &subscription_metrics,
         )?;
-        let (tx, tx_offset) = self.guard_tx(tx, <_>::default());
+
+        let (tx, tx_offset) = self.materialize_views_and_downgrade_tx(mut_tx, &queries, auth.caller)?;
 
         check_row_limit(
             &queries,
@@ -1015,13 +1034,61 @@ impl ModuleSubscriptions {
         }))
     }
 
-    /// Helper that starts a new read transaction, and guards it using
-    /// [`Self::guard_tx`] with the default configuration.
-    fn begin_tx(&self, workload: Workload) -> (ScopeGuard<TxId, impl FnOnce(TxId) + '_>, TransactionOffset) {
-        self.guard_tx(self.relational_db.begin_tx(workload), <_>::default())
+    /// Collect view ids from `view_collector` and decrement their subscriber count.
+    fn drop_views(
+        &self,
+        view_collector: &impl CollectViews,
+        sender: Identity,
+    ) -> Result<(TxGuard<impl FnOnce(TxId) + '_>, TransactionOffset), DBError> {
+        use IsolationLevel::*;
+        use Workload::*;
+        let mut tx = self.relational_db.begin_mut_tx(Serializable, Unsubscribe);
+        tx.dec_subscriber_count_for_views(view_collector, sender)?;
+        let (tx_data, tx_metrics_mut, tx) = tx.commit_downgrade(Unsubscribe);
+        let opts = GuardTxOptions::from_mut(tx_data, tx_metrics_mut);
+        Ok(self.guard_tx(tx, opts))
+    }
+
+    /// Downgrade this mutable `tx` after:
+    /// 1. Collecting view ids from `view_collector` and
+    /// 2. Decrementing their subscriber count
+    fn drop_views_and_downgrade_tx(
+        &self,
+        mut tx: MutTxId,
+        view_collector: &impl CollectViews,
+        sender: Identity,
+    ) -> Result<(TxGuard<impl FnOnce(TxId) + '_>, TransactionOffset), DBError> {
+        tx.dec_subscriber_count_for_views(view_collector, sender)?;
+        let (tx_data, tx_metrics_mut, tx) = tx.commit_downgrade(Workload::Subscribe);
+        let opts = GuardTxOptions::from_mut(tx_data, tx_metrics_mut);
+        Ok(self.guard_tx(tx, opts))
+    }
+
+    /// Downgrade this mutable `tx` after:
+    /// 1. Collecting view ids from `view_collector` and
+    /// 2. Materializing them if necessary
+    fn materialize_views_and_downgrade_tx(
+        &self,
+        mut tx: MutTxId,
+        view_collector: &impl CollectViews,
+        sender: Identity,
+    ) -> Result<(TxGuard<impl FnOnce(TxId) + '_>, TransactionOffset), DBError> {
+        tx.materialize_views(view_collector, sender)?;
+        let (tx_data, tx_metrics_mut, tx) = tx.commit_downgrade(Workload::Subscribe);
+        let opts = GuardTxOptions::from_mut(tx_data, tx_metrics_mut);
+        Ok(self.guard_tx(tx, opts))
+    }
+
+    /// Helper that starts a new mutable transaction, and guards it using
+    /// [`Self::guard_mut_tx`] with the default configuration.
+    fn begin_mut_tx(&self, workload: Workload) -> (MutTxGuard<impl FnOnce(MutTxId) + '_>, TransactionOffset) {
+        self.guard_mut_tx(
+            self.relational_db.begin_mut_tx(IsolationLevel::Serializable, workload),
+            <_>::default(),
+        )
     }
 
-    /// Helper wrapping `tx` in a scopegard, with a configurable drop fn.
+    /// Helper wrapping a [`TxId`] in a scopegard, with a configurable drop fn.
     ///
     /// By default, `tx` is released when the returned [`ScopeGuard`] is dropped,
     /// and reports the transaction metrics via [`RelationalDB::report_tx_metrics`].
@@ -1036,27 +1103,42 @@ impl ModuleSubscriptions {
     /// If another receiver of the transaction offset is needed, its sending
     /// side can be passed in as `extra_tx_offset_sender`. It will be sent the
     /// offset as well.
-    fn guard_tx(
-        &self,
-        tx: TxId,
-        GuardTxOptions {
-            extra_tx_offset_sender,
-            tx_data,
-            tx_metrics_mut,
-        }: GuardTxOptions,
-    ) -> (ScopeGuard<TxId, impl FnOnce(TxId) + '_>, TransactionOffset) {
+    fn guard_tx(&self, tx: TxId, opts: GuardTxOptions) -> (TxGuard<impl FnOnce(TxId) + '_>, TransactionOffset) {
         let (offset_tx, offset_rx) = oneshot::channel();
         let guard = scopeguard::guard(tx, |tx| {
             let (tx_offset, tx_metrics, reducer) = self.relational_db.release_tx(tx);
             log::trace!("read tx released with offset {tx_offset}");
             let _ = offset_tx.send(tx_offset);
-            if let Some(extra) = extra_tx_offset_sender {
+            if let Some(extra) = opts.extra_tx_offset_sender {
                 let _ = extra.send(tx_offset);
             }
             self.relational_db
-                .report_tx_metrics(reducer, tx_data, tx_metrics_mut, Some(tx_metrics));
+                .report_tx_metrics(reducer, opts.tx_data, opts.tx_metrics_mut, Some(tx_metrics));
         });
+        (guard, offset_rx)
+    }
 
+    /// The same as [`Self::guard_tx`] but for mutable transactions.
+    ///
+    /// By default, `tx` is committed when the returned [`ScopeGuard`] is dropped,
+    /// and reports the transaction metrics via [`RelationalDB::report_tx_metrics`].
+    fn guard_mut_tx(
+        &self,
+        tx: MutTxId,
+        opts: GuardTxOptions,
+    ) -> (MutTxGuard<impl FnOnce(MutTxId) + '_>, TransactionOffset) {
+        let (offset_tx, offset_rx) = oneshot::channel();
+        let guard = scopeguard::guard(tx, |tx| {
+            if let Ok(Some((tx_offset, tx_data, tx_metrics_mut, reducer))) = self.relational_db.commit_tx(tx) {
+                log::trace!("mutable tx committed with offset {tx_offset}");
+                let _ = offset_tx.send(tx_offset);
+                if let Some(extra) = opts.extra_tx_offset_sender {
+                    let _ = extra.send(tx_offset);
+                }
+                self.relational_db
+                    .report_tx_metrics(reducer, Some(Arc::new(tx_data)), Some(tx_metrics_mut), None);
+            }
+        });
         (guard, offset_rx)
     }
 }
@@ -1084,10 +1166,24 @@ impl GuardTxOptions {
             tx_metrics_mut: tx_metrics_mut.into(),
         }
     }
+
+    fn from_mut(tx_data: TxData, tx_metrics_mut: TxMetrics) -> Self {
+        Self {
+            extra_tx_offset_sender: None,
+            tx_data: Some(Arc::new(tx_data)),
+            tx_metrics_mut: tx_metrics_mut.into(),
+        }
+    }
 }
 
 pub struct WriteConflict;
 
+/// A [`ScopeGuard`] for [`TxId`]
+type TxGuard<F> = ScopeGuard<TxId, F>;
+
+/// A [`ScopeGuard`] for [`MutTxId`]
+type MutTxGuard<F> = ScopeGuard<MutTxId, F>;
+
 #[cfg(test)]
 mod tests {
     use super::{AssertTxFn, ModuleSubscriptions};

crates/core/src/subscription/module_subscription_manager.rs

@@ -22,9 +22,10 @@ use spacetimedb_client_api_messages::websocket::{
 use spacetimedb_data_structures::map::{Entry, IntMap};
 use spacetimedb_datastore::locking_tx_datastore::state_view::StateView;
 use spacetimedb_durability::TxOffset;
+use spacetimedb_expr::expr::CollectViews;
 use spacetimedb_lib::metrics::ExecutionMetrics;
 use spacetimedb_lib::{AlgebraicValue, ConnectionId, Identity, ProductValue};
-use spacetimedb_primitives::{ColId, IndexId, TableId};
+use spacetimedb_primitives::{ColId, IndexId, TableId, ViewDatabaseId};
 use spacetimedb_subscription::{JoinEdge, SubscriptionPlan, TableName};
 use std::collections::BTreeMap;
 use std::fmt::Debug;
@@ -53,6 +54,14 @@ pub struct Plan {
     plans: Vec<SubscriptionPlan>,
 }
 
+impl CollectViews for Plan {
+    fn collect_views(&self, views: &mut std::collections::HashSet<ViewDatabaseId>) {
+        for plan in &self.plans {
+            plan.collect_views(views);
+        }
+    }
+}
+
 impl Plan {
     /// Create a new subscription plan to be cached
     pub fn new(plans: Vec<SubscriptionPlan>, hash: QueryHash, text: String) -> Self {

crates/core/src/subscription/query.rs

@@ -5,6 +5,8 @@ use crate::sql::compiler::compile_sql;
 use crate::subscription::subscription::SupportedQuery;
 use once_cell::sync::Lazy;
 use regex::Regex;
+use spacetimedb_datastore::locking_tx_datastore::state_view::StateView;
+use spacetimedb_execution::Datastore;
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_subscription::SubscriptionPlan;
 use spacetimedb_vm::expr::{self, Crud, CrudExpr, QueryExpr};
@@ -93,7 +95,7 @@ pub fn compile_read_only_query(auth: &AuthCtx, tx: &Tx, input: &str) -> Result<P
 
 /// Compile a string into a single read-only query.
 /// This returns an error if the string has multiple queries or mutations.
-pub fn compile_query_with_hashes(
+pub fn compile_query_with_hashes<Tx: Datastore + StateView>(
     auth: &AuthCtx,
     tx: &Tx,
     input: &str,