Skip to content

Commit

Permalink
feedback
Browse files Browse the repository at this point in the history
  • Loading branch information
@lbrndnr
lbrndnr committed Nov 8, 2024
1 parent 72639de commit 3903f79
Show file tree
Hide file tree
Showing 2 changed files with 137 additions and 162 deletions.
205 changes: 106 additions & 99 deletions libbpf-rs/src/map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -369,7 +370,7 @@ fn lookup_batch_raw<M>(
elem_flags: MapFlags,
flags: MapFlags,
delete: bool,
) -> Result<BatchedMapIter<'_>>
) -> BatchedMapIter<'_>
where
M: MapCore + ?Sized,
{
Expand All @@ -382,21 +383,35 @@ where
..Default::default()
};

let key_size = match map.map_type() {
MapType::Hash | MapType::PercpuHash | MapType::LruHash | MapType::LruPercpuHash => {
map.key_size().max(4)
}
_ => map.key_size(),
// 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()
};

Ok(BatchedMapIter::new(
map.as_fd(),
count,
key_size,
map.value_size(),
opts,
delete,
))
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<M>(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)]
Expand Down Expand Up @@ -446,18 +461,7 @@ 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<Option<Vec<u8>>> {
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)
}
Expand All @@ -471,19 +475,8 @@ pub trait MapCore: Debug + AsFd + private::Sealed {
elem_flags: MapFlags,
flags: MapFlags,
) -> Result<BatchedMapIter<'_>> {
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(),
)));
}

lookup_batch_raw(self, count, elem_flags, flags, false)
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.
Expand All @@ -495,19 +488,8 @@ pub trait MapCore: Debug + AsFd + private::Sealed {
elem_flags: MapFlags,
flags: MapFlags,
) -> Result<BatchedMapIter<'_>> {
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(),
)));
}

lookup_batch_raw(self, count, elem_flags, flags, true)
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`.
Expand Down Expand Up @@ -1253,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 {
Expand Down Expand Up @@ -1370,16 +1360,16 @@ impl Iterator for MapKeyIter<'_> {
#[derive(Debug)]
pub struct BatchedMapIter<'map> {
map_fd: BorrowedFd<'map>,
count: u32,
key_size: u32,
value_size: u32,
delete: bool,
count: usize,
key_size: usize,
value_size: usize,
keys: Vec<u8>,
values: Vec<u8>,
prev: Option<Vec<u8>>,
next: Vec<u8>,
batch_opts: libbpf_sys::bpf_map_batch_opts,
done: bool,
delete: bool,
index: Option<usize>,
}

impl<'map> BatchedMapIter<'map> {
Expand All @@ -1393,82 +1383,99 @@ impl<'map> BatchedMapIter<'map> {
) -> Self {
Self {
map_fd,
count,
key_size,
value_size,
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,
done: false,
delete,
index: None,
}
}
}

impl Iterator for BatchedMapIter<'_> {
type Item = (Vec<Vec<u8>>, Vec<Vec<u8>>);

fn next(&mut self) -> Option<Self::Item> {
if self.done {
return 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;
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() as _,
self.keys.as_mut_ptr() as *mut c_void,
self.values.as_mut_ptr() as *mut c_void,
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() as _,
self.keys.as_mut_ptr() as *mut c_void,
self.values.as_mut_ptr() as *mut c_void,
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 ret == -14 || count == 0 {
None
} else {
self.prev = Some(self.next.clone());
if ret != 0 {
self.done = true;
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;
}
}
}

unsafe {
self.keys.set_len((self.key_size * count) as usize);
self.values.set_len((self.value_size * count) as usize);
}
self.prev = Some(self.next.clone());
self.index = Some(0);

let keys = self
.keys
.chunks_exact(self.key_size as usize)
.map(|c| c.to_vec())
.collect();
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<u8>, Vec<u8>);

let values = self
.values
.chunks_exact(self.value_size as usize)
.map(|c| c.to_vec())
.collect();
fn next(&mut self) -> Option<Self::Item> {
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,
};

Some((keys, values))
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))
}
}

Expand Down
Loading

0 comments on commit 3903f79

Please sign in to comment.