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>
---
 rust/helpers/slab.c            |  8 ++--
 rust/helpers/vmalloc.c         |  4 +-
 rust/kernel/alloc.rs           | 80 ++++++++++++++++++++++++++++++++--
 rust/kernel/alloc/allocator.rs | 42 ++++++++++--------
 4 files changed, 107 insertions(+), 27 deletions(-)

diff --git a/rust/helpers/slab.c b/rust/helpers/slab.c
index a842bfbddcba..8472370a4338 100644
--- a/rust/helpers/slab.c
+++ b/rust/helpers/slab.c
@@ -3,13 +3,13 @@
 #include <linux/slab.h>
 
 void * __must_check __realloc_size(2)
-rust_helper_krealloc(const void *objp, size_t new_size, gfp_t flags)
+rust_helper_krealloc_node(const void *objp, size_t new_size, gfp_t flags, int node)
 {
-	return krealloc(objp, new_size, flags);
+	return krealloc_node(objp, new_size, flags, node);
 }
 
 void * __must_check __realloc_size(2)
-rust_helper_kvrealloc(const void *p, size_t size, gfp_t flags)
+rust_helper_kvrealloc_node(const void *p, size_t size, gfp_t flags, int node)
 {
-	return kvrealloc(p, size, flags);
+	return kvrealloc_node(p, size, flags, node);
 }
diff --git a/rust/helpers/vmalloc.c b/rust/helpers/vmalloc.c
index 80d34501bbc0..62d30db9a1a6 100644
--- a/rust/helpers/vmalloc.c
+++ b/rust/helpers/vmalloc.c
@@ -3,7 +3,7 @@
 #include <linux/vmalloc.h>
 
 void * __must_check __realloc_size(2)
-rust_helper_vrealloc(const void *p, size_t size, gfp_t flags)
+rust_helper_vrealloc_node(const void *p, size_t size, gfp_t flags, int node)
 {
-	return vrealloc(p, size, flags);
+	return vrealloc_node(p, size, flags, node);
 }
diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index a2c49e5494d3..e886ab31108f 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -28,7 +28,9 @@
 /// Indicates an allocation error.
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub struct AllocError;
+
 use core::{alloc::Layout, ptr::NonNull};
+use crate::error::{code::EINVAL, Result};
 
 /// Flags to be used when allocating memory.
 ///
@@ -115,6 +117,31 @@ pub mod flags {
     pub const __GFP_NOWARN: Flags = Flags(bindings::__GFP_NOWARN);
 }
 
+/// 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);
+}
+
 /// 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> {
+        // 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 NUMA_NO_NODE if the caller doesn't care.
     ///
     /// If the requested size is zero, `realloc` behaves equivalent to `free`.
     ///
@@ -191,13 +249,29 @@ 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, numa::NUMA_NO_NODE) }
+    }
+
+
     /// Free an existing memory allocation.
     ///
     /// # Safety
diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
index aa2dfa9dca4c..2e86e9839a1b 100644
--- a/rust/kernel/alloc/allocator.rs
+++ b/rust/kernel/alloc/allocator.rs
@@ -13,7 +13,7 @@
 use core::ptr;
 use core::ptr::NonNull;
 
-use crate::alloc::{AllocError, Allocator};
+use crate::alloc::{AllocError, Allocator, NumaNode};
 use crate::bindings;
 use crate::pr_warn;
 
