reduce code duplication and join in reader threads cleanly

Author: cre4ture

tests/by-util/test_env.rs

@@ -250,13 +250,15 @@ fn test_fail_change_directory() {
 fn test_simulation_of_terminal_false() {
     let scene = TestScenario::new(util_name!());
 
-    let out = scene
-        .ucmd()
-        .arg("sh")
-        .arg("is_atty.sh")
-        .succeeds();
-    assert_eq!(String::from_utf8_lossy(out.stdout()), "stdin is not atty\nstdout is not atty\nstderr is not atty\n");
-    assert_eq!(String::from_utf8_lossy(out.stderr()), "This is an error message.\n");
+    let out = scene.ucmd().arg("sh").arg("is_atty.sh").succeeds();
+    assert_eq!(
+        String::from_utf8_lossy(out.stdout()),
+        "stdin is not atty\nstdout is not atty\nstderr is not atty\n"
+    );
+    assert_eq!(
+        String::from_utf8_lossy(out.stderr()),
+        "This is an error message.\n"
+    );
 }
 
 #[test]
@@ -269,8 +271,14 @@ fn test_simulation_of_terminal_true() {
         .arg("is_atty.sh")
         .terminal_simulation(true)
         .succeeds();
-    assert_eq!(String::from_utf8_lossy(out.stdout()), "stdin is atty\r\nstdout is atty\r\nstderr is atty\r\n");
-    assert_eq!(String::from_utf8_lossy(out.stderr()), "This is an error message.\r\n");
+    assert_eq!(
+        String::from_utf8_lossy(out.stdout()),
+        "stdin is atty\r\nstdout is atty\r\nstderr is atty\r\n"
+    );
+    assert_eq!(
+        String::from_utf8_lossy(out.stderr()),
+        "This is an error message.\r\n"
+    );
 }
 
 #[test]
@@ -288,23 +296,26 @@ fn test_simulation_of_terminal_pty_sends_eot_automatically() {
     assert_eq!(String::from_utf8_lossy(out.stderr()), "");
 }
 
-
 #[test]
 fn test_simulation_of_terminal_pty_pipes_into_data_and_sends_eot_automatically() {
     let scene = TestScenario::new(util_name!());
 
+    let message = "Hello stdin forwarding!";
+
     let mut cmd = scene.ucmd();
     cmd.args(&["cat", "-"]);
     cmd.terminal_simulation(true);
-    cmd.pipe_in("Hello stdin forwarding!");
+    cmd.pipe_in(message);
     let child = cmd.run_no_wait();
     let out = child.wait().unwrap();
 
-    assert_eq!(String::from_utf8_lossy(out.stdout()), "Hello stdin forwarding!\r\n");
+    assert_eq!(
+        String::from_utf8_lossy(out.stdout()),
+        format!("{}\r\n", message)
+    );
     assert_eq!(String::from_utf8_lossy(out.stderr()), "");
 }
 
-
 #[test]
 fn test_simulation_of_terminal_pty_write_in_data_and_sends_eot_automatically() {
     let scene = TestScenario::new(util_name!());
@@ -316,7 +327,9 @@ fn test_simulation_of_terminal_pty_write_in_data_and_sends_eot_automatically() {
     child.write_in("Hello stdin forwarding via write_in!");
     let out = child.wait().unwrap();
 
-    assert_eq!(String::from_utf8_lossy(out.stdout()), "Hello stdin forwarding via write_in!\r\n");
+    assert_eq!(
+        String::from_utf8_lossy(out.stdout()),
+        "Hello stdin forwarding via write_in!\r\n"
+    );
     assert_eq!(String::from_utf8_lossy(out.stderr()), "");
 }
-

tests/by-util/test_nohup.rs

@@ -41,9 +41,9 @@ fn test_nohup_multiple_args_and_flags() {
     target_vendor = "apple"
 ))]
 fn test_nohup_with_pseudo_terminal_emulation_on_stdin_stdout_stderr_get_replaced() {
-
     let ts = TestScenario::new(util_name!());
-    let result = ts.ucmd()
+    let result = ts
+        .ucmd()
         .terminal_simulation(true)
         .args(&["sh", "is_atty.sh"])
         .succeeds();
@@ -61,4 +61,3 @@ fn test_nohup_with_pseudo_terminal_emulation_on_stdin_stdout_stderr_get_replaced
         "stdin is not atty\nstdout is not atty\nstderr is not atty\n"
     );
 }
