Skip to content

Commit

Permalink
add: support for bpf_map_update_batch method
Browse files Browse the repository at this point in the history
Signed-off-by: yogaraj.s <yogaraj.s@zohocorp.com>
  • Loading branch information
yogaraj.s authored and danielocfb committed Dec 5, 2023
1 parent 4ef1505 commit 071d80d
Show file tree
Hide file tree
Showing 2 changed files with 138 additions and 0 deletions.
50 changes: 50 additions & 0 deletions libbpf-rs/src/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -765,6 +765,56 @@ impl MapHandle {
self.update_raw(key, value, flags)
}

/// Updates many elements in batch mode in the map
///
/// `keys` must have exactly [`MapHandle::key_size()` * count] elements. `value` must have exactly
/// [`MapHandle::key_size()` * count] elements
pub fn update_batch(
&self,
keys: &[u8],
values: &[u8],
count: u32,
elem_flags: MapFlags,
flags: MapFlags,
) -> Result<()> {
if keys.len() as u32 / count != self.key_size() || (keys.len() as u32) % count != 0 {
return Err(Error::with_invalid_data(format!(
"batch key_size {} != {} * {}",
keys.len(),
self.key_size(),
count
)));
};

if values.len() as u32 / count != self.value_size() || (values.len() as u32) % count != 0 {
return Err(Error::with_invalid_data(format!(
"batch value_size {} != {} * {}",
values.len(),
self.value_size(),
count
)));
}

let opts = libbpf_sys::bpf_map_batch_opts {
sz: mem::size_of::<libbpf_sys::bpf_map_batch_opts>() as _,
elem_flags: elem_flags.bits(),
flags: flags.bits(),
};

let mut count = count;
let ret = unsafe {
libbpf_sys::bpf_map_update_batch(
self.fd.as_raw_fd(),
keys.as_ptr() as *const c_void,
values.as_ptr() as *const c_void,
(&mut count) as *mut u32,
&opts as *const libbpf_sys::bpf_map_batch_opts,
)
};

util::parse_ret(ret)
}

/// Update an element in an per-cpu map with one value per cpu.
///
/// `key` must have exactly [`MapHandle::key_size()`] elements. `value` must have one
Expand Down
88 changes: 88 additions & 0 deletions libbpf-rs/tests/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,94 @@ fn test_sudo_object_map_key_value_size() {
.is_err());
}

#[test]
fn test_sudo_object_map_update_batch() {
bump_rlimit_mlock();

let mut obj = get_test_object("runqslower.bpf.o");
let start = obj.map_mut("start").expect("failed to find map");

let key1 = 1u32.to_ne_bytes();
let key2 = 2u32.to_ne_bytes();
let key3 = 3u32.to_ne_bytes();
let key4 = 4u32.to_ne_bytes();

let value1 = 369u64.to_ne_bytes();
let value2 = 258u64.to_ne_bytes();
let value3 = 147u64.to_ne_bytes();
let value4 = 159u64.to_ne_bytes();

let batch_key1 = key1.into_iter().chain(key2).collect::<Vec<_>>();
let batch_value1 = value1.into_iter().chain(value2).collect::<Vec<_>>();

let batch_key2 = key2.into_iter().chain(key3).chain(key4).collect::<Vec<_>>();
let batch_value2 = value2
.into_iter()
.chain(value3)
.chain(value4)
.collect::<Vec<_>>();

// Update batch with wrong key size
assert!(start
.update_batch(
&[1, 2, 3],
&batch_value1,
2,
MapFlags::ANY,
MapFlags::NO_EXIST
)
.is_err());

// Update batch with wrong value size
assert!(start
.update_batch(
&batch_key1,
&[1, 2, 3],
2,
MapFlags::ANY,
MapFlags::NO_EXIST
)
.is_err());

// Update batch with wrong count.
assert!(start
.update_batch(
&batch_key1,
&batch_value1,
1,
MapFlags::ANY,
MapFlags::NO_EXIST
)
.is_err());

// Update batch with 1 key.
assert!(start
.update_batch(&key1, &value1, 1, MapFlags::ANY, MapFlags::NO_EXIST)
.is_ok());

// Update batch with multiple keys.
assert!(start
.update_batch(
&batch_key2,
&batch_value2,
3,
MapFlags::ANY,
MapFlags::NO_EXIST
)
.is_ok());

// Update batch with existing keys.
assert!(start
.update_batch(
&batch_key2,
&batch_value2,
3,
MapFlags::NO_EXIST,
MapFlags::NO_EXIST
)
.is_err());
}

#[test]
fn test_sudo_object_map_delete_batch() {
bump_rlimit_mlock();
Expand Down

0 comments on commit 071d80d

Please sign in to comment.