diff --git a/libbpf-rs/src/map.rs b/libbpf-rs/src/map.rs index 81c8761c..55b561c7 100644 --- a/libbpf-rs/src/map.rs +++ b/libbpf-rs/src/map.rs @@ -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::() 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 diff --git a/libbpf-rs/tests/test.rs b/libbpf-rs/tests/test.rs index d606e703..c25948a0 100644 --- a/libbpf-rs/tests/test.rs +++ b/libbpf-rs/tests/test.rs @@ -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::>(); + let batch_value1 = value1.into_iter().chain(value2).collect::>(); + + let batch_key2 = key2.into_iter().chain(key3).chain(key4).collect::>(); + let batch_value2 = value2 + .into_iter() + .chain(value3) + .chain(value4) + .collect::>(); + + // 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();