-

tests/common/util.rs

@@ -3,7 +3,7 @@
 // For the full copyright and license information, please view the LICENSE
 // file that was distributed with this source code.
 
-//spell-checker: ignore (linux) rlimit prlimit coreutil ggroups uchild uncaptured scmd SHLVL canonicalized
+//spell-checker: ignore (linux) rlimit prlimit coreutil ggroups uchild uncaptured scmd SHLVL canonicalized openpty
 
 #![allow(dead_code)]
 
@@ -38,7 +38,7 @@ use std::rc::Rc;
 use std::sync::mpsc::{self, RecvTimeoutError};
 use std::thread::{sleep, JoinHandle};
 use std::time::{Duration, Instant};
-use std::{env, hint, thread, mem};
+use std::{env, hint, mem, thread};
 use tempfile::{Builder, TempDir};
 
 static TESTS_DIR: &str = "tests";
@@ -1429,6 +1429,32 @@ impl UCommand {
         }
     }
 
+    fn spawn_reader_thread(
+        &self,
+        captured_output: Option<CapturedOutput>,
+        pty_fd_master: OwnedFd,
+        name: String,
+    ) -> Option<CapturedOutput> {
+        if let Some(mut captured_output_i) = captured_output {
+            let fd = captured_output_i.try_clone().unwrap();
+
+            let handle = std::thread::Builder::new()
+                .name(name)
+                .spawn(move || {
+                    let mut buffer_out = String::default();
+                    Self::read_string_from_pty(pty_fd_master, &mut buffer_out);
+                    let mut out = std::fs::File::from(fd);
+                    out.write_all(buffer_out.as_bytes()).unwrap();
+                })
+                .unwrap();
+
+            captured_output_i.reader_thread_handle = Some(handle);
+            Some(captured_output_i)
+        } else {
+            None
+        }
+    }
+
     /// Build the `std::process::Command` and apply the defaults on fields which were not specified
     /// by the user.
     ///
