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

[Alice Ryhl <aliceryhl@google.com>]

On Sat, Jul 05, 2025 at 09:17:10AM +0200, Vitaly Wool wrote:
> 
> 
> > On Jul 4, 2025, at 5:23 PM, Danilo Krummrich <dakr@kernel.org> wrote:
> > 
> > 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.
> 
> I don’t have a strong preference either but if I modify alloc() and
> realloc(), I’ll have to change KVec and KBox in the same patch, would
> that be okay with you?

That's okay with me.

Alice