Use prlimit for linux gnu

Author: LMJW

src/sys/resource.rs

@@ -131,29 +131,42 @@ cfg_if! {
 ///
 /// [`Resource`]: enum.Resource.html
 pub fn getrlimit(resource: Resource) -> Result<(Option<rlim_t>, Option<rlim_t>)> {
-    let mut rlim = mem::MaybeUninit::<rlimit>::uninit();
-
-    #[cfg(all(target_os = "linux", target_env = "gnu"))]
-    let res =
-        unsafe { libc::getrlimit(resource as __rlimit_resource_t, rlim.as_mut_ptr() as *mut _) };
-    #[cfg(any(
-        target_os = "freebsd",
-        target_os = "openbsd",
-        target_os = "netbsd",
-        target_os = "macos",
-        target_os = "ios",
-        target_os = "android",
-        target_os = "dragonfly",
-        target_os = "bitrig",
-        target_os = "linux", // target_env != "gnu"
-    ))]
-    let res = unsafe { libc::getrlimit(resource as c_int, rlim.as_mut_ptr() as *mut _) };
+    let mut old_rlim = mem::MaybeUninit::<rlimit>::uninit();
+
+    cfg_if! {
+        if #[cfg(all(target_os = "linux", target_env = "gnu"))]{
+            // the below implementation is mimicing the similar implementation in golang
+            // https://go-review.googlesource.com/c/sys/+/230478/2/unix/syscall_linux_arm64.go#176
+            // seems for some of the architectures, we prefer to use prlimit instead of {g,s}etrlimit
+
+            let res = unsafe { libc::prlimit(0, resource as __rlimit_resource_t, std::ptr::null(), old_rlim.as_mut_ptr() as *mut _) };
+            if res == -1 {
+                // when error happens, the map will return an Err, the (None, None) is just make compiler
+                // happy, it will not go through
+                return Errno::result(res).map(|_|{ (None, None) });
+            }
+            let res = unsafe { libc::getrlimit(resource as __rlimit_resource_t, old_rlim.as_mut_ptr() as *mut _) };
+
+        }else if #[cfg(any(
+            target_os = "freebsd",
+            target_os = "openbsd",
+            target_os = "netbsd",
+            target_os = "macos",
+            target_os = "ios",
+            target_os = "android",
+            target_os = "dragonfly",
+            target_os = "bitrig",
+            target_os = "linux", // target_env != "gnu"
+        ))]{
+            let res = unsafe { libc::getrlimit(resource as c_int, old_rlim.as_mut_ptr() as *mut _) };
+        }
+    }
 
-    let rlim = unsafe { rlim.assume_init() };
+    let old_rlim = unsafe { old_rlim.assume_init() };
     Errno::result(res).map(|_| {
         (
-            Some(rlim.rlim_cur).filter(|x| *x != RLIM_INFINITY),
-            Some(rlim.rlim_max).filter(|x| *x != RLIM_INFINITY),
+            Some(old_rlim.rlim_cur).filter(|x| *x != RLIM_INFINITY),
+            Some(old_rlim.rlim_max).filter(|x| *x != RLIM_INFINITY),
         )
     })
 }
@@ -191,23 +204,36 @@ pub fn setrlimit(
     soft_limit: Option<rlim_t>,
     hard_limit: Option<rlim_t>,
 ) -> Result<()> {
-    let mut rlim = unsafe { mem::MaybeUninit::<rlimit>::uninit().assume_init() };
-    rlim.rlim_cur = soft_limit.unwrap_or(RLIM_INFINITY);
-    rlim.rlim_max = hard_limit.unwrap_or(RLIM_INFINITY);
+    let mut new_rlim = unsafe { mem::MaybeUninit::<rlimit>::uninit().assume_init() };
+    new_rlim.rlim_cur = soft_limit.unwrap_or(RLIM_INFINITY);
+    new_rlim.rlim_max = hard_limit.unwrap_or(RLIM_INFINITY);
+
+    cfg_if! {
+        if #[cfg(all(target_os = "linux", target_env = "gnu"))]{
+            // the below implementation is mimicing the similar implementation in golang
+            // https://go-review.googlesource.com/c/sys/+/230478/2/unix/syscall_linux_arm64.go#176
+            // seems for some of the architectures, we prefer to use prlimit instead of {g,s}etrlimit
+            let res = unsafe { libc::prlimit(0, resource as __rlimit_resource_t, &new_rlim as *const _, std::ptr::null_mut()) };
+            if res == -1 {
+                return Errno::result(res).map(|_| ());
+            }
+
+            let res = unsafe { libc::setrlimit(resource as __rlimit_resource_t, &new_rlim as *const _) };
+
+        }else if #[cfg(any(
+            target_os = "freebsd",
+            target_os = "openbsd",
+            target_os = "netbsd",
+            target_os = "macos",
+            target_os = "ios",
+            target_os = "android",
+            target_os = "dragonfly",
+            target_os = "bitrig",
+            target_os = "linux", // target_env != "gnu"
+        ))]{
+            let res = unsafe { libc::setrlimit(resource as c_int, &new_rlim as *const _) };
+        }
+    }
 
-    #[cfg(all(target_os = "linux", target_env = "gnu"))]
-    let res = unsafe { libc::setrlimit(resource as __rlimit_resource_t, &rlim as *const _) };
-    #[cfg(any(
-        target_os = "freebsd",
-        target_os = "openbsd",
-        target_os = "netbsd",
-        target_os = "macos",
-        target_os = "ios",
-        target_os = "android",
-        target_os = "dragonfly",
-        target_os = "bitrig",
-        target_os = "linux", // target_env != "gnu"
-    ))]
-    let res = unsafe { libc::setrlimit(resource as c_int, &rlim as *const _) };
     Errno::result(res).map(|_| ())
 }

test/test_resource.rs

@@ -33,11 +33,13 @@ pub fn test_resource_limits_nofile() {
 pub fn test_resource_limits_stack() {
     let (mut soft_limit, hard_limit) = getrlimit(Resource::RLIMIT_STACK).unwrap();
     let orig_limit = (soft_limit, hard_limit);
-
-    soft_limit = hard_limit.or(Some(10000));
+    eprintln!("\n{:?} --- {:?}", soft_limit, hard_limit);
+    soft_limit = Some(4194304);
     setrlimit(Resource::RLIMIT_STACK, soft_limit, hard_limit).unwrap();
 
     let limit2 = getrlimit(Resource::RLIMIT_STACK).unwrap();
+    eprintln!("\n{:?} --- {:?}", limit2.0, limit2.1);
+
     assert_eq!(soft_limit, limit2.0);
     assert_eq!(hard_limit, limit2.1);
 
@@ -46,9 +48,4 @@ pub fn test_resource_limits_stack() {
     let final_limit = getrlimit(Resource::RLIMIT_STACK).unwrap();
     assert_eq!(orig_limit.0, final_limit.0);
     assert_eq!(orig_limit.1, final_limit.1);
-
-    // assert_eq!(
-    //     format!("softlimit-{:?}, hardlimit-{:?}, limit2.0-{:?}, limit2.1-{:?}, final_limit.0-{:?}, final_limit.0-{:?}"
-    //     , soft_limit, hard_limit, limit2.0, limit2.1, final_limit.0, final_limit.1),
-    //     "");
 }