map: Add lookup_into() method

Allows looking up into a preallocated buffer.

Author: Jongy

libbpf-rs/CHANGELOG.md

@@ -12,6 +12,8 @@ Unreleased
 - Added `Program::{stdout,stderr}` for accessing BPF stdout and stderr
   streams
 - Added `OpenProgram::autoload` getter
+- Added `MapCore::lookup_into` for looking up values into preallocated
+  buffers
 
 
 0.26.0-beta.0

libbpf-rs/src/map.rs

@@ -311,8 +311,14 @@ where
     key.as_ptr() as *const c_void
 }
 
-/// Internal function to return a value from a map into a buffer of the given size.
-fn lookup_raw<M>(map: &M, key: &[u8], flags: MapFlags, out_size: usize) -> Result<Option<Vec<u8>>>
+/// Internal function to perform a map lookup and write the value into raw pointer.
+/// Returns `Ok(true)` if the key was found, `Ok(false)` if not found, or an error.
+fn lookup_raw<M>(
+    map: &M,
+    key: &[u8],
+    value: &mut [mem::MaybeUninit<u8>],
+    flags: MapFlags,
+) -> Result<bool>
 where
     M: MapCore + ?Sized,
 {
@@ -322,34 +328,65 @@ where
             key.len(),
             map.key_size()
         )));
-    };
+    }
 
-    let mut out: Vec<u8> = Vec::with_capacity(out_size);
+    // Make sure the internal users of this function pass the expected buffer size
+    debug_assert_eq!(
+        value.len(),
+        if map.map_type().is_percpu() {
+            percpu_buffer_size(map).unwrap()
+        } else {
+            map.value_size() as usize
+        }
+    );
 
     let ret = unsafe {
         libbpf_sys::bpf_map_lookup_elem_flags(
             map.as_fd().as_raw_fd(),
             map_key(map, key),
-            out.as_mut_ptr() as *mut c_void,
+            // TODO: Use `MaybeUninit::slice_as_mut_ptr` once stable.
+            value.as_mut_ptr().cast(),
             flags.bits(),
         )
     };
 
     if ret == 0 {
-        unsafe {
-            out.set_len(out_size);
-        }
-        Ok(Some(out))
+        Ok(true)
     } else {
         let err = io::Error::last_os_error();
         if err.kind() == io::ErrorKind::NotFound {
-            Ok(None)
+            Ok(false)
         } else {
             Err(Error::from(err))
         }
     }
 }
 
+/// Internal function to return a value from a map into a buffer of the given size.
+fn lookup_raw_vec<M>(
+    map: &M,
+    key: &[u8],
+    flags: MapFlags,
+    out_size: usize,
+) -> Result<Option<Vec<u8>>>
+where
+    M: MapCore + ?Sized,
+{
+    // Allocate without initializing (avoiding memset)
+    let mut out = Vec::with_capacity(out_size);
+
+    match lookup_raw(map, key, out.spare_capacity_mut(), flags)? {
+        true => {
+            // SAFETY: `lookup_raw` successfully filled the buffer
+            unsafe {
+                out.set_len(out_size);
+            }
+            Ok(Some(out))
+        }
+        false => Ok(None),
+    }
+}
+
 /// Internal function to update a map. This does not check the length of the
 /// supplied value.
 fn update_raw<M>(map: &M, key: &[u8], value: &[u8], flags: MapFlags) -> Result<()>
