libbpf-rs: add support for bpf_map_update_batch method

libbpf-rs/src/map.rs

@@ -766,6 +766,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

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::<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();