make JsInstance methods async

Author: Centril

crates/core/src/host/module_host.rs

@@ -11,6 +11,8 @@ use crate::error::DBError;
 use crate::estimation::estimate_rows_scanned;
 use crate::hash::Hash;
 use crate::host::scheduler::{handle_queued_call_reducer_params, QueueItem};
+use crate::host::v8::JsInstance;
+use crate::host::wasmtime::ModuleInstance;
 use crate::host::InvalidFunctionArguments;
 use crate::identity::Identity;
 use crate::messages::control_db::{Database, HostType};
@@ -386,76 +388,6 @@ impl Instance {
         }
     }
 
-    /// Update the module instance's database to match the schema of the module instance.
-    fn update_database(
-        &mut self,
-        program: Program,
-        old_module_info: Arc<ModuleInfo>,
-        policy: MigrationPolicy,
-    ) -> anyhow::Result<UpdateDatabaseResult> {
-        match self {
-            Instance::Wasm(inst) => inst.update_database(program, old_module_info, policy),
-            Instance::Js(inst) => inst.update_database(program, old_module_info, policy),
-        }
-    }
-
-    fn call_reducer(&mut self, tx: Option<MutTxId>, params: CallReducerParams) -> ReducerCallResult {
-        match self {
-            Instance::Wasm(inst) => inst.call_reducer(tx, params),
-            Instance::Js(inst) => inst.call_reducer(tx, params),
-        }
-    }
-
-    fn clear_all_clients(&self) -> anyhow::Result<()> {
-        match self {
-            Instance::Wasm(inst) => inst.clear_all_clients(),
-            Instance::Js(inst) => inst.clear_all_clients(),
-        }
-    }
-
-    fn call_identity_connected(
-        &mut self,
-        caller_auth: ConnectionAuthCtx,
-        caller_connection_id: ConnectionId,
-    ) -> Result<(), ClientConnectedError> {
-        match self {
-            Instance::Wasm(inst) => inst.call_identity_connected(caller_auth, caller_connection_id),
-            Instance::Js(inst) => inst.call_identity_connected(caller_auth, caller_connection_id),
-        }
-    }
-
-    fn call_identity_disconnected(
-        &mut self,
-        caller_identity: Identity,
-        caller_connection_id: ConnectionId,
-    ) -> Result<(), ReducerCallError> {
-        match self {
-            Instance::Wasm(inst) => inst.call_identity_disconnected(caller_identity, caller_connection_id),
-            Instance::Js(inst) => inst.call_identity_disconnected(caller_identity, caller_connection_id),
-        }
-    }
-
-    fn disconnect_client(&mut self, client_id: ClientActorId) -> Result<(), ReducerCallError> {
-        match self {
-            Instance::Wasm(inst) => inst.disconnect_client(client_id),
-            Instance::Js(inst) => inst.disconnect_client(client_id),
-        }
-    }
-
-    fn call_scheduled_reducer(&mut self, item: QueueItem) -> Result<ReducerCallResult, ReducerCallError> {
-        match self {
-            Instance::Wasm(inst) => inst.call_scheduled_reducer(item),
-            Instance::Js(inst) => inst.call_scheduled_reducer(item),
-        }
-    }
-
-    fn init_database(&mut self, program: Program) -> anyhow::Result<Option<ReducerCallResult>> {
-        match self {
-            Instance::Wasm(inst) => inst.init_database(program),
-            Instance::Js(inst) => inst.init_database(program),
-        }
-    }
-
     async fn call_procedure(&mut self, params: CallProcedureParams) -> Result<ProcedureCallResult, ProcedureCallError> {
         match self {
             Instance::Wasm(inst) => inst.call_procedure(params).await,
@@ -1050,79 +982,119 @@ impl ModuleHost {
         })
     }
 
-    async fn call_async_with_instance<Fun, Fut, R>(&self, label: &str, f: Fun) -> Result<R, NoSuchModule>
+    /// Run a function for this module which has access to the module instance.
+    async fn with_instance<'a, Guard, R, F>(
+        &'a self,
+        kind: &str,
+        label: &str,
+        timer: impl FnOnce(&str) -> Guard,
+        work: impl FnOnce(Guard, &'a SingleCoreExecutor, Instance) -> F,
+    ) -> Result<R, NoSuchModule>
     where
