|
| 1 | +//! Windows Job Object for reliable child process cleanup. |
| 2 | +//! |
| 3 | +//! This module provides a wrapper around Windows Job Objects with the |
| 4 | +//! `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE` flag set. When the job object handle |
| 5 | +//! is closed (including when the parent process exits or crashes), Windows |
| 6 | +//! automatically terminates all processes assigned to the job. |
| 7 | +//! |
| 8 | +//! This is more reliable than manual cleanup because it works even if: |
| 9 | +//! - The parent process crashes |
| 10 | +//! - The parent is killed via Task Manager |
| 11 | +//! - The RunEvent::Exit handler fails to run |
| 12 | +
|
| 13 | +use std::io::{Error, Result}; |
| 14 | +#[cfg(windows)] |
| 15 | +use std::sync::Mutex; |
| 16 | +use windows::Win32::Foundation::{CloseHandle, HANDLE}; |
| 17 | +use windows::Win32::System::JobObjects::{ |
| 18 | + AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, |
| 19 | + JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, |
| 20 | + SetInformationJobObject, |
| 21 | +}; |
| 22 | +use windows::Win32::System::Threading::{OpenProcess, PROCESS_SET_QUOTA, PROCESS_TERMINATE}; |
| 23 | + |
| 24 | +/// A Windows Job Object configured to kill all assigned processes when closed. |
| 25 | +/// |
| 26 | +/// When this struct is dropped or when the owning process exits (even abnormally), |
| 27 | +/// Windows will automatically terminate all processes that have been assigned to it. |
| 28 | +pub struct JobObject(HANDLE); |
| 29 | + |
| 30 | +// SAFETY: HANDLE is just a pointer-sized value, and Windows job objects |
| 31 | +// can be safely accessed from multiple threads. |
| 32 | +unsafe impl Send for JobObject {} |
| 33 | +unsafe impl Sync for JobObject {} |
| 34 | + |
| 35 | +impl JobObject { |
| 36 | + /// Creates a new anonymous job object with `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE` set. |
| 37 | + /// |
| 38 | + /// When the last handle to this job is closed (including on process exit), |
| 39 | + /// Windows will terminate all processes assigned to the job. |
| 40 | + pub fn new() -> Result<Self> { |
| 41 | + unsafe { |
| 42 | + // Create an anonymous job object |
| 43 | + let job = CreateJobObjectW(None, None).map_err(|e| Error::other(e.message()))?; |
| 44 | + |
| 45 | + // Configure the job to kill all processes when the handle is closed |
| 46 | + let mut info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default(); |
| 47 | + info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; |
| 48 | + |
| 49 | + SetInformationJobObject( |
| 50 | + job, |
| 51 | + JobObjectExtendedLimitInformation, |
| 52 | + &info as *const _ as *const std::ffi::c_void, |
| 53 | + std::mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32, |
| 54 | + ) |
| 55 | + .map_err(|e| Error::other(e.message()))?; |
| 56 | + |
| 57 | + Ok(Self(job)) |
| 58 | + } |
| 59 | + } |
| 60 | + |
| 61 | + /// Assigns a process to this job object by its process ID. |
| 62 | + /// |
| 63 | + /// Once assigned, the process will be terminated when this job object is dropped |
| 64 | + /// or when the owning process exits. |
| 65 | + /// |
| 66 | + /// # Arguments |
| 67 | + /// * `pid` - The process ID of the process to assign |
| 68 | + pub fn assign_pid(&self, pid: u32) -> Result<()> { |
| 69 | + unsafe { |
| 70 | + // Open a handle to the process with the minimum required permissions |
| 71 | + // PROCESS_SET_QUOTA and PROCESS_TERMINATE are required by AssignProcessToJobObject |
| 72 | + let process = OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, false, pid) |
| 73 | + .map_err(|e| Error::other(e.message()))?; |
| 74 | + |
| 75 | + // Assign the process to the job |
| 76 | + let result = AssignProcessToJobObject(self.0, process); |
| 77 | + |
| 78 | + // Close our handle to the process - the job object maintains its own reference |
| 79 | + let _ = CloseHandle(process); |
| 80 | + |
| 81 | + result.map_err(|e| Error::other(e.message())) |
| 82 | + } |
| 83 | + } |
| 84 | +} |
| 85 | + |
| 86 | +impl Drop for JobObject { |
| 87 | + fn drop(&mut self) { |
| 88 | + unsafe { |
| 89 | + // When this handle is closed and it's the last handle to the job, |
| 90 | + // Windows will terminate all processes in the job due to KILL_ON_JOB_CLOSE |
| 91 | + let _ = CloseHandle(self.0); |
| 92 | + } |
| 93 | + } |
| 94 | +} |
| 95 | + |
| 96 | +/// Holds the Windows Job Object that ensures child processes are killed when the app exits. |
| 97 | +/// On Windows, when the job object handle is closed (including on crash), all assigned |
| 98 | +/// processes are automatically terminated by the OS. |
| 99 | +#[cfg(windows)] |
| 100 | +pub struct JobObjectState { |
| 101 | + job: Mutex<Option<JobObject>>, |
| 102 | + error: Mutex<Option<String>>, |
| 103 | +} |
| 104 | + |
| 105 | +#[cfg(windows)] |
| 106 | +impl JobObjectState { |
| 107 | + pub fn new() -> Self { |
| 108 | + match JobObject::new() { |
| 109 | + Ok(job) => Self { |
| 110 | + job: Mutex::new(Some(job)), |
| 111 | + error: Mutex::new(None), |
| 112 | + }, |
| 113 | + Err(e) => { |
| 114 | + eprintln!("Failed to create job object: {e}"); |
| 115 | + Self { |
| 116 | + job: Mutex::new(None), |
| 117 | + error: Mutex::new(Some(format!("Failed to create job object: {e}"))), |
| 118 | + } |
| 119 | + } |
| 120 | + } |
| 121 | + } |
| 122 | + |
| 123 | + pub fn assign_pid(&self, pid: u32) { |
| 124 | + if let Some(job) = self.job.lock().unwrap().as_ref() { |
| 125 | + if let Err(e) = job.assign_pid(pid) { |
| 126 | + eprintln!("Failed to assign process {pid} to job object: {e}"); |
| 127 | + *self.error.lock().unwrap() = |
| 128 | + Some(format!("Failed to assign process to job object: {e}")); |
| 129 | + } else { |
| 130 | + println!("Assigned process {pid} to job object for automatic cleanup"); |
| 131 | + } |
| 132 | + } |
| 133 | + } |
| 134 | +} |
| 135 | + |
| 136 | +#[cfg(test)] |
| 137 | +mod tests { |
| 138 | + use super::*; |
| 139 | + |
| 140 | + #[test] |
| 141 | + fn test_job_object_creation() { |
| 142 | + let job = JobObject::new(); |
| 143 | + assert!(job.is_ok(), "Failed to create job object: {:?}", job.err()); |
| 144 | + } |
| 145 | +} |
0 commit comments