diff --git a/libbpf-rs/CHANGELOG.md b/libbpf-rs/CHANGELOG.md index 50c928f5..056a1269 100644 --- a/libbpf-rs/CHANGELOG.md +++ b/libbpf-rs/CHANGELOG.md @@ -2,6 +2,7 @@ ------ - Added `Program::attach_netfilter_with_opts` for attaching to netfilter hooks +- Added `Map::lookup_batch` and `Map::lookup_and_delete_batch` method 0.24.5 diff --git a/libbpf-rs/src/map.rs b/libbpf-rs/src/map.rs index 0fc4ca67..65f59e11 100644 --- a/libbpf-rs/src/map.rs +++ b/libbpf-rs/src/map.rs @@ -27,6 +27,7 @@ use bitflags::bitflags; use libbpf_sys::bpf_map_info; use libbpf_sys::bpf_obj_get_info_by_fd; +use crate::error; use crate::util; use crate::util::parse_ret_i32; use crate::util::validate_bpf_ret; @@ -42,7 +43,6 @@ pub type OpenMap<'obj> = OpenMapImpl<'obj>; /// A mutable parsed but not yet loaded BPF map. pub type OpenMapMut<'obj> = OpenMapImpl<'obj, Mut>; - /// Represents a parsed but not yet loaded BPF map. /// /// This object exposes operations that need to happen before the map is created. @@ -363,6 +363,57 @@ where util::parse_ret(ret) } +/// Internal function to batch lookup (and delete) elements from a map. +fn lookup_batch_raw( + map: &M, + count: u32, + elem_flags: MapFlags, + flags: MapFlags, + delete: bool, +) -> BatchedMapIter<'_> +where + M: MapCore + ?Sized, +{ + #[allow(clippy::needless_update)] + let opts = libbpf_sys::bpf_map_batch_opts { + sz: mem::size_of::() as _, + elem_flags: elem_flags.bits(), + flags: flags.bits(), + // bpf_map_batch_opts might have padding fields on some platform + ..Default::default() + }; + + // for maps of type BPF_MAP_TYPE_{HASH, PERCPU_HASH, LRU_HASH, LRU_PERCPU_HASH} + // the key size must be at least 4 bytes + let key_size = if map.map_type().is_hash_map() { + map.key_size().max(4) + } else { + map.key_size() + }; + + BatchedMapIter::new(map.as_fd(), count, key_size, map.value_size(), opts, delete) +} + +/// Intneral function that returns an error for per-cpu and bloom filter maps. +fn check_not_bloom_or_percpu(map: &M) -> Result<()> +where + M: MapCore + ?Sized, +{ + if map.map_type().is_bloom_filter() { + return Err(Error::with_invalid_data( + "lookup_bloom_filter() must be used for bloom filter maps", + )); + } + if map.map_type().is_percpu() { + return Err(Error::with_invalid_data(format!( + "lookup_percpu() must be used for per-cpu maps (type of the map is {:?})", + map.map_type(), + ))); + } + + Ok(()) +} + #[allow(clippy::wildcard_imports)] mod private { use super::*; @@ -410,22 +461,37 @@ pub trait MapCore: Debug + AsFd + private::Sealed { /// must be used. /// If the map is of type bloom_filter the function [`Self::lookup_bloom_filter()`] must be used fn lookup(&self, key: &[u8], flags: MapFlags) -> Result>> { - if self.map_type().is_bloom_filter() { - return Err(Error::with_invalid_data( - "lookup_bloom_filter() must be used for bloom filter maps", - )); - } - if self.map_type().is_percpu() { - return Err(Error::with_invalid_data(format!( - "lookup_percpu() must be used for per-cpu maps (type of the map is {:?})", - self.map_type(), - ))); - } - + check_not_bloom_or_percpu(self)?; let out_size = self.value_size() as usize; lookup_raw(self, key, flags, out_size) } + /// Returns many elements in batch mode from the map. + /// + /// `count` specifies the batch size. + fn lookup_batch( + &self, + count: u32, + elem_flags: MapFlags, + flags: MapFlags, + ) -> Result> { + check_not_bloom_or_percpu(self)?; + Ok(lookup_batch_raw(self, count, elem_flags, flags, false)) + } + + /// Returns many elements in batch mode from the map. + /// + /// `count` specifies the batch size. + fn lookup_and_delete_batch( + &self, + count: u32, + elem_flags: MapFlags, + flags: MapFlags, + ) -> Result> { + check_not_bloom_or_percpu(self)?; + Ok(lookup_batch_raw(self, count, elem_flags, flags, true)) + } + /// Returns if the given value is likely present in bloom_filter as `bool`. /// /// `value` must have exactly [`Self::value_size()`] elements. @@ -1169,6 +1235,14 @@ impl MapType { ) } + /// Returns if the map is of one of the hashmap types. + pub fn is_hash_map(&self) -> bool { + matches!( + self, + MapType::Hash | MapType::PercpuHash | MapType::LruHash | MapType::LruPercpuHash + ) + } + /// Returns if the map is keyless map type as per documentation of libbpf /// Keyless map types are: Queues, Stacks and Bloom Filters fn is_keyless(&self) -> bool { @@ -1282,6 +1356,129 @@ impl Iterator for MapKeyIter<'_> { } } +/// An iterator over batches of key value pairs of a BPF map. +#[derive(Debug)] +pub struct BatchedMapIter<'map> { + map_fd: BorrowedFd<'map>, + delete: bool, + count: usize, + key_size: usize, + value_size: usize, + keys: Vec, + values: Vec, + prev: Option>, + next: Vec, + batch_opts: libbpf_sys::bpf_map_batch_opts, + index: Option, +} + +impl<'map> BatchedMapIter<'map> { + fn new( + map_fd: BorrowedFd<'map>, + count: u32, + key_size: u32, + value_size: u32, + batch_opts: libbpf_sys::bpf_map_batch_opts, + delete: bool, + ) -> Self { + Self { + map_fd, + delete, + count: count as usize, + key_size: key_size as usize, + value_size: value_size as usize, + keys: vec![0; (count * key_size) as usize], + values: vec![0; (count * value_size) as usize], + prev: None, + next: vec![0; key_size as usize], + batch_opts, + index: None, + } + } + + fn lookup_next_batch(&mut self) { + let prev = self.prev.as_ref().map_or(ptr::null(), |p| p.as_ptr()); + let mut count = self.count as u32; + + let ret = unsafe { + if self.delete { + libbpf_sys::bpf_map_lookup_and_delete_batch( + self.map_fd.as_raw_fd(), + prev as _, + self.next.as_mut_ptr().cast(), + self.keys.as_mut_ptr().cast(), + self.values.as_mut_ptr().cast(), + (&mut count) as *mut u32, + &self.batch_opts as *const libbpf_sys::bpf_map_batch_opts, + ) + } else { + libbpf_sys::bpf_map_lookup_batch( + self.map_fd.as_raw_fd(), + prev as _, + self.next.as_mut_ptr().cast(), + self.keys.as_mut_ptr().cast(), + self.values.as_mut_ptr().cast(), + (&mut count) as *mut u32, + &self.batch_opts as *const libbpf_sys::bpf_map_batch_opts, + ) + } + }; + + if let Err(e) = util::parse_ret(ret) { + match e.kind() { + // in this case we can trust the returned count value + error::ErrorKind::NotFound => {} + // retry with same input arguments + error::ErrorKind::Interrupted => { + return self.lookup_next_batch(); + } + _ => { + self.index = None; + return; + } + } + } + + self.prev = Some(self.next.clone()); + self.index = Some(0); + + unsafe { + self.keys.set_len(self.key_size * count as usize); + self.values.set_len(self.value_size * count as usize); + } + } +} + +impl Iterator for BatchedMapIter<'_> { + type Item = (Vec, Vec); + + fn next(&mut self) -> Option { + let load_next_batch = match self.index { + Some(index) => { + let batch_finished = index * self.key_size >= self.keys.len(); + let last_batch = self.keys.len() < self.key_size * self.count; + batch_finished && !last_batch + } + None => true, + }; + + if load_next_batch { + self.lookup_next_batch(); + } + + let index = self.index?; + let key = self.keys.chunks_exact(self.key_size).nth(index)?.to_vec(); + let val = self + .values + .chunks_exact(self.value_size) + .nth(index)? + .to_vec(); + + self.index = Some(index + 1); + Some((key, val)) + } +} + /// A convenience wrapper for [`bpf_map_info`][libbpf_sys::bpf_map_info]. It /// provides the ability to retrieve the details of a certain map. #[derive(Debug)] diff --git a/libbpf-rs/tests/test.rs b/libbpf-rs/tests/test.rs index 305f7c8e..c2dc6249 100644 --- a/libbpf-rs/tests/test.rs +++ b/libbpf-rs/tests/test.rs @@ -3,6 +3,7 @@ mod common; +use std::collections::HashMap; use std::collections::HashSet; use std::env::current_exe; use std::ffi::c_int; @@ -58,7 +59,6 @@ use crate::common::get_test_object_path; use crate::common::open_test_object; use crate::common::with_ringbuffer; - #[tag(root)] #[test] fn test_object_build_and_load() { @@ -253,6 +253,68 @@ fn test_object_map_update_batch() { .is_err()); } +#[tag(root)] +#[test] +fn test_object_map_lookup_batch() { + bump_rlimit_mlock(); + + let mut obj = get_test_object("runqslower.bpf.o"); + let start = get_map_mut(&mut obj, "start"); + let data = HashMap::from([ + (1u32, 9999u64), + (2u32, 42u64), + (3u32, 18u64), + (4u32, 1337u64), + ]); + + for (key, val) in data.iter() { + assert!(start + .update(&key.to_ne_bytes(), &val.to_ne_bytes(), MapFlags::ANY) + .is_ok()); + } + + let elems = start + .lookup_batch(2, MapFlags::ANY, MapFlags::ANY) + .expect("failed to lookup batch") + .collect::>(); + assert_eq!(elems.len(), 4); + + for (key, val) in elems.into_iter() { + let key = u32::from_ne_bytes(key.try_into().unwrap()); + let val = u64::from_ne_bytes(val.try_into().unwrap()); + assert_eq!(val, data[&key]); + } + + // test lookup with batch size larger than the number of keys + let elems = start + .lookup_batch(5, MapFlags::ANY, MapFlags::ANY) + .expect("failed to lookup batch") + .collect::>(); + assert_eq!(elems.len(), 4); + + for (key, val) in elems.into_iter() { + let key = u32::from_ne_bytes(key.try_into().unwrap()); + let val = u64::from_ne_bytes(val.try_into().unwrap()); + assert_eq!(val, data[&key]); + } + + // test lookup and delete with batch size that does not divide total count + let elems = start + .lookup_and_delete_batch(3, MapFlags::ANY, MapFlags::ANY) + .expect("failed to lookup batch") + .collect::>(); + assert_eq!(elems.len(), 4); + + for (key, val) in elems.into_iter() { + let key = u32::from_ne_bytes(key.try_into().unwrap()); + let val = u64::from_ne_bytes(val.try_into().unwrap()); + assert_eq!(val, data[&key]); + } + + // Map should be empty now. + assert!(start.keys().collect::>().is_empty()) +} + #[tag(root)] #[test] fn test_object_map_delete_batch() {