@@ -480,7 +517,40 @@ pub trait MapCore: Debug + AsFd + private::Sealed {
     fn lookup(&self, key: &[u8], flags: MapFlags) -> Result<Option<Vec<u8>>> {
         check_not_bloom_or_percpu(self)?;
         let out_size = self.value_size() as usize;
-        lookup_raw(self, key, flags, out_size)
+        lookup_raw_vec(self, key, flags, out_size)
+    }
+
+    /// Looks up a map value into a pre-allocated buffer, avoiding allocation.
+    ///
+    /// This method provides a zero-allocation alternative to [`Self::lookup()`].
+    ///
+    /// `key` must have exactly [`Self::key_size()`] elements.
+    /// `value` must have exactly [`Self::value_size()`] elements.
+    ///
+    /// Returns `Ok(true)` if the key was found and the buffer was filled,
+    /// `Ok(false)` if the key was not found, or an error.
+    ///
+    /// If the map is one of the per-cpu data structures, this function cannot be used.
+    /// If the map is of type `bloom_filter`, this function cannot be used.
+    fn lookup_into(&self, key: &[u8], value: &mut [u8], flags: MapFlags) -> Result<bool> {
+        check_not_bloom_or_percpu(self)?;
+
+        if value.len() != self.value_size() as usize {
+            return Err(Error::with_invalid_data(format!(
+                "value buffer size {} != {}",
+                value.len(),
+                self.value_size()
+            )));
+        }
+
+        // SAFETY: `u8` and `MaybeUninit<u8>` have the same in-memory representation.
+        let value = unsafe {
+            slice::from_raw_parts_mut::<mem::MaybeUninit<u8>>(
+                value.as_mut_ptr().cast(),
+                value.len(),
+            )
+        };
+        lookup_raw(self, key, value, flags)
     }
 
     /// Returns many elements in batch mode from the map.
@@ -548,7 +618,7 @@ pub trait MapCore: Debug + AsFd + private::Sealed {
         let aligned_val_size = percpu_aligned_value_size(self);
         let out_size = percpu_buffer_size(self)?;
 
-        let raw_res = lookup_raw(self, key, flags, out_size)?;
+        let raw_res = lookup_raw_vec(self, key, flags, out_size)?;
         if let Some(raw_vals) = raw_res {
             let mut out = Vec::new();
             for chunk in raw_vals.chunks_exact(aligned_val_size) {

libbpf-rs/tests/test.rs

@@ -667,6 +667,90 @@ fn test_object_map_mutation() {
         .is_none());
 }
 
+#[tag(root)]
+#[test]
+fn test_object_map_lookup_into() {
+    let mut obj = get_test_object("runqslower.bpf.o");
+    let start = get_map_mut(&mut obj, "start");
+
+    // Insert a test value
+    start
+        .update(&[1, 2, 3, 4], &[1, 2, 3, 4, 5, 6, 7, 8], MapFlags::empty())
+        .expect("failed to write");
+
+    // Test successful lookup with pre-allocated buffer
+    let mut value = [0u8; 8];
+    let found = start
+        .lookup_into(&[1, 2, 3, 4], &mut value, MapFlags::empty())
+        .expect("failed to lookup_into");
+
+    assert!(found, "key should be found");
+    assert_eq!(value, [1, 2, 3, 4, 5, 6, 7, 8]);
+
+    // Test lookup of non-existent key
+    let mut value2 = [0u8; 8];
+    let found2 = start
+        .lookup_into(&[5, 6, 7, 8], &mut value2, MapFlags::empty())
+        .expect("failed to lookup_into for non-existent key");
+
+    assert!(!found2, "key should not be found");
+    // Buffer should remain unchanged when key is not found
+    assert_eq!(value2, [0u8; 8]);
+}
+
+#[tag(root)]
+#[test]
+fn test_object_map_lookup_into_wrong_size() {
+    let mut obj = get_test_object("runqslower.bpf.o");
+    let start = get_map_mut(&mut obj, "start");
+
+    // Insert a test value
+    start
+        .update(&[1, 2, 3, 4], &[1, 2, 3, 4, 5, 6, 7, 8], MapFlags::empty())
+        .expect("failed to write");
+
+    // Test with wrong buffer size (too small)
+    let mut value_small = [0u8; 4];
+    let result = start.lookup_into(&[1, 2, 3, 4], &mut value_small, MapFlags::empty());
+    assert!(result.is_err(), "should fail with wrong buffer size");
+
+    // Test with wrong buffer size (too large)
+    let mut value_large = [0u8; 16];
+    let result = start.lookup_into(&[1, 2, 3, 4], &mut value_large, MapFlags::empty());
+    assert!(result.is_err(), "should fail with wrong buffer size");
+}
+
+#[tag(root)]
+#[test]
+fn test_object_map_lookup_into_consistency() {
+    let mut obj = get_test_object("runqslower.bpf.o");
+    let start = get_map_mut(&mut obj, "start");
+
+    // Insert a test value
+    let test_value = [10, 20, 30, 40, 50, 60, 70, 80];
+    start
+        .update(&[1, 2, 3, 4], &test_value, MapFlags::empty())
+        .expect("failed to write");
+
+    // Compare results from lookup() and lookup_into()
+    let lookup_result = start
+        .lookup(&[1, 2, 3, 4], MapFlags::empty())
+        .expect("failed to lookup")
+        .expect("key not found");
+
+    let mut value_buffer = [0u8; 8];
+    let found = start
+        .lookup_into(&[1, 2, 3, 4], &mut value_buffer, MapFlags::empty())
+        .expect("failed to lookup_into");
+
+    assert!(found, "key should be found");
+    assert_eq!(
+        lookup_result.as_slice(),
+        &value_buffer,
+        "lookup() and lookup_into() should return the same value"
+    );
+}
+
 #[tag(root)]
 #[test]
 fn test_object_map_lookup_flags() {