add deferred command in execution context and update run method

1 files changed, 91 insertions, 46 deletions

diff --git a/src/bootstrap/src/utils/execution_context.rs b/src/bootstrap/src/utils/execution_context.rs
index a5e1e9bcc07..778ef1b0fc1 100644
--- a/src/bootstrap/src/utils/execution_context.rs
+++ b/src/bootstrap/src/utils/execution_context.rs
@@ -3,6 +3,7 @@
 //! This module provides the [`ExecutionContext`] type, which holds global configuration
 //! relevant during the execution of commands in bootstrap. This includes dry-run
 //! mode, verbosity level, and behavior on failure.
+use std::process::Child;
 use std::sync::{Arc, Mutex};
 
 use crate::core::config::DryRun;
@@ -80,15 +81,16 @@ impl ExecutionContext {
     /// Note: Ideally, you should use one of the BootstrapCommand::run* functions to
     /// execute commands. They internally call this method.
     #[track_caller]
-    pub fn run(
+    pub fn start<'a>(
         &self,
-        command: &mut BootstrapCommand,
+        command: &'a mut BootstrapCommand,
         stdout: OutputMode,
         stderr: OutputMode,
-    ) -> CommandOutput {
+    ) -> DeferredCommand<'a> {
         command.mark_as_executed();
+
         if self.dry_run() && !command.run_always {
-            return CommandOutput::default();
+            return DeferredCommand::new(None, stdout, stderr, command, Arc::new(self.clone()));
         }
 
         #[cfg(feature = "tracing")]
@@ -105,7 +107,75 @@ impl ExecutionContext {
         cmd.stdout(stdout.stdio());
         cmd.stderr(stderr.stdio());
 
-        let output = cmd.output();
+        let child = cmd.spawn().unwrap();
+
+        DeferredCommand::new(Some(child), stdout, stderr, command, Arc::new(self.clone()))
+    }
+
+    /// Execute a command and return its output.
+    /// Note: Ideally, you should use one of the BootstrapCommand::run* functions to
+    /// execute commands. They internally call this method.
+    #[track_caller]
+    pub fn run(
+        &self,
+        command: &mut BootstrapCommand,
+        stdout: OutputMode,
+        stderr: OutputMode,
+    ) -> CommandOutput {
+        self.start(command, stdout, stderr).wait_for_output()
+    }
+
+    fn fail(&self, message: &str, output: CommandOutput) -> ! {
+        if self.is_verbose() {
+            println!("{message}");
+        } else {
+            let (stdout, stderr) = (output.stdout_if_present(), output.stderr_if_present());
+            // If the command captures output, the user would not see any indication that
+            // it has failed. In this case, print a more verbose error, since to provide more
+            // context.
+            if stdout.is_some() || stderr.is_some() {
+                if let Some(stdout) = output.stdout_if_present().take_if(|s| !s.trim().is_empty()) {
+                    println!("STDOUT:\n{stdout}\n");
+                }
+                if let Some(stderr) = output.stderr_if_present().take_if(|s| !s.trim().is_empty()) {
+                    println!("STDERR:\n{stderr}\n");
+                }
+                println!("Command has failed. Rerun with -v to see more details.");
+            } else {
+                println!("Command has failed. Rerun with -v to see more details.");
+            }
+        }
+        exit!(1);
+    }
+}
+
+pub struct DeferredCommand<'a> {
+    process: Option<Child>,
+    command: &'a mut BootstrapCommand,
+    stdout: OutputMode,
+    stderr: OutputMode,
+    exec_ctx: Arc<ExecutionContext>,
+}
+
+impl<'a> DeferredCommand<'a> {
+    pub fn new(
+        child: Option<Child>,
+        stdout: OutputMode,
+        stderr: OutputMode,
+        command: &'a mut BootstrapCommand,
+        exec_ctx: Arc<ExecutionContext>,
+    ) -> Self {
+        DeferredCommand { process: child, stdout, stderr, command, exec_ctx }
+    }
+
+    pub fn wait_for_output(mut self) -> CommandOutput {
+        if self.process.is_none() {
+            return CommandOutput::default();
+        }
+        let output = self.process.take().unwrap().wait_with_output();
+
+        let created_at = self.command.get_created_location();
+        let executed_at = std::panic::Location::caller();
 
         use std::fmt::Write;
 
@@ -113,30 +183,31 @@ impl ExecutionContext {
         let output: CommandOutput = match output {
             // Command has succeeded
             Ok(output) if output.status.success() => {
-                CommandOutput::from_output(output, stdout, stderr)
+                CommandOutput::from_output(output, self.stdout, self.stderr)
             }
             // Command has started, but then it failed
             Ok(output) => {
                 writeln!(
                     message,
                     r#"
-Command {command:?} did not execute successfully.
+Command {:?} did not execute successfully.
 Expected success, got {}
 Created at: {created_at}
 Executed at: {executed_at}"#,
-                    output.status,
+                    self.command, output.status,
                 )
                 .unwrap();
 
-                let output: CommandOutput = CommandOutput::from_output(output, stdout, stderr);
+                let output: CommandOutput =
+                    CommandOutput::from_output(output, self.stdout, self.stderr);
 
                 // If the output mode is OutputMode::Capture, we can now print the output.
                 // If it is OutputMode::Print, then the output has already been printed to
                 // stdout/stderr, and we thus don't have anything captured to print anyway.
-                if stdout.captures() {
+                if self.stdout.captures() {
                     writeln!(message, "\nSTDOUT ----\n{}", output.stdout().trim()).unwrap();
                 }
-                if stderr.captures() {
+                if self.stderr.captures() {
                     writeln!(message, "\nSTDERR ----\n{}", output.stderr().trim()).unwrap();
                 }
                 output
@@ -145,52 +216,26 @@ Executed at: {executed_at}"#,
             Err(e) => {
                 writeln!(
                     message,
-                    "\n\nCommand {command:?} did not execute successfully.\
-            \nIt was not possible to execute the command: {e:?}"
+                    "\n\nCommand {:?} did not execute successfully.\
+            \nIt was not possible to execute the command: {e:?}",
+                    self.command
                 )
                 .unwrap();
-                CommandOutput::did_not_start(stdout, stderr)
-            }
-        };
-
-        let fail = |message: &str, output: CommandOutput| -> ! {
-            if self.is_verbose() {
-                println!("{message}");
-            } else {
-                let (stdout, stderr) = (output.stdout_if_present(), output.stderr_if_present());
-                // If the command captures output, the user would not see any indication that
-                // it has failed. In this case, print a more verbose error, since to provide more
-                // context.
-                if stdout.is_some() || stderr.is_some() {
-                    if let Some(stdout) =
-                        output.stdout_if_present().take_if(|s| !s.trim().is_empty())
-                    {
-                        println!("STDOUT:\n{stdout}\n");
-                    }
-                    if let Some(stderr) =
-                        output.stderr_if_present().take_if(|s| !s.trim().is_empty())
-                    {
-                        println!("STDERR:\n{stderr}\n");
-                    }
-                    println!("Command {command:?} has failed. Rerun with -v to see more details.");
-                } else {
-                    println!("Command has failed. Rerun with -v to see more details.");
-                }
+                CommandOutput::did_not_start(self.stdout, self.stderr)
             }
-            exit!(1);
         };
 
         if !output.is_success() {
-            match command.failure_behavior {
+            match self.command.failure_behavior {
                 BehaviorOnFailure::DelayFail => {
-                    if self.fail_fast {
-                        fail(&message, output);
+                    if self.exec_ctx.fail_fast {
+                        self.exec_ctx.fail(&message, output);
                     }
 
-                    self.add_to_delay_failure(message);
+                    self.exec_ctx.add_to_delay_failure(message);
                 }
                 BehaviorOnFailure::Exit => {
-                    fail(&message, output);
+                    self.exec_ctx.fail(&message, output);
                 }
                 BehaviorOnFailure::Ignore => {
                     // If failures are allowed, either the error has been printed already