From d134c1d89afabdc6f9e2718df5f186e9039e70c1 Mon Sep 17 00:00:00 2001
From: Boshen <boshenc@gmail.com>
Date: Sun, 19 Mar 2023 11:32:46 +0800
Subject: [PATCH] refactor(oxc_cli): clean up code around spawning tasks

---
 crates/oxc_cli/src/lint/mod.rs    |  1 +
 crates/oxc_cli/src/lint/runner.rs | 74 ++++++++++++++++---------------
 2 files changed, 40 insertions(+), 35 deletions(-)

diff --git a/crates/oxc_cli/src/lint/mod.rs b/crates/oxc_cli/src/lint/mod.rs
index d04a0ac8c..f03c34418 100644
--- a/crates/oxc_cli/src/lint/mod.rs
+++ b/crates/oxc_cli/src/lint/mod.rs
@@ -6,6 +6,7 @@ use clap::ArgMatches;
 
 pub use self::{command::lint_command, runner::LintRunner};
 
+#[derive(Debug)]
 pub struct LintOptions {
     pub paths: Vec<PathBuf>,
     pub fix: bool,
diff --git a/crates/oxc_cli/src/lint/runner.rs b/crates/oxc_cli/src/lint/runner.rs
index 58d1ec1f9..581b7f002 100644
--- a/crates/oxc_cli/src/lint/runner.rs
+++ b/crates/oxc_cli/src/lint/runner.rs
@@ -3,7 +3,10 @@ use std::{
     io::{BufWriter, Write},
     path::{Path, PathBuf},
     rc::Rc,
-    sync::{mpsc, Arc},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        mpsc, Arc,
+    },
 };
 
 use miette::NamedSource;
@@ -34,18 +37,15 @@ impl LintRunner {
     pub fn run(&self) -> CliRunResult {
         let now = std::time::Instant::now();
 
-        let mut number_of_files = 0;
-        let mut number_of_warnings = 0;
-        let mut number_of_diagnostics = 0;
-
+        let number_of_files = Arc::new(AtomicUsize::new(0));
         let (tx_error, rx_error) = mpsc::channel::<(PathBuf, Vec<Error>)>();
 
-        self.process_paths(&mut number_of_files, tx_error);
-        self.process_diagnostics(&mut number_of_warnings, &mut number_of_diagnostics, &rx_error);
+        self.process_paths(&number_of_files, tx_error);
+        let (number_of_warnings, number_of_diagnostics) = self.process_diagnostics(&rx_error);
 
         CliRunResult::LintResult {
             duration: now.elapsed(),
-            number_of_files,
+            number_of_files: number_of_files.load(Ordering::Relaxed),
             number_of_diagnostics,
             number_of_warnings,
             max_warnings_exceeded: self
@@ -57,46 +57,49 @@ impl LintRunner {
 
     fn process_paths(
         &self,
-        number_of_files: &mut usize,
+        number_of_files: &Arc<AtomicUsize>,
         tx_error: mpsc::Sender<(PathBuf, Vec<Error>)>,
     ) {
         let (tx_path, rx_path) = mpsc::channel::<Box<Path>>();
-        let options = &self.options;
-        let fix = options.fix;
-        rayon::join(
-            move || {
-                Walk::new(options).iter().for_each(|path| {
-                    *number_of_files += 1;
-                    tx_path.send(path).unwrap();
+
+        let walk = Walk::new(&self.options);
+        let number_of_files = Arc::clone(number_of_files);
+        rayon::spawn(move || {
+            let mut count = 0;
+            walk.iter().for_each(|path| {
+                count += 1;
+                tx_path.send(path).unwrap();
+            });
+            number_of_files.store(count, Ordering::Relaxed);
+        });
+
+        let fix = self.options.fix;
+        rayon::spawn(move || {
+            while let Ok(path) = rx_path.recv() {
+                let tx_error = tx_error.clone();
+                rayon::spawn(move || {
+                    if let Some(diagnostics) = Self::lint_path(&path, fix) {
+                        tx_error.send(diagnostics).unwrap();
+                    }
+                    drop(tx_error);
                 });
-            },
-            move || {
-                while let Ok(path) = rx_path.recv() {
-                    let tx_error = tx_error.clone();
-                    rayon::spawn(move || {
-                        if let Some(diagnostics) = Self::lint_path(&path, fix) {
-                            tx_error.send(diagnostics).unwrap();
-                        }
-                        drop(tx_error);
-                    });
-                }
-            },
-        );
+            }
+        });
     }
 
     fn process_diagnostics(
         &self,
-        number_of_warnings: &mut usize,
-        number_of_diagnostics: &mut usize,
         rx_error: &mpsc::Receiver<(PathBuf, Vec<Error>)>,
-    ) {
+    ) -> (usize, usize) {
+        let mut number_of_warnings = 0;
+        let mut number_of_diagnostics = 0;
         let mut buf_writer = BufWriter::new(std::io::stdout());
 
         while let Ok((path, diagnostics)) = rx_error.recv() {
-            *number_of_diagnostics += diagnostics.len();
+            number_of_diagnostics += diagnostics.len();
             for diagnostic in diagnostics {
                 if diagnostic.severity() == Some(Severity::Warning) {
-                    *number_of_warnings += 1;
+                    number_of_warnings += 1;
                     // The --quiet flag follows ESLint's --quiet behavior as documented here: https://eslint.org/docs/latest/use/command-line-interface#--quiet
                     // Note that it does not disable ALL diagnostics, only Warning diagnostics
                     if self.options.quiet {
@@ -104,7 +107,7 @@ impl LintRunner {
                     }
 
                     if let Some(max_warnings) = self.options.max_warnings {
-                        if *number_of_warnings > max_warnings {
+                        if number_of_warnings > max_warnings {
                             continue;
                         }
                     }
@@ -122,6 +125,7 @@ impl LintRunner {
         }
 
         buf_writer.flush().unwrap();
+        (number_of_warnings, number_of_diagnostics)
     }
 
     fn lint_path(path: &Path, fix: bool) -> Option<(PathBuf, Vec<Error>)> {