Move locking operations onto the module's thread (#2866)

Author: jsdt

crates/core/src/client/client_connection.rs

@@ -530,12 +530,13 @@ impl ClientConnection {
         timer: Instant,
     ) -> Result<Option<ExecutionMetrics>, DBError> {
         let me = self.clone();
-        asyncify(move || {
-            me.module
-                .subscriptions()
-                .add_single_subscription(me.sender, subscription, timer, None)
-        })
-        .await
+        self.module
+            .on_module_thread("subscribe_single", move || {
+                me.module
+                    .subscriptions()
+                    .add_single_subscription(me.sender, subscription, timer, None)
+            })
+            .await?
     }
 
     pub async fn unsubscribe(&self, request: Unsubscribe, timer: Instant) -> Result<Option<ExecutionMetrics>, DBError> {
@@ -554,12 +555,13 @@ impl ClientConnection {
         timer: Instant,
     ) -> Result<Option<ExecutionMetrics>, DBError> {
         let me = self.clone();
-        asyncify(move || {
-            me.module
-                .subscriptions()
-                .add_multi_subscription(me.sender, request, timer, None)
-        })
-        .await
+        self.module
+            .on_module_thread("subscribe_multi", move || {
+                me.module
+                    .subscriptions()
+                    .add_multi_subscription(me.sender, request, timer, None)
+            })
+            .await?
     }
 
     pub async fn unsubscribe_multi(
@@ -568,12 +570,13 @@ impl ClientConnection {
         timer: Instant,
     ) -> Result<Option<ExecutionMetrics>, DBError> {
         let me = self.clone();
-        asyncify(move || {
-            me.module
-                .subscriptions()
-                .remove_multi_subscription(me.sender, request, timer)
-        })
-        .await
+        self.module
+            .on_module_thread("unsubscribe_multi", move || {
+                me.module
+                    .subscriptions()
+                    .remove_multi_subscription(me.sender, request, timer)
+            })
+            .await?
     }
 
     pub async fn subscribe(&self, subscription: Subscribe, timer: Instant) -> Result<ExecutionMetrics, DBError> {

crates/core/src/host/host_controller.rs

@@ -334,7 +334,9 @@ impl HostController {
             warn!("database operation panicked");
             on_panic();
         });
-        let result = asyncify(move || f(&module.replica_ctx().relational_db)).await;
+
+        let db = module.replica_ctx().relational_db.clone();
+        let result = module.on_module_thread("using_database", move || f(&db)).await?;
         Ok(result)
     }
 
@@ -643,6 +645,7 @@ async fn make_replica_ctx(
             let Some(subscriptions) = downgraded.upgrade() else {
                 break;
             };
+            // This should happen on the module thread, but we haven't created the module yet.
             asyncify(move || subscriptions.write().remove_dropped_clients()).await
         }
     });