base: Add call_with_extended_max_files()

Add a utility function `base::call_with_extended_max_files()` that executes
a given callback with extended file counts and restore the original file limit.

This utility will be used by the following CL for pmem-ext2.

BUG=b:329359333
TEST=presubmit

Change-Id: Id85781e4e4cd5542d4da630eb48779947b098490
Reviewed-on: https://chromium-review.googlesource.com/c/crosvm/crosvm/+/5872164
Reviewed-by: Daniel Verkamp <dverkamp@chromium.org>
Commit-Queue: Keiichi Watanabe <keiichiw@chromium.org>
This commit is contained in:
Keiichi Watanabe 2024-09-18 19:20:01 +09:00 committed by crosvm LUCI
parent 56ada36fc0
commit 6327cc1cae

View file

@ -569,6 +569,43 @@ pub fn max_open_files() -> Result<libc::rlimit64> {
}
}
/// Executes the given callback with extended soft limit of max number of open files. After the
/// callback executed, restore the limit.
pub fn call_with_extended_max_files<T, E>(
callback: impl FnOnce() -> std::result::Result<T, E>,
) -> Result<std::result::Result<T, E>> {
let cur_limit = max_open_files()?;
let new_limit = libc::rlimit64 {
rlim_cur: cur_limit.rlim_max,
..cur_limit
};
let needs_extension = cur_limit.rlim_cur < new_limit.rlim_cur;
if needs_extension {
set_max_open_files(new_limit)?;
}
let r = callback();
// Restore the soft limit.
if needs_extension {
set_max_open_files(cur_limit)?;
}
Ok(r)
}
/// Set the soft and hard limits of max number of open files to the given value.
fn set_max_open_files(limit: libc::rlimit64) -> Result<()> {
// SAFETY: RLIMIT_NOFILE is known only to read a buffer of size rlimit64, and we have always
// rlimit64 allocated.
let res = unsafe { libc::setrlimit64(libc::RLIMIT_NOFILE, &limit) };
if res == 0 {
Ok(())
} else {
errno_result()
}
}
/// Moves the requested PID/TID to a particular cgroup
pub fn move_to_cgroup(cgroup_path: PathBuf, id_to_write: Pid, cgroup_file: &str) -> Result<()> {
use std::io::Write;