@@ -58,18 +58,20 @@ fn aligned_size(new_layout: Layout) -> usize {
 ///
 /// One of the following: `krealloc`, `vrealloc`, `kvrealloc`.
 struct ReallocFunc(
-    unsafe extern "C" fn(*const crate::ffi::c_void, usize, u32) -> *mut crate::ffi::c_void,
+    unsafe extern "C" fn(
+        *const crate::ffi::c_void, usize,  u32, crate::ffi::c_int,
+    ) -> *mut crate::ffi::c_void,
 );
 
 impl ReallocFunc {
-    // INVARIANT: `krealloc` satisfies the type invariants.
-    const KREALLOC: Self = Self(bindings::krealloc);
+    // INVARIANT: `krealloc_node` satisfies the type invariants.
+    const KREALLOC: Self = Self(bindings::krealloc_node);
 
-    // INVARIANT: `vrealloc` satisfies the type invariants.
-    const VREALLOC: Self = Self(bindings::vrealloc);
+    // INVARIANT: `vrealloc_node` satisfies the type invariants.
+    const VREALLOC: Self = Self(bindings::vrealloc_node);
 
-    // INVARIANT: `kvrealloc` satisfies the type invariants.
-    const KVREALLOC: Self = Self(bindings::kvrealloc);
+    // INVARIANT: `kvrealloc_node` satisfies the type invariants.
+    const KVREALLOC: Self = Self(bindings::kvrealloc_node);
 
     /// # Safety
     ///
@@ -87,6 +89,7 @@ unsafe fn call(
         layout: Layout,
         old_layout: Layout,
         flags: Flags,
+        nid: NumaNode,
     ) -> Result<NonNull<[u8]>, AllocError> {
         let size = aligned_size(layout);
         let ptr = match ptr {
@@ -110,7 +113,7 @@ unsafe fn call(
         // - Those functions provide the guarantees of this function.
         let raw_ptr = unsafe {
             // If `size == 0` and `ptr != NULL` the memory behind the pointer is freed.
-            self.0(ptr.cast(), size, flags.0).cast()
+            self.0(ptr.cast(), size, flags.0, nid.0).cast()
         };
 
         let ptr = if size == 0 {
@@ -123,34 +126,36 @@ unsafe fn call(
     }
 }
 
-// SAFETY: `realloc` delegates to `ReallocFunc::call`, which guarantees that
+// SAFETY: `realloc_node` delegates to `ReallocFunc::call`, which guarantees that
 // - memory remains valid until it is explicitly freed,
 // - passing a pointer to a valid memory allocation is OK,
 // - `realloc` satisfies the guarantees, since `ReallocFunc::call` has the same.
 unsafe impl Allocator for Kmalloc {
     #[inline]
-    unsafe fn realloc(
+    unsafe fn realloc_node(
         ptr: Option<NonNull<u8>>,
         layout: Layout,
         old_layout: Layout,
         flags: Flags,
+        nid: NumaNode,
     ) -> Result<NonNull<[u8]>, AllocError> {
         // SAFETY: `ReallocFunc::call` has the same safety requirements as `Allocator::realloc`.
-        unsafe { ReallocFunc::KREALLOC.call(ptr, layout, old_layout, flags) }
+        unsafe { ReallocFunc::KREALLOC.call(ptr, layout, old_layout, flags, nid) }
     }
 }
 
-// SAFETY: `realloc` delegates to `ReallocFunc::call`, which guarantees that
+// SAFETY: `realloc_node` delegates to `ReallocFunc::call`, which guarantees that
 // - memory remains valid until it is explicitly freed,
 // - passing a pointer to a valid memory allocation is OK,
 // - `realloc` satisfies the guarantees, since `ReallocFunc::call` has the same.
 unsafe impl Allocator for Vmalloc {
     #[inline]
-    unsafe fn realloc(
+    unsafe fn realloc_node(
         ptr: Option<NonNull<u8>>,
         layout: Layout,
         old_layout: Layout,
         flags: Flags,
+        nid: NumaNode,
     ) -> Result<NonNull<[u8]>, AllocError> {
         // TODO: Support alignments larger than PAGE_SIZE.
         if layout.align() > bindings::PAGE_SIZE {
@@ -160,21 +165,22 @@ unsafe fn realloc(
 
         // SAFETY: If not `None`, `ptr` is guaranteed to point to valid memory, which was previously
         // allocated with this `Allocator`.
-        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags) }
+        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags, nid) }
     }
 }
 
-// SAFETY: `realloc` delegates to `ReallocFunc::call`, which guarantees that
+// SAFETY: `realloc_node` delegates to `ReallocFunc::call`, which guarantees that
 // - memory remains valid until it is explicitly freed,
 // - passing a pointer to a valid memory allocation is OK,
 // - `realloc` satisfies the guarantees, since `ReallocFunc::call` has the same.
 unsafe impl Allocator for KVmalloc {
     #[inline]
-    unsafe fn realloc(
+    unsafe fn realloc_node(
         ptr: Option<NonNull<u8>>,
         layout: Layout,
         old_layout: Layout,
         flags: Flags,
+        nid: NumaNode,
     ) -> Result<NonNull<[u8]>, AllocError> {
         // TODO: Support alignments larger than PAGE_SIZE.
         if layout.align() > bindings::PAGE_SIZE {
@@ -184,6 +190,6 @@ unsafe fn realloc(
 
         // SAFETY: If not `None`, `ptr` is guaranteed to point to valid memory, which was previously
         // allocated with this `Allocator`.
-        unsafe { ReallocFunc::KVREALLOC.call(ptr, layout, old_layout, flags) }
+        unsafe { ReallocFunc::KVREALLOC.call(ptr, layout, old_layout, flags, nid) }
     }
 }
-- 
2.39.2

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

> On Jul 1, 2025, at 12:34 PM, Alice Ryhl <aliceryhl@google.com> wrote:
> 
> 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.

This is an elegant solution indeed but I think that keeping the existing approach goes better with the overall kernel trend of having better NUMA support. My point is, if we add NodeAllocator as a super-trait and in a foreseeable future all the Rust allocators will want/be required to support NUMA (which is likely to happen), we’ll have to “flatten” the traits and effectively go back to the approach expressed in this patch.

>> +/// 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.

Thanks, noted.

> 
>> /// 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?
> 
> 
Will do, thanks.

~Vitaly

On Tue, Jul 1, 2025 at 1:19 PM Vitaly Wool <vitaly.wool@konsulko.se> wrote:
>
>
>
> > On Jul 1, 2025, at 12:34 PM, Alice Ryhl <aliceryhl@google.com> wrote:
> >
> > 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.
>
> This is an elegant solution indeed but I think that keeping the existing approach goes better with the overall kernel trend of having better NUMA support. My point is, if we add NodeAllocator as a super-trait and in a foreseeable future all the Rust allocators will want/be required to support NUMA (which is likely to happen), we’ll have to “flatten” the traits and effectively go back to the approach expressed in this patch.

If we are not going to have allocators without numa support, then what
you did is reasonable. Though in that case I would consider just
changing the existing methods instead of having methods both with and
without a numa node argument.

Alice