From e815fa80697249afef6c25fb3652929749089890 Mon Sep 17 00:00:00 2001
From: Michael Alyn Miller <malyn@strangeGizmo.com>
Date: Fri, 18 Nov 2022 21:06:34 -0800
Subject: [PATCH] Wait for all processes in the process group

This avoids zombie processes that would otherwise be created if the child
process exits after spawning its own children, but without waiting for
those children to exit. This change ensures that `wait` reaps all of
those zombie processes, while still returning the exit code from the
original process that was spawned (Command Group's direct child).
---
 src/stdlib/child/unix.rs | 52 +++++++++++++++++++++++++++++++++-------
 src/tokio/child/unix.rs  | 49 ++++++++++++++++++++++++++++++-------
 2 files changed, 84 insertions(+), 17 deletions(-)

diff --git a/src/stdlib/child/unix.rs b/src/stdlib/child/unix.rs
index 15cf29c..54b1730 100644
--- a/src/stdlib/child/unix.rs
+++ b/src/stdlib/child/unix.rs
@@ -67,15 +67,49 @@ impl ChildImp {
 	fn wait_imp(&mut self, flag: WaitPidFlag) -> Result<Option<ExitStatus>> {
 		let negpid = Pid::from_raw(-self.pgid.as_raw());
 
-		// we can't use the safe wrapper directly because it doesn't return the raw status, and we
-		// need it to convert to the std's ExitStatus.
-		let mut status: i32 = 0;
-		match unsafe { libc::waitpid(negpid.into(), &mut status as *mut libc::c_int, flag.bits()) }
-		{
-			0 => Ok(None),
-			res => Errno::result(res)
-				.map_err(Error::from)
-				.map(|_| Some(ExitStatus::from_raw(status))),
+		// Wait for processes in a loop until every process in this
+		// process group has exited (this ensures that we reap any
+		// zombies that may have been created if the parent exited after
+		// spawning children, but didn't wait for those children to
+		// exit).
+		let mut parent_exit_status: Option<ExitStatus> = None;
+		loop {
+			// we can't use the safe wrapper directly because it doesn't
+			// return the raw status, and we need it to convert to the
+			// std's ExitStatus.
+			let mut status: i32 = 0;
+			match unsafe {
+				libc::waitpid(negpid.into(), &mut status as *mut libc::c_int, flag.bits())
+			} {
+				0 => {
+					// Zero should only happen if WNOHANG was passed in,
+					// and means that no processes have yet to exit.
+					return Ok(None);
+				}
+				-1 => {
+					match Errno::last() {
+						Errno::ECHILD => {
+							// No more children to reap; this is a
+							// graceful exit.
+							return Ok(parent_exit_status);
+						}
+						errno => {
+							return Err(Error::from(errno));
+						}
+					}
+				}
+				pid => {
+					// *A* process exited. Was it the parent process
+					// that we started? If so, collect the exit signal,
+					// otherwise we reaped a zombie process and should
+					// continue in the loop.
+					if self.pgid.as_raw() == pid {
+						parent_exit_status = Some(ExitStatus::from_raw(status));
+					} else {
+						// Reaped a zombie child; keep looping.
+					}
+				}
+			};
 		}
 	}
 
diff --git a/src/tokio/child/unix.rs b/src/tokio/child/unix.rs
index 8f7ca2f..5c23ad4 100644
--- a/src/tokio/child/unix.rs
+++ b/src/tokio/child/unix.rs
@@ -70,14 +70,47 @@ impl ChildImp {
 	}
 
 	fn wait_imp(pgid: i32, flag: WaitPidFlag) -> Result<Option<ExitStatus>> {
-		// we can't use the safe wrapper directly because it doesn't return the raw status, and we
-		// need it to convert to the std's ExitStatus.
-		let mut status: i32 = 0;
-		match unsafe { libc::waitpid(-pgid, &mut status as *mut libc::c_int, flag.bits()) } {
-			0 => Ok(None),
-			res => Errno::result(res)
-				.map_err(Error::from)
-				.map(|_| Some(ExitStatus::from_raw(status))),
+		// Wait for processes in a loop until every process in this
+		// process group has exited (this ensures that we reap any
+		// zombies that may have been created if the parent exited after
+		// spawning children, but didn't wait for those children to
+		// exit).
+		let mut parent_exit_status: Option<ExitStatus> = None;
+		loop {
+			// we can't use the safe wrapper directly because it doesn't
+			// return the raw status, and we need it to convert to the
+			// std's ExitStatus.
+			let mut status: i32 = 0;
+			match unsafe { libc::waitpid(-pgid, &mut status as *mut libc::c_int, flag.bits()) } {
+				0 => {
+					// Zero should only happen if WNOHANG was passed in,
+					// and means that no processes have yet to exit.
+					return Ok(None);
+				}
+				-1 => {
+					match Errno::last() {
+						Errno::ECHILD => {
+							// No more children to reap; this is a
+							// graceful exit.
+							return Ok(parent_exit_status);
+						}
+						errno => {
+							return Err(Error::from(errno));
+						}
+					}
+				}
+				pid => {
+					// *A* process exited. Was it the parent process
+					// that we started? If so, collect the exit signal,
+					// otherwise we reaped a zombie process and should
+					// continue in the loop.
+					if pgid == pid {
+						parent_exit_status = Some(ExitStatus::from_raw(status));
+					} else {
+						// Reaped a zombie child; keep looping.
+					}
+				}
+			};
 		}
 	}