Re: [PATCH v10 3/4] rust: add support for NUMA ids in allocations

[Danilo Krummrich <dakr@kernel.org>]

On Wed, Jul 02, 2025 at 06:09:10PM +0200, Vitaly Wool wrote:
>  /// The kernel's [`Allocator`] trait.
>  ///
>  /// An implementation of [`Allocator`] can allocate, re-allocate and free memory buffers described
> @@ -148,7 +172,7 @@ pub unsafe trait Allocator {
>      ///
>      /// When the return value is `Ok(ptr)`, then `ptr` is
>      /// - valid for reads and writes for `layout.size()` bytes, until it is passed to
> -    ///   [`Allocator::free`] or [`Allocator::realloc`],
> +    ///   [`Allocator::free`], [`Allocator::realloc`] or [`Allocator::realloc_node`],
>      /// - aligned to `layout.align()`,
>      ///
>      /// Additionally, `Flags` are honored as documented in
> @@ -159,7 +183,41 @@ fn alloc(layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError> {
>          unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags) }
>      }
>  
> -    /// Re-allocate an existing memory allocation to satisfy the requested `layout`.
> +    /// Allocate memory based on `layout`, `flags` and `nid`.
> +    ///
> +    /// 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_node` when called with `None`.
> +    ///
> +    /// # Guarantees
> +    ///
> +    /// When the return value is `Ok(ptr)`, then `ptr` is
> +    /// - valid for reads and writes for `layout.size()` bytes, until it is passed to
> +    ///   [`Allocator::free`], [`Allocator::realloc`] or [`Allocator::realloc_node`],
> +    /// - aligned to `layout.align()`,
> +    ///
> +    /// Additionally, `Flags` are honored as documented in
> +    /// <https://docs.kernel.org/core-api/mm-api.html#mm-api-gfp-flags>.
> +    fn alloc_node(
> +        layout: Layout,
> +        flags: Flags,
> +        nid: NumaNode,
> +    ) -> Result<NonNull<[u8]>, AllocError> {
> +        // SAFETY: Passing `None` to `realloc_node` is valid by its safety requirements and
> +        // asks for a new memory allocation.
> +        unsafe { Self::realloc_node(None, layout, Layout::new::<()>(), flags, nid) }
> +    }
> +
> +    /// Re-allocate an existing memory allocation to satisfy the requested `layout` and
> +    /// a specific NUMA node request to allocate the memory for.
> +    ///
> +    /// Systems employing a Non Uniform Memory Access (NUMA) architecture contain collections of
> +    /// hardware resources including processors, memory, and I/O buses, that comprise what is
> +    /// commonly known as a NUMA node.
> +    ///
> +    /// `nid` stands for NUMA id, i. e. NUMA node identifier, which is a non-negative
> +    /// integer if a node needs to be specified, or NO_NODE if the caller doesn't care.
>      ///
>      /// If the requested size is zero, `realloc` behaves equivalent to `free`.
>      ///
> @@ -191,13 +249,27 @@ fn alloc(layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError> {
>      ///   and old size, i.e. `ret_ptr[0..min(layout.size(), old_layout.size())] ==
>      ///   p[0..min(layout.size(), old_layout.size())]`.
>      /// - when the return value is `Err(AllocError)`, then `ptr` is still valid.
> -    unsafe fn realloc(
> +    unsafe fn realloc_node(
>          ptr: Option<NonNull<u8>>,
>          layout: Layout,
>          old_layout: Layout,
>          flags: Flags,
> +        nid: NumaNode,
>      ) -> Result<NonNull<[u8]>, AllocError>;
>  
> +    /// Re-allocate an existing memory allocation to satisfy the requested `layout`. This
> +    /// function works exactly as realloc_node() but it doesn't give the ability to specify
> +    /// the NUMA node in the call.
> +    unsafe fn realloc(
> +        ptr: Option<NonNull<u8>>,
> +        layout: Layout,
> +        old_layout: Layout,
> +        flags: Flags,
> +    ) -> Result<NonNull<[u8]>, AllocError> {
> +        // SAFETY: guaranteed by realloc_node()
> +        unsafe { Self::realloc_node(ptr, layout, old_layout, flags, NumaNode::NO_NODE) }
> +    }

I think Alice suggested to just drop alloc_node() and realloc_node() and make
alloc() and realloc() always take a NumaNode argument.

I don't have a strong preference, but keeping only alloc() and realloc() for
seems indeed simpler, so let's remove the _node() variants.

Regardless, please note that realloc() as you implemented it above misses the
safaety requirement, which should just mention that the safety requirements are
identical to realloc_node().

The safety comment on the subsequent call to realloc_node() would then be
justified with realloc() having identical safety requirements as realloc_node().