@@ -5,22 +5,37 @@ use std::future::Future;
55use std:: io:: { Read , Write } ;
66use std:: process:: Command ;
77use std:: process:: ExitStatus ;
8+ use std:: sync:: { Arc , Mutex } ;
89use std:: thread;
910
11+ struct CmdRunnerOptions < F > {
12+ on_process_spawned : Option < F > ,
13+ capture_stdout : bool ,
14+ }
15+
16+ impl < F > Default for CmdRunnerOptions < F > {
17+ fn default ( ) -> Self {
18+ Self {
19+ on_process_spawned : None ,
20+ capture_stdout : false ,
21+ }
22+ }
23+ }
24+
1025/// Run a command and log its output to stdout and stderr
1126///
1227/// # Arguments
1328/// - `cmd`: The command to run.
14- /// - `cb `: A callback function that takes the process ID and returns a result.
29+ /// - `options `: Configuration options for the runner (e.g. capture output, run callback)
1530///
1631/// # Returns
17- ///
18- /// The exit status of the command.
19- ///
20- pub async fn run_command_with_log_pipe_and_callback < F , Fut > (
32+ /// A tuple containing:
33+ /// - `ExitStatus`: The exit status of the executed command
34+ /// - `Option<String>`: Captured stdout if `capture_stdout` was true, otherwise None
35+ async fn run_command_with_log_pipe_and_options < F , Fut > (
2136 mut cmd : Command ,
22- cb : F ,
23- ) -> Result < ExitStatus >
37+ options : CmdRunnerOptions < F > ,
38+ ) -> Result < ( ExitStatus , Option < String > ) >
2439where
2540 F : FnOnce ( u32 ) -> Fut ,
2641 Fut : Future < Output = anyhow:: Result < ( ) > > ,
@@ -29,14 +44,23 @@ where
2944 mut reader : impl Read ,
3045 mut writer : impl Write ,
3146 log_prefix : Option < & str > ,
47+ captured_output : Option < Arc < Mutex < Vec < u8 > > > > ,
3248 ) -> Result < ( ) > {
3349 let prefix = log_prefix. unwrap_or ( "" ) ;
3450 let mut buffer = [ 0 ; 1024 ] ;
51+ let mut capture_guard = captured_output
52+ . as_ref ( )
53+ . map ( |capture| capture. lock ( ) . unwrap ( ) ) ;
3554 loop {
3655 let bytes_read = reader. read ( & mut buffer) ?;
3756 if bytes_read == 0 {
3857 break ;
3958 }
59+
60+ if let Some ( ref mut output) = capture_guard {
61+ output. extend_from_slice ( & buffer[ ..bytes_read] ) ;
62+ }
63+
4064 suspend_progress_bar ( || {
4165 writer. write_all ( & buffer[ ..bytes_read] ) . unwrap ( ) ;
4266 trace ! (
@@ -57,19 +81,76 @@ where
5781 . context ( "failed to spawn the process" ) ?;
5882 let stdout = process. stdout . take ( ) . expect ( "unable to get stdout" ) ;
5983 let stderr = process. stderr . take ( ) . expect ( "unable to get stderr" ) ;
60- thread:: spawn ( move || {
61- log_tee ( stdout, std:: io:: stdout ( ) , None ) . unwrap ( ) ;
62- } ) ;
6384
64- thread:: spawn ( move || {
65- log_tee ( stderr, std:: io:: stderr ( ) , Some ( "[stderr]" ) ) . unwrap ( ) ;
66- } ) ;
85+ let captured_stdout = if options. capture_stdout {
86+ Some ( Arc :: new ( Mutex :: new ( Vec :: new ( ) ) ) )
87+ } else {
88+ None
89+ } ;
90+ let ( stdout_handle, stderr_handle) = {
91+ let stdout_capture = captured_stdout. clone ( ) ;
92+ let stdout_handle = thread:: spawn ( move || {
93+ log_tee ( stdout, std:: io:: stdout ( ) , None , stdout_capture) . unwrap ( ) ;
94+ } ) ;
95+ let stderr_handle = thread:: spawn ( move || {
96+ log_tee ( stderr, std:: io:: stderr ( ) , Some ( "[stderr]" ) , None ) . unwrap ( ) ;
97+ } ) ;
6798
68- cb ( process. id ( ) ) . await ?;
99+ ( stdout_handle, stderr_handle)
100+ } ;
69101
70- process. wait ( ) . context ( "failed to wait for the process" )
102+ if let Some ( cb) = options. on_process_spawned {
103+ cb ( process. id ( ) ) . await ?;
104+ }
105+
106+ let exit_status = process. wait ( ) . context ( "failed to wait for the process" ) ?;
107+ let _ = ( stdout_handle. join ( ) . unwrap ( ) , stderr_handle. join ( ) . unwrap ( ) ) ;
108+
109+ let stdout_output = captured_stdout
110+ . map ( |capture| String :: from_utf8_lossy ( & capture. lock ( ) . unwrap ( ) ) . to_string ( ) ) ;
111+ Ok ( ( exit_status, stdout_output) )
112+ }
113+
114+ pub async fn run_command_with_log_pipe_and_callback < F , Fut > (
115+ cmd : Command ,
116+ cb : F ,
117+ ) -> Result < ( ExitStatus , Option < String > ) >
118+ where
119+ F : FnOnce ( u32 ) -> Fut ,
120+ Fut : Future < Output = anyhow:: Result < ( ) > > ,
121+ {
122+ run_command_with_log_pipe_and_options (
123+ cmd,
124+ CmdRunnerOptions {
125+ on_process_spawned : Some ( cb) ,
126+ capture_stdout : false ,
127+ } ,
128+ )
129+ . await
71130}
72131
73132pub async fn run_command_with_log_pipe ( cmd : Command ) -> Result < ExitStatus > {
74- run_command_with_log_pipe_and_callback ( cmd, async |_| Ok ( ( ) ) ) . await
133+ let ( exit_status, _) = run_command_with_log_pipe_and_options (
134+ cmd,
135+ CmdRunnerOptions :: < fn ( u32 ) -> futures:: future:: Ready < anyhow:: Result < ( ) > > > {
136+ on_process_spawned : None ,
137+ capture_stdout : false ,
138+ } ,
139+ )
140+ . await ?;
141+ Ok ( exit_status)
142+ }
143+
144+ pub async fn run_command_with_log_pipe_capture_stdout (
145+ cmd : Command ,
146+ ) -> Result < ( ExitStatus , String ) > {
147+ let ( exit_status, stdout) = run_command_with_log_pipe_and_options (
148+ cmd,
149+ CmdRunnerOptions :: < fn ( u32 ) -> futures:: future:: Ready < anyhow:: Result < ( ) > > > {
150+ on_process_spawned : None ,
151+ capture_stdout : true ,
152+ } ,
153+ )
154+ . await ?;
155+ Ok ( ( exit_status, stdout. unwrap_or_default ( ) ) )
75156}
0 commit comments