add trap handling for reducers, pooled instances for procedures

Author: joshua-spacetime

crates/core/src/host/module_host.rs

@@ -359,7 +359,8 @@ struct WasmtimeModuleHost {
 
 struct V8ModuleHost {
     module: super::v8::JsModule,
-    instance: super::v8::JsInstance,
+    reducer_lane: super::v8::JsReducerLane,
+    procedure_instances: ModuleInstanceManager<super::v8::JsModule>,
 }
 
 /// A module; used as a bound on `InstanceManager`.
@@ -407,7 +408,7 @@ impl GenericModule for super::v8::JsModule {
 
 impl GenericModuleInstance for super::v8::JsInstance {
     fn trapped(&self) -> bool {
-        false
+        self.trapped()
     }
 }
 
@@ -615,6 +616,7 @@ pub fn call_identity_connected(
     }
 }
 
+#[derive(Clone)]
 pub struct CallReducerParams {
     pub timestamp: Timestamp,
     pub caller_identity: Identity,
@@ -801,6 +803,23 @@ impl<M: GenericModule> ModuleInstanceManager<M> {
         }
     }
 
+    fn new_empty(module: M, database_identity: Identity) -> Self {
+        let host_type = module.host_type();
+        let create_instance_time_metric = CreateInstanceTimeMetric {
+            metric: WORKER_METRICS
+                .module_create_instance_time_seconds
+                .with_label_values(&database_identity, &host_type),
+            host_type,
+            database_identity,
+        };
+
+        Self {
+            instances: Mutex::new(VecDeque::new()),
+            module,
+            create_instance_time_metric,
+        }
+    }
+
     async fn with_instance<R>(&self, f: impl AsyncFnOnce(M::Instance) -> (R, M::Instance)) -> R {
         let inst = self.get_instance().await;
         let (res, inst) = f(inst).await;
@@ -1029,9 +1048,12 @@ impl ModuleHost {
             }
             ModuleWithInstance::Js { module, init_inst } => {
                 info = module.info();
+                let reducer_lane = super::v8::JsReducerLane::new(module.clone(), init_inst);
+                let procedure_instances = ModuleInstanceManager::new_empty(module.clone(), database_identity);
                 Arc::new(ModuleHostInner::Js(V8ModuleHost {
                     module,
-                    instance: init_inst,
+                    reducer_lane,
+                    procedure_instances,
                 }))
             }
         };
@@ -1100,13 +1122,13 @@ impl ModuleHost {
                     })
                     .await
             }
-            ModuleHostInner::Js(V8ModuleHost { instance, .. }) => {
-                instance
+            ModuleHostInner::Js(V8ModuleHost { reducer_lane, .. }) => {
+                reducer_lane
                     .run_on_thread(async move || {
                         drop(timer_guard);
                         f().await
                     })
-                    .await
+                    .await?
             }
         })
     }
@@ -1145,7 +1167,7 @@ impl ModuleHost {
         arg: A,
         timer: impl FnOnce(&str) -> Guard,
         work_wasm: impl AsyncFnOnce(Guard, &SingleCoreExecutor, Box<ModuleInstance>, A) -> (R, Box<ModuleInstance>),
-        work_js: impl AsyncFnOnce(Guard, &JsInstance, A) -> R,
+        work_js: impl AsyncFnOnce(Guard, &super::v8::JsReducerLane, A) -> R,
     ) -> Result<R, NoSuchModule> {
         self.guard_closed()?;
         let timer_guard = timer(label);
@@ -1172,7 +1194,7 @@ impl ModuleHost {
                     .with_instance(async |inst| work_wasm(timer_guard, executor, inst, arg).await)
                     .await
             }
-            ModuleHostInner::Js(V8ModuleHost { instance, .. }) => work_js(timer_guard, instance, arg).await,
+            ModuleHostInner::Js(V8ModuleHost { reducer_lane, .. }) => work_js(timer_guard, reducer_lane, arg).await,
         })
     }
 
@@ -1185,7 +1207,7 @@ impl ModuleHost {
         label: &str,
         arg: A,
         wasm: impl AsyncFnOnce(A, &mut ModuleInstance) -> R + Send + 'static,
-        js: impl AsyncFnOnce(A, &JsInstance) -> R,
+        js: impl AsyncFnOnce(A, &super::v8::JsReducerLane) -> R,
     ) -> Result<R, NoSuchModule>
     where
         R: Send + 'static,
@@ -1217,6 +1239,35 @@ impl ModuleHost {
         .await
     }
 