-        Fun: (FnOnce(Instance) -> Fut) + Send + 'static,
-        Fut: Future<Output = (R, Instance)> + Send + 'static,
-        R: Send + 'static,
+        F: Future<Output = (R, Instance)>,
     {
         self.guard_closed()?;
-        let timer_guard = self.start_call_timer(label);
+        let timer_guard = timer(label);
 
+        // Operations on module instances (e.g. calling reducers) is blocking,
+        // partially because the computation can potentially take a long time
+        // and partially because interacting with the database requires taking
+        // a blocking lock. So, we run `f` on a dedicated thread with `self.executor`.
+        // This will bubble up any panic that may occur.
+
+        // If a function call panics, we **must** ensure to call `self.on_panic`
+        // so that the module is discarded by the host controller.
         scopeguard::defer_on_unwind!({
-            log::warn!("procedure {label} panicked");
+            log::warn!("{kind} {label} panicked");
             (self.on_panic)();
         });
 
         // TODO: should we be calling and/or `await`-ing `get_instance` within the below `run_job`?
         // Unclear how much overhead this call can have.
-        let instance = self.instance_manager.lock().await.get_instance().await;
+        let inst = self.instance_manager.lock().await.get_instance().await;
 
-        let (res, instance) = self
-            .executor
-            .run_job(async move {
-                drop(timer_guard);
-                f(instance).await
-            })
-            .await;
+        let (res, inst) = work(timer_guard, &self.executor, inst).await;
 
-        self.instance_manager.lock().await.return_instance(instance);
+        self.instance_manager.lock().await.return_instance(inst);
 
         Ok(res)
     }
 
-    /// Run a function on the JobThread for this module which has access to the module instance.
-    async fn call<F, R>(&self, label: &str, f: F) -> Result<R, NoSuchModule>
+    async fn call_async_with_instance<Fun, Fut, R>(&self, label: &str, work: Fun) -> Result<R, NoSuchModule>
     where
-        F: FnOnce(&mut Instance) -> R + Send + 'static,
+        Fun: (FnOnce(Instance) -> Fut) + Send + 'static,
+        Fut: Future<Output = (R, Instance)> + Send + 'static,
         R: Send + 'static,
     {
-        self.guard_closed()?;
-        let timer_guard = self.start_call_timer(label);
-
-        // Operations on module instances (e.g. calling reducers) is blocking,
-        // partially because the computation can potentially take a long time
-        // and partially because interacting with the database requires taking
-        // a blocking lock. So, we run `f` on a dedicated thread with `self.executor`.
-        // This will bubble up any panic that may occur.
-
-        // If a reducer call panics, we **must** ensure to call `self.on_panic`
-        // so that the module is discarded by the host controller.
-        scopeguard::defer_on_unwind!({
-            log::warn!("reducer {label} panicked");
-            (self.on_panic)();
-        });
-
-        let mut instance = self.instance_manager.lock().await.get_instance().await;
-
-        let (res, instance) = self
-            .executor
-            .run_sync_job(move || {
-                drop(timer_guard);
-                let res = f(&mut instance);
-                (res, instance)
-            })
-            .await;
-
-        self.instance_manager.lock().await.return_instance(instance);
+        self.with_instance(
+            "procedure",
+            label,
+            |l| self.start_call_timer(l),
+            |timer_guard, executor, inst| {
+                executor.run_job(async move {
+                    drop(timer_guard);
+                    work(inst).await
+                })
+            },
+        )
+        .await
+    }
 
