Re: [PATCH v10 2/2] rust: xarray: Add an abstraction for XArray
From: Tamir Duberstein
Date: Tue Dec 03 2024 - 10:03:00 EST
On Tue, Dec 3, 2024 at 7:30 AM Alice Ryhl <aliceryhl@xxxxxxxxxx> wrote:
>
> On Wed, Nov 20, 2024 at 12:48 PM Tamir Duberstein <tamird@xxxxxxxxx> 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@xxxxxxxxxx>
> > Inspired-by: Asahi Lina <lina@xxxxxxxxxxxxx>
> > Signed-off-by: Tamir Duberstein <tamird@xxxxxxxxx>
> > ---
> > 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 | 266 ++++++++++++++++++++++++++++++++++++++++
> > 6 files changed, 307 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)
> > + }
>
> Is this used anywhere?
It was used in an older version to fill reservations which were
expected to not need to allocate memory. Technically it is public API
so calling `store` on an occupied entry with `Flags::empty()` would be
a valid use.
> > 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..acbac787dc9a38684538697d53f590880fa9903a
> > --- /dev/null
> > +++ b/rust/kernel/xarray.rs
> > @@ -0,0 +1,266 @@
> > +// SPDX-License-Identifier: GPL-2.0
> > +
> > +//! XArray abstraction.
> > +//!
> > +//! C header: [`include/linux/xarray.h`](srctree/include/linux/xarray.h)
> > +
> > +use core::pin::Pin;
>
> Could be merged with the imports below.
Will do.
> > +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};
> > +use macros::{pin_data, pinned_drop};
>
> I think these are in crate::prelude.
I prefer to be explicit, unless there's guidance on this somewhere?
> > +/// 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@xxxxxxxxxxx/ is applied.
> > + const MIN: core::ffi::c_ulong = core::ffi::c_ulong::MIN;
> > + const MAX: core::ffi::c_ulong = core::ffi::c_ulong::MAX;
>
> Isn't MIN just zero?
I liked the symmetry, but I could change it if you feel strongly.
> > + let mut index = MIN;
> > +
> > + // 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()))
> > + }
> > +
> > + /// 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.
> > + //
> > + // SAFETY: 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@xxxxxxxxxxx/ 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();
>
> Two safety comments?
There are two properties that must be upheld. How would you like to
see it formatted?
> > + // SAFETY: `ptr` is either NULL or came from `T::into_foreign`.
> > + unsafe { T::try_from_foreign(ptr) }
> > + }
> > +
> > + /// Stores an entry in the array.
> > + ///
> > + /// On success, returns the entry which was previously at the given index.
> > + ///
> > + /// On failure, returns the entry which was attempted to be stored.
>
> I'd like to see documentation about the gfp flags. This may unlock the
> spinlock temporarily if GFP_KERNEL is used.
Will add the language from the C documentation: "May drop the lock if
needed to allocate memory, and then reacquire it afterwards."
> > + pub fn store(
> > + &mut self,
> > + index: usize,
> > + value: T,
> > + gfp: alloc::Flags,
> > + ) -> Result<Option<T>, (T, Error)> {
> > + 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`.
> > + Ok(unsafe { T::try_from_foreign(old) })
>
> It can't be XA_ZERO_ENTRY?
No it can't. XA_ZERO_ENTRY is never returned from the "normal" API.
XA_ZERO_ENTRY presents as NULL.
> > + }
> > + 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,
> > +// `T` is `Send` because XArray is thread-safe and all mutation operations are internally locked.
> > +unsafe impl<T: ForeignOwnable + Send + Sync> Sync for XArray<T> {}
> >
> > --
> > 2.47.0
> >