+    async fn with_js_procedure_instance<R>(
+        &self,
+        label: &str,
+        f: impl AsyncFnOnce(&JsInstance) -> R,
+    ) -> Result<R, NoSuchModule> {
+        self.guard_closed()?;
+        let timer_guard = self.start_call_timer(label);
+
+        scopeguard::defer_on_unwind!({
+            log::warn!("procedure {label} panicked");
+            (self.on_panic)();
+        });
+
+        Ok(match &*self.inner {
+            ModuleHostInner::Wasm(_) => unreachable!("WASM procedures should not use the JS procedure instance path"),
+            ModuleHostInner::Js(V8ModuleHost {
+                procedure_instances, ..
+            }) => {
+                procedure_instances
+                    .with_instance(async |inst| {
+                        drop(timer_guard);
+                        let res = f(&inst).await;
+                        (res, inst)
+                    })
+                    .await
+            }
+        })
+    }
+
     pub async fn disconnect_client(&self, client_id: ClientActorId) {
         log::trace!("disconnecting client {client_id}");
         if let Err(e) = self
@@ -1833,26 +1884,62 @@ impl ModuleHost {
         name: &str,
         params: CallProcedureParams,
     ) -> Result<CallProcedureReturn, NoSuchModule> {
-        self.call(
-            name,
-            params,
-            async move |params, inst| inst.call_procedure(params).await,
-            async move |params, inst| inst.call_procedure(params).await,
-        )
-        .await
+        match &*self.inner {
+            ModuleHostInner::Wasm(_) => {
+                self.call(
+                    name,
+                    params,
+                    async move |params, inst| inst.call_procedure(params).await,
+                    async move |_params, _inst| unreachable!("JS procedure lane is not used for WASM modules"),
+                )
+                .await
+            }
+            ModuleHostInner::Js(_) => {
+                self.with_js_procedure_instance(name, async move |inst| inst.call_procedure(params).await)
+                    .await
+            }
+        }
     }
 
     pub(super) async fn call_scheduled_function(
         &self,
         params: ScheduledFunctionParams,
     ) -> Result<CallScheduledFunctionResult, NoSuchModule> {
-        self.call(
-            "unknown scheduled function",
-            params,
-            async move |params, inst| inst.call_scheduled_function(params).await,
-            async move |params, inst| inst.call_scheduled_function(params).await,
-        )
-        .await
+        match &*self.inner {
+            ModuleHostInner::Wasm(_) => {
+                self.call(
+                    "unknown scheduled function",
+                    params,
+                    async move |params, inst| inst.call_scheduled_function(params).await,
+                    async move |_params, _inst| unreachable!("JS scheduled-function lane is not used for WASM modules"),
+                )
+                .await
+            }
+            ModuleHostInner::Js(V8ModuleHost { module, .. }) => {
+                let use_procedure_lane =
+                    match params.uses_procedure_lane(&self.info, module.replica_ctx().relational_db.as_ref()) {
+                        Ok(use_procedure_lane) => use_procedure_lane,
+                        Err(err) => {
+                            log::error!("failed to classify scheduled JS function; routing to procedure lane: {err:#}");
+                            true
+                        }
+                    };
+                if use_procedure_lane {
+                    self.with_js_procedure_instance("unknown scheduled function", async move |inst| {
+                        inst.call_scheduled_function(params).await
+                    })
+                    .await
+                } else {
+                    self.call(
+                        "unknown scheduled function",
+                        params,
+                        async move |params, inst| inst.call_scheduled_function(params).await,
+                        async move |params, inst| inst.call_scheduled_function(params).await,
+                    )
+                    .await
+                }
+            }
+        }
     }
 
     /// Materializes the views return by the `view_collector`, if not already materialized,

crates/core/src/host/scheduler.rs

@@ -268,16 +268,26 @@ struct SchedulerActor {
     module_host: WeakModuleHost,
 }
 
+#[derive(Clone)]
 enum QueueItem {
     Id { id: ScheduledFunctionId, at: Timestamp },
     VolatileNonatomicImmediate { function_name: String, args: FunctionArgs },
 }
 
+#[derive(Clone)]
 pub(crate) struct ScheduledFunctionParams(QueueItem);
 
 #[cfg(target_pointer_width = "64")]
 spacetimedb_table::static_assert_size!(QueueItem, 64);
 
+impl ScheduledFunctionParams {
+    pub(crate) fn uses_procedure_lane(&self, module: &ModuleInfo, db: &RelationalDB) -> anyhow::Result<bool> {
+        let tx = db.begin_mut_tx(IsolationLevel::Serializable, Workload::Internal);
+        let params = call_params_for_queued_item(module, db, &tx, self.0.clone())?;
+        Ok(matches!(params, Some((_timestamp, _instant, CallParams::Procedure(_)))))
+    }
+}
+
 impl SchedulerActor {
     async fn run(mut self) {
         loop {