diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..6936990 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +/target +**/*.rs.bk +Cargo.lock diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..2ee0fb4 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "win32job" +version = "0.1.0" +authors = ["Ohad Ravid "] +edition = "2018" + +[dependencies] +winapi = { version = "0.3", features = ["handleapi", "winbase", "psapi", "processthreadsapi", "jobapi2"] } +thiserror = "1.0" \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..3d00a15 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,172 @@ +//! # wun32job-rs +//! +//! A safe API for Windows' job objects, which can be used to set various limits to +//! processes associated with them. +//! See also https://docs.microsoft.com/en-us/windows/win32/api/jobapi2/nf-jobapi2-createjobobjectw +//! +//! # Using the low level API +//! +//! The most basic API is getting and raw `JOBOBJECT_BASIC_LIMIT_INFORMATION`, modify it directly +//! and set it back to the job. +//! +//! It's important to remeber to set the needed `LimitFlags` for each limit used. +//! +//! ```edition2018 +//! # use win32job::*; +//! # fn main() -> Result<(), JobError> { +//! use winapi::um::winnt::JOB_OBJECT_LIMIT_WORKINGSET; +//! let job = Job::create()?; +//! let mut info = job.basic_limit_info()?; +//! +//! info.MinimumWorkingSetSize = 1 * 1024 * 1024; +//! info.MaximumWorkingSetSize = 4 * 1024 * 1024; +//! info.LimitFlags |= JOB_OBJECT_LIMIT_WORKINGSET; +//! +//! job.set_basic_limit_info(&mut info)?; +//! job.assign_current_process()?; +//! # Ok(()) +//! # } +//! ``` +mod utils; + +use std::{io, mem, ptr}; +use thiserror::Error; +use winapi::shared::minwindef::*; +use winapi::um::handleapi::*; +use winapi::um::jobapi2::*; +use winapi::um::winnt::*; + +pub use crate::utils::{get_current_process, get_process_memory_info}; + +#[derive(Error, Debug)] +pub enum JobError { + #[error("Failed to create job")] + CreateFailed(io::Error), + #[error("Failed to assign job")] + AssignFailed(io::Error), + #[error("Failed to set info for job")] + SetInfoFailed(io::Error), + #[error("Failed to get info for job")] + GetInfoFailed(io::Error), +} + +#[derive(Debug)] +pub struct Job { + handle: HANDLE, +} + +impl Job { + pub fn create() -> Result { + let job_handle = unsafe { CreateJobObjectW(ptr::null_mut(), ptr::null()) }; + + if job_handle.is_null() { + return Err(JobError::CreateFailed(io::Error::last_os_error())); + } + + Ok(Job { handle: job_handle }) + } + + pub fn basic_limit_info(&self) -> Result { + let mut info: JOBOBJECT_BASIC_LIMIT_INFORMATION = unsafe { mem::zeroed() }; + let return_value = unsafe { + QueryInformationJobObject( + self.handle, + JobObjectBasicLimitInformation, + &mut info as *mut _ as LPVOID, + mem::size_of_val(&info) as DWORD, + 0 as *mut _, + ) + }; + + if return_value == 0 { + Err(JobError::GetInfoFailed(io::Error::last_os_error())) + } else { + Ok(info) + } + } + + pub fn set_basic_limit_info( + &self, + basic_info: &mut JOBOBJECT_BASIC_LIMIT_INFORMATION, + ) -> Result<(), JobError> { + let return_value = unsafe { + SetInformationJobObject( + self.handle, + JobObjectBasicLimitInformation, + basic_info as *mut _ as LPVOID, + mem::size_of_val(basic_info) as DWORD, + ) + }; + + if return_value == 0 { + Err(JobError::SetInfoFailed(io::Error::last_os_error())) + } else { + Ok(()) + } + } + + pub fn assign_process(&self, proc_handle: HANDLE) -> Result<(), JobError> { + let return_value = unsafe { AssignProcessToJobObject(self.handle, proc_handle) }; + + if return_value == 0 { + Err(JobError::AssignFailed(io::Error::last_os_error())) + } else { + Ok(()) + } + } + + pub fn assign_current_process(&self) -> Result<(), JobError> { + let current_proc_handle = get_current_process(); + + self.assign_process(current_proc_handle) + } +} + +impl Drop for Job { + fn drop(&mut self) { + unsafe { + CloseHandle(self.handle); + } + } +} + +#[cfg(test)] +mod tests { + use crate::utils::{get_current_process, get_process_memory_info}; + use crate::Job; + use winapi::um::winnt::JOB_OBJECT_LIMIT_WORKINGSET; + + #[test] + fn it_works() { + let job = Job::create().unwrap(); + + let mut info = job.basic_limit_info().unwrap(); + + assert_eq!(info.LimitFlags, 0); + + // This is the default. + assert_eq!(info.SchedulingClass, 5); + + info.SchedulingClass = 3; + info.MinimumWorkingSetSize = 1 * 1024 * 1024; + info.MaximumWorkingSetSize = 4 * 1024 * 1024; + + info.LimitFlags |= JOB_OBJECT_LIMIT_WORKINGSET; + + job.set_basic_limit_info(&mut info).unwrap(); + + let test_vec_size = 8 * 1024 * 1024; + let mut big_vec: Vec = Vec::with_capacity(test_vec_size); + big_vec.resize_with(test_vec_size, || 1); + + let memory_info = get_process_memory_info(get_current_process()).unwrap(); + + assert!(memory_info.WorkingSetSize >= info.MaximumWorkingSetSize); + + job.assign_current_process().unwrap(); + + let memory_info = get_process_memory_info(get_current_process()).unwrap(); + + assert!(memory_info.WorkingSetSize <= info.MaximumWorkingSetSize); + } +} diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..6f20a68 --- /dev/null +++ b/src/utils.rs @@ -0,0 +1,31 @@ +use std::{io, mem}; +use winapi::um::psapi::PROCESS_MEMORY_COUNTERS; +use winapi::um::winnt::*; +use winapi::um::{processthreadsapi, psapi}; + +/// Return a pseudo handle to the current process. +/// See also https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-getcurrentprocess +pub fn get_current_process() -> HANDLE { + unsafe { processthreadsapi::GetCurrentProcess() } +} + +/// Retrieves information about the memory usage of the specified process. +/// See also https://docs.microsoft.com/en-us/windows/win32/api/psapi/nf-psapi-getprocessmemoryinfo +pub fn get_process_memory_info( + process_handle: HANDLE, +) -> Result { + let mut counters: PROCESS_MEMORY_COUNTERS = unsafe { mem::zeroed() }; + let return_value = unsafe { + psapi::GetProcessMemoryInfo( + process_handle, + &mut counters as *mut _, + mem::size_of::() as u32, + ) + }; + + if return_value == 0 { + Err(io::Error::last_os_error()) + } else { + Ok(counters) + } +}