diff --git a/Cargo.lock b/Cargo.lock index 045ce9e..1880414 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -817,6 +817,7 @@ dependencies = [ "serde", "serde_json", "tokio", + "windows-sys", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 2474dd2..274e31f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,8 @@ macmon = { version = "0.10", default-features = false } [target.'cfg(windows)'.dependencies] # ADLX, the AMD driver's library, is loaded at run time: no AMD driver, no source. libloading = "0.8" +# Already in the lockfile through tokio. +windows-sys = { version = "0.61", features = ["Win32_System_JobObjects", "Win32_System_Threading"] } [dev-dependencies] serde_json = "1" diff --git a/src/sampler.rs b/src/sampler.rs index c62ca76..aafc357 100644 --- a/src/sampler.rs +++ b/src/sampler.rs @@ -9,6 +9,9 @@ //! reports power (Jetson's integrated GPU answers `[N/A]`) stops the sampler //! after its first round, and so does a macOS VM, whose GPU reads 0 W. +#[cfg(windows)] +mod job; + use anyhow::Result; use tokio::io::{AsyncBufRead, Lines}; use tokio::sync::mpsc; @@ -155,8 +158,11 @@ mod nvidia { if !listing.status.success() || gpus == 0 { return Ok(false); } - // `kill_on_drop` ends nvidia-smi with the plugin when it returns; a - // plugin killed outright can leave it looping on Windows. + // `kill_on_drop` ends nvidia-smi with the plugin when it returns. On + // Windows a plugin killed outright would leave it looping, so it also + // runs in a job that Windows closes with the plugin. + #[cfg(windows)] + let job = super::job::KillOnClose::new(); let mut child = Command::new("nvidia-smi") .args([ "--query-gpu=index,power.draw", @@ -169,6 +175,22 @@ mod nvidia { .kill_on_drop(true) .spawn() .context("start nvidia-smi power sampling")?; + // Held until sampling ends, since closing the job ends nvidia-smi. + // Without it the readings still work; only that cleanup is lost. A + // plugin killed between spawn and assignment still leaves one behind: + // starting it suspended would close that gap, which tokio's Command + // does not offer. + #[cfg(windows)] + let _job = match job.and_then(|job| { + let process = child.raw_handle().ok_or(std::io::ErrorKind::NotFound)?; + job.assign(process).map(|()| job) + }) { + Ok(job) => Some(job), + Err(error) => { + eprintln!("powermetrics-plugin: nvidia-smi may outlive the plugin: {error}"); + None + } + }; let stdout = child .stdout .take() diff --git a/src/sampler/job.rs b/src/sampler/job.rs new file mode 100644 index 0000000..bd2cd59 --- /dev/null +++ b/src/sampler/job.rs @@ -0,0 +1,84 @@ +//! A Windows job object that ends the processes in it once nothing holds it. +//! +//! `kill_on_drop` only runs when the plugin unwinds. A plugin ended outright +//! (the host killed, a console closed) leaves `nvidia-smi --loop-ms` running +//! for good on Windows, where a write to the closed pipe does not stop it. The +//! job's handle lives in the plugin process, so Windows closes it whichever +//! way the plugin ends, and the job then ends `nvidia-smi` too. + +use std::io; +use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle}; +use windows_sys::Win32::System::JobObjects::{ + AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, + SetInformationJobObject, +}; + +pub(super) struct KillOnClose(OwnedHandle); + +impl KillOnClose { + pub(super) fn new() -> io::Result { + // SAFETY: no attributes and no name ask for an anonymous job; the + // handle is checked before it is owned. + let raw = unsafe { CreateJobObjectW(std::ptr::null(), std::ptr::null()) }; + if raw.is_null() { + return Err(io::Error::last_os_error()); + } + // SAFETY: `raw` is a valid job handle that nothing else owns. + let job = Self(unsafe { OwnedHandle::from_raw_handle(raw) }); + let mut limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default(); + limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; + // SAFETY: the pointer and size describe `limits`, which outlives the call. + let set = unsafe { + SetInformationJobObject( + job.0.as_raw_handle(), + JobObjectExtendedLimitInformation, + (&raw const limits).cast(), + size_of::() as u32, + ) + }; + if set == 0 { + return Err(io::Error::last_os_error()); + } + Ok(job) + } + + /// Puts the process behind `process` in the job. + pub(super) fn assign(&self, process: std::os::windows::io::RawHandle) -> io::Result<()> { + // SAFETY: both handles are open for the duration of the call. + if unsafe { AssignProcessToJobObject(self.0.as_raw_handle(), process) } == 0 { + return Err(io::Error::last_os_error()); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::process::{Command, Stdio}; + use std::time::{Duration, Instant}; + + #[test] + fn closing_the_job_ends_the_process_in_it() { + // A child that would otherwise run for a minute. + let mut child = Command::new("ping") + .args(["-n", "60", "127.0.0.1"]) + .stdout(Stdio::null()) + .spawn() + .unwrap(); + let job = KillOnClose::new().unwrap(); + job.assign(child.as_raw_handle()).unwrap(); + assert!( + child.try_wait().unwrap().is_none(), + "still running in the job" + ); + + drop(job); + let deadline = Instant::now() + Duration::from_secs(10); + while child.try_wait().unwrap().is_none() { + assert!(Instant::now() < deadline, "the child outlived its job"); + std::thread::sleep(Duration::from_millis(50)); + } + } +}