use windows::Win32::System::{
JobObjects::{
JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JOB_OBJECT_LIMIT_AFFINITY,
JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, JOB_OBJECT_LIMIT_PRIORITY_CLASS,
JOB_OBJECT_LIMIT_SCHEDULING_CLASS, JOB_OBJECT_LIMIT_WORKINGSET,
},
Threading::{
ABOVE_NORMAL_PRIORITY_CLASS, BELOW_NORMAL_PRIORITY_CLASS, HIGH_PRIORITY_CLASS,
IDLE_PRIORITY_CLASS, NORMAL_PRIORITY_CLASS, REALTIME_PRIORITY_CLASS,
},
};
#[derive(Debug)]
pub struct ExtendedLimitInfo(pub(crate) JOBOBJECT_EXTENDED_LIMIT_INFORMATION);
#[derive(Debug, Clone, Copy)]
#[repr(u32)]
pub enum PriorityClass {
Normal = NORMAL_PRIORITY_CLASS.0,
Idle = IDLE_PRIORITY_CLASS.0,
High = HIGH_PRIORITY_CLASS.0,
Realtime = REALTIME_PRIORITY_CLASS.0,
BelowNormal = BELOW_NORMAL_PRIORITY_CLASS.0,
AboveNormal = ABOVE_NORMAL_PRIORITY_CLASS.0,
}
impl Default for ExtendedLimitInfo {
fn default() -> Self {
Self::new()
}
}
impl ExtendedLimitInfo {
pub fn new() -> Self {
let inner = Default::default();
ExtendedLimitInfo(inner)
}
pub fn limit_working_memory(&mut self, min: usize, max: usize) -> &mut Self {
self.0.BasicLimitInformation.MinimumWorkingSetSize = min;
self.0.BasicLimitInformation.MaximumWorkingSetSize = max;
self.0.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_WORKINGSET;
self
}
pub fn limit_kill_on_job_close(&mut self) -> &mut Self {
self.0.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
self
}
pub fn limit_priority_class(&mut self, priority_class: PriorityClass) -> &mut Self {
self.0.BasicLimitInformation.PriorityClass = priority_class as u32;
self.0.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_PRIORITY_CLASS;
self
}
pub fn limit_scheduling_class(&mut self, scheduling_class: u8) -> &mut Self {
self.0.BasicLimitInformation.SchedulingClass = scheduling_class as u32;
self.0.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_SCHEDULING_CLASS;
self
}
pub fn limit_affinity(&mut self, affinity: usize) -> &mut Self {
self.0.BasicLimitInformation.Affinity = affinity;
self.0.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_AFFINITY;
self
}
pub fn clear_limits(&mut self) -> &mut Self {
self.0.BasicLimitInformation.LimitFlags.0 = 0;
self
}
}
#[cfg(test)]
mod tests {
use crate::utils::{get_current_process, get_process_affinity_mask, get_process_memory_info};
use crate::{Job, PriorityClass};
use rusty_fork::rusty_fork_test;
rusty_fork_test! {
#[test]
fn working_mem_limits() {
let job = Job::create().unwrap();
let mut info = job.query_extended_limit_info().unwrap();
let min = 1 * 1024 * 1024;
let max = 4 * 1024 * 1024;
info.limit_working_memory(min, max);
job.set_extended_limit_info(&mut info).unwrap();
job.assign_current_process().unwrap();
let test_vec_size = max * 4;
let mut big_vec: Vec<u8> = 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.working_set_size <= max * 2);
info.clear_limits();
job.set_extended_limit_info(&mut info).unwrap();
}
}
rusty_fork_test! {
#[test]
fn kill_on_job_close_limits() {
let job = Job::create().unwrap();
let mut info = job.query_extended_limit_info().unwrap();
info.limit_kill_on_job_close();
job.set_extended_limit_info(&mut info).unwrap();
job.assign_current_process().unwrap();
drop(job);
panic!();
}
}
rusty_fork_test! {
#[test]
fn priority_class_limits() {
let job = Job::create().unwrap();
let mut info = job.query_extended_limit_info().unwrap();
info.limit_priority_class(PriorityClass::BelowNormal);
job.set_extended_limit_info(&mut info).unwrap();
let info = job.query_extended_limit_info().unwrap();
assert_eq!(info.0.BasicLimitInformation.PriorityClass, PriorityClass::BelowNormal as u32);
}
}
rusty_fork_test! {
#[test]
fn scheduling_class_limits() {
let job = Job::create().unwrap();
let mut info = job.query_extended_limit_info().unwrap();
info.limit_scheduling_class(1);
job.set_extended_limit_info(&mut info).unwrap();
let info = job.query_extended_limit_info().unwrap();
assert_eq!(info.0.BasicLimitInformation.SchedulingClass, 1);
}
}
rusty_fork_test! {
#[test]
fn affinity_limits() {
let job = Job::create().unwrap();
let mut info = job.query_extended_limit_info().unwrap();
info.limit_affinity(1);
job.set_extended_limit_info(&mut info).unwrap();
let (proc_affinity, _) = get_process_affinity_mask(get_current_process()).unwrap();
assert_ne!(proc_affinity, 1);
job.assign_current_process().unwrap();
let (proc_affinity, _) = get_process_affinity_mask(get_current_process()).unwrap();
assert_eq!(proc_affinity, 1);
}
}
}