Raise RLIMIT_MEMLOCK on kernel version < 5.11.0

See https://github.com/cilium/ebpf/blob/main/docs/ebpf/concepts/rlimit.md.

aya/src/maps/mod.rs

@@ -50,7 +50,7 @@
 use std::{
     borrow::Borrow,
     ffi::CString,
-    fmt, io,
+    io,
     marker::PhantomData,
     mem,
     ops::Deref,
@@ -60,8 +60,6 @@ use std::{
 };
 
 use aya_obj::{EbpfSectionKind, InvalidTypeBinding, generated::bpf_map_type, parse_map_info};
-use libc::{RLIM_INFINITY, RLIMIT_MEMLOCK, getrlimit, rlim_t, rlimit};
-use log::warn;
 use thiserror::Error;
 
 use crate::{
@@ -71,7 +69,7 @@ use crate::{
         SyscallError, bpf_create_map, bpf_get_object, bpf_map_freeze, bpf_map_get_fd_by_id,
         bpf_map_get_next_key, bpf_map_update_elem_ptr, bpf_pin_object,
     },
-    util::{KernelVersion, nr_cpus},
+    util::nr_cpus,
 };
 
 pub mod array;
@@ -235,40 +233,6 @@ impl AsFd for MapFd {
     }
 }
 
-/// Raises a warning about rlimit. Should be used only if creating a map was not
-/// successful.
-fn maybe_warn_rlimit() {
-    let mut limit = mem::MaybeUninit::<rlimit>::uninit();
-    let ret = unsafe { getrlimit(RLIMIT_MEMLOCK, limit.as_mut_ptr()) };
-    if ret == 0 {
-        let limit = unsafe { limit.assume_init() };
-
-        if limit.rlim_cur == RLIM_INFINITY {
-            return;
-        }
-        struct HumanSize(rlim_t);
-
-        impl fmt::Display for HumanSize {
-            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-                let &Self(size) = self;
-                if size < 1024 {
-                    write!(f, "{size} bytes")
-                } else if size < 1024 * 1024 {
-                    write!(f, "{} KiB", size / 1024)
-                } else {
-                    write!(f, "{} MiB", size / 1024 / 1024)
-                }
-            }
-        }
-        warn!(
-            "RLIMIT_MEMLOCK value is {}, not RLIM_INFINITY; if experiencing problems with creating \
-            maps, try raising RLIMIT_MEMLOCK either to RLIM_INFINITY or to a higher value sufficient \
-            for the size of your maps",
-            HumanSize(limit.rlim_cur)
-        );
-    }
-}
-
 /// eBPF map types.
 #[derive(Debug)]
 pub enum Map {
@@ -576,16 +540,11 @@ impl MapData {
             }
         };
 
