From 6327cc1caecb8f2895d2be9547b6855096dbc99f Mon Sep 17 00:00:00 2001
From: Keiichi Watanabe <keiichiw@google.com>
Date: Wed, 18 Sep 2024 19:20:01 +0900
Subject: [PATCH] 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>
---
 base/src/sys/linux/mod.rs | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/base/src/sys/linux/mod.rs b/base/src/sys/linux/mod.rs
index dbd582b3d7..ea185c8161 100644
--- a/base/src/sys/linux/mod.rs
+++ b/base/src/sys/linux/mod.rs
@@ -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;