Use Set (instead of Map) to hold links

This uses hashbrown instead of std because the latter relies on Borrow
which requires a reference; hashbrown's Equivalent is more flexible.

aya/Cargo.toml

@@ -18,6 +18,7 @@ async-io = { workspace = true, optional = true }
 aya-obj = { path = "../aya-obj", version = "^0.2.1", features = ["std"] }
 bitflags = { workspace = true }
 bytes = { workspace = true }
+hashbrown = { workspace = true }
 libc = { workspace = true }
 log = { workspace = true }
 object = { workspace = true, features = ["elf", "read_core", "std", "write"] }

aya/src/programs/cgroup_device.rs

@@ -5,8 +5,8 @@ use std::os::fd::AsFd;
 use crate::{
     generated::{bpf_attach_type::BPF_CGROUP_DEVICE, bpf_prog_type::BPF_PROG_TYPE_CGROUP_DEVICE},
     programs::{
-        bpf_prog_get_fd_by_id, define_link_wrapper, load_program, query, CgroupAttachMode, FdLink,
-        Link, ProgAttachLink, ProgramData, ProgramError, ProgramFd,
+        bpf_prog_get_fd_by_id, define_link_wrapper, id_as_key, load_program, query,
+        CgroupAttachMode, FdLink, Link, ProgAttachLink, ProgramData, ProgramError, ProgramFd,
     },
     sys::{bpf_link_create, LinkTarget, ProgQueryTarget, SyscallError},
     util::KernelVersion,
@@ -153,6 +153,8 @@ impl Link for CgroupDeviceLinkInner {
     }
 }
 
+id_as_key!(CgroupDeviceLinkInner, CgroupDeviceLinkIdInner);
+
 define_link_wrapper!(
     /// The link used by [CgroupDevice] programs.
     CgroupDeviceLink,

aya/src/programs/cgroup_skb.rs

@@ -8,8 +8,8 @@ use crate::{
         bpf_prog_type::BPF_PROG_TYPE_CGROUP_SKB,
     },
     programs::{
-        define_link_wrapper, load_program, CgroupAttachMode, FdLink, Link, ProgAttachLink,
-        ProgramData, ProgramError,
+        define_link_wrapper, id_as_key, load_program, CgroupAttachMode, FdLink, Link,
+        ProgAttachLink, ProgramData, ProgramError,
     },
     sys::{bpf_link_create, LinkTarget, SyscallError},
     util::KernelVersion,
@@ -172,6 +172,8 @@ impl Link for CgroupSkbLinkInner {
     }
 }
 
+id_as_key!(CgroupSkbLinkInner, CgroupSkbLinkIdInner);
+
 define_link_wrapper!(
     /// The link used by [CgroupSkb] programs.
     CgroupSkbLink,

aya/src/programs/cgroup_sock.rs

@@ -7,8 +7,8 @@ pub use aya_obj::programs::CgroupSockAttachType;
 use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_CGROUP_SOCK,
     programs::{
-        define_link_wrapper, load_program, CgroupAttachMode, FdLink, Link, ProgAttachLink,
-        ProgramData, ProgramError,
+        define_link_wrapper, id_as_key, load_program, CgroupAttachMode, FdLink, Link,
+        ProgAttachLink, ProgramData, ProgramError,
     },
     sys::{bpf_link_create, LinkTarget, SyscallError},
     util::KernelVersion,
@@ -147,6 +147,8 @@ impl Link for CgroupSockLinkInner {
     }
 }
 
+id_as_key!(CgroupSockLinkInner, CgroupSockLinkIdInner);
+
 define_link_wrapper!(
     /// The link used by [CgroupSock] programs.
     CgroupSockLink,

aya/src/programs/cgroup_sock_addr.rs

@@ -7,8 +7,8 @@ pub use aya_obj::programs::CgroupSockAddrAttachType;
 use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_CGROUP_SOCK_ADDR,
     programs::{
-        define_link_wrapper, load_program, CgroupAttachMode, FdLink, Link, ProgAttachLink,
-        ProgramData, ProgramError,
+        define_link_wrapper, id_as_key, load_program, CgroupAttachMode, FdLink, Link,
+        ProgAttachLink, ProgramData, ProgramError,
     },
     sys::{bpf_link_create, LinkTarget, SyscallError},
     util::KernelVersion,
@@ -148,6 +148,8 @@ impl Link for CgroupSockAddrLinkInner {
     }
 }
 
