aya: add Bpf::many_programs_mut

Store programs in a `hashbrown::HashMap` and expose `get_many_mut`. We
can revisit this dependency when
rust-lang/rust#97601 is resolved.

Author: tamird

aya/Cargo.toml

@@ -16,6 +16,7 @@ async-io = { workspace = true, optional = true }
 aya-obj = { workspace = true, features = ["std"] }
 bitflags = { workspace = true }
 bytes = { workspace = true }
+hashbrown = { workspace = true, default-features = true }
 lazy_static = { workspace = true }
 libc = { workspace = true }
 log = { workspace = true }

aya/src/bpf.rs

@@ -768,7 +768,9 @@ impl<'a> Default for BpfLoader<'a> {
 #[derive(Debug)]
 pub struct Bpf {
     maps: HashMap<String, Map>,
-    programs: HashMap<String, Program>,
+    // TODO(https://github.com/rust-lang/rust/issues/97601): use `HashMap` when `get_many_mut` is
+    // stabilized.
+    programs: hashbrown::HashMap<String, Program>,
 }
 
 impl Bpf {
@@ -914,6 +916,31 @@ impl Bpf {
         self.programs.get_mut(name)
     }
 
+    /// Attempts to get mutable references to `N` programs in the map at once.
+    ///
+    /// Returns an array of length `N` with the results of each query. For soundness, at most one
+    /// mutable reference will be returned to any value. `None` will be returned if any of the
+    /// keys are duplicates or missing.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # let mut bpf = aya::Bpf::load(&[])?;
+    /// use aya::programs::UProbe;
+    ///
+    /// let [ssl_read, ssl_write ]= bpf.many_programs_mut(["SSL_read", "SSL_write"]).unwrap();
+    /// let ssl_read: &mut UProbe = ssl_read.try_into()?;
+    /// let ssl_write: &mut UProbe = ssl_write.try_into()?;
+    /// ssl_read.load()?;
+    /// ssl_read.attach(Some("SSL_read"), 0, "libssl", None)?;
+    /// ssl_write.load()?;
+    /// ssl_write.attach(Some("SSL_write"), 0, "libssl", None)?;
+    /// # Ok::<(), aya::BpfError>(())
+    /// ```
+    pub fn many_programs_mut<const N: usize>(&mut self, names: [&str; N]) -> Option<[&mut Program; N]> {
+        self.programs.get_many_mut(names)
+    }
+
     /// An iterator over all the programs.
     ///
     /// # Examples

aya/src/programs/extension.rs

@@ -37,22 +37,13 @@ pub enum ExtensionError {
 /// use aya::{BpfLoader, programs::{Xdp, XdpFlags, Extension}};
 ///
 /// let mut bpf = BpfLoader::new().extension("extension").load_file("app.o")?;
-/// let mut prog = None;
-/// let mut ext = None;
-/// for (name, program) in bpf.programs_mut() {
-///     match name {
-///         "main" => prog = Some(program),
-///         "extension" => ext = Some(program),
-///         _ => {},
-///     }
-/// }
-///
-/// let prog: &mut Xdp = prog.unwrap().try_into()?;
+/// let [prog, ext] = bpf.many_programs_mut(["main", "extension"]).unwrap();
+/// let prog: &mut Xdp = prog.try_into()?;
 /// prog.load()?;
 /// prog.attach("eth0", XdpFlags::default())?;
 ///
 /// let prog_fd = prog.fd().unwrap();
-/// let ext: &mut Extension = ext.unwrap().try_into()?;
+/// let ext: &mut Extension = ext.try_into()?;
 /// ext.load(prog_fd, "function_to_replace")?;
 /// ext.attach()?;
 /// Ok::<(), aya::BpfError>(())