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

[Alice Ryhl <aliceryhl@google.com>]

On Tue, Jul 01, 2025 at 12:16:40AM +0200, Vitaly Wool wrote:
> Add a new type to support specifying NUMA identifiers in Rust
> allocators and extend the allocators to have NUMA id as a
> parameter. Thus, modify ReallocFunc to use the new extended realloc
> primitives from the C side of the kernel (i. e.
> k[v]realloc_node_align/vrealloc_node_align) and add the new function
> alloc_node to the Allocator trait while keeping the existing one
> (alloc) for backward compatibility.
> 
> This will allow to specify node to use for allocation of e. g.
> {KV}Box, as well as for future NUMA aware users of the API.
> 
> Signed-off-by: Vitaly Wool <vitaly.wool@konsulko.se>

My main feedback is that we should consider introducing a new trait
instead of modifying Allocator. What we could do is have a NodeAllocator
trait that is a super-trait of Allocator and has additional methods with
a node parameter.

A sketch:

pub unsafe trait NodeAllocator: Allocator {
    fn alloc_node(layout: Layout, flags: Flags, nid: NumaNode)
                -> Result<NonNull<[u8]>, AllocError>;

    unsafe fn realloc_node(
        ptr: Option<NonNull<u8>>,
        layout: Layout,
        old_layout: Layout,
        flags: Flags,
        nid: NumaNode,
    ) -> Result<NonNull<[u8]>, AllocError>;
}

By doing this, it's possible to have allocators that do not support
specifying the numa node which only implement Allocator, and to have
other allocators that implement both Allocator and NumaAllocator where
you are able to specify the node.

If all allocators in the kernel support numa nodes, then you can ignore
this.
> +/// Non Uniform Memory Access (NUMA) node identifier
> +#[derive(Clone, Copy, PartialEq)]
> +pub struct NumaNode(i32);
> +
> +impl NumaNode {
> +    /// create a new NUMA node identifer (non-negative integer)
> +    /// returns EINVAL if a negative id or an id exceeding MAX_NUMNODES is specified
> +    pub fn new(node: i32) -> Result<Self> {
> +        // SAFETY: MAX_NUMNODES never exceeds 2**10 because NODES_SHIFT is 0..10
> +        if node < 0 || node >= bindings::MAX_NUMNODES as i32 {
> +            return Err(EINVAL);
> +        }
> +        Ok(Self(node))
> +    }
> +}
> +
> +/// Specify necessary constant to pass the information to Allocator that the caller doesn't care
> +/// about the NUMA node to allocate memory from.
> +pub mod numa {
> +    use super::NumaNode;
> +
> +    /// No preference for NUMA node
> +    pub const NUMA_NO_NODE: NumaNode = NumaNode(bindings::NUMA_NO_NODE);
> +}

Instead of using a module, you can make it an associated constant of the
struct.

impl NumaNode {
    pub const NO_NODE: NumaNode = NumaNode(bindings::NUMA_NO_NODE);
}

This way you can access the constant as NumaNode::NO_NODE.

>  /// The kernel's [`Allocator`] trait.
>  ///
>  /// An implementation of [`Allocator`] can allocate, re-allocate and free memory buffers described
> @@ -148,7 +175,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 +186,38 @@ 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> {

I don't think this is how rustfmt would format this. Can you run rustfmt
on your patch?

Alice