+id_as_key!(CgroupSockAddrLinkInner, CgroupSockAddrLinkIdInner);
+
 define_link_wrapper!(
     /// The link used by [CgroupSockAddr] programs.
     CgroupSockAddrLink,

aya/src/programs/cgroup_sockopt.rs

@@ -7,8 +7,8 @@ pub use aya_obj::programs::CgroupSockoptAttachType;
 use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_CGROUP_SOCKOPT,
     programs::{
-        define_link_wrapper, load_program, CgroupAttachMode, FdLink, Link, ProgAttachLink,
-        ProgramData, ProgramError,
+        define_link_wrapper, id_as_key, load_program, CgroupAttachMode, FdLink, Link,
+        ProgAttachLink, ProgramData, ProgramError,
     },
     sys::{bpf_link_create, LinkTarget, SyscallError},
     util::KernelVersion,
@@ -147,6 +147,8 @@ impl Link for CgroupSockoptLinkInner {
     }
 }
 
+id_as_key!(CgroupSockoptLinkInner, CgroupSockoptLinkIdInner);
+
 define_link_wrapper!(
     /// The link used by [CgroupSockopt] programs.
     CgroupSockoptLink,

aya/src/programs/cgroup_sysctl.rs

@@ -5,8 +5,8 @@ use std::{hash::Hash, os::fd::AsFd};
 use crate::{
     generated::{bpf_attach_type::BPF_CGROUP_SYSCTL, bpf_prog_type::BPF_PROG_TYPE_CGROUP_SYSCTL},
     programs::{
-        define_link_wrapper, load_program, CgroupAttachMode, FdLink, Link, ProgAttachLink,
-        ProgramData, ProgramError,
+        define_link_wrapper, id_as_key, load_program, CgroupAttachMode, FdLink, Link,
+        ProgAttachLink, ProgramData, ProgramError,
     },
     sys::{bpf_link_create, LinkTarget, SyscallError},
     util::KernelVersion,
@@ -128,6 +128,8 @@ impl Link for CgroupSysctlLinkInner {
     }
 }
 
+id_as_key!(CgroupSysctlLinkInner, CgroupSysctlLinkIdInner);
+
 define_link_wrapper!(
     /// The link used by [CgroupSysctl] programs.
     CgroupSysctlLink,

aya/src/programs/links.rs

@@ -1,12 +1,12 @@
 //! Program links.
 use std::{
-    collections::{hash_map::Entry, HashMap},
     ffi::CString,
     io,
     os::fd::{AsFd as _, AsRawFd as _, BorrowedFd, RawFd},
     path::{Path, PathBuf},
 };
 
+use hashbrown::hash_set::{Entry, HashSet};
 use thiserror::Error;
 
 use crate::{
@@ -20,9 +20,9 @@ use crate::{
 };
 
 /// A Link.
-pub trait Link: std::fmt::Debug + 'static {
+pub trait Link: std::fmt::Debug + Eq + std::hash::Hash + 'static {
     /// Unique Id
-    type Id: std::fmt::Debug + std::hash::Hash + Eq + PartialEq;
+    type Id: std::fmt::Debug + Eq + std::hash::Hash + hashbrown::Equivalent<Self>;
 
     /// Returns the link id
     fn id(&self) -> Self::Id;
@@ -56,48 +56,50 @@ impl From<CgroupAttachMode> for u32 {
 }
 
 #[derive(Debug)]
-pub(crate) struct LinkMap<T: Link> {
-    links: HashMap<T::Id, T>,
+pub(crate) struct Links<T: Link> {
+    links: HashSet<T>,
 }
 
-impl<T: Link> LinkMap<T> {
+impl<T> Links<T>
+where
+    T: Eq + std::hash::Hash + Link,
+    T::Id: hashbrown::Equivalent<T> + Eq + std::hash::Hash,
+{
     pub(crate) fn new() -> Self {
         Self {
-            links: HashMap::new(),
+            links: Default::default(),
         }
     }
 
     pub(crate) fn insert(&mut self, link: T) -> Result<T::Id, ProgramError> {
-        let id = link.id();
-
-        match self.links.entry(link.id()) {
-            Entry::Occupied(_) => return Err(ProgramError::AlreadyAttached),
-            Entry::Vacant(e) => e.insert(link),
-        };
-
-        Ok(id)
+        match self.links.entry(link) {
+            Entry::Occupied(_entry) => Err(ProgramError::AlreadyAttached),
+            Entry::Vacant(entry) => Ok(entry.insert().get().id()),
+        }
     }
 
     pub(crate) fn remove(&mut self, link_id: T::Id) -> Result<(), ProgramError> {
         self.links
-            .remove(&link_id)
+            .take(&link_id)
             .ok_or(ProgramError::NotAttached)?
             .detach()
     }
 
+    pub(crate) fn forget(&mut self, link_id: T::Id) -> Result<T, ProgramError> {
+        self.links.take(&link_id).ok_or(ProgramError::NotAttached)
+    }
+}
+
+impl<T: Link> Links<T> {
     pub(crate) fn remove_all(&mut self) -> Result<(), ProgramError> {
-        for (_, link) in self.links.drain() {
+        for link in self.links.drain() {
             link.detach()?;
         }
         Ok(())
     }
-
-    pub(crate) fn forget(&mut self, link_id: T::Id) -> Result<T, ProgramError> {
-        self.links.remove(&link_id).ok_or(ProgramError::NotAttached)
-    }
 }
 
-impl<T: Link> Drop for LinkMap<T> {
+impl<T: Link> Drop for Links<T> {
     fn drop(&mut self) {
         let _ = self.remove_all();
     }
@@ -206,6 +208,8 @@ impl Link for FdLink {
     }
 }
 
+id_as_key!(FdLink, FdLinkId);
+
 impl From<PinnedLink> for FdLink {
     fn from(p: PinnedLink) -> Self {
         p.inner
@@ -321,6 +325,40 @@ impl Link for ProgAttachLink {
     }
 }
 
+id_as_key!(ProgAttachLink, ProgAttachLinkId);
+
+macro_rules! id_as_key {
+    ($wrapper:ident, $wrapper_id:ident) => {
+        impl PartialEq for $wrapper {
+            fn eq(&self, other: &Self) -> bool {
+                use $crate::programs::links::Link as _;
+
+                self.id() == other.id()
+            }
+        }
+
+        impl Eq for $wrapper {}
+
+        impl std::hash::Hash for $wrapper {
+            fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+                use $crate::programs::links::Link as _;
+
+                self.id().hash(state)
+            }
+        }
+
+        impl hashbrown::Equivalent<$wrapper> for $wrapper_id {
+            fn equivalent(&self, key: &$wrapper) -> bool {
+                use $crate::programs::links::Link as _;
+
+                *self == key.id()
+            }
+        }
+    };
+}
+
+pub(crate) use id_as_key;
+
 macro_rules! define_link_wrapper {
     (#[$doc1:meta] $wrapper:ident, #[$doc2:meta] $wrapper_id:ident, $base:ident, $base_id:ident, $program:ident,) => {
         #[$doc2]
@@ -350,7 +388,7 @@ macro_rules! define_link_wrapper {
 
         impl Drop for $wrapper {
             fn drop(&mut self) {
-                use crate::programs::links::Link;
+                use $crate::programs::links::Link as _;
 
                 if let Some(base) = self.0.take() {
                     let _ = base.detach();
@@ -370,6 +408,8 @@ macro_rules! define_link_wrapper {
             }
         }
 
+        $crate::programs::links::id_as_key!($wrapper, $wrapper_id);
+
         impl From<$base> for $wrapper {
             fn from(b: $base) -> $wrapper {
                 $wrapper(Some(b))
@@ -540,7 +580,7 @@ mod tests {
     use assert_matches::assert_matches;
     use tempfile::tempdir;
 
-    use super::{FdLink, Link, LinkMap};
+    use super::{FdLink, Link, Links};
     use crate::{
         generated::{BPF_F_ALLOW_MULTI, BPF_F_ALLOW_OVERRIDE},
         programs::{CgroupAttachMode, ProgramError},
@@ -578,9 +618,11 @@ mod tests {
         }
     }
 
+    id_as_key!(TestLink, TestLinkId);
+
     #[test]
     fn test_link_map() {
-        let mut links = LinkMap::new();
+        let mut links = Links::new();
         let l1 = TestLink::new(1, 2);
         let l1_detached = Rc::clone(&l1.detached);
         let l2 = TestLink::new(1, 3);
@@ -603,7 +645,7 @@ mod tests {
 
     #[test]
     fn test_already_attached() {
-        let mut links = LinkMap::new();
+        let mut links = Links::new();
 
         links.insert(TestLink::new(1, 2)).unwrap();
         assert_matches!(
@@ -614,7 +656,7 @@ mod tests {
 
     #[test]
     fn test_not_attached() {
-        let mut links = LinkMap::new();
+        let mut links = Links::new();
 
         let l1 = TestLink::new(1, 2);
         let l1_id1 = l1.id();
@@ -632,7 +674,7 @@ mod tests {
         let l2_detached = Rc::clone(&l2.detached);
 
         {
-            let mut links = LinkMap::new();
+            let mut links = Links::new();
             let id1 = links.insert(l1).unwrap();
             links.insert(l2).unwrap();
             // manually remove one link
@@ -653,7 +695,7 @@ mod tests {
         let l2_detached = Rc::clone(&l2.detached);
 
         let owned_l1 = {
-            let mut links = LinkMap::new();
+            let mut links = Links::new();
             let id1 = links.insert(l1).unwrap();
             links.insert(l2).unwrap();
             // manually forget one link

aya/src/programs/lirc_mode2.rs

@@ -4,8 +4,8 @@ use std::os::fd::{AsFd, AsRawFd as _, RawFd};
 use crate::{
     generated::{bpf_attach_type::BPF_LIRC_MODE2, bpf_prog_type::BPF_PROG_TYPE_LIRC_MODE2},
     programs::{
-        load_program, query, CgroupAttachMode, Link, ProgramData, ProgramError, ProgramFd,
-        ProgramInfo,
+        id_as_key, load_program, query, CgroupAttachMode, Link, ProgramData, ProgramError,
+        ProgramFd, ProgramInfo,
     },
     sys::{bpf_prog_attach, bpf_prog_detach, bpf_prog_get_fd_by_id, ProgQueryTarget},
 };
@@ -153,3 +153,5 @@ impl Link for LircLink {
             .map_err(Into::into)
     }
 }
+
+id_as_key!(LircLink, LircLinkId);

aya/src/programs/mod.rs

@@ -487,7 +487,7 @@ pub(crate) struct ProgramData<T: Link> {
     pub(crate) name: Option<String>,
     pub(crate) obj: Option<(obj::Program, obj::Function)>,
     pub(crate) fd: Option<ProgramFd>,
-    pub(crate) links: LinkMap<T>,
+    pub(crate) links: Links<T>,
     pub(crate) expected_attach_type: Option<bpf_attach_type>,
     pub(crate) attach_btf_obj_fd: Option<crate::MockableFd>,
     pub(crate) attach_btf_id: Option<u32>,
@@ -509,7 +509,7 @@ impl<T: Link> ProgramData<T> {
             name,
             obj: Some(obj),
             fd: None,
-            links: LinkMap::new(),
+            links: Links::new(),
             expected_attach_type: None,
             attach_btf_obj_fd: None,
             attach_btf_id: None,
@@ -541,7 +541,7 @@ impl<T: Link> ProgramData<T> {
             name,
             obj: None,
             fd: Some(ProgramFd(fd)),
-            links: LinkMap::new(),
+            links: Links::new(),
             expected_attach_type: None,
             attach_btf_obj_fd,
             attach_btf_id,