diff --git a/Cargo.lock b/Cargo.lock index 21fd2010fa3b27..a416cd2a4cd3ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -19769,6 +19769,7 @@ dependencies = [ "util_macros", "walkdir", "which 6.0.3", + "windows 0.61.3", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 037afc59bfe16c..72167caf83661a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -860,6 +860,7 @@ features = [ "Win32_System_Diagnostics_Debug", "Win32_System_DataExchange", "Win32_System_IO", + "Win32_System_JobObjects", "Win32_System_LibraryLoader", "Win32_System_Memory", "Win32_System_Ole", diff --git a/crates/dap/src/transport.rs b/crates/dap/src/transport.rs index ff04c414517dcd..726c62689298d9 100644 --- a/crates/dap/src/transport.rs +++ b/crates/dap/src/transport.rs @@ -617,8 +617,8 @@ impl Transport for TcpTransport { if has_process { let status = process.lock().as_mut().unwrap().try_status(); if let Ok(Some(_)) = status { - let process = process.lock().take().unwrap().into_inner(); - let output = process.output().await?; + let child = process.lock().take().unwrap(); + let output = child.output().await?; let output = if output.stderr.is_empty() { String::from_utf8_lossy(&output.stdout).to_string() } else { diff --git a/crates/util/Cargo.toml b/crates/util/Cargo.toml index 13f40811225561..340faa25222c14 100644 --- a/crates/util/Cargo.toml +++ b/crates/util/Cargo.toml @@ -60,6 +60,7 @@ mach2.workspace = true [target.'cfg(windows)'.dependencies] tendril = "0.4.3" +windows.workspace = true [dev-dependencies] rand.workspace = true diff --git a/crates/util/src/process.rs b/crates/util/src/process.rs index eaf543dbd817ba..93cbe65ce58486 100644 --- a/crates/util/src/process.rs +++ b/crates/util/src/process.rs @@ -2,9 +2,17 @@ use anyhow::{Context as _, Result}; use std::process::Stdio; /// A wrapper around `smol::process::Child` that ensures all subprocesses -/// are killed when the process is terminated by using process groups. +/// are killed when the process is terminated: on Unix by using process +/// groups, and on Windows by using job objects. +/// +/// On Windows, dropping this struct closes the job object handle, which +/// terminates all processes in the job. This also applies when the Zed +/// process exits for any reason (including crashes), since the OS closes +/// its handles, so spawned process trees can never outlive Zed. pub struct Child { process: smol::process::Child, + #[cfg(windows)] + job: Option, } impl std::ops::Deref for Child { @@ -52,8 +60,6 @@ impl Child { stdout: Stdio, stderr: Stdio, ) -> Result { - // TODO(windows): create a job object and add the child process handle to it, - // see https://learn.microsoft.com/en-us/windows/win32/procthread/job-objects let mut command = smol::process::Command::from(command); let process = command .stdin(stdin) @@ -67,11 +73,42 @@ impl Child { ) })?; - Ok(Self { process }) + // Assign the child to a job object configured to kill the entire + // process tree when the last job handle is closed, so descendants + // (e.g. node workers and MCP servers spawned by agent servers) are + // reaped even if the direct child doesn't clean them up. Any process + // the child spawns after this assignment is automatically part of the + // job. + // + // There is a small race: descendants the child spawns between the + // `spawn()` call returning and the assignment below escape the job. + // Closing it fully would require creating the process suspended + // (`CREATE_SUSPENDED`), assigning it, then resuming it, which the + // std/smol process APIs don't support without reimplementing process + // creation. The window is microseconds, and the children we care + // about (`npx`, `node`, etc.) take far longer to load their runtime + // and spawn anything, so in practice nothing escapes. + let job = windows_job::JobObject::new() + .and_then(|job| { + job.assign_process(process.id())?; + Ok(job) + }) + .map_err(|error| { + log::error!("failed to assign spawned process to a job object: {error:#}"); + }) + .ok(); + + Ok(Self { process, job }) } - pub fn into_inner(self) -> smol::process::Child { - self.process + /// Consumes the child, draining its stdout/stderr and waiting for it to + /// exit, then returns the collected output. + pub async fn output(self) -> Result { + // NOTE: Keep `self` alive across this await, do not destructure it to + // pull `process` out first. On Windows that drops the job object early, + // which triggers `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE` and kills the + // child before `output()` finishes collecting its stdout/stderr. + Ok(self.process.output().await?) } #[cfg(not(windows))] @@ -85,8 +122,175 @@ impl Child { #[cfg(windows)] pub fn kill(&mut self) -> Result<()> { - // TODO(windows): terminate the job object in kill - self.process.kill()?; - Ok(()) + if let Some(job) = &self.job { + job.terminate() + } else { + self.process.kill()?; + Ok(()) + } + } +} + +#[cfg(windows)] +mod windows_job { + use crate::ResultExt as _; + use anyhow::{Context as _, Result}; + use windows::Win32::{ + Foundation::{CloseHandle, HANDLE}, + System::{ + JobObjects::{ + AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, + SetInformationJobObject, TerminateJobObject, + }, + Threading::{OpenProcess, PROCESS_SET_QUOTA, PROCESS_TERMINATE}, + }, + }; + + /// A Win32 job object configured with `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE`: + /// all processes assigned to the job (and their descendants) are terminated + /// when the last handle to the job is closed, which happens when this struct + /// is dropped, or when the OS closes the owning process's handles after it + /// exits for any reason. + pub(crate) struct JobObject(HANDLE); + + // SAFETY: Job object handles can be used from any thread. + unsafe impl Send for JobObject {} + unsafe impl Sync for JobObject {} + + impl JobObject { + pub(crate) fn new() -> Result { + unsafe { + let job = + Self(CreateJobObjectW(None, None).context("failed to create job object")?); + let mut info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default(); + info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; + SetInformationJobObject( + job.0, + JobObjectExtendedLimitInformation, + &info as *const _ as *const _, + size_of::() as u32, + ) + .context("failed to set job object limits")?; + Ok(job) + } + } + + pub(crate) fn assign_process(&self, pid: u32) -> Result<()> { + unsafe { + let process = OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, false, pid) + .context("failed to open process")?; + let result = AssignProcessToJobObject(self.0, process) + .context("failed to assign process to job object"); + CloseHandle(process).log_err(); + result + } + } + + pub(crate) fn terminate(&self) -> Result<()> { + unsafe { TerminateJobObject(self.0, 1).context("failed to terminate job object") } + } + } + + impl Drop for JobObject { + fn drop(&mut self) { + unsafe { + CloseHandle(self.0).log_err(); + } + } + } +} + +#[cfg(all(test, windows))] +mod windows_tests { + use super::*; + use std::time::{Duration, Instant}; + + /// Spawns a process tree `powershell -> ping` via `Child::spawn` and + /// returns the `Child` along with the pid of the grandchild (`ping`). + fn spawn_process_tree(temp_dir: &std::path::Path) -> (Child, u32) { + let pid_file = temp_dir.join("grandchild_pid"); + let mut command = std::process::Command::new("powershell.exe"); + command.args(["-NoProfile", "-Command"]).arg(format!( + "$p = Start-Process -FilePath ping.exe -ArgumentList @('-n','60','127.0.0.1') -PassThru -WindowStyle Hidden; \ + Set-Content -LiteralPath '{}' -Value $p.Id; \ + Wait-Process -Id $p.Id", + pid_file.display() + )); + let child = Child::spawn(command, Stdio::null(), Stdio::null(), Stdio::null()) + .expect("failed to spawn powershell"); + + let deadline = Instant::now() + Duration::from_secs(5); + let grandchild_pid = loop { + if let Ok(contents) = std::fs::read_to_string(&pid_file) + && let Ok(pid) = contents.trim().parse::() + { + break pid; + } + assert!( + Instant::now() < deadline, + "timed out waiting for grandchild pid file" + ); + std::thread::sleep(Duration::from_millis(50)); + }; + assert!( + process_is_alive(grandchild_pid), + "grandchild should be alive after spawning" + ); + (child, grandchild_pid) + } + + fn process_is_alive(pid: u32) -> bool { + use windows::Win32::{ + Foundation::{CloseHandle, STILL_ACTIVE}, + System::Threading::{ + GetExitCodeProcess, OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION, + }, + }; + + unsafe { + let Ok(handle) = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, pid) else { + return false; + }; + let mut exit_code = 0u32; + let alive = GetExitCodeProcess(handle, &mut exit_code).is_ok() + && exit_code == STILL_ACTIVE.0 as u32; + CloseHandle(handle).expect("failed to close process handle"); + alive + } + } + + fn assert_process_exits(pid: u32, message: &str) { + let deadline = Instant::now() + Duration::from_secs(2); + while process_is_alive(pid) { + assert!(Instant::now() < deadline, "{message} (pid {pid})"); + std::thread::sleep(Duration::from_millis(100)); + } + } + + #[test] + fn test_kill_terminates_grandchildren() { + let temp_dir = tempfile::tempdir().unwrap(); + let (mut child, grandchild_pid) = spawn_process_tree(temp_dir.path()); + + child.kill().expect("failed to kill child"); + + assert_process_exits( + grandchild_pid, + "grandchild should be terminated after killing the child", + ); + } + + #[test] + fn test_drop_terminates_grandchildren() { + let temp_dir = tempfile::tempdir().unwrap(); + let (child, grandchild_pid) = spawn_process_tree(temp_dir.path()); + + drop(child); + + assert_process_exits( + grandchild_pid, + "grandchild should be terminated after dropping the child", + ); } }