feat: add `Ebpf::maps_disjoint_mut`

Implemented using `HashMap::get_disjoint_mut` introduced in Rust 1.86.0.

aya/src/bpf.rs

@@ -1055,6 +1055,35 @@ impl Ebpf {
         self.maps.iter_mut().map(|(name, map)| (name.as_str(), map))
     }
 
+    /// Attempts to get mutable references to `N` maps at once.
+    ///
+    /// Returns an array of length `N` with the results of each query, in the same order
+    /// as the requested map names. For soundness, at most one mutable reference will be
+    /// returned to any map. `None` will be used if a map with the given name is missing.
+    ///
+    /// This method performs a check to ensure that there are no duplicate map names,
+    /// which currently has a time-complexity of *O(n²)*. Be careful when passing a large
+    /// number of names.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any names are duplicated.
+    ///
+    /// # Examples
+    /// ```no_run
+    /// # let mut bpf = aya::Ebpf::load(&[])?;
+    /// match bpf.maps_disjoint_mut(["MAP1", "MAP2"]) {
+    ///     [Some(m1), Some(m2)] => println!("Got MAP1 and MAP2"),
+    ///     [Some(m1), None] => println!("Got only MAP1"),
+    ///     [None, Some(m2)] => println!("Got only MAP2"),
+    ///     [None, None] => println!("No maps"),
+    /// }
+    /// # Ok::<(), aya::EbpfError>(())
+    /// ```
+    pub fn maps_disjoint_mut<const N: usize>(&mut self, names: [&str; N]) -> [Option<&mut Map>; N] {
+        self.maps.get_disjoint_mut(names)
+    }
+
     /// Returns a reference to the program with the given name.
     ///
     /// You can use this to inspect a program and its properties. To load and attach a program, use

test/integration-ebpf/src/map_test.rs

@@ -3,9 +3,9 @@
 #![expect(unused_crate_dependencies, reason = "used in other bins")]
 
 use aya_ebpf::{
-    macros::{map, socket_filter},
+    macros::{map, socket_filter, uprobe},
     maps::{Array, HashMap},
-    programs::SkBuffContext,
+    programs::{ProbeContext, SkBuffContext},
 };
 #[cfg(not(test))]
 extern crate ebpf_panic;
@@ -35,3 +35,20 @@ fn simple_prog(_ctx: SkBuffContext) -> i64 {
 
     0
 }
+
+#[uprobe]
+fn simple_prog_mut(_ctx: ProbeContext) -> i64 {
+    if let Some(array_value) = FOO.get_ptr_mut(0) {
+        unsafe {
+            *array_value += 1;
+        }
+    }
+
+    if let Some(map_value) = BAZ.get_ptr_mut(0) {
+        unsafe {
+            *map_value += 1;
+        }
+    }
+
+    0
+}

test/integration-test/src/tests.rs

@@ -10,6 +10,7 @@ mod load;
 mod log;
 mod lsm;
 mod map_pin;
+mod maps_disjoint;
 mod raw_tracepoint;
 mod rbpf;
 mod relocations;

test/integration-test/src/tests/maps_disjoint.rs

@@ -0,0 +1,44 @@
+use aya::{
+    Ebpf,
+    maps::{Array, HashMap},
+    programs::UProbe,
+};
+
+#[unsafe(no_mangle)]
+#[inline(never)]
+extern "C" fn trigger_ebpf_program_maps_disjoint() {
+    core::hint::black_box(trigger_ebpf_program_maps_disjoint);
+}
+
+#[test_log::test]
+fn test_maps_disjoint() {
+    let mut bpf: Ebpf = Ebpf::load(crate::MAP_TEST).unwrap();
+    let prog: &mut UProbe = bpf
+        .program_mut("simple_prog_mut")
+        .unwrap()
+        .try_into()
+        .unwrap();
+
+    prog.load().unwrap();
+    prog.attach(
+        "trigger_ebpf_program_maps_disjoint",
+        "/proc/self/exe",
+        None,
+        None,
+    )
+    .unwrap();
+
+    let [foo, bar, baz] = bpf.maps_disjoint_mut(["FOO", "BAR", "BAZ"]);
+
+    let mut foo: Array<_, u32> = Array::try_from(foo.unwrap()).unwrap();
+    let mut bar: HashMap<_, u32, u8> = HashMap::try_from(bar.unwrap()).unwrap();
+    assert!(baz.is_none());
+
+    foo.set(0, 5, 0).unwrap();
+    bar.insert(0, 10, 0).unwrap();
+
+    trigger_ebpf_program_maps_disjoint();
+
+    assert_eq!(foo.get(&0, 0).unwrap(), 6);
+    assert_eq!(bar.get(&0, 0).unwrap(), 11);
+}

xtask/public-api/aya.txt

@@ -10629,6 +10629,7 @@ pub fn aya::Ebpf::load_file<P: core::convert::AsRef<std::path::Path>>(path: P) -
 pub fn aya::Ebpf::map(&self, name: &str) -> core::option::Option<&aya::maps::Map>
 pub fn aya::Ebpf::map_mut(&mut self, name: &str) -> core::option::Option<&mut aya::maps::Map>
 pub fn aya::Ebpf::maps(&self) -> impl core::iter::traits::iterator::Iterator<Item = (&str, &aya::maps::Map)>
+pub fn aya::Ebpf::maps_disjoint_mut<const N: usize>(&mut self, names: [&str; N]) -> [core::option::Option<&mut aya::maps::Map>; N]
 pub fn aya::Ebpf::maps_mut(&mut self) -> impl core::iter::traits::iterator::Iterator<Item = (&str, &mut aya::maps::Map)>
 pub fn aya::Ebpf::program(&self, name: &str) -> core::option::Option<&aya::programs::Program>
 pub fn aya::Ebpf::program_mut(&mut self, name: &str) -> core::option::Option<&mut aya::programs::Program>