Add user ring buffer support to libbpf-rs

This change adds support for user ring buffers to libbpf-rs. The user
ring buffer is a type of BPF map that allows user space to write data
to a ring buffer that can be read by a BPF program.

The motivation for this change is to improve the performance of the
`scx_rustland` scheduler in the sched_ext project. We have verified the
change by running scx_rustland with the new user ring buffer support.

The `UserRingBuffer` struct is a wrapper around the `Map` type and
provides a new set of functions for interacting with the ring buffer.
The API closely resembles the libbpf API for user ring buffers. The
struct was intentionally made !Send to avoid introducing locking
overhead unless necessary. It is expected that the user will handle
synchronization if needed.

The `UserRingBufferSample` struct contains a mutable reference to the
sample in the ring buffer memory. It is the interface for writing data
to the ring buffer. To write to the sample, dereference with `as_mut()`
to get a mutable reference to the raw byte slice.

The `user_ring_buffer__discard` function was not exposed because the
implementation handles this internally to simplify the API.

The `user_ring_buffer__reserve_blocking` function is not yet implemented
and will be added in the future.

Link: https://lore.kernel.org/bpf/[email protected]/
Link: https://github.com/sched-ext/scx/tree/main/scheds/rust/scx_rustland
Signed-off-by: Jose Fernandez <[email protected]>
Reviewed-by: Andrea Righi <[email protected]>

Author: jfernandez

libbpf-rs/Cargo.toml

@@ -25,7 +25,7 @@ vendored = ["libbpf-sys/vendored"]
 
 [dependencies]
 bitflags = "2.0"
-libbpf-sys = { version = "1.4", default-features = false }
+libbpf-sys = { version = "1.4.1", default-features = false }
 libc = "0.2"
 num_enum = "0.5"
 strum_macros = "0.24"

libbpf-rs/src/lib.rs

@@ -92,6 +92,7 @@ pub mod query;
 mod ringbuf;
 mod skeleton;
 mod tc;
+mod user_ringbuf;
 mod util;
 mod xdp;
 
@@ -144,6 +145,7 @@ pub use crate::tc::TC_H_INGRESS;
 pub use crate::tc::TC_H_MIN_EGRESS;
 pub use crate::tc::TC_H_MIN_INGRESS;
 pub use crate::tc::TC_INGRESS;
+pub use crate::user_ringbuf::UserRingBuffer;
 pub use crate::util::num_possible_cpus;
 pub use crate::xdp::Xdp;
 pub use crate::xdp::XdpFlags;

libbpf-rs/src/map.rs