@@ -1448,7 +1474,14 @@ impl UCommand {
     /// * `stderr_to_stdout`: `false`
     /// * `bytes_into_stdin`: `None`
     /// * `limits`: `None`.
-    fn build(&mut self) -> (Command, Option<CapturedOutput>, Option<CapturedOutput>, Option<OwnedFd>) {
+    fn build(
+        &mut self,
+    ) -> (
+        Command,
+        Option<CapturedOutput>,
+        Option<CapturedOutput>,
+        Option<OwnedFd>,
+    ) {
         if self.bin_path.is_some() {
             if let Some(util_name) = &self.util_name {
                 self.args.push_front(util_name.into());
@@ -1562,7 +1595,6 @@ impl UCommand {
         };
 
         if self.terminal_simulation {
-
             let terminal_size = libc::winsize {
                 ws_col: 80,
                 ws_row: 30,
@@ -1585,41 +1617,12 @@ impl UCommand {
 
             stdin_pty = Some(pi_master);
 
-            if let Some(mut captured_stdout_i) = captured_stdout {
-                let fd = captured_stdout_i.try_clone().unwrap();
-
-                std::thread::Builder::new()
-                    .name("fwd_stdout".to_string())
-                    .spawn(move ||{
-                        let mut buffer_out = String::default();
-                        Self::read_string_from_pty(po_master, &mut buffer_out);
-                        let mut out = std::fs::File::from(fd);
-                        out.write_all(buffer_out.as_bytes()).unwrap();
-                        //std::mem::forget(_pi_master); // _pi_master needs to be kept alive otherwise stdin.is_terminal() fails
-                    }).unwrap();
-
-                captured_stdout = Some(captured_stdout_i);
-            }
-
-            if let Some(mut captured_stderr_i) = captured_stderr {
-                let fd = captured_stderr_i.try_clone().unwrap();
-
-                std::thread::Builder::new()
-                    .name("fwd_stderr".to_string())
-                    .spawn(move ||{
-                        let mut buffer_out = String::default();
-                        Self::read_string_from_pty(pe_master, &mut buffer_out);
-                        let mut out = std::fs::File::from(fd);
-                        out.write_all(buffer_out.as_bytes()).unwrap();
-                    }).unwrap();
-
-                captured_stderr = Some(captured_stderr_i);
-            }
+            captured_stdout =
+                self.spawn_reader_thread(captured_stdout, po_master, "stdout_reader".to_string());
+            captured_stderr =
+                self.spawn_reader_thread(captured_stderr, pe_master, "stderr_reader".to_string());
 
-            command
-                .stdin(pi_slave)
-                .stdout(po_slave)
-                .stderr(pe_slave);
+            command.stdin(pi_slave).stdout(po_slave).stderr(pe_slave);
         }
 
         (command, captured_stdout, captured_stderr, stdin_pty)
@@ -1631,10 +1634,7 @@ impl UCommand {
         assert!(!self.has_run, "{}", ALREADY_RUN);
         self.has_run = true;
 
-        let (mut command,
-            captured_stdout,
-            captured_stderr,
-            stdin_pty) = self.build();
+        let (mut command, captured_stdout, captured_stderr, stdin_pty) = self.build();
         log_info("run", self.to_string());
 
         let child = command.spawn().unwrap();
@@ -1715,6 +1715,7 @@ impl std::fmt::Display for UCommand {
 struct CapturedOutput {
     current_file: File,
     output: tempfile::NamedTempFile, // drop last
+    reader_thread_handle: Option<thread::JoinHandle<()>>,
 }
 
 impl CapturedOutput {
@@ -1723,6 +1724,7 @@ impl CapturedOutput {
         Self {
             current_file: output.reopen().unwrap(),
             output,
+            reader_thread_handle: None,
         }
     }
 
@@ -1799,6 +1801,7 @@ impl Default for CapturedOutput {
         Self {
             current_file: file.reopen().unwrap(),
             output: file,
+            reader_thread_handle: None,
         }
     }
 }
@@ -1946,7 +1949,7 @@ impl UChild {
         child: Child,
         captured_stdout: Option<CapturedOutput>,
         captured_stderr: Option<CapturedOutput>,
-        stdin_pty: Option<OwnedFd>
+        stdin_pty: Option<OwnedFd>,
     ) -> Self {
         Self {
             raw: child,
@@ -2095,7 +2098,6 @@ impl UChild {
     /// error.
     #[deprecated = "Please use wait() -> io::Result<CmdResult> instead."]
     pub fn wait_with_output(mut self) -> io::Result<Output> {
-
         // some apps do not stop execution until their stdin gets closed.
         // to prevent a endless waiting here, we close the stdin.
         self.join(); // ensure that all pending async input is piped in
@@ -2107,7 +2109,8 @@ impl UChild {
             let (sender, receiver) = mpsc::channel();
             let handle = thread::Builder::new()
                 .name("wait_with_output".to_string())
-                .spawn(move || sender.send(child.wait_with_output())).unwrap();
+                .spawn(move || sender.send(child.wait_with_output()))
+                .unwrap();
 
             match receiver.recv_timeout(timeout) {
                 Ok(result) => {
@@ -2139,9 +2142,17 @@ impl UChild {
         };
 
         if let Some(stdout) = self.captured_stdout.as_mut() {
+            stdout
+                .reader_thread_handle
+                .take()
+                .map(|handle| handle.join().unwrap());
             output.stdout = stdout.output_bytes();
         }
         if let Some(stderr) = self.captured_stderr.as_mut() {
+            stderr
+                .reader_thread_handle
+                .take()
+                .map(|handle| handle.join().unwrap());
             output.stderr = stderr.output_bytes();
         }
 
@@ -2346,7 +2357,6 @@ impl UChild {
     ///
     /// [`JoinHandle`]: std::thread::JoinHandle
     pub fn pipe_in<T: Into<Vec<u8>>>(&mut self, content: T) -> &mut Self {
-
         let ignore_stdin_write_error = self.ignore_stdin_write_error;
         let mut content: Vec<u8> = content.into();
         if self.stdin_pty.is_some() {
@@ -2369,7 +2379,8 @@ impl UChild {
                     Ok(()) | Err(_) => Ok(()),
                 };
                 result
-            }).unwrap();
+            })
+            .unwrap();
 
         self.join_handle = Some(join_handle);
         self