Re: [PATCH v4 01/28] rust: alloc: add `Allocator` trait

[Benno Lossin <benno.lossin@proton.me>]

On 05.08.24 17:19, Danilo Krummrich wrote:
> Add a kernel specific `Allocator` trait, that in contrast to the one in
> Rust's core library doesn't require unstable features and supports GFP
> flags.
> 
> Subsequent patches add the following trait implementors: `Kmalloc`,
> `Vmalloc` and `KVmalloc`.
> 
> Signed-off-by: Danilo Krummrich <dakr@kernel.org>
> ---
>  rust/kernel/alloc.rs | 79 ++++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 79 insertions(+)
> 
> diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
> index 1966bd407017..8a71a589469d 100644
> --- a/rust/kernel/alloc.rs
> +++ b/rust/kernel/alloc.rs
> @@ -11,6 +11,7 @@
>  /// Indicates an allocation error.
>  #[derive(Copy, Clone, PartialEq, Eq, Debug)]
>  pub struct AllocError;
> +use core::{alloc::Layout, ptr::NonNull};
> 
>  /// Flags to be used when allocating memory.
>  ///
> @@ -86,3 +87,81 @@ pub mod flags {
>      /// small allocations.
>      pub const GFP_NOWAIT: Flags = Flags(bindings::GFP_NOWAIT);
>  }
> +
> +/// The kernel's [`Allocator`] trait.
> +///
> +/// An implementation of [`Allocator`] can allocate, re-allocate and free memory buffer described
> +/// via [`Layout`].
> +///
> +/// [`Allocator`] is designed to be implemented as a ZST; [`Allocator`] functions do not operate on
> +/// an object instance.

This will prevent us from implementing arena-type allocators [^1]. Do we
want/need those?
I have heard that some people use them in embedded systems, but I can't 
say for sure. But this is a rather big design decision, so we should
discuss it now.

[^1]: For those who don't know what I mean by that here is a quick
sketch (without handling flags and optimizations):

    pub struct ArenaAlloc<const SIZE: usize> {
        memory: Opaque<[u8; SIZE]>,
        head: Cell<usize>,
    }

    impl<const SIZE: usize> ArenaAlloc<SIZE> {
        pub fn new() -> Self {
            Self {
                memory: Opaque::uninit(),
                head: 0,
            }
        }
    }

    impl<const SIZE: usize> Allocator for ArenaAlloc<SIZE> {
        fn alloc(&self, layout: Layout, _flags: Flags) -> Result<NonNull<u8>, AllocError> {
            let head = self.head.get();
            if head + layout.size() >= SIZE {
                return Err(AllocError);
            }
            let ptr = self.memory.get();
            let ptr = ptr.cast::<u8>();
            let ptr = unsafe { ptr.add(head) };
            self.head.set(head + layout.size());
            unsafe { NonNull::new_unchecked(ptr) }
        }

        unsafe fn realloc(
            &self,
            ptr: Option<NonNull<u8>>,
            old_layout: Layout, // Note that we also need `old_layout`!
            layout: Layout,
            flags: Flags
        ) -> Result<NonNull<u8>, AllocError> {
            let new = self.alloc(layout, flags)?;
            let Some(ptr) = ptr else { return Ok(new); };
            unsafe { core::ptr::copy_nonoverlapping(ptr.as_ptr(), new.as_ptr(), old_layout.size()) };
            self.free(ptr);
            Ok(new)
        }

        fn free(&self, ptr: NonNull<u8>) { /* noop */ }
    }

> +///
> +/// In order to be able to support `#[derive(SmartPointer)]` later on, we need to avoid a design
> +/// that requires an `Allocator` to be instantiated, hence its functions must not contain any kind
> +/// of `self` parameter.

Ah I see, so since `#[derive(SmartPointer)]` needs `Box` to only consist
of one non ZST field... I skimmed the RFC discussion and it seems like a
problem that *might* be solved in the future, but probably not in the
(very) near future. I guess this is just a bullet that we have to bite.
We can always have an `ArenaBox` that can deal with that (although
without `DispatchFromDyn`).
We should revisit this when `#[derive(SmartPointer)]` becomes advanced
enough.

> +///
> +/// # Safety
> +///
> +/// Memory returned from an allocator must point to a valid memory buffer and remain valid until
> +/// it is explicitly freed.
> +///
> +/// Any pointer to a memory buffer which is currently allocated must be valid to be passed to any
> +/// other [`Allocator`] function of the same type. The same applies for a NULL pointer.
> +///
> +/// If `realloc` is called with:
> +///   - a size of zero, the given memory allocation, if any, must be freed
> +///   - a NULL pointer, a new memory allocation must be created
> +pub unsafe trait Allocator {
> +    /// Allocate memory based on `layout` and `flags`.
> +    ///
> +    /// On success, returns a buffer represented as `NonNull<[u8]>` that satisfies the layout
> +    /// constraints (i.e. minimum size and alignment as specified by `layout`).
> +    ///
> +    /// This function is equivalent to `realloc` when called with a NULL pointer.
> +    fn alloc(layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError> {
> +        // SAFETY: Passing a NULL pointer to `realloc` is valid by it's safety requirements and asks
> +        // for a new memory allocation.
> +        unsafe { Self::realloc(None, layout, flags) }
> +    }
> +
> +    /// Re-allocate an existing memory allocation to satisfy the requested `layout`. If the
> +    /// requested size is zero, `realloc` behaves equivalent to `free`.
> +    ///
> +    /// If the requested size is larger than the size of the existing allocation, a successful call
> +    /// to `realloc` guarantees that the new or grown buffer has at least `Layout::size` bytes, but
> +    /// may also be larger.
> +    ///
> +    /// If the requested size is smaller than the size of the existing allocation, `realloc` may or
> +    /// may not shrink the buffer; this is implementation specific to the allocator.
> +    ///
> +    /// On allocation failure, the existing buffer, if any, remains valid.
> +    ///
> +    /// The buffer is represented as `NonNull<[u8]>`.
> +    ///
> +    /// # Safety
> +    ///
> +    /// `Some(ptr)` must point to an existing and valid memory allocation created by this allocator

This is the wrong way around, `ptr: Option<NonNull<u8>>`, so
`Some(ptr): Option<Option<NonNull<u8>>>`. Instead I would write
"If `ptr = Some(p)`, then `p` must point to...".

> +    /// instance. The alignment encoded in `layout` must be smaller than or equal to the alignment
> +    /// requested in the previous `alloc` or `realloc` call of the same allocation.
> +    ///
> +    /// Additionally, `ptr` is allowed to be `None`; in this case a new memory allocation is
> +    /// created.
> +    ///
> +    unsafe fn realloc(
> +        ptr: Option<NonNull<u8>>,
> +        layout: Layout,
> +        flags: Flags,
> +    ) -> Result<NonNull<[u8]>, AllocError>;
> +
> +    /// Free an existing memory allocation.
> +    ///
> +    /// # Safety
> +    ///
> +    /// `ptr` must point to an existing and valid memory allocation created by this `Allocator`
> +    /// instance.

Additionally, you need "The memory allocation at `ptr` must never again
be read from or written to.".

---
Cheers,
Benno

> +    unsafe fn free(ptr: NonNull<u8>) {
> +        // SAFETY: `ptr` is guaranteed to be previously allocated with this `Allocator` or NULL.
> +        // Calling `realloc` with a buffer size of zero, frees the buffer `ptr` points to.
> +        let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), Flags(0)) };
> +    }
> +}
> --
> 2.45.2
>