refact: thread pool (#936)

* refactor: minimum tokio thread pool

* refact: use a static rayon thread pool

* refact: import style

* refact: code style

* refact: import style

crates/core/Cargo.toml

@@ -83,7 +83,7 @@ serde_json            = { workspace = true }
 serde_yaml            = "0.9.22"
 svgr-rs               = "0.1.3"
 thiserror             = "1.0.43"
-tokio                 = { version = "1", features = ["full"] }
+tokio                 = { version = "1", features = ["rt", "sync"] }
 tokio-tungstenite     = "0.19.0"
 toml                  = "0.7.6"
 tracing               = "0.1.37"

crates/mako/src/ast.rs

@@ -311,16 +311,14 @@ mod tests {
     use std::path::PathBuf;
     use std::sync::Arc;
 
-    use mako_core::tokio;
-
     use super::build_js_ast;
     use crate::assert_debug_snapshot;
     use crate::ast::js_ast_to_code;
     use crate::compiler::Context;
     use crate::test_helper::create_mock_module;
 
-    #[tokio::test(flavor = "multi_thread")]
-    async fn test_chinese_ascii() {
+    #[test]
+    fn test_chinese_ascii() {
         let module = create_mock_module(
             PathBuf::from("/path/to/test"),
             r#"

crates/mako/src/ast_2/tests.rs

@@ -70,11 +70,10 @@ impl TestUtils {
     }
 
     pub fn gen_css_ast(content: String) -> TestUtils {
-        let test_utils = TestUtils::new(TestUtilsOpts {
+        TestUtils::new(TestUtilsOpts {
             file: Some("test.css".to_string()),
             content: Some(content),
-        });
-        test_utils
+        })
     }
 
     pub fn gen_js_ast(content: String) -> TestUtils {

crates/mako/src/build.rs

@@ -1,4 +1,5 @@
 use std::collections::HashSet;
+use std::sync::mpsc::channel;
 use std::sync::Arc;
 
 use mako_core::anyhow;
@@ -14,8 +15,8 @@ use crate::load::Load;
 use crate::module::{Module, ModuleAst, ModuleId, ModuleInfo};
 use crate::parse::Parse;
 use crate::resolve::ResolverResource;
+use crate::thread_pool;
 use crate::transform::Transform;
-use crate::util::create_thread_pool;
 
 #[derive(Debug, Error)]
 pub enum BuildError {
@@ -29,11 +30,12 @@ impl Compiler {
             return Ok(HashSet::new());
         }
 
-        let (pool, rs, rr) = create_thread_pool::<Result<Module>>();
+        let (rs, rr) = channel::<Result<Module>>();
+
         let build_with_pool = |file: File, parent_resource: Option<ResolverResource>| {
             let rs = rs.clone();
             let context = self.context.clone();
-            pool.spawn(move || {
+            thread_pool::spawn(move || {
                 let result = Self::build_module(&file, parent_resource, context.clone());
                 let result = Self::handle_build_result(result, &file, context);
                 rs.send(result).unwrap();

crates/mako/src/generate_chunks.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 use std::path::Path;
+use std::sync::mpsc::channel;
 use std::sync::Arc;
 use std::time::Instant;
 use std::vec;
@@ -15,8 +16,8 @@ use crate::chunk::{Chunk, ChunkType};
 use crate::chunk_pot::ChunkPot;
 use crate::compiler::{Compiler, Context};
 use crate::module::{ModuleAst, ModuleId};
+use crate::thread_pool;
 use crate::transform_in_generate::transform_css_generate;
-use crate::util::create_thread_pool;
 
 #[derive(Clone)]
 pub enum ChunkFileType {
@@ -139,7 +140,9 @@ impl Compiler {
         mako_core::mako_profile_function!();
         let chunk_graph = self.context.chunk_graph.read().unwrap();
         let chunks = chunk_graph.get_chunks();
-        let (pool, rs, rr) = create_thread_pool::<Result<Vec<ChunkFile>>>();
+
+        let (rs, rr) = channel::<Result<Vec<ChunkFile>>>();
+
         for chunk in chunks.iter() {
             if !matches!(
                 chunk.chunk_type,
@@ -148,7 +151,7 @@ impl Compiler {
                 let rs = rs.clone();
                 let context = self.context.clone();
                 let chunk_id = chunk.id.clone();
-                pool.spawn(move || {
+                thread_pool::spawn(move || {
                     let chunk_graph = context.chunk_graph.read().unwrap();
                     let module_graph = context.module_graph.read().unwrap();
                     let chunk = chunk_graph.chunk(&chunk_id).unwrap();

crates/mako/src/lib.rs

@@ -37,6 +37,7 @@ mod swc_helpers;
 mod targets;
 #[cfg(test)]
 mod test_helper;
+mod thread_pool;
 mod transform;
 mod transform_in_generate;
 mod transformers;

crates/mako/src/main.rs

@@ -51,6 +51,7 @@ mod swc_helpers;
 mod targets;
 #[cfg(test)]
 mod test_helper;
+mod thread_pool;
 mod transform;
 mod transform_in_generate;
 mod transformers;
@@ -72,8 +73,16 @@ static GLOBAL: mimalloc_rust::GlobalMiMalloc = mimalloc_rust::GlobalMiMalloc;
 #[global_allocator]
 static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
-#[tokio::main]
-async fn main() -> Result<()> {
+fn main() -> Result<()> {
+    let fut = async { run().await };
+
+    tokio::runtime::Builder::new_current_thread()
+        .build()
+        .expect("Failed to create tokio runtime.")
+        .block_on(fut)
+}
+
+async fn run() -> Result<()> {
     // logger
     init_logger();
 

crates/mako/src/thread_pool.rs

@@ -0,0 +1,16 @@
+use mako_core::lazy_static::lazy_static;
+use mako_core::rayon::{ThreadPool, ThreadPoolBuilder};
+
+lazy_static! {
+    static ref THREAD_POOL: ThreadPool = ThreadPoolBuilder::new()
+        .thread_name(|i| format!("rayon thread {}", i))
+        .build()
+        .expect("failed to create rayon thread pool.");
+}
+
+pub fn spawn<F>(func: F)
+where
+    F: FnOnce() + Send + 'static,
+{
+    THREAD_POOL.spawn(func)
+}

crates/mako/src/transform_in_generate.rs

@@ -1,4 +1,5 @@
 use std::collections::{HashMap, HashSet};
+use std::sync::mpsc::channel;
 use std::sync::Arc;
 use std::time::Instant;
 
@@ -24,15 +25,14 @@ use crate::compiler::{Compiler, Context};
 use crate::config::OutputMode;
 use crate::module::{Dependency, ModuleAst, ModuleId, ResolveType};
 use crate::swc_helpers::SwcHelpers;
-use crate::targets;
 use crate::transformers::transform_async_module::AsyncModule;
 use crate::transformers::transform_css_handler::CssHandler;
 use crate::transformers::transform_dep_replacer::{DepReplacer, DependenciesToReplace};
 use crate::transformers::transform_dynamic_import::DynamicImport;
 use crate::transformers::transform_mako_require::MakoRequire;
 use crate::transformers::transform_meta_url_replacer::MetaUrlReplacer;
 use crate::transformers::transform_optimize_define_utils::OptimizeDefineUtils;
-use crate::util::create_thread_pool;
+use crate::{targets, thread_pool};
 
 impl Compiler {
     pub fn transform_all(&self) -> Result<()> {
@@ -93,14 +93,15 @@ pub fn transform_modules_in_thread(
     async_deps_by_module_id: HashMap<ModuleId, Vec<Dependency>>,
 ) -> Result<()> {
     mako_core::mako_profile_function!();
-    let (pool, rs, rr) =
-        create_thread_pool::<Result<(ModuleId, ModuleAst, Option<HashSet<String>>)>>();
+
+    let (rs, rr) = channel::<Result<(ModuleId, ModuleAst, Option<HashSet<String>>)>>();
+
     for module_id in module_ids {
         let context = context.clone();
         let rs = rs.clone();
         let module_id = module_id.clone();
         let async_deps = async_deps_by_module_id.get(&module_id).unwrap().clone();
-        pool.spawn(move || {
+        thread_pool::spawn(move || {
             let module_graph = context.module_graph.read().unwrap();
             let deps = module_graph.get_dependencies(&module_id);
             let mut resolved_deps: HashMap<String, (String, String)> = deps

crates/mako/src/tree_shaking/module_side_effects_flag.rs

@@ -87,8 +87,6 @@ fn match_glob_pattern(pattern: &str, path: &str) -> bool {
 
 #[cfg(test)]
 mod tests {
-    use mako_core::tokio;
-
     use super::match_glob_pattern;
     use crate::test_helper::{get_module, setup_compiler};
 
@@ -105,8 +103,8 @@ mod tests {
         ));
     }
 
-    #[tokio::test(flavor = "multi_thread")]
-    async fn test_side_effects_flag() {
+    #[test]
+    fn test_side_effects_flag() {
         let compiler = setup_compiler("test/build/side-effects-flag", false);
         compiler.compile().unwrap();
         let foo = get_module(&compiler, "node_modules/foo/index.ts");

crates/mako/src/util.rs

@@ -1,22 +1,13 @@
 use std::path::PathBuf;
-use std::sync::mpsc::{channel, Receiver, Sender};
-use std::sync::Arc;
 
 use mako_core::anyhow::{anyhow, Result};
 use mako_core::base64::engine::general_purpose;
 use mako_core::base64::Engine;
 use mako_core::merge_source_map::sourcemap::SourceMap;
 use mako_core::merge_source_map::{merge, MergeOptions};
 use mako_core::pathdiff::diff_paths;
-use mako_core::rayon::{ThreadPool, ThreadPoolBuilder};
 use mako_core::regex::Regex;
 
-pub fn create_thread_pool<T>() -> (Arc<ThreadPool>, Sender<T>, Receiver<T>) {
-    let pool = Arc::new(ThreadPoolBuilder::new().build().unwrap());
-    let (rs, rr) = channel();
-    (pool, rs, rr)
-}
-
 pub fn base64_decode(bytes: &[u8]) -> Vec<u8> {
     general_purpose::STANDARD.decode(bytes).unwrap()
 }

crates/mako/src/visitors/css_dep_analyzer.rs

@@ -88,10 +88,10 @@ mod tests {
 
     #[test]
     fn test_remote() {
-        assert_eq!(run(r#"@import url(https://a.com/a.css);"#).is_empty(), true);
-        assert_eq!(run(r#"@import url(http://a.com/a.css);"#).is_empty(), true);
-        assert_eq!(run(r#"@import url(data://a.com/a.css);"#).is_empty(), true);
-        assert_eq!(run(r#"@import url(//a.com/a.css);"#).is_empty(), true);
+        assert!(run(r#"@import url(https://a.com/a.css);"#).is_empty());
+        assert!(run(r#"@import url(http://a.com/a.css);"#).is_empty());
+        assert!(run(r#"@import url(data://a.com/a.css);"#).is_empty());
+        assert!(run(r#"@import url(//a.com/a.css);"#).is_empty());
     }
 
     #[test]

crates/mako/src/visitors/js_dep_analyzer.rs

@@ -177,11 +177,8 @@ mod tests {
     #[test]
     fn test_require() {
         assert_eq!(run(r#"require('a');"#), vec!["a"]);
-        assert_eq!(
-            run(r#"const require = 'a'; require('a');"#).is_empty(),
-            true
-        );
-        assert_eq!(run(r#"require(a);"#).is_empty(), true);
+        assert!(run(r#"const require = 'a'; require('a');"#).is_empty());
+        assert!(run(r#"require(a);"#).is_empty());
     }
 
     #[test]
@@ -191,24 +188,15 @@ mod tests {
             vec!["a"]
         );
         // Worker is defined
-        assert_eq!(
-            run(r#"const Worker = 1;new Worker(new URL('a', import.meta.url));"#).is_empty(),
-            true
-        );
+        assert!(run(r#"const Worker = 1;new Worker(new URL('a', import.meta.url));"#).is_empty());
         // URL is defined
-        assert_eq!(
-            run(r#"const URL = 1;new Worker(new URL('a', import.meta.url));"#).is_empty(),
-            true
-        );
+        assert!(run(r#"const URL = 1;new Worker(new URL('a', import.meta.url));"#).is_empty());
         // no import.meta.url
-        assert_eq!(run(r#"new Worker(new URL('a'));"#).is_empty(), true);
+        assert!(run(r#"new Worker(new URL('a'));"#).is_empty());
         // no new URL
-        assert_eq!(run(r#"new Worker('a');"#).is_empty(), true);
+        assert!(run(r#"new Worker('a');"#).is_empty());
         // ignore remote
-        assert_eq!(
-            run(r#"new Worker(new URL('https://a', import.meta.url));"#).is_empty(),
-            true
-        );
+        assert!(run(r#"new Worker(new URL('https://a', import.meta.url));"#).is_empty());
     }
 
     #[test]

crates/mako/src/watch.rs

@@ -1,12 +1,12 @@
 use std::collections::{HashMap, HashSet};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::time::Instant;
 
 use mako_core::anyhow::{self, Ok};
 use mako_core::colored::Colorize;
 use mako_core::notify::{self, EventKind, Watcher as NotifyWatcher};
 use mako_core::notify_debouncer_full::DebouncedEvent;
-use mako_core::tokio::time::Instant;
 use mako_core::tracing::debug;
 
 use crate::compiler::Compiler;