Avoid reallocation of multi-cpu vectors during iteration.

This commit is contained in:
Herbert Wolverson 2023-04-19 16:38:39 +00:00
parent cf7ccb7954
commit 0c62268ed6

View File

@ -29,7 +29,9 @@ struct BpfMapIterator<KEY, VALUE> {
unsafe impl<KEY, VALUE> Sync for BpfMapIterator<KEY, VALUE> {} unsafe impl<KEY, VALUE> Sync for BpfMapIterator<KEY, VALUE> {}
unsafe impl<KEY, VALUE> Send for BpfMapIterator<KEY, VALUE> {} unsafe impl<KEY, VALUE> Send for BpfMapIterator<KEY, VALUE> {}
impl<KEY, VALUE> BpfMapIterator<KEY, VALUE> { impl<KEY, VALUE> BpfMapIterator<KEY, VALUE>
where VALUE: FromBytes + Debug + Clone + Default
{
/// Create a new link to an eBPF map, that *must* have an iterator /// Create a new link to an eBPF map, that *must* have an iterator
/// function defined in the eBPF program - and exposed in the /// function defined in the eBPF program - and exposed in the
/// skeleton. /// skeleton.
@ -107,10 +109,13 @@ pub(crate) struct BpfMapIter<K, V> {
buffer: Vec<u8>, buffer: Vec<u8>,
index: usize, index: usize,
_phantom: PhantomData<(K, V)>, _phantom: PhantomData<(K, V)>,
num_cpus: u32, num_cpus: usize,
row_data: Vec<V>,
} }
impl<K, V> BpfMapIter<K, V> { impl<K, V> BpfMapIter<K, V>
where V: FromBytes + Debug + Clone + Default
{
const KEY_SIZE: usize = std::mem::size_of::<K>(); const KEY_SIZE: usize = std::mem::size_of::<K>();
const VALUE_SIZE: usize = std::mem::size_of::<V>(); const VALUE_SIZE: usize = std::mem::size_of::<V>();
const TOTAL_SIZE: usize = Self::KEY_SIZE + Self::VALUE_SIZE; const TOTAL_SIZE: usize = Self::KEY_SIZE + Self::VALUE_SIZE;
@ -120,7 +125,7 @@ impl<K, V> BpfMapIter<K, V> {
/// throughout. /// throughout.
fn new(buffer: Vec<u8>) -> Self { fn new(buffer: Vec<u8>) -> Self {
let first_four : [u8; 4] = [buffer[0], buffer[1], buffer[2], buffer[3]]; let first_four : [u8; 4] = [buffer[0], buffer[1], buffer[2], buffer[3]];
let num_cpus = u32::from_ne_bytes(first_four); let num_cpus = u32::from_ne_bytes(first_four) as usize;
//println!("CPUs: {num_cpus}"); //println!("CPUs: {num_cpus}");
Self { Self {
@ -128,14 +133,15 @@ impl<K, V> BpfMapIter<K, V> {
index: std::mem::size_of::<i32>(), index: std::mem::size_of::<i32>(),
_phantom: PhantomData, _phantom: PhantomData,
num_cpus, num_cpus,
row_data: vec![V::default(); num_cpus],
} }
} }
} }
impl<K, V> Iterator for BpfMapIter<K, V> impl<K, V> Iterator for BpfMapIter<K, V>
where where
K: FromBytes + Debug, K: FromBytes + Debug + Clone,
V: FromBytes + Debug, V: FromBytes + Debug + Clone + Default,
{ {
type Item = (K, Vec<V>); type Item = (K, Vec<V>);
@ -143,17 +149,16 @@ where
if self.index + Self::TOTAL_SIZE <= self.buffer.len() { if self.index + Self::TOTAL_SIZE <= self.buffer.len() {
let key = K::read_from(&self.buffer[self.index..self.index + Self::KEY_SIZE]); let key = K::read_from(&self.buffer[self.index..self.index + Self::KEY_SIZE]);
self.index += Self::KEY_SIZE; self.index += Self::KEY_SIZE;
let mut vals = Vec::new(); for cpu in 0..self.num_cpus {
for _ in 0..self.num_cpus {
let value = V::read_from( let value = V::read_from(
&self.buffer &self.buffer
[self.index ..self.index + Self::VALUE_SIZE], [self.index ..self.index + Self::VALUE_SIZE],
); );
vals.push(value.unwrap()); self.row_data[cpu] = value.unwrap();
self.index += Self::VALUE_SIZE; self.index += Self::VALUE_SIZE;
} }
//println!("{key:?} {vals:?}"); //println!("{key:?} {vals:?}");
Some((key.unwrap(), vals)) Some((key.unwrap(), self.row_data.clone()))
} else { } else {
None None
} }