@@ -1022,6 +1022,7 @@ pub enum MapType {
     InodeStorage,
     TaskStorage,
     BloomFilter,
+    UserRingBuf,
     /// We choose to specify our own "unknown" type here b/c it's really up to the kernel
     /// to decide if it wants to reject the map. If it accepts it, it just means whoever
     /// using this library is a bit out of date.

libbpf-rs/src/user_ringbuf.rs

@@ -0,0 +1,170 @@
+use libc::E2BIG;
+use libc::ENOSPC;
+use std::ops::Deref;
+use std::ops::DerefMut;
+use std::os::fd::AsFd;
+use std::os::fd::AsRawFd;
+use std::os::raw::c_uint;
+use std::os::raw::c_void;
+use std::ptr::null_mut;
+use std::ptr::NonNull;
+use std::slice::from_raw_parts;
+use std::slice::from_raw_parts_mut;
+
+use crate::AsRawLibbpf;
+use crate::Error;
+use crate::MapHandle;
+use crate::MapType;
+use crate::Result;
+
+/// A mutable reference to sample from a [`UserRingBuffer`].
+///
+/// To write to the sample, dereference with `as_mut()` to get a mutable
+/// reference to the raw byte slice. You may find libraries such as
+/// [`plain`](https://crates.io/crates/plain) helpful to convert between raw
+/// bytes and structs.
+#[derive(Debug)]
+pub struct UserRingBufferSample<'slf> {
+    // A pointer to an 8-byte aligned reserved region of the user ring buffer
+    ptr: NonNull<c_void>,
+
+    // The size of the sample in bytes.
+    size: usize,
+
+    // Reference to the owning ring buffer. This is used to discard the sample
+    // if it is not submitted before being dropped.
+    rb: &'slf UserRingBuffer,
+
+    // Track whether the sample has been submitted.
+    submitted: bool,
+}
+
+impl Deref for UserRingBufferSample<'_> {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        unsafe { from_raw_parts(self.ptr.as_ptr() as *const u8, self.size) }
+    }
+}
+
+impl DerefMut for UserRingBufferSample<'_> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        unsafe { from_raw_parts_mut(self.ptr.as_ptr() as *mut u8, self.size) }
+    }
+}
+
+impl Drop for UserRingBufferSample<'_> {
+    fn drop(&mut self) {
+        // If the sample has not been submitted, explicitly discard it.
+        // This is necessary to avoid leaking ring buffer memory.
+        if !self.submitted {
+            unsafe {
+                libbpf_sys::user_ring_buffer__discard(self.rb.ptr.as_ptr(), self.ptr.as_ptr());
+            }
+        }
+    }
+}
+
+/// Represents a user ring buffer. This is a special kind of map that is used to
+/// transfer data between user space and kernel space.
+#[derive(Debug)]
+pub struct UserRingBuffer {
+    // A non-null pointer to the underlying user ring buffer.
+    ptr: NonNull<libbpf_sys::user_ring_buffer>,
+}
+
+impl UserRingBuffer {
+    /// Create a new user ring buffer from a map.
+    ///
+    /// # Errors
+    /// * If the map is not a user ring buffer.
+    /// * If the underlying libbpf function fails.
+    pub fn new(map: &MapHandle) -> Result<Self> {
+        if map.map_type() != MapType::UserRingBuf {
+            return Err(Error::with_invalid_data("must use a UserRingBuf map"));
+        }
+
+        let fd = map.as_fd();
+        let raw_ptr = unsafe { libbpf_sys::user_ring_buffer__new(fd.as_raw_fd(), null_mut()) };
+
+        let ptr = NonNull::new(raw_ptr).ok_or_else(|| {
+            // Safely get the last OS error after a failed call to user_ring_buffer__new
+            let errno = unsafe { *libc::__errno_location() };
+            Error::from_raw_os_error(errno)
+        })?;
+
+        Ok(UserRingBuffer { ptr })
+    }
+
+    /// Reserve a sample in the user ring buffer.
+    ///
+    /// Returns a [`UserRingBufferSample`](UserRingBufferSample<'slf>)
+    /// that contains a mutable reference to sample that can be written to.
+    /// The sample must be submitted via [`UserRingBuffer::submit`] before it is
+    /// dropped.
+    ///
+    /// # Parameters
+    /// * `size` - The size of the sample in bytes.
+    ///
+    /// This function is *not* thread-safe. It is necessary to synchronize
+    /// amongst multiple producers when invoking this function.
+    pub fn reserve(&self, size: usize) -> Result<UserRingBufferSample<'_>> {
+        let sample_ptr =
+            unsafe { libbpf_sys::user_ring_buffer__reserve(self.ptr.as_ptr(), size as c_uint) };
+
+        let ptr = NonNull::new(sample_ptr).ok_or_else(|| {
+            // Fetch the current value of errno to determine the type of error.
+            let errno = unsafe { *libc::__errno_location() };
+            match errno {
+                E2BIG => Error::with_invalid_data("requested size is too large"),
+                ENOSPC => Error::with_invalid_data("not enough space in the ring buffer"),
+                _ => Error::from_raw_os_error(errno),
+            }
+        })?;
+
+        Ok(UserRingBufferSample {
+            ptr,
+            size,
+            submitted: false,
+            rb: self,
+        })
+    }
+
+    /// Submit a sample to the user ring buffer.
+    ///
+    /// This function takes ownership of the sample and submits it to the ring
+    /// buffer. After submission, the consumer will be able to read the sample
+    /// from the ring buffer.
+    ///
+    /// This function is thread-safe. It is *not* necessary to synchronize
+    /// amongst multiple producers when invoking this function.
+    pub fn submit(&self, mut sample: UserRingBufferSample<'_>) -> Result<()> {
+        unsafe {
+            libbpf_sys::user_ring_buffer__submit(self.ptr.as_ptr(), sample.ptr.as_ptr());
+        }
+
+        sample.submitted = true;
+
+        // The libbpf API does not return an error code, so we cannot determine
+        // if the submission was successful. Return a `Result` to enable future
+        // validation while maintaing backwards compatibility.
+        Ok(())
+    }
+}
+
+impl AsRawLibbpf for UserRingBuffer {
+    type LibbpfType = libbpf_sys::user_ring_buffer;
+
+    /// Retrieve the underlying [`libbpf_sys::user_ring_buffer`].
+    fn as_libbpf_object(&self) -> NonNull<Self::LibbpfType> {
+        self.ptr
+    }
+}
+
+impl Drop for UserRingBuffer {
+    fn drop(&mut self) {
+        unsafe {
+            libbpf_sys::user_ring_buffer__free(self.ptr.as_ptr());
+        }
+    }
+}

libbpf-rs/tests/bin/src/user_ringbuf.bpf.c

@@ -0,0 +1,46 @@
+// SPDX-License-Identifier: GPL-2.0
+// Copyright (c) 2024 Jose Fernandez
+
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+struct {
+	__uint(type, BPF_MAP_TYPE_USER_RINGBUF);
+	__uint(max_entries, 4096 /* one page */);
+} user_ringbuf SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_HASH);
+	__type(key, u32);
+	__type(value, u32);
+	__uint(max_entries, 100);
+} samples SEC(".maps");
+
+struct my_struct_t {
+	u32 key;
+	u32 value;
+};
+
+static long user_ringbuf_callback(struct bpf_dynptr *dynptr, void *context)
+{
+	const struct my_struct_t *data;
+
+	data = bpf_dynptr_data(dynptr, 0, sizeof(*data));
+	if (!data)
+		return 0;
+
+	bpf_map_update_elem(&samples, &data->key, &data->value, BPF_ANY);
+
+	return 0;
+}
+
+SEC("tp/syscalls/sys_enter_getpid")
+int handle__sys_enter_getpid(void *ctx)
+{
+	bpf_user_ringbuf_drain(&user_ringbuf, user_ringbuf_callback, NULL, 0);
+
+	return 0;
+}
+
+char LICENSE[] SEC("license") = "GPL";

