use std::fs::File;
use std::sync::atomic::{AtomicUsize, Ordering};
pub fn set_priority(nice: Option<i32>) {
if let Some(n) = nice {
unsafe {
libc::setpriority(libc::PRIO_PROCESS, 0, n);
}
}
}
#[cfg(target_os = "linux")]
pub fn set_io_idle() {
const IOPRIO_WHO_PROCESS: libc::c_long = 1;
const IOPRIO_CLASS_IDLE: libc::c_long = 3;
const IOPRIO_CLASS_SHIFT: libc::c_long = 13;
unsafe {
libc::syscall(
libc::SYS_ioprio_set,
IOPRIO_WHO_PROCESS,
0 as libc::c_long,
IOPRIO_CLASS_IDLE << IOPRIO_CLASS_SHIFT,
);
}
}
#[cfg(target_os = "macos")]
pub fn set_io_idle() {
const IOPOL_TYPE_DISK: libc::c_int = 0;
const IOPOL_SCOPE_PROCESS: libc::c_int = 0;
const IOPOL_THROTTLE: libc::c_int = 3;
unsafe extern "C" {
fn setiopolicy_np(
iotype: libc::c_int,
scope: libc::c_int,
policy: libc::c_int,
) -> libc::c_int;
}
unsafe {
setiopolicy_np(IOPOL_TYPE_DISK, IOPOL_SCOPE_PROCESS, IOPOL_THROTTLE);
}
}
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
pub fn set_io_idle() {}
#[cfg(target_os = "linux")]
pub fn prepare(f: &File, _drop_cache: bool) {
use std::os::unix::io::AsRawFd;
unsafe {
libc::posix_fadvise(f.as_raw_fd(), 0, 0, libc::POSIX_FADV_SEQUENTIAL);
}
}
#[cfg(target_os = "macos")]
pub fn prepare(f: &File, drop_cache: bool) {
use std::os::unix::io::AsRawFd;
if drop_cache {
unsafe {
libc::fcntl(f.as_raw_fd(), libc::F_NOCACHE, 1);
}
}
}
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
pub fn prepare(_f: &File, _drop_cache: bool) {}
#[cfg(target_os = "linux")]
pub fn finish(f: &File, drop_cache: bool) {
use std::os::unix::io::AsRawFd;
if drop_cache {
unsafe {
libc::posix_fadvise(f.as_raw_fd(), 0, 0, libc::POSIX_FADV_DONTNEED);
}
}
}
#[cfg(not(target_os = "linux"))]
pub fn finish(_f: &File, _drop_cache: bool) {
}
#[cfg(any(target_os = "linux", target_os = "macos"))]
pub fn load_average_1m() -> Option<f64> {
let mut avg = [0f64; 1];
let got = unsafe { libc::getloadavg(avg.as_mut_ptr(), 1) };
(got == 1).then_some(avg[0])
}
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
pub fn load_average_1m() -> Option<f64> {
None
}
const LOAD_CHECK: std::time::Duration = std::time::Duration::from_secs(5);
pub fn starting_gate(jobs: usize, limit: f64, load: f64) -> usize {
if load > limit { 1 } else { jobs.max(1) }
}
pub fn next_gate(gate: usize, jobs: usize, load: f64, limit: f64) -> usize {
let others = load - gate as f64;
if others > limit { 1 } else { (gate + 1).min(jobs).max(1) }
}
pub fn load_gate(jobs: usize, limit: Option<f64>) -> Option<std::sync::Arc<AtomicUsize>> {
let limit = limit?;
if jobs <= 1 {
return None;
}
let gate = std::sync::Arc::new(AtomicUsize::new(starting_gate(jobs, limit, load_average_1m()?)));
let held = std::sync::Arc::downgrade(&gate);
let _ = std::thread::Builder::new()
.name("trg-load".into())
.spawn(move || {
loop {
std::thread::sleep(LOAD_CHECK);
let Some(gate) = held.upgrade() else { return };
if let Some(load) = load_average_1m() {
let now = gate.load(Ordering::Relaxed);
gate.store(next_gate(now, jobs, load, limit), Ordering::Relaxed);
}
}
});
Some(gate)
}