use libc::{getrlimit, rlim_t, rlimit, setrlimit, RLIMIT_NOFILE};
use std::mem::MaybeUninit;
use std::{cmp, fs, io};
const DEFAULT_NOFILE: rlim_t = 1_000_000;
fn get_max_nofile() -> Result<rlim_t, String> {
let path = "/proc/sys/fs/nr_open";
let max_str = fs::read_to_string(path).map_err(|error| format!("Reading {path}: {error:?}"))?;
max_str
.trim()
.parse()
.map_err(|error| format!("Parsing {path}: {error:?}"))
}
fn get_nofile_limits() -> Result<rlimit, String> {
let mut limits = MaybeUninit::<rlimit>::zeroed();
let ret = unsafe { getrlimit(RLIMIT_NOFILE, limits.as_mut_ptr()) };
if ret != 0 {
return Err(format!("getrlimit: {}", io::Error::last_os_error()));
}
Ok(unsafe { limits.assume_init() })
}
fn setup_rlimit_nofile_to(nofile: rlim_t) -> Result<(), String> {
let rlimit = rlimit {
rlim_cur: nofile,
rlim_max: nofile,
};
let ret = unsafe { setrlimit(RLIMIT_NOFILE, &rlimit) };
if ret < 0 {
Err(format!(
"Failed to increase the limit: {:?}",
io::Error::last_os_error()
))
} else {
Ok(())
}
}
pub fn setup_rlimit_nofile(nofile: Option<u64>) -> Result<u64, String> {
let max_nofile = get_max_nofile()?;
let rlimit { rlim_cur, rlim_max } = get_nofile_limits()?;
let target_limit = if let Some(nofile) = nofile {
if nofile == 0 {
return Ok(rlim_cur); }
nofile
} else {
if DEFAULT_NOFILE <= rlim_cur {
return Ok(rlim_cur); }
cmp::min(DEFAULT_NOFILE, max_nofile)
};
if target_limit > max_nofile {
return Err(format!("It cannot be increased above {max_nofile}"));
}
let new_limit = if let Err(error) = setup_rlimit_nofile_to(target_limit) {
if nofile.is_some() {
return Err(error);
} else {
warn!(
"Failure when trying to set the limit to {target_limit}, \
the hard limit ({rlim_max}) of open file descriptors is used instead."
);
setup_rlimit_nofile_to(rlim_max).map_err(|error| {
format!("Cannot increase the soft limit to the hard limit: {error}")
})?;
rlim_max
}
} else {
target_limit
};
Ok(new_limit)
}