-        Ok(res)
+    /// Run a function for this module which has access to the module instance.
+    ///
+    /// For WASM, the function is run on the module's JobThread.
+    /// For V8/JS, the function is run in the current task.
+    async fn call<A, R, JF>(
+        &self,
+        label: &str,
+        arg: A,
+        wasm: impl FnOnce(A, &mut ModuleInstance) -> R + Send + 'static,
+        js: impl FnOnce(A, Box<JsInstance>) -> JF,
+    ) -> Result<R, NoSuchModule>
+    where
+        JF: Future<Output = (R, Box<JsInstance>)>,
+        R: Send + 'static,
+        A: Send + 'static,
+    {
+        self.with_instance(
+            "reducer",
+            label,
+            |l| self.start_call_timer(l),
+            // Operations on module instances (e.g. calling reducers) is blocking,
+            // partially because the computation can potentially take a long time
+            // and partially because interacting with the database requires taking a blocking lock.
+            // So, we run `work` on a dedicated thread with `self.executor`.
+            // This will bubble up any panic that may occur.
+            |timer_guard, executor, inst| async move {
+                match inst {
+                    Instance::Wasm(mut inst) => {
+                        executor
+                            .run_sync_job(move || {
+                                drop(timer_guard);
+                                (wasm(arg, &mut inst), Instance::Wasm(inst))
+                            })
+                            .await
+                    }
+                    Instance::Js(inst) => {
+                        drop(timer_guard);
+                        let (res, inst) = js(arg, inst).await;
+                        (res, Instance::Js(inst))
+                    }
+                }
+            },
+        )
+        .await
     }
 
     pub async fn disconnect_client(&self, client_id: ClientActorId) {
         log::trace!("disconnecting client {client_id}");
         if let Err(e) = self
-            .call("disconnect_client", move |inst| inst.disconnect_client(client_id))
+            .call(
+                "disconnect_client",
+                client_id,
+                |client_id, inst| inst.disconnect_client(client_id),
+                |client_id, inst| inst.disconnect_client(client_id),
+            )
             .await
         {
             log::error!("Error from client_disconnected transaction: {e}");
@@ -1165,9 +1137,12 @@ impl ModuleHost {
         caller_auth: ConnectionAuthCtx,
         caller_connection_id: ConnectionId,
     ) -> Result<(), ClientConnectedError> {
-        self.call("call_identity_connected", move |inst| {
-            inst.call_identity_connected(caller_auth, caller_connection_id)
-        })
+        self.call(
+            "call_identity_connected",
+            (caller_auth, caller_connection_id),
+            |(a, b), inst| inst.call_identity_connected(a, b),
+            |(a, b), inst| inst.call_identity_connected(a, b),
+        )
         .await
         .map_err(ReducerCallError::from)?
     }
@@ -1305,16 +1280,24 @@ impl ModuleHost {
         caller_identity: Identity,
         caller_connection_id: ConnectionId,
     ) -> Result<(), ReducerCallError> {
-        self.call("call_identity_disconnected", move |inst| {
-            inst.call_identity_disconnected(caller_identity, caller_connection_id)
-        })
+        self.call(
+            "call_identity_disconnected",
+            (caller_identity, caller_connection_id),
+            |(a, b), inst| inst.call_identity_disconnected(a, b),
+            |(a, b), inst| inst.call_identity_disconnected(a, b),
+        )
         .await?
     }
 
     /// Empty the system tables tracking clients without running any lifecycle reducers.
     pub async fn clear_all_clients(&self) -> anyhow::Result<()> {
-        self.call("clear_all_clients", move |inst| inst.clear_all_clients())
-            .await?
+        self.call(
+            "clear_all_clients",
+            (),
+            |_, inst| inst.clear_all_clients(),
+            |_, inst| inst.clear_all_clients(),
+        )
+        .await?
     }
 
     fn call_reducer_params(
@@ -1369,9 +1352,12 @@ impl ModuleHost {
         };
 
         Ok(self
-            .call(&reducer_def.name, move |inst| {
-                inst.call_reducer(None, call_reducer_params)
-            })
+            .call(
+                &reducer_def.name,
+                call_reducer_params,
+                |p, inst| inst.call_reducer(None, p),
+                |p, inst| inst.call_reducer(None, p),
+            )
             .await?)
     }
 
@@ -1492,9 +1478,12 @@ impl ModuleHost {
     // because their reducer arguments are stored in the database and need to be fetched
     // within the same transaction as the reducer call.
     pub(crate) async fn call_scheduled_reducer(&self, item: QueueItem) -> Result<ReducerCallResult, ReducerCallError> {
-        self.call(SCHEDULED_REDUCER, move |inst: &mut Instance| {
-            inst.call_scheduled_reducer(item)
-        })
+        self.call(
+            SCHEDULED_REDUCER,
+            item,
+            |item, inst| inst.call_scheduled_reducer(item),
+            |item, inst| inst.call_scheduled_reducer(item),
+        )
         .await?
     }
 
@@ -1503,9 +1492,14 @@ impl ModuleHost {
     }
 
     pub async fn init_database(&self, program: Program) -> Result<Option<ReducerCallResult>, InitDatabaseError> {
-        self.call("<init_database>", move |inst| inst.init_database(program))
-            .await?
-            .map_err(InitDatabaseError::Other)
+        self.call(
+            "<init_database>",
+            program,
+            |p, inst| inst.init_database(p),
+            |p, inst| inst.init_database(p),
+        )
+        .await?
+        .map_err(InitDatabaseError::Other)
     }
 
     pub async fn update_database(
@@ -1514,9 +1508,12 @@ impl ModuleHost {
         old_module_info: Arc<ModuleInfo>,
         policy: MigrationPolicy,
     ) -> Result<UpdateDatabaseResult, anyhow::Error> {
-        self.call("<update_database>", move |inst| {
-            inst.update_database(program, old_module_info, policy)
-        })
+        self.call(
+            "<update_database>",
+            (program, old_module_info, policy),
+            |(a, b, c), inst| inst.update_database(a, b, c),
+            |(a, b, c), inst| inst.update_database(a, b, c),
+        )
         .await?
     }