Re: [PATCH v6] rust: xarray: Add an abstraction for XArray

["Andreas Hindborg (Samsung)" <nmi@metaspace.dk>]

Boqun Feng <boqun.feng@gmail.com> writes:

> On Tue, Jan 16, 2024 at 12:06:03PM -0300, Maíra Canal wrote:
>> From: Asahi Lina <lina@asahilina.net>
>> 
> [...]
>> v5 -> v6: https://lore.kernel.org/rust-for-linux/20231201195300.1329092-1-mcanal@igalia.com/T/
>> 
>> - Update constants to the new format (RUST_CONST_HELPER)
>> - Add invariant for `self.0` being a pointer derived from `T::from_foreign` (Benno Lossin)
>> - Fix the place of the INVARIANT comments (Benno Lossin)
>> - Use the Pin-Init API (Benno Lossin)
>> - Remove PhantomPinned from XArray (Benno Lossin)
>> - Add new requirements for calling `xa_unlock()` (Benno Lossin)
>> - Improve SAFETY comments (Benno Lossin)
>> - Split unsafe block (Benno Lossin)
>> - s/alloc_limits_opt/insert_between (Benno Lossin)
>> - Specify the target type of the cast (Andreas Hindborg/Trevor Gross)
>> - Guarantee that T is 4-byte aligned (Andreas Hindborg)
>> - Add code examples in the code (Boqun Feng)
>
> Thanks!
>
>> 
>> [1] https://github.com/mairacanal/linux/pull/11
>> 
>>  rust/bindings/bindings_helper.h |  17 ++
>>  rust/helpers.c                  |  37 +++
>>  rust/kernel/lib.rs              |   1 +
>>  rust/kernel/xarray.rs           | 383 ++++++++++++++++++++++++++++++++
>>  4 files changed, 438 insertions(+)
>>  create mode 100644 rust/kernel/xarray.rs
>> 
>
> [...]
>
>> +
>> +/// Wrapper for a value owned by the `XArray` which holds the `XArray` lock until dropped.
>> +///
>
> Although the comment here says that `Guard` represents a *locked* slot
> in the xarray, but..
>
>> +/// INVARIANT: `Guard` holds a reference (`self.0`) to the underlying value owned by XArray.
>> +/// You can use the `into_foreign` method to obtain a pointer to the foreign representation
>> +/// of the owned value, which is valid for the lifetime of the Guard.
>
> nothing about the xa lock is mentioned here in the type invariants. I
> think it's better to say:
>
> /// INVARIANT: `Guard` holds a reference (`self.0`) to the underlying
> /// value owned by the `XArray` (`self.1`) with its lock held...
>
> and then at the Guard::drop..
>
>> +pub struct Guard<'a, T: ForeignOwnable>(NonNull<T>, &'a XArray<T>);
>> +
>> +impl<'a, T: ForeignOwnable> Guard<'a, T> {
>> +    /// Borrow the underlying value wrapped by the `Guard`.
>> +    ///
>> +    /// Returns a `T::Borrowed` type for the owned `ForeignOwnable` type.
>> +    pub fn borrow(&self) -> T::Borrowed<'_> {
>> +        // SAFETY: The value is owned by the `XArray`, the lifetime it is borrowed for must not
>> +        // outlive the `XArray` itself, nor the Guard that holds the lock ensuring the value
>> +        // remains in the `XArray`.
>> +        //
>> +        // By the type invariant of `Guard`, we can guarantee that `Guard` holds this reference
>> +        // (`self.0`).
>> +        unsafe { T::borrow(self.0.as_ptr().cast()) }
>> +    }
>> +}
>> +
>> +// Convenience impl for `ForeignOwnable` types whose `Borrowed`
>> +// form implements Deref.
>> +impl<'a, T: ForeignOwnable> Deref for Guard<'a, T>
>> +where
>> +    T::Borrowed<'a>: Deref,
>> +    for<'b> T::Borrowed<'b>: Into<&'b <T::Borrowed<'a> as Deref>::Target>,
>> +{
>> +    type Target = <T::Borrowed<'a> as Deref>::Target;
>> +
>> +    fn deref(&self) -> &Self::Target {
>> +        self.borrow().into()
>> +    }
>> +}
>> +
>> +impl<'a, T: ForeignOwnable> Drop for Guard<'a, T> {
>> +    fn drop(&mut self) {
>> +        // SAFETY: By the type invariant, we guarantee that we have a reference
>> +        // that owns the C xarray object.
>
> you can (and should) also mention, "by the type invariant, `Guard` is
> the owner of the xarray lock, so it's safe to unlock". Justifying the
> ownership of lock at unlock point is safety related, since otherwise a
> data race may happen.
>
>> +        unsafe { bindings::xa_unlock(self.1.xa.get()) };
>> +    }
>> +}
>> +
>> +/// Represents a reserved slot in an `XArray`, which does not yet have a value but has an assigned
>> +/// index and may not be allocated by any other user. If the Reservation is dropped without
>> +/// being filled, the entry is marked as available again.
>> +///
>> +/// Users must ensure that reserved slots are not filled by other mechanisms, or otherwise their
>> +/// contents may be dropped and replaced (which will print a warning).
>
> I feel like this is a bit weird: instead of dropping and replacing, I
> would expect we do nothing if someone steals our unused reservation ;-)
> But let's see whether there would be a real user complain ;-)
>
>> +pub struct Reservation<'a, T: ForeignOwnable>(&'a XArray<T>, usize, PhantomData<T>);
>> +
>> +impl<'a, T: ForeignOwnable> Reservation<'a, T> {
>> +    /// Stores a value into the reserved slot.
>> +    pub fn store(self, value: T) -> Result<usize> {
>> +        if self.0.replace(self.1, value)?.is_some() {
>> +            crate::pr_err!("XArray: Reservation stored but the entry already had data!\n");
>> +            // Consider it a success anyway, not much we can do
>> +        }
>> +        let index = self.1;
>> +        // The reservation is now fulfilled, so do not run our destructor.
>> +        core::mem::forget(self);
>> +        Ok(index)
>> +    }
>> +
>> +    /// Returns the index of this reservation.
>> +    pub fn index(&self) -> usize {
>> +        self.1
>> +    }
>> +}
>> +
>> +impl<'a, T: ForeignOwnable> Drop for Reservation<'a, T> {
>> +    fn drop(&mut self) {
>> +        if self.0.remove(self.1).is_some() {
>> +            crate::pr_err!("XArray: Reservation dropped but the entry was not empty!\n");
>> +        }
>> +    }
>> +}
>> +
>> +/// An array which efficiently maps sparse integer indices to owned objects.
>> +///
>> +/// This is similar to a `Vec<Option<T>>`, but more efficient when there are holes in the
>> +/// index space, and can be efficiently grown.
>> +///
>> +/// This structure is expected to often be used with an inner type that
>> +/// can be efficiently cloned, such as an `Arc<T>`.
>> +///
>> +/// INVARIANT: All pointers stored in the array are pointers obtained by
>> +/// calling `T::into_foreign`.
>
> ... or null pointers (considering reserve() cases).
>
>> +#[pin_data(PinnedDrop)]
>> +pub struct XArray<T: ForeignOwnable> {
>> +    #[pin]
>> +    xa: Opaque<bindings::xarray>,
>> +    _p: PhantomData<T>,
>> +}
>> +
>> +///
>> +/// # Examples
>> +///
>> +/// ```rust
>> +/// # use kernel::xarray;
>> +/// # use kernel::prelude::*;
>> +/// # use kernel::sync::Arc;
>> +///
>> +/// # struct Foo {
>> +/// #     a: u32,
>> +/// #     b: u32,
>> +/// # }
>> +///
>> +/// let arr = Box::pin_init(xarray::XArray::<Arc<Foo>>::new(xarray::flags::ALLOC1))
>> +///                                                    .expect("Unable to allocate XArray");
>> +///
>> +/// let foo = Arc::try_new(Foo { a : 1, b: 2 }).expect("Unable to allocate Foo");
>> +/// let index = arr.alloc(foo).expect("Error allocating Index");
>> +///
>> +/// if let Some(guard) = arr.get(index) {
>> +///     assert_eq!(guard.borrow().a, 1);
>> +///     assert_eq!(guard.borrow().b, 2);
>> +/// } else {
>> +///     pr_info!("No value found in index {}", index);
>> +/// }
>> +///
>> +/// let foo = Arc::try_new(Foo { a : 3, b: 4 }).expect("Unable to allocate Foo");
>> +/// let index = arr.alloc(foo).expect("Error allocating Index");
>> +///
>> +/// if let Some(removed_data) = arr.remove(index) {
>> +///     assert_eq!(removed_data.a, 3);
>> +///     assert_eq!(removed_data.b, 4);
>> +/// } else {
>> +///     pr_info!("No value found in index {}", index);
>> +/// }
>
> Again, thanks!
>
>> +/// ```
>> +impl<T: ForeignOwnable> XArray<T> {
>> +    /// Creates a new `XArray` with the given flags.
>> +    pub fn new(flags: Flags) -> impl PinInit<Self> {
>> +        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,
>> +        })
>> +    }
>> +
>> +    /// Replaces an entry with a new value, returning the old value (if any).
>> +    pub fn replace(&self, index: usize, value: T) -> Result<Option<T>> {
>> +        let new = value.into_foreign();
>> +
>> +        build_assert!(mem::align_of::<T>() <= 4);
>
> I know Benno suggested this, but this is not quite right ;-) Sorry I
> didn't mention this earlier. Here T is a ForeignOwnable type, e.g.
> Box<_>, Arc<_>, while they may be 4-byte aligned but the return pointers
> of into_foreign() may not. What we need here is that `new` is at least 4
> bytes aligned. Yes, this is a boring difference, but it's a difference..

Also, this line would assert that the alignment of `T` is less than or
equal to 4 bytes? So a 2 byte aligned type would be accepted, which is
not what we want. 4, 8, 12, etc. is fine.

Instead of this, we can check that the last two bits of `new` is zero.
But that would not be const 😬

This is the relevant part of documentation:

> Normal pointers may be stored in the XArray directly.  They must be 4-byte
> aligned, which is true for any pointer returned from kmalloc() and
> alloc_page().  It isn't true for arbitrary user-space pointers,
> nor for function pointers.  You can store pointers to statically allocated
> objects, as long as those objects have an alignment of at least 4.

BR Andreas