-        let fd = bpf_create_map(&c_name, &obj, btf_fd).map_err(|io_error| {
-            if !KernelVersion::at_least(5, 11, 0) {
-                maybe_warn_rlimit();
-            }
-
-            MapError::CreateError {
+        let fd =
+            bpf_create_map(&c_name, &obj, btf_fd).map_err(|io_error| MapError::CreateError {
                 name: name.into(),
                 io_error,
-            }
-        })?;
+            })?;
         Ok(Self {
             obj,
             fd: MapFd::from_fd(fd),

aya/src/sys/bpf.rs

@@ -1,7 +1,7 @@
 use std::{
     cmp,
     ffi::{CStr, CString, c_char},
-    io, iter,
+    fmt, io, iter,
     mem::{self, MaybeUninit},
     os::fd::{AsFd as _, AsRawFd as _, BorrowedFd, FromRawFd as _, RawFd},
 };
@@ -22,7 +22,11 @@ use aya_obj::{
     },
     maps::{LegacyMap, bpf_map_def},
 };
-use libc::{ENOENT, ENOSPC};
+use libc::{
+    EBADF, ENOENT, ENOSPC, EPERM, RLIM_INFINITY, RLIMIT_MEMLOCK, getrlimit, rlim_t, rlimit,
+    setrlimit,
+};
+use log::warn;
 
 use crate::{
     Btf, Pod, VerifierLogLevel,
@@ -99,8 +103,7 @@ pub(crate) fn bpf_create_map(
             .copy_from_slice(unsafe { mem::transmute::<&[u8], &[c_char]>(&name_bytes[..len]) });
     }
 
-    // SAFETY: BPF_MAP_CREATE returns a new file descriptor.
-    unsafe { fd_sys_bpf(bpf_cmd::BPF_MAP_CREATE, &mut attr) }
+    bpf_map_create(&mut attr)
 }
 
 pub(crate) fn bpf_pin_object(fd: BorrowedFd<'_>, path: &CStr) -> io::Result<()> {
@@ -695,6 +698,65 @@ pub(super) unsafe fn fd_sys_bpf(
     Ok(unsafe { crate::MockableFd::from_raw_fd(fd) })
 }
 
+static RAISE_MEMLIMIT: std::sync::Once = std::sync::Once::new();
+fn with_raised_rlimit_retry<T, F: FnMut() -> io::Result<T>>(mut op: F) -> io::Result<T> {
+    let mut result = op();
+    if matches!(result.as_ref(), Err(err) if err.raw_os_error() == Some(EPERM)) {
+        RAISE_MEMLIMIT.call_once(|| {
+            if KernelVersion::at_least(5, 11, 0) {
+                return;
+            }
+            let mut limit = mem::MaybeUninit::<rlimit>::uninit();
+            let ret = unsafe { getrlimit(RLIMIT_MEMLOCK, limit.as_mut_ptr()) };
+            if ret != 0 {
+                warn!("getrlimit(RLIMIT_MEMLOCK) failed: {ret}");
+                return;
+            }
+            let rlimit {
+                rlim_cur,
+                rlim_max: _,
+            } = unsafe { limit.assume_init() };
+
+            if rlim_cur == RLIM_INFINITY {
+                return;
+            }
+            let limit = rlimit {
+                rlim_cur: RLIM_INFINITY,
+                rlim_max: RLIM_INFINITY,
+            };
+            let ret = unsafe { setrlimit(RLIMIT_MEMLOCK, &limit) };
+            if ret != 0 {
+                struct HumanSize(rlim_t);
+
+                impl fmt::Display for HumanSize {
+                    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                        let &Self(size) = self;
+                        if size < 1024 {
+                            write!(f, "{size} bytes")
+                        } else if size < 1024 * 1024 {
+                            write!(f, "{} KiB", size / 1024)
+                        } else {
+                            write!(f, "{} MiB", size / 1024 / 1024)
+                        }
+                    }
+                }
+                warn!(
+                    "setrlimit(RLIMIT_MEMLOCK, RLIM_INFINITIY) failed: {ret}; current value is {}",
+                    HumanSize(rlim_cur)
+                );
+            }
+        });
+        // Retry after raising the limit.
+        result = op();
+    }
+    result
+}
+
+pub(super) fn bpf_map_create(attr: &mut bpf_attr) -> io::Result<crate::MockableFd> {
+    // SAFETY: BPF_MAP_CREATE returns a new file descriptor.
+    with_raised_rlimit_retry(|| unsafe { fd_sys_bpf(bpf_cmd::BPF_MAP_CREATE, attr) })
+}
+
 pub(crate) fn bpf_btf_get_fd_by_id(id: u32) -> Result<crate::MockableFd, SyscallError> {
     let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
     attr.__bindgen_anon_6.__bindgen_anon_1.btf_id = id;
@@ -849,7 +911,7 @@ pub(crate) fn is_perf_link_supported() -> bool {
                 None,
             );
             // Returns EINVAL if unsupported. EBADF if supported.
-            matches!(link, Err(err) if err.raw_os_error() == Some(libc::EBADF))
+            matches!(link, Err(err) if err.raw_os_error() == Some(EBADF))
         } else {
             false
         }
@@ -944,9 +1006,7 @@ pub(crate) fn is_prog_id_supported(map_type: bpf_map_type) -> bool {
     u.max_entries = 1;
     u.map_flags = 0;
 
-    // SAFETY: BPF_MAP_CREATE returns a new file descriptor.
-    let fd = unsafe { fd_sys_bpf(bpf_cmd::BPF_MAP_CREATE, &mut attr) };
-    fd.is_ok()
+    bpf_map_create(&mut attr).is_ok()
 }
 
 pub(crate) fn is_btf_supported() -> bool {
@@ -1107,7 +1167,7 @@ pub(crate) fn is_btf_type_tag_supported() -> bool {
 
 pub(super) fn bpf_prog_load(attr: &mut bpf_attr) -> io::Result<crate::MockableFd> {
     // SAFETY: BPF_PROG_LOAD returns a new file descriptor.
-    unsafe { fd_sys_bpf(bpf_cmd::BPF_PROG_LOAD, attr) }
+    with_raised_rlimit_retry(|| unsafe { fd_sys_bpf(bpf_cmd::BPF_PROG_LOAD, attr) })
 }
 
 fn sys_bpf(cmd: bpf_cmd, attr: &mut bpf_attr) -> io::Result<i64> {
@@ -1225,7 +1285,7 @@ pub(crate) fn retry_with_verifier_logs<T>(
 
 #[cfg(test)]
 mod tests {
-    use libc::{EBADF, EINVAL};
+    use libc::EINVAL;
 
     use super::*;
     use crate::sys::override_syscall;

aya/src/sys/feature_probe.rs

@@ -10,7 +10,7 @@ use aya_obj::{
 };
 use libc::{E2BIG, EBADF, EINVAL};
 
-use super::{SyscallError, bpf_prog_load, fd_sys_bpf, unit_sys_bpf, with_trivial_prog};
+use super::{SyscallError, bpf_map_create, bpf_prog_load, unit_sys_bpf, with_trivial_prog};
 use crate::{
     MockableFd,
     maps::MapType,
@@ -278,13 +278,10 @@ pub fn is_map_supported(map_type: MapType) -> Result<bool, SyscallError> {
             u_map.key_size = 1;
             u_map.value_size = 1;
             u_map.max_entries = 1;
-            // SAFETY: BPF_MAP_CREATE returns a new file descriptor.
-            inner_map_fd = unsafe { fd_sys_bpf(bpf_cmd::BPF_MAP_CREATE, &mut attr_map) }.map_err(
-                |io_error| SyscallError {
-                    call: "bpf_map_create",
-                    io_error,
-                },
-            )?;
+            inner_map_fd = bpf_map_create(&mut attr_map).map_err(|io_error| SyscallError {
+                call: "bpf_map_create",
+                io_error,
+            })?;
 
             u.inner_map_fd = inner_map_fd.as_raw_fd() as u32;
         }
@@ -299,8 +296,8 @@ pub fn is_map_supported(map_type: MapType) -> Result<bool, SyscallError> {
     }
 
     // SAFETY: BPF_MAP_CREATE returns a new file descriptor.
-    let io_error = match unsafe { fd_sys_bpf(bpf_cmd::BPF_MAP_CREATE, &mut attr) } {
-        Ok(_) => return Ok(true),
+    let io_error = match bpf_map_create(&mut attr) {
+        Ok(_fd) => return Ok(true),
         Err(io_error) => io_error,
     };