Skip to main content

soma_codemode/pool/
job_guard.rs

1#[cfg(not(windows))]
2#[derive(Debug, Default)]
3pub struct JobGuard;
4
5#[cfg(not(windows))]
6impl JobGuard {
7    #[must_use]
8    pub fn new(_pid: Option<u32>) -> Self {
9        Self
10    }
11}
12
13#[cfg(windows)]
14#[derive(Debug)]
15pub struct JobGuard {
16    job: windows_sys::Win32::Foundation::HANDLE,
17}
18
19#[cfg(windows)]
20impl JobGuard {
21    #[must_use]
22    pub fn new(pid: Option<u32>) -> Self {
23        use windows_sys::Win32::System::JobObjects::{
24            AssignProcessToJobObject, CreateJobObjectW, JobObjectExtendedLimitInformation,
25            SetInformationJobObject, JOBOBJECT_EXTENDED_LIMIT_INFORMATION,
26            JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
27        };
28        use windows_sys::Win32::System::Threading::{
29            OpenProcess, PROCESS_SET_QUOTA, PROCESS_TERMINATE,
30        };
31
32        let Some(pid) = pid else {
33            return Self {
34                job: std::ptr::null_mut(),
35            };
36        };
37        unsafe {
38            let job = CreateJobObjectW(std::ptr::null(), std::ptr::null());
39            if job.is_null() {
40                return Self { job };
41            }
42            let mut info: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = std::mem::zeroed();
43            info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
44            let _ = SetInformationJobObject(
45                job,
46                JobObjectExtendedLimitInformation,
47                &mut info as *mut _ as *mut _,
48                std::mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
49            );
50            let process = OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, 0, pid);
51            if !process.is_null() {
52                let _ = AssignProcessToJobObject(job, process);
53                windows_sys::Win32::Foundation::CloseHandle(process);
54            }
55            Self { job }
56        }
57    }
58}
59
60#[cfg(windows)]
61impl Drop for JobGuard {
62    fn drop(&mut self) {
63        if !self.job.is_null() {
64            unsafe {
65                windows_sys::Win32::Foundation::CloseHandle(self.job);
66            }
67        }
68    }
69}