Refactor: Extract new fn run_commands_in_parallel (#1549)

Simplify `compile_objects` function and prepare for parallel is_flags_supported checks

Author: NobodyXu

src/command_helpers.rs

@@ -119,7 +119,7 @@ impl StderrForwarder {
         }
     }
 
-    fn forward_available(&mut self) -> bool {
+    pub(crate) fn forward_available(&mut self) -> bool {
         if let Some((stderr, buffer)) = self.inner.as_mut() {
             loop {
                 // For non-blocking we check to see if there is data available, so we should try to
@@ -226,7 +226,7 @@ impl StderrForwarder {
     }
 
     #[cfg(feature = "parallel")]
-    fn forward_all(&mut self) {
+    pub(crate) fn forward_all(&mut self) {
         while !self.forward_available() {}
     }
 
@@ -424,38 +424,3 @@ pub(crate) fn command_add_output_file(cmd: &mut Command, dst: &Path, args: CmdAd
         cmd.arg("-o").arg(dst);
     }
 }
-
-#[cfg(feature = "parallel")]
-pub(crate) fn try_wait_on_child(
-    cmd: &Command,
-    child: &mut Child,
-    stdout: &mut dyn io::Write,
-    stderr_forwarder: &mut StderrForwarder,
-) -> Result<Option<()>, Error> {
-    stderr_forwarder.forward_available();
-
-    match child.try_wait() {
-        Ok(Some(status)) => {
-            stderr_forwarder.forward_all();
-
-            let _ = writeln!(stdout, "{}", status);
-
-            if status.success() {
-                Ok(Some(()))
-            } else {
-                Err(Error::new(
-                    ErrorKind::ToolExecError,
-                    format!("command did not execute successfully (status code {status}): {cmd:?}"),
-                ))
-            }
-        }
-        Ok(None) => Ok(None),
-        Err(e) => {
-            stderr_forwarder.forward_all();
-            Err(Error::new(
-                ErrorKind::ToolExecError,
-                format!("failed to wait on spawned child process `{cmd:?}`: {e}"),
-            ))
-        }
-    }
-}

src/lib.rs

@@ -248,8 +248,6 @@ use std::fmt::{self, Display};
 use std::fs;
 use std::io::{self, Write};
 use std::path::{Component, Path, PathBuf};
-#[cfg(feature = "parallel")]
-use std::process::Child;
 use std::process::{Command, Stdio};
 use std::sync::{
     atomic::{AtomicU8, Ordering::Relaxed},
@@ -260,8 +258,10 @@ use shlex::Shlex;
 
 #[cfg(feature = "parallel")]
 mod parallel;
+
 mod target;
 use self::target::*;
+
 /// A helper module to looking for windows-specific tools:
 /// 1. On Windows host, probe the Windows Registry if needed;
 /// 2. On non-Windows host, check specified environment variables.
@@ -1755,144 +1755,17 @@ impl Build {
         Ok(objects.into_iter().map(|v| v.dst).collect())
     }
 
-    #[cfg(feature = "parallel")]
     fn compile_objects(&self, objs: &[Object]) -> Result<(), Error> {
-        use std::cell::Cell;
-
-        use parallel::async_executor::{block_on, YieldOnce};
-
         check_disabled()?;
 
-        if objs.len() <= 1 {
-            return self.compile_objects_sequential(objs);
-        }
-
-        // Limit our parallelism globally with a jobserver.
-        let mut tokens = parallel::job_token::ActiveJobTokenServer::new();
-
-        // When compiling objects in parallel we do a few dirty tricks to speed
-        // things up:
-        //
-        // * First is that we use the `jobserver` crate to limit the parallelism
-        //   of this build script. The `jobserver` crate will use a jobserver
-        //   configured by Cargo for build scripts to ensure that parallelism is
-        //   coordinated across C compilations and Rust compilations. Before we
-        //   compile anything we make sure to wait until we acquire a token.
-        //
-        //   Note that this jobserver is cached globally so we only used one per
-        //   process and only worry about creating it once.
-        //
-        // * Next we use spawn the process to actually compile objects in
-        //   parallel after we've acquired a token to perform some work
-        //
-        // With all that in mind we compile all objects in a loop here, after we
-        // acquire the appropriate tokens, Once all objects have been compiled
-        // we wait on all the processes and propagate the results of compilation.
-
-        let pendings =
-            Cell::new(Vec::<(Command, KillOnDrop, parallel::job_token::JobToken)>::new());
-        let is_disconnected = Cell::new(false);
-        let has_made_progress = Cell::new(false);
-
-        let wait_future = async {
-            let mut error = None;
-            // Buffer the stdout
-            let mut stdout = io::BufWriter::with_capacity(128, io::stdout());
-
-            loop {
-                // If the other end of the pipe is already disconnected, then we're not gonna get any new jobs,
-                // so it doesn't make sense to reuse the tokens; in fact,
-                // releasing them as soon as possible (once we know that the other end is disconnected) is beneficial.
-                // Imagine that the last file built takes an hour to finish; in this scenario,
-                // by not releasing the tokens before that last file is done we would effectively block other processes from
-                // starting sooner - even though we only need one token for that last file, not N others that were acquired.
-
-                let mut pendings_is_empty = false;
-
-                cell_update(&pendings, |mut pendings| {
-                    // Try waiting on them.
-                    pendings.retain_mut(|(cmd, child, _token)| {
-                        match try_wait_on_child(cmd, &mut child.0, &mut stdout, &mut child.1) {
-                            Ok(Some(())) => {
-                                // Task done, remove the entry
-                                has_made_progress.set(true);
-                                false
-                            }
-                            Ok(None) => true, // Task still not finished, keep the entry
-                            Err(err) => {
-                                // Task fail, remove the entry.
-                                // Since we can only return one error, log the error to make
-                                // sure users always see all the compilation failures.
-                                has_made_progress.set(true);
-
-                                if self.cargo_output.warnings {
-                                    let _ = writeln!(stdout, "cargo:warning={}", err);
-                                }
-                                error = Some(err);
-
-                                false
-                            }
-                        }
-                    });
-                    pendings_is_empty = pendings.is_empty();
-                    pendings
-                });
-
-                if pendings_is_empty && is_disconnected.get() {
-                    break if let Some(err) = error {
-                        Err(err)
-                    } else {
-                        Ok(())
-                    };
-                }
-
-                YieldOnce::default().await;
-            }
-        };
-        let spawn_future = async {
-            for obj in objs {
-                let mut cmd = self.create_compile_object_cmd(obj)?;
-                let token = tokens.acquire().await?;
-                let mut child = spawn(&mut cmd, &self.cargo_output)?;
-                let mut stderr_forwarder = StderrForwarder::new(&mut child);
-                stderr_forwarder.set_non_blocking()?;
-
-                cell_update(&pendings, |mut pendings| {
-                    pendings.push((cmd, KillOnDrop(child, stderr_forwarder), token));
-                    pendings
-                });
-
-                has_made_progress.set(true);
-            }
-            is_disconnected.set(true);
-
-            Ok::<_, Error>(())
-        };
-
-        return block_on(wait_future, spawn_future, &has_made_progress);
-
-        struct KillOnDrop(Child, StderrForwarder);
-
-        impl Drop for KillOnDrop {
-            fn drop(&mut self) {
-                let child = &mut self.0;
-
-                child.kill().ok();
-            }
-        }
-
-        fn cell_update<T, F>(cell: &Cell<T>, f: F)
-        where
-            T: Default,
-            F: FnOnce(T) -> T,
-        {
-            let old = cell.take();
-            let new = f(old);
-            cell.set(new);
+        #[cfg(feature = "parallel")]
+        if objs.len() > 1 {
+            return parallel::run_commands_in_parallel(
+                &self.cargo_output,
+                &mut objs.iter().map(|obj| self.create_compile_object_cmd(obj)),
+            );
         }
-    }
 
-    fn compile_objects_sequential(&self, objs: &[Object]) -> Result<(), Error> {
         for obj in objs {
             let mut cmd = self.create_compile_object_cmd(obj)?;
             run(&mut cmd, &self.cargo_output)?;
@@ -1901,13 +1774,6 @@ impl Build {
         Ok(())
     }
 
-    #[cfg(not(feature = "parallel"))]
-    fn compile_objects(&self, objs: &[Object]) -> Result<(), Error> {
-        check_disabled()?;
-
-        self.compile_objects_sequential(objs)
-    }
-
     fn create_compile_object_cmd(&self, obj: &Object) -> Result<Command, Error> {
         let asm_ext = AsmFileExt::from_path(&obj.src);
         let is_asm = asm_ext.is_some();