Use spawn_blocking for the fallback commit in identity_dis/connected (#2730)

Author: coolreader18

crates/core/src/client/client_connection.rs

@@ -9,6 +9,7 @@ use crate::error::DBError;
 use crate::host::module_host::ClientConnectedError;
 use crate::host::{ModuleHost, NoSuchModule, ReducerArgs, ReducerCallError, ReducerCallResult};
 use crate::messages::websocket::Subscribe;
+use crate::util::asyncify;
 use crate::util::prometheus_handle::IntGaugeExt;
 use crate::worker_metrics::WORKER_METRICS;
 use bytes::Bytes;
@@ -302,24 +303,22 @@ impl ClientConnection {
 
     pub async fn subscribe_single(&self, subscription: SubscribeSingle, timer: Instant) -> Result<(), DBError> {
         let me = self.clone();
-        tokio::task::spawn_blocking(move || {
+        asyncify(move || {
             me.module
                 .subscriptions()
                 .add_single_subscription(me.sender, subscription, timer, None)
         })
         .await
-        .unwrap() // TODO: is unwrapping right here?
     }
 
     pub async fn unsubscribe(&self, request: Unsubscribe, timer: Instant) -> Result<(), DBError> {
         let me = self.clone();
-        tokio::task::spawn_blocking(move || {
+        asyncify(move || {
             me.module
                 .subscriptions()
                 .remove_single_subscription(me.sender, request, timer)
         })
         .await
-        .unwrap() // TODO: is unwrapping right here?
     }
 
     pub async fn subscribe_multi(
@@ -328,13 +327,12 @@ impl ClientConnection {
         timer: Instant,
     ) -> Result<Option<ExecutionMetrics>, DBError> {
         let me = self.clone();
-        tokio::task::spawn_blocking(move || {
+        asyncify(move || {
             me.module
                 .subscriptions()
                 .add_multi_subscription(me.sender, request, timer, None)
         })
         .await
-        .unwrap() // TODO: is unwrapping right here?
     }
 
     pub async fn unsubscribe_multi(
@@ -343,24 +341,22 @@ impl ClientConnection {
         timer: Instant,
     ) -> Result<Option<ExecutionMetrics>, DBError> {
         let me = self.clone();
-        tokio::task::spawn_blocking(move || {
+        asyncify(move || {
             me.module
                 .subscriptions()
                 .remove_multi_subscription(me.sender, request, timer)
         })
         .await
-        .unwrap() // TODO: is unwrapping right here?
     }
 
     pub async fn subscribe(&self, subscription: Subscribe, timer: Instant) -> Result<ExecutionMetrics, DBError> {
         let me = self.clone();
-        tokio::task::spawn_blocking(move || {
+        asyncify(move || {
             me.module
                 .subscriptions()
                 .add_legacy_subscriber(me.sender, subscription, timer, None)
         })
         .await
-        .unwrap()
     }
 
     pub fn one_off_query_json(&self, query: &str, message_id: &[u8], timer: Instant) -> Result<(), anyhow::Error> {

crates/core/src/database_logger.rs

@@ -6,6 +6,8 @@ use tokio::sync::broadcast;
 
 use spacetimedb_paths::server::{ModuleLogPath, ModuleLogsDir};
 
+use crate::util::asyncify;
+
 pub struct DatabaseLogger {
     inner: Mutex<DatabaseLoggerInner>,
     pub tx: broadcast::Sender<bytes::Bytes>,
@@ -176,22 +178,20 @@ impl DatabaseLogger {
             }
             // if there's none for today, read the directory and
             let logs_dir = path.popped();
-            return tokio::task::spawn_blocking(move || match logs_dir.most_recent()? {
+            return asyncify(move || match logs_dir.most_recent()? {
                 Some(newest_log_file) => std::fs::read_to_string(newest_log_file),
                 None => Ok(String::new()),
             })
             .await
-            .unwrap()
             .expect("couldn't read log file");
         };
 
         if num_lines == 0 {
             return String::new();
         }
 
-        tokio::task::spawn_blocking(move || read_latest_lines(logs_dir, num_lines))
+        asyncify(move || read_latest_lines(logs_dir, num_lines))
             .await
-            .unwrap()
             .expect("couldn't read log file")
     }
 

crates/core/src/db/relational_db.rs

@@ -19,7 +19,7 @@ use crate::db::datastore::system_tables::{StModuleRow, WASM_MODULE};
 use crate::error::{DBError, DatabaseError, TableError};
 use crate::execution_context::{ReducerContext, Workload};
 use crate::messages::control_db::HostType;
-use crate::util::spawn_rayon;
+use crate::util::{asyncify, spawn_rayon};
 use anyhow::{anyhow, Context};
 use fs2::FileExt;
 use futures::channel::mpsc;
@@ -157,15 +157,14 @@ impl SnapshotWorkerActor {
         let start_time = std::time::Instant::now();
         let committed_state = self.committed_state.clone();
         let snapshot_repo = self.repo.clone();
-        let res = tokio::task::spawn_blocking(move || {
+        let res = asyncify(move || {
             Locking::take_snapshot_internal(&committed_state, &snapshot_repo).inspect(|opts| {
                 if let Some(opts) = opts {
                     Locking::compress_older_snapshot_internal(&snapshot_repo, opts.0);
                 }
             })
         })
-        .await
-        .unwrap();
+        .await;
         match res {
             Err(e) => {
                 log::error!(

crates/core/src/host/host_controller.rs

@@ -13,7 +13,7 @@ use crate::module_host_context::ModuleCreationContext;
 use crate::replica_context::ReplicaContext;
 use crate::subscription::module_subscription_actor::ModuleSubscriptions;
 use crate::subscription::module_subscription_manager::SubscriptionManager;
-use crate::util::spawn_rayon;
+use crate::util::{asyncify, spawn_rayon};
 use anyhow::{anyhow, ensure, Context};
 use async_trait::async_trait;
 use durability::{Durability, EmptyHistory};
@@ -282,13 +282,11 @@ impl HostController {
         trace!("using database {}/{}", database.database_identity, replica_id);
         let module = self.get_or_launch_module_host(database, replica_id).await?;
         let on_panic = self.unregister_fn(replica_id);
-        let result = tokio::task::spawn_blocking(move || f(&module.replica_ctx().relational_db))
-            .await
-            .unwrap_or_else(|e| {
-                warn!("database operation panicked");
-                on_panic();
-                std::panic::resume_unwind(e.into_panic())
-            });
+        scopeguard::defer_on_unwind!({
+            warn!("database operation panicked");
+            on_panic();
+        });
+        let result = asyncify(move || f(&module.replica_ctx().relational_db)).await;
         Ok(result)
     }
 
@@ -539,11 +537,7 @@ async fn make_replica_ctx(
             let Some(subscriptions) = downgraded.upgrade() else {
                 break;
             };
-            tokio::task::spawn_blocking(move || {
-                subscriptions.write().remove_dropped_clients();
-            })
-            .await
-            .unwrap();
+            asyncify(move || subscriptions.write().remove_dropped_clients()).await
         }
     });
 

crates/core/src/host/module_host.rs

@@ -15,6 +15,7 @@ use crate::sql::ast::SchemaViewer;
 use crate::subscription::module_subscription_actor::ModuleSubscriptions;
 use crate::subscription::tx::DeltaTx;
 use crate::subscription::{execute_plan, record_exec_metrics};
+use crate::util::asyncify;
 use crate::util::lending_pool::{LendingPool, LentResource, PoolClosed};
 use crate::vm::check_row_limit;
 use crate::worker_metrics::WORKER_METRICS;
@@ -535,10 +536,7 @@ impl ModuleHost {
     pub async fn disconnect_client(&self, client_id: ClientActorId) {
         log::trace!("disconnecting client {}", client_id);
         let this = self.clone();
-        let _ = tokio::task::spawn_blocking(move || {
-            this.subscriptions().remove_subscriber(client_id);
-        })
-        .await;
+        asyncify(move || this.subscriptions().remove_subscriber(client_id)).await;
         // ignore NoSuchModule; if the module's already closed, that's fine
         if let Err(e) = self
             .call_identity_disconnected(client_id.identity, client_id.connection_id)
@@ -622,18 +620,18 @@ impl ModuleHost {
                 timestamp: Timestamp::now(),
                 arg_bsatn: Bytes::new(),
             });
-            self.inner
-                .replica_ctx()
-                .relational_db
-                .with_auto_commit(workload, |mut_tx| {
+
+            let stdb = self.inner.replica_ctx().relational_db.clone();
+            asyncify(move || {
+                stdb.with_auto_commit(workload, |mut_tx| {
                     mut_tx.insert_st_client(caller_identity, caller_connection_id)
                 })
-                .inspect_err(|e| {
-                    log::error!(
-                        "`call_identity_connected`: fallback transaction to insert into `st_client` failed: {e:#?}"
-                    );
-                })
-                .map_err(Into::into)
+            })
+            .await
+            .inspect_err(|e| {
+                log::error!("`call_identity_connected`: fallback transaction to insert into `st_client` failed: {e:#?}")
+            })
+            .map_err(Into::into)
         }
     }
 
@@ -660,7 +658,7 @@ impl ModuleHost {
         let reducer_lookup = self.info.module_def.lifecycle_reducer(Lifecycle::OnDisconnect);
 
         // A fallback transaction that deletes the client from `st_client`.
-        let fallback = || {
+        let fallback = || async {
             let reducer_name = reducer_lookup
                 .as_ref()
                 .map(|(_, def)| &*def.name)
@@ -673,24 +671,24 @@ impl ModuleHost {
                 timestamp: Timestamp::now(),
                 arg_bsatn: Bytes::new(),
             });
-            self.inner
-                .replica_ctx()
-                .relational_db
-                .with_auto_commit(workload, |mut_tx| {
-                    mut_tx.delete_st_client(caller_identity, caller_connection_id, self.info.database_identity)
-                })
-                .inspect_err(|e| {
-                    log::error!(
-                        "`call_identity_disconnected`: fallback transaction to delete from `st_client` failed: {e}"
-                    );
-                })
-                .map_err(|err| {
-                    InvalidReducerArguments {
-                        err: err.into(),
-                        reducer: reducer_name.into(),
-                    }
-                    .into()
+            let stdb = self.inner.replica_ctx().relational_db.clone();
+            let database_identity = self.info.database_identity;
+            asyncify(move || {
+                stdb.with_auto_commit(workload, |mut_tx| {
+                    mut_tx.delete_st_client(caller_identity, caller_connection_id, database_identity)
                 })
+            })
+            .await
+            .map_err(|err| {
+                log::error!(
+                    "`call_identity_disconnected`: fallback transaction to delete from `st_client` failed: {err}"
+                );
+                InvalidReducerArguments {
+                    err: err.into(),
+                    reducer: reducer_name.into(),
+                }
+                .into()
+            })
         };
 
         if let Some((reducer_id, reducer_def)) = reducer_lookup {
@@ -720,12 +718,12 @@ impl ModuleHost {
             match result {
                 Err(e) => {
                     log::error!("call_reducer_inner of client_disconnected failed: {e:#?}");
-                    fallback()
+                    fallback().await
                 }
                 Ok(ReducerCallResult {
                     outcome: ReducerOutcome::Failed(_) | ReducerOutcome::BudgetExceeded,
                     ..
-                }) => fallback(),
+                }) => fallback().await,
 
                 // If it succeeded, as mentioend above, `st_client` is already updated.
                 Ok(ReducerCallResult {
@@ -736,7 +734,7 @@ impl ModuleHost {
         } else {
             // The module doesn't define a `client_disconnected` reducer.
             // Commit a transaction to update `st_clients`.
-            fallback()
+            fallback().await
         }
     }
 

crates/core/src/util/mod.rs

@@ -32,6 +32,23 @@ pub fn spawn_rayon<R: Send + 'static>(f: impl FnOnce() -> R + Send + 'static) ->
     rx.map(|res| res.unwrap().unwrap_or_else(|err| std::panic::resume_unwind(err)))
 }
 
+/// Ergonomic wrapper for `tokio::task::spawn_blocking(f).await`.
+///
+/// If `f` panics, it will be bubbled up to the calling task.
+pub async fn asyncify<F, R>(f: F) -> R
+where
+    F: FnOnce() -> R + Send + 'static,
+    R: Send + 'static,
+{
+    tokio::task::spawn_blocking(f)
+        .await
+        .unwrap_or_else(|e| match e.try_into_panic() {
+            Ok(panic_payload) => std::panic::resume_unwind(panic_payload),
+            // the only other variant is cancelled, which shouldn't happen because we don't cancel it.
+            Err(e) => panic!("Unexpected JoinError: {e}"),
+        })
+}
+
 /// Await `fut`, while also polling `also`.
 pub async fn also_poll<Fut: Future>(fut: Fut, also: impl Future<Output = ()>) -> Fut::Output {
     let mut also = pin!(also.fuse());