Materialize views for subscriptions

Author: joshua-spacetime

crates/core/src/client/client_connection.rs

@@ -865,10 +865,11 @@ impl ClientConnection {
     ) -> Result<Option<ExecutionMetrics>, DBError> {
         let me = self.clone();
         self.module()
-            .on_module_thread("subscribe_single", move || {
-                me.module()
-                    .subscriptions()
-                    .add_single_subscription(me.sender, subscription, timer, None)
+            .on_module_thread_async("subscribe_single", async move || {
+                let host = me.module();
+                host.subscriptions()
+                    .add_single_subscription(Some(&host), me.sender, subscription, timer, None)
+                    .await
             })
             .await?
     }
@@ -890,10 +891,11 @@ impl ClientConnection {
     ) -> Result<Option<ExecutionMetrics>, DBError> {
         let me = self.clone();
         self.module()
-            .on_module_thread("subscribe_multi", move || {
-                me.module()
-                    .subscriptions()
-                    .add_multi_subscription(me.sender, request, timer, None)
+            .on_module_thread_async("subscribe_multi", async move || {
+                let host = me.module();
+                host.subscriptions()
+                    .add_multi_subscription(Some(&host), me.sender, request, timer, None)
+                    .await
             })
             .await?
     }
@@ -915,12 +917,14 @@ impl ClientConnection {
 
     pub async fn subscribe(&self, subscription: Subscribe, timer: Instant) -> Result<ExecutionMetrics, DBError> {
         let me = self.clone();
-        asyncify(move || {
-            me.module()
-                .subscriptions()
-                .add_legacy_subscriber(me.sender, subscription, timer, None)
-        })
-        .await
+        self.module()
+            .on_module_thread_async("subscribe", async move || {
+                let host = me.module();
+                host.subscriptions()
+                    .add_legacy_subscriber(Some(&host), me.sender, subscription, timer, None)
+                    .await
+            })
+            .await?
     }
 
     pub async fn one_off_query_json(

crates/core/src/host/module_host.rs

@@ -39,8 +39,7 @@ use spacetimedb_data_structures::error_stream::ErrorStream;
 use spacetimedb_data_structures::map::{HashCollectionExt as _, IntMap};
 use spacetimedb_datastore::error::DatastoreError;
 use spacetimedb_datastore::execution_context::{ExecutionContext, ReducerContext, Workload, WorkloadType};
-use spacetimedb_datastore::locking_tx_datastore::datastore::TxMetrics;
-use spacetimedb_datastore::locking_tx_datastore::{MutTxId, TxId};
+use spacetimedb_datastore::locking_tx_datastore::MutTxId;
 use spacetimedb_datastore::system_tables::{ST_CLIENT_ID, ST_CONNECTION_CREDENTIALS_ID, ST_VIEW_SUB_ID};
 use spacetimedb_datastore::traits::{IsolationLevel, Program, TxData};
 use spacetimedb_durability::DurableOffset;
@@ -1120,7 +1119,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.unsubscribe_views(caller_identity) {
                     log::error!("`call_identity_disconnected`: failed to delete client view data: {err}");
                 }
             }
@@ -1489,16 +1488,19 @@ impl ModuleHost {
         .await?
     }
 
-    /// Downgrade this mutable `tx` after:
-    /// 1. Collecting view ids from `view_collector` and
-    /// 2. Materializing them if necessary
-    pub async fn materialize_views_and_downgrade_tx(
+    /// Materializes the views return by the `view_collector`, if not already materialized,
+    /// and updates `st_view_sub` accordingly.
+    ///
+    /// Passing [`Workload::Sql`] will update `st_view_sub.last_called`.
+    /// Passing [`Workload::Subscribe`] will also increment `st_view_sub.num_subscribers`,
+    /// in addition to updating `st_view_sub.last_called`.
+    pub async fn materialize_views(
         &self,
         mut tx: MutTxId,
         view_collector: &impl CollectViews,
         sender: Identity,
         workload: Workload,
-    ) -> Result<(TxData, TxMetrics, TxId), ViewCallError> {
+    ) -> Result<MutTxId, ViewCallError> {
         use FunctionArgs::*;
         let mut view_ids = HashSet::new();
         view_collector.collect_views(&mut view_ids);
@@ -1507,9 +1509,16 @@ impl ModuleHost {
             if !tx.is_view_materialized(view_id, ArgId::SENTINEL, sender)? {
                 tx = self.call_view(tx, &name, Nullary, sender, None).await?.tx;
             }
-            tx.st_view_sub_update_or_insert_last_called(view_id, ArgId::SENTINEL, sender)?;
+            // If this is a sql call, we only update this view's "last called" timestamp
+            if let Workload::Sql = workload {
+                tx.update_view_timestamp(view_id, ArgId::SENTINEL, sender)?;
+            }
+            // If this is a subscribe call, we also increment this view's subscriber count
+            if let Workload::Subscribe = workload {
+                tx.subscribe_view(view_id, ArgId::SENTINEL, sender)?;
+            }
         }
-        Ok(tx.commit_downgrade(workload))
+        Ok(tx)
     }
 
     pub async fn call_view(

crates/core/src/sql/execute.rs

@@ -194,24 +194,23 @@ pub async fn run(
     module: Option<&ModuleHost>,
     head: &mut Vec<(Box<str>, AlgebraicType)>,
 ) -> Result<SqlResult, DBError> {
-    let module = module
-        .as_ref()
-        .ok_or_else(|| anyhow!("Cannot execute views without module context"))?;
-
-    let mut metrics = ExecutionMetrics::default();
-
     // We parse the sql statement in a mutable transaction.
     // If it turns out to be a query, we downgrade the tx.
     let (tx, stmt) = db.with_auto_rollback(db.begin_mut_tx(IsolationLevel::Serializable, Workload::Sql), |tx| {
         compile_sql_stmt(sql_text, &SchemaViewer::new(tx, &auth), &auth)
     })?;
 
+    let mut metrics = ExecutionMetrics::default();
+
     match stmt {
         Statement::Select(stmt) => {
             // Materialize views and downgrade to a read-only transaction
-            let (tx_data, tx_metrics_mut, tx) = module
-                .materialize_views_and_downgrade_tx(tx, &stmt, auth.caller, Workload::Sql)
-                .await?;
+            let tx = match module {
+                Some(module) => module.materialize_views(tx, &stmt, auth.caller, Workload::Sql).await?,
+                None => tx,
+            };
+
+            let (tx_data, tx_metrics_mut, tx) = tx.commit_downgrade(Workload::Sql);
 
             let (tx_offset_send, tx_offset) = oneshot::channel();
             // Release the tx on drop, so that we record metrics