libbpf-rs/tests/bin/user_ringbuf.bpf.o

@@ -37,6 +37,7 @@ use libbpf_rs::ProgramType;
 use libbpf_rs::TracepointOpts;
 use libbpf_rs::UprobeOpts;
 use libbpf_rs::UsdtOpts;
+use libbpf_rs::UserRingBuffer;
 use plain::Plain;
 use probe::probe;
 use scopeguard::defer;
@@ -1244,6 +1245,102 @@ fn test_sudo_object_ringbuf_with_closed_map() {
     });
 }
 
+#[test]
+fn test_sudo_object_user_ringbuf() {
+    #[repr(C)]
+    struct MyStruct {
+        key: u32,
+        value: u32,
+    }
+
+    unsafe impl Plain for MyStruct {}
+
+    bump_rlimit_mlock();
+
+    let mut obj = get_test_object("user_ringbuf.bpf.o");
+    let prog = obj
+        .prog_mut("handle__sys_enter_getpid")
+        .expect("failed to find program");
+    let _link = prog.attach().expect("failed to attach prog");
+    let urb_map = obj
+        .map("user_ringbuf")
+        .expect("failed to find user ringbuf map");
+    let user_ringbuf = UserRingBuffer::new(urb_map).expect("failed to create user ringbuf");
+    let mut urb_sample = user_ringbuf
+        .reserve(size_of::<MyStruct>())
+        .expect("failed to reserve space");
+    let bytes = urb_sample.as_mut();
+    let my_struct = plain::from_mut_bytes::<MyStruct>(bytes).expect("failed to convert bytes");
+    my_struct.key = 42;
+    my_struct.value = 1337;
+    user_ringbuf
+        .submit(urb_sample)
+        .expect("failed to submit sample");
+
+    // Trigger BPF program.
+    let _pid = unsafe { libc::getpid() };
+
+    // At this point, the BPF program should have run and consumed the sample in
+    // the user ring buffer, and stored the key/value in the samples map.
+    let samples_map = obj.map("samples").expect("failed to find map");
+    let key: u32 = 42;
+    let value: u32 = 1337;
+    let res = samples_map
+        .lookup(&key.to_ne_bytes(), MapFlags::ANY)
+        .expect("failed to lookup")
+        .expect("failed to find value for key");
+
+    // The value in the samples map should be the same as the value we submitted
+    assert_eq!(res.len(), size_of::<u32>());
+    let mut array = [0; size_of::<u32>()];
+    array.copy_from_slice(&res[..]);
+    assert_eq!(u32::from_ne_bytes(array), value);
+}
+
+#[test]
+fn test_sudo_object_user_ringbuf_reservation_too_big() {
+    bump_rlimit_mlock();
+
+    let mut obj = get_test_object("user_ringbuf.bpf.o");
+    let prog = obj
+        .prog_mut("handle__sys_enter_getpid")
+        .expect("failed to find program");
+    let _link = prog.attach().expect("failed to attach prog");
+    let urb_map = obj
+        .map("user_ringbuf")
+        .expect("failed to find user ringbuf map");
+    let user_ringbuf = UserRingBuffer::new(urb_map).expect("failed to create user ringbuf");
+    let err = user_ringbuf.reserve(1024 * 1024).unwrap_err();
+    assert!(
+        err.to_string().contains("requested size is too large"),
+        "{err:#}"
+    );
+}
+
+#[test]
+fn test_sudo_object_user_ringbuf_not_enough_space() {
+    bump_rlimit_mlock();
+
+    let mut obj = get_test_object("user_ringbuf.bpf.o");
+    let prog = obj
+        .prog_mut("handle__sys_enter_getpid")
+        .expect("failed to find program");
+    let _link = prog.attach().expect("failed to attach prog");
+    let urb_map = obj
+        .map("user_ringbuf")
+        .expect("failed to find user ringbuf map");
+    let user_ringbuf = UserRingBuffer::new(urb_map).expect("failed to create user ringbuf");
+    let _ = user_ringbuf
+        .reserve(1024 * 3)
+        .expect("failed to reserve space");
+    let err = user_ringbuf.reserve(1024 * 3).unwrap_err();
+    assert!(
+        err.to_string()
+            .contains("not enough space in the ring buffer"),
+        "{err:#}"
+    );
+}
+
 #[test]
 fn test_sudo_object_task_iter() {
     bump_rlimit_mlock();