diff --git a/crates/mako/src/compiler.rs b/crates/mako/src/compiler.rs index 3d9117de7..77274c22d 100644 --- a/crates/mako/src/compiler.rs +++ b/crates/mako/src/compiler.rs @@ -1,7 +1,6 @@ use std::collections::HashMap; use std::fs; use std::path::PathBuf; -use std::sync::mpsc::channel; use std::sync::{Arc, Mutex, RwLock}; use std::time::{Instant, UNIX_EPOCH}; @@ -377,14 +376,9 @@ impl Compiler { } let result = { mako_core::mako_profile_scope!("Generate Stage"); - let (rs, rr) = channel::>(); // need to put all rayon parallel iterators run in the existed scope, or else rayon // will create a new thread pool for those parallel iterators - thread_pool::install(|| { - let res = self.generate(); - rs.send(res).unwrap(); - }); - rr.recv().unwrap() + thread_pool::scope(|_| self.generate()) }; let t_compiler_duration = t_compiler.elapsed(); if result.is_ok() { diff --git a/crates/mako/src/thread_pool.rs b/crates/mako/src/thread_pool.rs index abad56b0d..131ac1b87 100644 --- a/crates/mako/src/thread_pool.rs +++ b/crates/mako/src/thread_pool.rs @@ -1,6 +1,6 @@ use std::sync::OnceLock; -use mako_core::rayon::{ThreadPool, ThreadPoolBuilder}; +use mako_core::rayon::{Scope, ThreadPool, ThreadPoolBuilder}; static THREAD_POOL: OnceLock = OnceLock::new(); @@ -18,10 +18,10 @@ where THREAD_POOL.get_or_init(build_rayon_thread_pool).spawn(func) } -pub fn install(op: OP) -> R +pub fn scope<'scope, OP, R>(op: OP) -> R where - OP: FnOnce() -> R + Send, + OP: FnOnce(&Scope<'scope>) -> R + Send, R: Send, { - THREAD_POOL.get_or_init(build_rayon_thread_pool).install(op) + THREAD_POOL.get_or_init(build_rayon_thread_pool).scope(op) }