diff --git a/tokio/src/process/unix/mod.rs b/tokio/src/process/unix/mod.rs index c9d1035f53d..2b82c33f146 100644 --- a/tokio/src/process/unix/mod.rs +++ b/tokio/src/process/unix/mod.rs @@ -106,7 +106,7 @@ impl OrphanQueue for GlobalOrphanQueue { pub(crate) enum Child { SignalReaper(Reaper), #[cfg(all(target_os = "linux", feature = "rt"))] - PidfdReaper(pidfd_reaper::PidfdReaper), + PidfdReaper(pidfd_reaper::PidfdReaper), } impl fmt::Debug for Child { @@ -115,21 +115,95 @@ impl fmt::Debug for Child { } } +pub(crate) struct ChildDropGuard { + child: Option, + kill_on_drop: bool, +} + +impl Drop for ChildDropGuard { + fn drop(&mut self) { + if let Some(child) = &mut self.child { + if self.kill_on_drop { + drop(child.kill()); + } + } + } +} + +impl ChildDropGuard { + pub(crate) fn new(child: StdChild) -> ChildDropGuard { + return ChildDropGuard { + child: Some(child), + kill_on_drop: true, + }; + } + + pub(crate) fn extract(mut self) -> StdChild { + self.child.take().unwrap() + } + + pub(crate) fn dont_kill_on_drop(&mut self) { + self.kill_on_drop = false; + } + + pub(crate) fn take_stdin(&mut self) -> Option { + self.child + .as_mut() + .expect("child has gone away") + .stdin + .take() + } + pub(crate) fn take_stdout(&mut self) -> Option { + self.child + .as_mut() + .expect("child has gone away") + .stdout + .take() + } + pub(crate) fn take_stderr(&mut self) -> Option { + self.child + .as_mut() + .expect("child has gone away") + .stderr + .take() + } + + pub(crate) fn inner_mut(&mut self) -> &mut StdChild { + self.child.as_mut().expect("child has gone away") + } +} + +impl Wait for ChildDropGuard { + fn id(&self) -> u32 { + self.child.as_ref().expect("child has gone away").id() + } + fn try_wait(&mut self) -> io::Result> { + self.child.as_mut().expect("child has gone away").try_wait() + } +} + +impl OrphanQueue for GlobalOrphanQueue { + fn push_orphan(&self, orphan: ChildDropGuard) { + get_orphan_queue().push_orphan(orphan.extract()); + } +} + pub(crate) fn spawn_child(cmd: &mut std::process::Command) -> io::Result { - let mut child = cmd.spawn()?; - let stdin = child.stdin.take().map(stdio).transpose()?; - let stdout = child.stdout.take().map(stdio).transpose()?; - let stderr = child.stderr.take().map(stdio).transpose()?; + let mut child = ChildDropGuard::new(cmd.spawn()?); + let stdin = child.take_stdin().map(stdio).transpose()?; + let stdout = child.take_stdout().map(stdio).transpose()?; + let stderr = child.take_stderr().map(stdio).transpose()?; #[cfg(all(target_os = "linux", feature = "rt"))] match pidfd_reaper::PidfdReaper::new(child, GlobalOrphanQueue) { - Ok(pidfd_reaper) => { + Ok(mut pidfd_reaper) => { + pidfd_reaper.inner_mut().dont_kill_on_drop(); return Ok(SpawnedChild { child: Child::PidfdReaper(pidfd_reaper), stdin, stdout, stderr, - }) + }); } Err((Some(err), _child)) => return Err(err), Err((None, child_returned)) => child = child_returned, @@ -138,7 +212,7 @@ pub(crate) fn spawn_child(cmd: &mut std::process::Command) -> io::Result signal_reaper.inner_mut(), #[cfg(all(target_os = "linux", feature = "rt"))] - Self::PidfdReaper(pidfd_reaper) => pidfd_reaper.inner_mut(), + Self::PidfdReaper(pidfd_reaper) => pidfd_reaper.inner_mut().inner_mut(), } } diff --git a/tokio/src/process/windows.rs b/tokio/src/process/windows.rs index 792a9c9610b..6774ee34792 100644 --- a/tokio/src/process/windows.rs +++ b/tokio/src/process/windows.rs @@ -66,15 +66,59 @@ struct Waiting { unsafe impl Sync for Waiting {} unsafe impl Send for Waiting {} +struct DroppableChild { + child: Option, +} + +impl Drop for DroppableChild { + fn drop(&mut self) { + if let Some(child) = &mut self.child { + drop(child.kill()); + } + } +} + +impl DroppableChild { + pub(crate) fn new(child: StdChild) -> DroppableChild { + return DroppableChild { child: Some(child) }; + } + + pub(crate) fn take_stdin(&mut self) -> Option { + self.child + .as_mut() + .expect("child has gone away") + .stdin + .take() + } + pub(crate) fn take_stdout(&mut self) -> Option { + self.child + .as_mut() + .expect("child has gone away") + .stdout + .take() + } + pub(crate) fn take_stderr(&mut self) -> Option { + self.child + .as_mut() + .expect("child has gone away") + .stderr + .take() + } + + pub(crate) fn take(mut self) -> StdChild { + self.child.take().expect("child has gone away") + } +} + pub(crate) fn spawn_child(cmd: &mut StdCommand) -> io::Result { - let mut child = cmd.spawn()?; - let stdin = child.stdin.take().map(stdio).transpose()?; - let stdout = child.stdout.take().map(stdio).transpose()?; - let stderr = child.stderr.take().map(stdio).transpose()?; + let mut child = DroppableChild::new(cmd.spawn()?); + let stdin = child.take_stdin().map(stdio).transpose()?; + let stdout = child.take_stdout().map(stdio).transpose()?; + let stderr = child.take_stderr().map(stdio).transpose()?; Ok(SpawnedChild { child: Child { - child, + child: child.take(), waiting: None, }, stdin,