Skip to content

Commit ddd9c71

Browse files
authored
feat(desktop): Tie desktop & CLI to the same Windows JobObject (anomalyco#8153)
1 parent 520a814 commit ddd9c71

4 files changed

Lines changed: 168 additions & 0 deletions

File tree

packages/desktop/src-tauri/Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/desktop/src-tauri/Cargo.toml

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,3 +44,11 @@ uuid = { version = "1.19.0", features = ["v4"] }
4444
[target.'cfg(target_os = "linux")'.dependencies]
4545
gtk = "0.18.2"
4646
webkit2gtk = "=2.0.1"
47+
48+
[target.'cfg(windows)'.dependencies]
49+
windows = { version = "0.61", features = [
50+
"Win32_Foundation",
51+
"Win32_System_JobObjects",
52+
"Win32_System_Threading",
53+
"Win32_Security"
54+
] }
Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
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+
}

packages/desktop/src-tauri/src/lib.rs

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
mod cli;
2+
#[cfg(windows)]
3+
mod job_object;
24
mod window_customizer;
35

46
use cli::{install_cli, sync_cli};
57
use futures::FutureExt;
68
use futures::future;
9+
#[cfg(windows)]
10+
use job_object::*;
711
use std::{
812
collections::VecDeque,
913
net::TcpListener,
@@ -251,6 +255,9 @@ pub fn run() {
251255
// Initialize log state
252256
app.manage(LogState(Arc::new(Mutex::new(VecDeque::new()))));
253257

258+
#[cfg(windows)]
259+
app.manage(JobObjectState::new());
260+
254261
let primary_monitor = app.primary_monitor().ok().flatten();
255262
let size = primary_monitor
256263
.map(|m| m.size().to_logical(m.scale_factor()))
@@ -303,7 +310,14 @@ pub fn run() {
303310

304311
let res = match setup_server_connection(&app, custom_url).await {
305312
Ok((child, url)) => {
313+
#[cfg(windows)]
314+
if let Some(child) = &child {
315+
let job_state = app.state::<JobObjectState>();
316+
job_state.assign_pid(child.pid());
317+
}
318+
306319
app.state::<ServerState>().set_child(child);
320+
307321
Ok(url)
308322
}
309323
Err(e) => Err(e),

0 commit comments

Comments
 (0)