Message-ID: <CAH5fLgjPir8LfzfouBd3PYBvfCkWgQEw+im-=Vo7z8kBmFLtrw@mail.gmail.com>
Date: Wed, 11 Dec 2024 16:03:45 +0100
From: Alice Ryhl <aliceryhl@...gle.com>
To: Tamir Duberstein <tamird@...il.com>
Cc: Danilo Krummrich <dakr@...nel.org>, Miguel Ojeda <ojeda@...nel.org>, 
	Alex Gaynor <alex.gaynor@...il.com>, Boqun Feng <boqun.feng@...il.com>, 
	Gary Guo <gary@...yguo.net>, Björn Roy Baron <bjorn3_gh@...tonmail.com>, 
	Benno Lossin <benno.lossin@...ton.me>, Andreas Hindborg <a.hindborg@...nel.org>, 
	Trevor Gross <tmgross@...ch.edu>, Maíra Canal <mcanal@...lia.com>, 
	Asahi Lina <lina@...hilina.net>, rust-for-linux@...r.kernel.org, 
	linux-kernel@...r.kernel.org
Subject: Re: [PATCH v11 2/2] rust: xarray: Add an abstraction for XArray

On Tue, Dec 3, 2024 at 11:44 PM Tamir Duberstein <tamird@...il.com> wrote:
>
> `XArray` is an efficient sparse array of pointers. Add a Rust
> abstraction for this type.
>
> This implementation bounds the element type on `ForeignOwnable` and
> requires explicit locking for all operations. Future work may leverage
> RCU to enable lockless operation.
>
> Inspired-by: Maíra Canal <mcanal@...lia.com>
> Inspired-by: Asahi Lina <lina@...hilina.net>
> Signed-off-by: Tamir Duberstein <tamird@...il.com>
> ---
>  rust/bindings/bindings_helper.h |   6 +
>  rust/helpers/helpers.c          |   1 +
>  rust/helpers/xarray.c           |  28 +++++
>  rust/kernel/alloc.rs            |   5 +
>  rust/kernel/lib.rs              |   1 +
>  rust/kernel/xarray.rs           | 270 ++++++++++++++++++++++++++++++++++++++++
>  6 files changed, 311 insertions(+)
>
> diff --git a/rust/bindings/bindings_helper.h b/rust/bindings/bindings_helper.h
> index 5c4dfe22f41a5a106330e8c43ffbd342c69c4e0b..9f39d673b240281aed2759b5bd076aa43fb54951 100644
> --- a/rust/bindings/bindings_helper.h
> +++ b/rust/bindings/bindings_helper.h
> @@ -30,6 +30,7 @@
>  #include <linux/tracepoint.h>
>  #include <linux/wait.h>
>  #include <linux/workqueue.h>
> +#include <linux/xarray.h>
>  #include <trace/events/rust_sample.h>
>
>  /* `bindgen` gets confused at certain things. */
> @@ -43,3 +44,8 @@ const gfp_t RUST_CONST_HELPER___GFP_ZERO = __GFP_ZERO;
>  const gfp_t RUST_CONST_HELPER___GFP_HIGHMEM = ___GFP_HIGHMEM;
>  const gfp_t RUST_CONST_HELPER___GFP_NOWARN = ___GFP_NOWARN;
>  const blk_features_t RUST_CONST_HELPER_BLK_FEAT_ROTATIONAL = BLK_FEAT_ROTATIONAL;
> +
> +const xa_mark_t RUST_CONST_HELPER_XA_PRESENT = XA_PRESENT;
> +
> +const gfp_t RUST_CONST_HELPER_XA_FLAGS_ALLOC = XA_FLAGS_ALLOC;
> +const gfp_t RUST_CONST_HELPER_XA_FLAGS_ALLOC1 = XA_FLAGS_ALLOC1;
> diff --git a/rust/helpers/helpers.c b/rust/helpers/helpers.c
> index dcf827a61b52e71e46fd5378878602eef5e538b6..ff28340e78c53c79baf18e2927cc90350d8ab513 100644
> --- a/rust/helpers/helpers.c
> +++ b/rust/helpers/helpers.c
> @@ -30,3 +30,4 @@
>  #include "vmalloc.c"
>  #include "wait.c"
>  #include "workqueue.c"
> +#include "xarray.c"
> diff --git a/rust/helpers/xarray.c b/rust/helpers/xarray.c
> new file mode 100644
> index 0000000000000000000000000000000000000000..60b299f11451d2c4a75e50e25dec4dac13f143f4
> --- /dev/null
> +++ b/rust/helpers/xarray.c
> @@ -0,0 +1,28 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +#include <linux/xarray.h>
> +
> +int rust_helper_xa_err(void *entry)
> +{
> +       return xa_err(entry);
> +}
> +
> +void rust_helper_xa_init_flags(struct xarray *xa, gfp_t flags)
> +{
> +       return xa_init_flags(xa, flags);
> +}
> +
> +int rust_helper_xa_trylock(struct xarray *xa)
> +{
> +       return xa_trylock(xa);
> +}
> +
> +void rust_helper_xa_lock(struct xarray *xa)
> +{
> +       return xa_lock(xa);
> +}
> +
> +void rust_helper_xa_unlock(struct xarray *xa)
> +{
> +       return xa_unlock(xa);
> +}
> diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
> index f2f7f3a53d298cf899e062346202ba3285ce3676..be9f164ece2e0fe71143e0201247d2b70c193c51 100644
> --- a/rust/kernel/alloc.rs
> +++ b/rust/kernel/alloc.rs
> @@ -39,6 +39,11 @@
>  pub struct Flags(u32);
>
>  impl Flags {
> +    /// Get a flags value with all bits unset.
> +    pub fn empty() -> Self {
> +        Self(0)
> +    }
> +
>      /// Get the raw representation of this flag.
>      pub(crate) fn as_raw(self) -> u32 {
>          self.0
> diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
> index e1065a7551a39e68d6379031d80d4be336e652a3..9ca524b15920c525c7db419e60dec4c43522751d 100644
> --- a/rust/kernel/lib.rs
> +++ b/rust/kernel/lib.rs
> @@ -68,6 +68,7 @@
>  pub mod types;
>  pub mod uaccess;
>  pub mod workqueue;
> +pub mod xarray;
>
>  #[doc(hidden)]
>  pub use bindings;
> diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
> new file mode 100644
> index 0000000000000000000000000000000000000000..39163ea037cb393a7c32a40b3e6539be33986b45
> --- /dev/null
> +++ b/rust/kernel/xarray.rs
> @@ -0,0 +1,270 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +//! XArray abstraction.
> +//!
> +//! C header: [`include/linux/xarray.h`](srctree/include/linux/xarray.h)
> +
> +use crate::{
> +    alloc, bindings, build_assert, build_error,
> +    error::{Error, Result},
> +    init::PinInit,
> +    pin_init,
> +    types::{ForeignOwnable, NotThreadSafe, Opaque},
> +};
> +use core::{iter, marker::PhantomData, mem, pin::Pin};
> +use macros::{pin_data, pinned_drop};
> +
> +/// An array which efficiently maps sparse integer indices to owned objects.
> +///
> +/// This is similar to a [`crate::alloc::kvec::Vec<Option<T>>`], but more efficient when there are
> +/// holes in the index space, and can be efficiently grown.
> +///
> +/// # Invariants
> +///
> +/// `self.xa` is always an initialized and valid [`bindings::xarray`] whose entries are either
> +/// `XA_ZERO_ENTRY` or came from `T::into_foreign`.
> +///
> +/// # Examples
> +///
> +/// ```rust
> +/// use kernel::alloc::KBox;
> +/// use kernel::xarray::{AllocKind, XArray};
> +///
> +/// let xa = KBox::pin_init(XArray::new(AllocKind::Alloc1), GFP_KERNEL)?;
> +///
> +/// let dead = KBox::new(0xdead, GFP_KERNEL)?;
> +/// let beef = KBox::new(0xbeef, GFP_KERNEL)?;
> +///
> +/// let mut guard = xa.lock();
> +///
> +/// assert_eq!(guard.get(0), None);
> +///
> +/// assert_eq!(guard.store(0, dead, GFP_KERNEL).unwrap().as_deref(), None);
> +/// assert_eq!(guard.get(0).copied(), Some(0xdead));
> +///
> +/// *guard.get_mut(0).unwrap() = 0xffff;
> +/// assert_eq!(guard.get(0).copied(), Some(0xffff));
> +///
> +/// assert_eq!(guard.store(0, beef, GFP_KERNEL).unwrap().as_deref().copied(), Some(0xffff));
> +/// assert_eq!(guard.get(0).copied(), Some(0xbeef));
> +///
> +/// guard.remove(0);
> +/// assert_eq!(guard.get(0), None);
> +///
> +/// # Ok::<(), Error>(())
> +/// ```
> +#[pin_data(PinnedDrop)]
> +pub struct XArray<T: ForeignOwnable> {
> +    #[pin]
> +    xa: Opaque<bindings::xarray>,
> +    _p: PhantomData<T>,
> +}
> +
> +#[pinned_drop]
> +impl<T: ForeignOwnable> PinnedDrop for XArray<T> {
> +    fn drop(self: Pin<&mut Self>) {
> +        self.iter().for_each(|ptr| {
> +            let ptr = ptr.as_ptr();
> +            // SAFETY: `ptr` came from `T::into_foreign`.
> +            //
> +            // INVARIANT: we own the only reference to the array which is being dropped so the
> +            // broken invariant is not observable on function exit.
> +            drop(unsafe { T::from_foreign(ptr) })
> +        });
> +
> +        // SAFETY: `self.xa` is always valid by the type invariant.
> +        unsafe { bindings::xa_destroy(self.xa.get()) };
> +    }
> +}
> +
> +/// Flags passed to [`XArray::new`] to configure the array's allocation tracking behavior.
> +pub enum AllocKind {
> +    /// Consider the first element to be at index 0.
> +    Alloc,
> +    /// Consider the first element to be at index 1.
> +    Alloc1,
> +}
> +
> +impl<T: ForeignOwnable> XArray<T> {
> +    /// Creates a new [`XArray`].
> +    pub fn new(kind: AllocKind) -> impl PinInit<Self> {
> +        let flags = match kind {
> +            AllocKind::Alloc => bindings::XA_FLAGS_ALLOC,
> +            AllocKind::Alloc1 => bindings::XA_FLAGS_ALLOC1,
> +        };
> +        pin_init!(Self {
> +            // SAFETY: `xa` is valid while the closure is called.
> +            xa <- Opaque::ffi_init(|xa| unsafe {
> +                bindings::xa_init_flags(xa, flags)
> +            }),
> +            _p: PhantomData,
> +        })
> +    }
> +
> +    fn iter(&self) -> impl Iterator<Item = core::ptr::NonNull<T::PointedTo>> + '_ {
> +        // TODO: Remove when https://lore.kernel.org/all/20240913213041.395655-5-gary@garyguo.net/ is applied.
> +        const MAX: core::ffi::c_ulong = core::ffi::c_ulong::MAX;

I think you can use kernel::ffi::c_ulong already. Enough things were
merged in 6.13 for that to work. If you import kernel::ffi::c_ulong at
the top of this file, then you can just do c_ulong::MAX in the
function calls below.

> +        let mut index = 0;
> +
> +        // SAFETY: `self.xa` is always valid by the type invariant.
> +        iter::once(unsafe {
> +            bindings::xa_find(self.xa.get(), &mut index, MAX, bindings::XA_PRESENT)
> +        })
> +        .chain(iter::from_fn(move || {
> +            // SAFETY: `self.xa` is always valid by the type invariant.
> +            Some(unsafe {
> +                bindings::xa_find_after(self.xa.get(), &mut index, MAX, bindings::XA_PRESENT)
> +            })
> +        }))
> +        .map_while(|ptr| core::ptr::NonNull::new(ptr.cast()))

You use core::ptr::NonNull in many places. Consider importing it.

> +    }
> +
> +    /// Attempts to lock the [`XArray`] for exclusive access.
> +    pub fn try_lock(&self) -> Option<Guard<'_, T>> {
> +        // SAFETY: `self.xa` is always valid by the type invariant.
> +        (unsafe { bindings::xa_trylock(self.xa.get()) } != 0).then(|| Guard {
> +            xa: self,
> +            _not_send: NotThreadSafe,
> +        })
> +    }
> +
> +    /// Locks the [`XArray`] for exclusive access.
> +    pub fn lock(&self) -> Guard<'_, T> {
> +        // SAFETY: `self.xa` is always valid by the type invariant.
> +        unsafe { bindings::xa_lock(self.xa.get()) };
> +
> +        Guard {
> +            xa: self,
> +            _not_send: NotThreadSafe,
> +        }
> +    }
> +}
> +
> +/// A lock guard.
> +///
> +/// The lock is unlocked when the guard goes out of scope.
> +#[must_use = "the lock unlocks immediately when the guard is unused"]
> +pub struct Guard<'a, T: ForeignOwnable> {
> +    xa: &'a XArray<T>,
> +    _not_send: NotThreadSafe,
> +}
> +
> +impl<T: ForeignOwnable> Drop for Guard<'_, T> {
> +    fn drop(&mut self) {
> +        // SAFETY:
> +        // - `self.xa.xa` is always valid by the type invariant.
> +        // - The caller holds the lock, so it is safe to unlock it.
> +        unsafe { bindings::xa_unlock(self.xa.xa.get()) };
> +    }
> +}
> +
> +// TODO: Remove when https://lore.kernel.org/all/20240913213041.395655-5-gary@garyguo.net/ is applied.
> +fn to_index(i: usize) -> core::ffi::c_ulong {
> +    i.try_into()
> +        .unwrap_or_else(|_| build_error!("cannot convert usize to c_ulong"))
> +}
> +
> +impl<'a, T: ForeignOwnable> Guard<'a, T> {
> +    fn load<F, U>(&self, index: usize, f: F) -> Option<U>
> +    where
> +        F: FnOnce(core::ptr::NonNull<T::PointedTo>) -> U,
> +    {
> +        // SAFETY: `self.xa.xa` is always valid by the type invariant.
> +        let ptr = unsafe { bindings::xa_load(self.xa.xa.get(), to_index(index)) };
> +        let ptr = core::ptr::NonNull::new(ptr.cast())?;
> +        Some(f(ptr))
> +    }
> +
> +    /// Loads an entry from the array.
> +    ///
> +    /// Returns the entry at the given index.
> +    pub fn get(&self, index: usize) -> Option<T::Borrowed<'_>> {
> +        self.load(index, |ptr| {
> +            // SAFETY: `ptr` came from `T::into_foreign`.
> +            unsafe { T::borrow(ptr.as_ptr()) }
> +        })
> +    }
> +
> +    /// Loads an entry from the array.
> +    ///
> +    /// Returns the entry at the given index.
> +    pub fn get_mut(&mut self, index: usize) -> Option<T::BorrowedMut<'_>> {
> +        self.load(index, |ptr| {
> +            // SAFETY: `ptr` came from `T::into_foreign`.
> +            unsafe { T::borrow_mut(ptr.as_ptr()) }
> +        })
> +    }
> +
> +    /// Erases an entry from the array.
> +    ///
> +    /// Returns the entry which was previously at the given index.
> +    pub fn remove(&mut self, index: usize) -> Option<T> {
> +        // SAFETY: `self.xa.xa` is always valid by the type invariant.
> +        //
> +        // SAFETY: The caller holds the lock.
> +        let ptr = unsafe { bindings::__xa_erase(self.xa.xa.get(), to_index(index)) }.cast();
> +        // SAFETY: `ptr` is either NULL or came from `T::into_foreign`.
> +        unsafe { T::try_from_foreign(ptr) }
> +    }
> +
> +    /// Stores an entry in the array.
> +    ///
> +    /// May drop the lock if needed to allocate memory, and then reacquire it afterwards.
> +    ///
> +    /// On success, returns the entry which was previously at the given index.
> +    ///
> +    /// On failure, returns the entry which was attempted to be stored.
> +    pub fn store(
> +        &mut self,
> +        index: usize,
> +        value: T,
> +        gfp: alloc::Flags,
> +    ) -> Result<Option<T>, (T, Error)> {

We can see in your examples that this return type is inconvenient.
Perhaps it would be better to make a new error type containing a T and
an Error, and implement From so that the question mark can convert
directly to Error (throwing away the T).

> +        build_assert!(
> +            mem::align_of::<T::PointedTo>() >= 4,
> +            "pointers stored in XArray must be 4-byte aligned"
> +        );
> +        let new = value.into_foreign();
> +
> +        let old = {
> +            let new = new.cast();
> +            // SAFETY: `self.xa.xa` is always valid by the type invariant.
> +            //
> +            // SAFETY: The caller holds the lock.
> +            //
> +            // INVARIANT: `new` came from `T::into_foreign`.
> +            unsafe { bindings::__xa_store(self.xa.xa.get(), to_index(index), new, gfp.as_raw()) }
> +        };
> +
> +        // SAFETY: `__xa_store` returns the old entry at this index on success or `xa_err` if an
> +        // error happened.
> +        match unsafe { bindings::xa_err(old) } {
> +            0 => {
> +                let old = old.cast();
> +                // SAFETY: `ptr` is either NULL or came from `T::into_foreign`.
> +                //
> +                // NB: `XA_ZERO_ENTRY` is never returned by functions belonging to the Normal XArray
> +                // API; such entries present as `NULL`.
> +                Ok(unsafe { T::try_from_foreign(old) })
> +            }
> +            errno => {
> +                // SAFETY: `new` came from `T::into_foreign` and `__xa_store` does not take
> +                // ownership of the value on error.
> +                let value = unsafe { T::from_foreign(new) };
> +                Err((value, Error::from_errno(errno)))
> +            }
> +        }
> +    }
> +}
> +
> +// SAFETY: It is safe to send `XArray<T>` to another thread when the underlying `T` is `Send`
> +// because XArray is thread-safe and all mutation operations are synchronized.
> +unsafe impl<T: ForeignOwnable + Send> Send for XArray<T> {}
> +
> +// SAFETY: It is safe to send `&XArray<T>` to another thread when the underlying `T` is `Sync`
> +// because it effectively means sharing `&T` (which is safe because `T` is `Sync`); additionally, it
> +// needs `T` to be `Send` because any thread that has a `&XArray<T>` may lock it and get a
> +// `Guard<T>` on that thread, so the thread may ultimately access `T` using a mutable reference, for
> +// example, using `get_mut` or `remove`.
> +unsafe impl<T: ForeignOwnable + Send + Sync> Sync for XArray<T> {}

I don't think Sync is needed due to the spinlock.

Alice

Powered by blists - more mailing lists