Parallelize QuerySet::eval (#891)

* Parallelize QuerySet::eval

* Reduce number of arguments for make_actor

crates/core/src/host/host_controller.rs

@@ -2,7 +2,8 @@ use crate::energy::{EnergyMonitor, EnergyQuanta, NullEnergyMonitor};
 use crate::hash::hash_bytes;
 use crate::host;
 use crate::messages::control_db::HostType;
-use crate::module_host_context::ModuleHostContext;
+use crate::module_host_context::{ModuleCreationContext, ModuleHostContext};
+use crate::util::spawn_rayon;
 use anyhow::Context;
 use parking_lot::Mutex;
 use serde::Serialize;
@@ -12,77 +13,13 @@ use std::sync::Arc;
 use std::time::{Duration, Instant};
 
 use super::module_host::{Catalog, EntityDef, EventStatus, ModuleHost, NoSuchModule, UpdateDatabaseResult};
-use super::scheduler::SchedulerStarter;
 use super::ReducerArgs;
 
 pub struct HostController {
     modules: Mutex<HashMap<u64, ModuleHost>>,
-    threadpool: Arc<HostThreadpool>,
     pub energy_monitor: Arc<dyn EnergyMonitor>,
 }
 
-pub struct HostThreadpool {
-    inner: rayon_core::ThreadPool,
-}
-
-/// A Rayon [spawn_handler](https://docs.rs/rustc-rayon-core/latest/rayon_core/struct.ThreadPoolBuilder.html#method.spawn_handler)
-/// which enters the given Tokio runtime at thread startup,
-/// so that the Rayon workers can send along async channels.
-///
-/// Other than entering the `rt`, this spawn handler behaves identitically to the default Rayon spawn handler,
-/// as documented in
-/// https://docs.rs/rustc-rayon-core/0.5.0/rayon_core/struct.ThreadPoolBuilder.html#method.spawn_handler
-///
-/// Having Rayon threads block on async operations is a code smell.
-/// We need to be careful that the Rayon threads never actually block,
-/// i.e. that every async operation they invoke immediately completes.
-/// I (pgoldman 2024-02-22) believe that our Rayon threads only ever send to unbounded channels,
-/// and therefore never wait.
-fn thread_spawn_handler(rt: tokio::runtime::Handle) -> impl FnMut(rayon::ThreadBuilder) -> Result<(), std::io::Error> {
-    move |thread| {
-        let rt = rt.clone();
-        let mut builder = std::thread::Builder::new();
-        if let Some(name) = thread.name() {
-            builder = builder.name(name.to_owned());
-        }
-        if let Some(stack_size) = thread.stack_size() {
-            builder = builder.stack_size(stack_size);
-        }
-        builder.spawn(move || {
-            let _rt_guard = rt.enter();
-            thread.run()
-        })?;
-        Ok(())
-    }
-}
-
-impl HostThreadpool {
-    fn new() -> Self {
-        let inner = rayon_core::ThreadPoolBuilder::new()
-            .thread_name(|_idx| "rayon-worker".to_string())
-            .spawn_handler(thread_spawn_handler(tokio::runtime::Handle::current()))
-            // TODO(perf, pgoldman 2024-02-22):
-            // in the case where we have many modules running many reducers,
-            // we'll wind up with Rayon threads competing with each other and with Tokio threads
-            // for CPU time.
-            //
-            // We should investigate creating two separate CPU pools,
-            // possibly via https://docs.rs/nix/latest/nix/sched/fn.sched_setaffinity.html,
-            // and restricting Tokio threads to one CPU pool
-            // and Rayon threads to the other.
-            // Then we should give Tokio and Rayon each a number of worker threads
-            // equal to the size of their pool.
-            .num_threads(std::thread::available_parallelism().unwrap().get())
-            .build()
-            .unwrap();
-        Self { inner }
-    }
-
-    pub fn spawn(&self, f: impl FnOnce() + Send + 'static) {
-        self.inner.spawn(f)
-    }
-}
-
 #[derive(PartialEq, Eq, Hash, Copy, Clone, Serialize, Debug)]
 pub enum DescribedEntityType {
     Table,
@@ -163,7 +100,6 @@ impl HostController {
     pub fn new(energy_monitor: Arc<impl EnergyMonitor>) -> Self {
         Self {
             modules: Mutex::new(HashMap::new()),
-            threadpool: Arc::new(HostThreadpool::new()),
             energy_monitor,
         }
     }
@@ -233,36 +169,38 @@ impl HostController {
     /// impact the logic of applications. The idea being that if SpacetimeDB is a distributed operating
     /// system, the applications will expect to be called when they are scheduled to be called regardless
     /// of whether the OS has been restarted.
-    pub async fn spawn_module_host(&self, module_host_context: ModuleHostContext) -> Result<ModuleHost, anyhow::Error> {
-        let key = module_host_context.dbic.database_instance_id;
-
-        let (module_host, start_scheduler) = self.make_module_host(module_host_context)?;
+    pub async fn spawn_module_host(&self, mhc: ModuleHostContext) -> Result<ModuleHost, anyhow::Error> {
+        let key = mhc.dbic.database_instance_id;
+
+        let mcc = ModuleCreationContext {
+            dbic: mhc.dbic,
+            scheduler: mhc.scheduler,
+            program_hash: hash_bytes(&mhc.program_bytes),
+            program_bytes: mhc.program_bytes,
+            energy_monitor: self.energy_monitor.clone(),
+        };
+        let module_host = spawn_rayon(move || Self::make_module_host(mhc.host_type, mcc)).await?;
 
         let old_module = self.modules.lock().insert(key, module_host.clone());
         if let Some(old_module) = old_module {
             old_module.exit().await
         }
         module_host.start();
-        start_scheduler.start(&module_host)?;
+        mhc.scheduler_starter.start(&module_host)?;
 
         Ok(module_host)
     }
 
-    fn make_module_host(&self, mhc: ModuleHostContext) -> anyhow::Result<(ModuleHost, SchedulerStarter)> {
-        let module_hash = hash_bytes(&mhc.program_bytes);
-        let (threadpool, energy_monitor) = (self.threadpool.clone(), self.energy_monitor.clone());
-        let module_host = match mhc.host_type {
+    fn make_module_host(host_type: HostType, mcc: ModuleCreationContext) -> anyhow::Result<ModuleHost> {
+        let module_host = match host_type {
             HostType::Wasm => {
-                // make_actor with block_in_place since it's going to take some time to compute.
                 let start = Instant::now();
-                let actor = tokio::task::block_in_place(|| {
-                    host::wasmtime::make_actor(mhc.dbic, module_hash, &mhc.program_bytes, mhc.scheduler, energy_monitor)
-                })?;
+                let actor = host::wasmtime::make_actor(mcc)?;
                 log::trace!("wasmtime::make_actor blocked for {:?}", start.elapsed());
-                ModuleHost::new(threadpool, actor)
+                ModuleHost::new(actor)
             }
         };
-        Ok((module_host, mhc.scheduler_starter))
+        Ok(module_host)
     }
 
     /// Determine if the module host described by [`ModuleHostContext`] is

crates/core/src/host/module_host.rs

@@ -6,9 +6,7 @@ use std::time::Duration;
 use base64::{engine::general_purpose::STANDARD as BASE_64_STD, Engine as _};
 use futures::{Future, FutureExt};
 use indexmap::IndexMap;
-use tokio::sync::oneshot;
 
-use super::host_controller::HostThreadpool;
 use super::{ArgsTuple, InvalidReducerArguments, ReducerArgs, ReducerCallResult, ReducerId, Timestamp};
 use crate::client::{ClientActorId, ClientConnectionSender};
 use crate::database_logger::LogLevel;
@@ -315,7 +313,7 @@ impl fmt::Debug for ModuleHost {
 
 #[async_trait::async_trait]
 trait DynModuleHost: Send + Sync + 'static {
-    async fn get_instance(&self, db: Address) -> Result<(&HostThreadpool, Box<dyn ModuleInstance>), NoSuchModule>;
+    async fn get_instance(&self, db: Address) -> Result<Box<dyn ModuleInstance>, NoSuchModule>;
     fn inject_logs(&self, log_level: LogLevel, message: &str);
     fn one_off_query(
         &self,
@@ -330,15 +328,14 @@ trait DynModuleHost: Send + Sync + 'static {
 
 struct HostControllerActor<T: Module> {
     module: Arc<T>,
-    threadpool: Arc<HostThreadpool>,
     instance_pool: LendingPool<T::Instance>,
     start: NotifyOnce,
 }
 
 impl<T: Module> HostControllerActor<T> {
     fn spinup_new_instance(&self) {
         let (module, instance_pool) = (self.module.clone(), self.instance_pool.clone());
-        self.threadpool.spawn(move || {
+        rayon::spawn(move || {
             let instance = module.create_instance();
             match instance_pool.add(instance) {
                 Ok(()) => {}
@@ -361,7 +358,7 @@ async fn select_first<A: Future, B: Future<Output = ()>>(fut_a: A, fut_b: B) ->
 
 #[async_trait::async_trait]
 impl<T: Module> DynModuleHost for HostControllerActor<T> {
-    async fn get_instance(&self, db: Address) -> Result<(&HostThreadpool, Box<dyn ModuleInstance>), NoSuchModule> {
+    async fn get_instance(&self, db: Address) -> Result<Box<dyn ModuleInstance>, NoSuchModule> {
         self.start.notified().await;
         // in the future we should do something like in the else branch here -- add more instances based on load.
         // we need to do write-skew retries first - right now there's only ever once instance per module.
@@ -379,11 +376,10 @@ impl<T: Module> DynModuleHost for HostControllerActor<T> {
             .await
             .map_err(|_| NoSuchModule)?
         };
-        let inst = AutoReplacingModuleInstance {
+        Ok(Box::new(AutoReplacingModuleInstance {
             inst,
             module: self.module.clone(),
-        };
-        Ok((&self.threadpool, Box::new(inst)))
+        }))
     }
 
     fn inject_logs(&self, log_level: LogLevel, message: &str) {
@@ -459,13 +455,12 @@ pub enum InitDatabaseError {
 }
 
 impl ModuleHost {
-    pub fn new(threadpool: Arc<HostThreadpool>, mut module: impl Module) -> Self {
+    pub fn new(mut module: impl Module) -> Self {
         let info = module.info();
         let instance_pool = LendingPool::new();
         instance_pool.add_multiple(module.initial_instances()).unwrap();
         let inner = Arc::new(HostControllerActor {
             module: Arc::new(module),
-            threadpool,
             instance_pool,
             start: NotifyOnce::new(),
         });
@@ -491,13 +486,12 @@ impl ModuleHost {
         F: FnOnce(&mut dyn ModuleInstance) -> R + Send + 'static,
         R: Send + 'static,
     {
-        let (threadpool, mut inst) = self.inner.get_instance(self.info.address).await?;
+        let mut inst = self.inner.get_instance(self.info.address).await?;
 
-        let (tx, rx) = oneshot::channel();
-        threadpool.spawn(move || {
-            let _ = tx.send(f(&mut *inst));
-        });
-        Ok(rx.await.expect("instance panicked"))
+        let result = tokio::task::spawn_blocking(move || f(&mut *inst))
+            .await
+            .unwrap_or_else(|e| std::panic::resume_unwind(e.into_panic()));
+        Ok(result)
     }
 
     pub async fn disconnect_client(&self, client_id: ClientActorId) {

crates/core/src/host/wasm_common/module_host_actor.rs

@@ -15,7 +15,6 @@ use crate::db::datastore::locking_tx_datastore::MutTxId;
 use crate::db::datastore::traits::IsolationLevel;
 use crate::energy::{EnergyMonitor, EnergyQuanta, ReducerBudget, ReducerFingerprint};
 use crate::execution_context::ExecutionContext;
-use crate::hash::Hash;
 use crate::host::instance_env::InstanceEnv;
 use crate::host::module_host::{
     CallReducerParams, DatabaseUpdate, EventStatus, Module, ModuleEvent, ModuleFunctionCall, ModuleInfo,
@@ -24,6 +23,7 @@ use crate::host::module_host::{
 use crate::host::{ArgsTuple, EntityDef, ReducerCallResult, ReducerId, ReducerOutcome, Scheduler, Timestamp};
 use crate::identity::Identity;
 use crate::messages::control_db::Database;
+use crate::module_host_context::ModuleCreationContext;
 use crate::sql;
 use crate::subscription::module_subscription_actor::ModuleSubscriptions;
 use crate::util::const_unwrap;
@@ -53,7 +53,7 @@ pub trait WasmInstance: Send + Sync + 'static {
 
     fn instance_env(&self) -> &InstanceEnv;
 
-    type Trap;
+    type Trap: Send;
 
     fn call_reducer(&mut self, op: ReducerOp<'_>, budget: ReducerBudget) -> ExecuteResult<Self::Trap>;
 
@@ -123,13 +123,14 @@ pub enum DescribeError {
 }
 
 impl<T: WasmModule> WasmModuleHostActor<T> {
-    pub fn new(
-        database_instance_context: Arc<DatabaseInstanceContext>,
-        module_hash: Hash,
-        module: T,
-        scheduler: Scheduler,
-        energy_monitor: Arc<dyn EnergyMonitor>,
-    ) -> Result<Self, InitializationError> {
+    pub fn new(mcc: ModuleCreationContext, module: T) -> Result<Self, InitializationError> {
+        let ModuleCreationContext {
+            dbic: database_instance_context,
+            scheduler,
+            program_bytes: _,
+            program_hash: module_hash,
+            energy_monitor,
+        } = mcc;
         log::trace!(
             "Making new module host actor for database {}",
             database_instance_context.address
@@ -533,7 +534,9 @@ impl<T: WasmInstance> WasmModuleInstance<T> {
         )
         .entered();
 
-        let (tx, result) = tx_slot.set(tx, || self.instance.call_reducer(op, budget));
+        // run the call_reducer call in rayon. it's important that we don't acquire a lock inside a rayon task,
+        // as that can lead to deadlock.
+        let (tx, result) = rayon::scope(|_| tx_slot.set(tx, || self.instance.call_reducer(op, budget)));
 
         let ExecuteResult {
             energy,

crates/core/src/host/wasmtime/mod.rs

@@ -1,14 +1,12 @@
 use std::borrow::Cow;
-use std::sync::Arc;
 
 use anyhow::Context;
 use once_cell::sync::Lazy;
 use wasmtime::{Engine, Linker, Module, StoreContext, StoreContextMut};
 
-use crate::database_instance_context::DatabaseInstanceContext;
-use crate::energy::{EnergyMonitor, EnergyQuanta, ReducerBudget};
+use crate::energy::{EnergyQuanta, ReducerBudget};
 use crate::error::NodesError;
-use crate::hash::Hash;
+use crate::module_host_context::ModuleCreationContext;
 use crate::stdb_path;
 
 mod wasm_instance_env;
@@ -20,7 +18,6 @@ use self::wasm_instance_env::WasmInstanceEnv;
 
 use super::wasm_common::module_host_actor::InitializationError;
 use super::wasm_common::{abi, module_host_actor::WasmModuleHostActor, ModuleCreationError};
-use super::Scheduler;
 
 static ENGINE: Lazy<Engine> = Lazy::new(|| {
     let mut config = wasmtime::Config::new();
@@ -55,14 +52,8 @@ static LINKER: Lazy<Linker<WasmInstanceEnv>> = Lazy::new(|| {
     linker
 });
 
-pub fn make_actor(
-    dbic: Arc<DatabaseInstanceContext>,
-    module_hash: Hash,
-    program_bytes: &[u8],
-    scheduler: Scheduler,
-    energy_monitor: Arc<dyn EnergyMonitor>,
-) -> Result<impl super::module_host::Module, ModuleCreationError> {
-    let module = Module::new(&ENGINE, program_bytes).map_err(ModuleCreationError::WasmCompileError)?;
+pub fn make_actor(mcc: ModuleCreationContext) -> Result<impl super::module_host::Module, ModuleCreationError> {
+    let module = Module::new(&ENGINE, &mcc.program_bytes).map_err(ModuleCreationError::WasmCompileError)?;
 
     let func_imports = module
         .imports()
@@ -77,7 +68,7 @@ pub fn make_actor(
 
     let module = WasmtimeModule::new(module);
 
-    WasmModuleHostActor::new(dbic, module_hash, module, scheduler, energy_monitor).map_err(Into::into)
+    WasmModuleHostActor::new(mcc, module).map_err(Into::into)
 }
 
 #[derive(Debug, derive_more::From)]

crates/core/src/module_host_context.rs

@@ -1,4 +1,7 @@
+use spacetimedb_lib::Hash;
+
 use crate::database_instance_context::DatabaseInstanceContext;
+use crate::energy::EnergyMonitor;
 use crate::host::scheduler::{Scheduler, SchedulerStarter};
 use crate::messages::control_db::HostType;
 use crate::util::AnyBytes;
@@ -11,3 +14,11 @@ pub struct ModuleHostContext {
     pub host_type: HostType,
     pub program_bytes: AnyBytes,
 }
+
+pub struct ModuleCreationContext {
+    pub dbic: Arc<DatabaseInstanceContext>,
+    pub scheduler: Scheduler,
+    pub program_bytes: AnyBytes,
+    pub program_hash: Hash,
+    pub energy_monitor: Arc<dyn EnergyMonitor>,
+}