feat: add support for multithreading (#330)

hougesen · Jun 25, 2024 · c22a4b6 · c22a4b6
1 parent b2e6f2f
commit c22a4b6
Show file tree

Hide file tree

Showing 8 changed files with 160 additions and 36 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -38,6 +38,7 @@ serde = { version = "1.0.203", features = ["derive"] }
 serde_json = "1.0.118"
 tempfile = "3.10.1"
 terminal_size = "0.3.0"
+threadpool = "1.8.1"
 which = "6.0.1"
 
 [dev-dependencies]

diff --git a/codegen/Cargo.lock b/codegen/Cargo.lock
diff --git a/src/cli.rs b/src/cli.rs
@@ -45,6 +45,10 @@ pub struct FormatCommandArguments {
 
     #[arg(long, value_enum)]
     pub log_level: Option<LogLevel>,
+
+    /// Amount of threads to use. Defaults to 0 (auto).
+    #[arg(long)]
+    pub threads: Option<usize>,
 }
 
 #[derive(clap::ValueEnum, Clone, Copy, PartialEq, Eq, Debug)]

diff --git a/src/commands/format.rs b/src/commands/format.rs
@@ -1,10 +1,31 @@
-use std::env::current_dir;
+use std::{
+    env::current_dir,
+    sync::{atomic::AtomicU32, Arc},
+};
 
 use clap::builder::OsStr;
 use mdsf::{cli::FormatCommandArguments, config::MdsfConfig, error::MdsfError, handle_file};
+use threadpool::ThreadPool;
 
 const MDSF_IGNORE_FILE_NAME: &str = ".mdsfignore";
 
+#[inline]
+fn detemine_threads_to_use(argument: &Option<usize>) -> usize {
+    if let Some(thread_arg) = argument {
+        if thread_arg > &0 {
+            return thread_arg.to_owned();
+        }
+    }
+
+    if let Ok(available_threads) = std::thread::available_parallelism().map(usize::from) {
+        if available_threads > 0 {
+            return available_threads;
+        }
+    }
+
+    1
+}
+
 pub fn run(args: FormatCommandArguments, dry_run: bool) -> Result<(), MdsfError> {
     mdsf::DEBUG.swap(args.debug, core::sync::atomic::Ordering::Relaxed);
 
@@ -18,10 +39,14 @@ pub fn run(args: FormatCommandArguments, dry_run: bool) -> Result<(), MdsfError>
 
     mdsf::runners::set_javascript_runtime(conf.javascript_runtime);
 
-    let mut changed_file_count = 0;
+    let changed_file_count = Arc::new(AtomicU32::new(0));
 
     if args.path.is_file() {
-        changed_file_count += u32::from(handle_file(&conf, &args.path, dry_run)?);
+        let was_formatted = handle_file(&conf, &args.path, dry_run);
+
+        if was_formatted {
+            changed_file_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+        }
     } else if args.path.is_dir() {
         let mut walk_builder = ignore::WalkBuilder::new(args.path);
 
@@ -33,19 +58,39 @@ pub fn run(args: FormatCommandArguments, dry_run: bool) -> Result<(), MdsfError>
 
         let md_ext = OsStr::from("md");
 
+        let thread_count = detemine_threads_to_use(&args.threads).max(1);
+
+        let pool = ThreadPool::new(thread_count);
+
+        let shared_config = Arc::new(conf);
+
         for entry in walk_builder.build().flatten() {
-            let file_path = entry.path();
+            let file_path = entry.path().to_path_buf();
 
             if file_path.extension() == Some(&md_ext) {
-                changed_file_count += u32::from(handle_file(&conf, file_path, dry_run)?);
+                let config_ref = shared_config.clone();
+
+                let changed_file_count_ref = changed_file_count.clone();
+
+                pool.execute(move || {
+                    let was_formatted = handle_file(&config_ref, &file_path, dry_run);
+
+                    if was_formatted {
+                        changed_file_count_ref.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+                    }
+                });
             }
         }
+
+        pool.join();
     } else {
         return Err(MdsfError::FileNotFound(args.path));
     }
 
-    if dry_run && changed_file_count > 0 {
-        Err(MdsfError::CheckModeChanges(changed_file_count))
+    let total_changed_files = changed_file_count.load(std::sync::atomic::Ordering::SeqCst);
+
+    if dry_run && total_changed_files > 0 {
+        Err(MdsfError::CheckModeChanges(total_changed_files))
     } else {
         Ok(())
     }

diff --git a/src/formatters/mod.rs b/src/formatters/mod.rs
@@ -1230,6 +1230,7 @@ impl MdsfFormatter<Tooling> {
                 if let Err(e) = &r {
                     if let MdsfError::MissingBinary(binary) = e {
                         print_binary_not_in_path(
+                            info.filename,
                             if &formatter_name == binary {
                                 formatter_name
                             } else {

diff --git a/src/lib.rs b/src/lib.rs
@@ -1,8 +1,9 @@
 use core::sync::atomic::AtomicBool;
 
+use terminal::{print_error_reading_file, print_error_saving_file};
+
 use crate::{
     config::MdsfConfig,
-    error::MdsfError,
     formatters::format_snippet,
     parser::{parse_generic_codeblock, parse_go_codeblock},
     terminal::{
@@ -111,34 +112,44 @@ fn format_file(config: &MdsfConfig, filename: &std::path::Path, input: &str) ->
 }
 
 #[inline]
-pub fn handle_file(
-    config: &MdsfConfig,
-    path: &std::path::Path,
-    dry_run: bool,
-) -> Result<bool, MdsfError> {
+pub fn handle_file(config: &MdsfConfig, path: &std::path::Path, dry_run: bool) -> bool {
     let time = std::time::Instant::now();
 
-    let input = std::fs::read_to_string(path)?;
+    match std::fs::read_to_string(path) {
+        Ok(input) => {
+            if !input.is_empty() {
+                let (modified, output) = format_file(config, path, &input);
 
-    if !input.is_empty() {
-        let (modified, output) = format_file(config, path, &input);
+                if modified && output != input {
+                    if dry_run {
+                        print_changed_file_error(path);
+                    } else {
+                        let write_result = std::fs::write(path, output);
 
-        if modified && output != input {
-            if dry_run {
-                print_changed_file_error(path);
-            } else {
-                std::fs::write(path, output)?;
+                        if let Err(write_error) = write_result {
+                            print_error_saving_file(path, &write_error);
+
+                            return false;
+                        }
+                    }
+
+                    print_changed_file(path, time.elapsed());
+
+                    return true;
+                }
             }
 
-            print_changed_file(path, time.elapsed());
+            print_unchanged_file(path, time.elapsed());
 
-            return Ok(true);
+            false
         }
-    }
 
-    print_unchanged_file(path, time.elapsed());
+        Err(error) => {
+            print_error_reading_file(path, &error);
 
-    Ok(false)
+            false
+        }
+    }
 }
 
 pub struct LineInfo<'a> {
@@ -201,7 +212,7 @@ fn add(a: i32, b: i32) -> i32 {
             let file =
                 setup_snippet(input, language_to_ext("markdown")).expect("it to create a file");
 
-            handle_file(&config, file.path(), false).expect("it to be a success");
+            assert!(handle_file(&config, file.path(), false));
 
             let output = std::fs::read_to_string(file.path()).expect("it to return the string");
 
@@ -941,7 +952,7 @@ fn add(a: i32, b: i32) i32 {
             let file =
                 setup_snippet(input, language_to_ext("markdown")).expect("it to create a file");
 
-            handle_file(&config, file.path(), false).expect("it to be a success");
+            assert!(handle_file(&config, file.path(), false));
 
             let output = std::fs::read_to_string(file.path()).expect("it to return the string");
 
@@ -1002,7 +1013,7 @@ type Whatever struct {
                 let file =
                     setup_snippet(input, language_to_ext("markdown")).expect("it to create a file");
 
-                handle_file(&config, file.path(), false).expect("it to be a success");
+                assert!(handle_file(&config, file.path(), false));
 
                 let output = std::fs::read_to_string(file.path()).expect("it to return the string");
 
@@ -1031,7 +1042,7 @@ type Whatever struct {
                 let file =
                     setup_snippet(input, language_to_ext("markdown")).expect("it to create a file");
 
-                handle_file(&config, file.path(), false).expect("it to be a success");
+                assert!(handle_file(&config, file.path(), false));
 
                 let output = std::fs::read_to_string(file.path()).expect("it to return the string");
 
@@ -1088,7 +1099,7 @@ type Whatever struct {
                 let file =
                     setup_snippet(input, language_to_ext("markdown")).expect("it to create a file");
 
-                handle_file(&config, file.path(), false).expect("it to be a success");
+                assert!(handle_file(&config, file.path(), false));
 
                 let output = std::fs::read_to_string(file.path()).expect("it to return the string");
 
@@ -1118,7 +1129,7 @@ type Whatever struct {
                 let file =
                     setup_snippet(input, language_to_ext("markdown")).expect("it to create a file");
 
-                handle_file(&config, file.path(), false).expect("it to be a success");
+                assert!(handle_file(&config, file.path(), false));
 
                 let output = std::fs::read_to_string(file.path()).expect("it to return the string");
 
@@ -1205,7 +1216,7 @@ func add(a int, b int) int {
                 let file =
                     setup_snippet(input, language_to_ext("markdown")).expect("it to create a file");
 
-                handle_file(&config, file.path(), false).expect("it to be a success");
+                assert!(handle_file(&config, file.path(), false));
 
                 let output = std::fs::read_to_string(file.path()).expect("it to return the string");
 
@@ -1234,7 +1245,7 @@ func add(a int, b int) int {
                 let file =
                     setup_snippet(input, language_to_ext("markdown")).expect("it to create a file");
 
-                handle_file(&config, file.path(), false).expect("it to be a success");
+                assert!(handle_file(&config, file.path(), false));
 
                 let output = std::fs::read_to_string(file.path()).expect("it to return the string");