Add support for specifying NUMA ids in Rust allocators as an Option
(i. e. providing `None` as nid corresponds to NUMA_NO_NODE). To do
this, modify ReallocFunc to use the new extended realloc primitives
from the C side of the kernel (i. e. k[v]realloc_node/vrealloc_node)
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 | 13 +++++++++++++
rust/helpers/vmalloc.c | 6 ++++++
rust/kernel/alloc.rs | 28 ++++++++++++++++++++++++++--
rust/kernel/alloc/allocator.rs | 33 ++++++++++++++++++++++-----------
rust/kernel/alloc/kvec.rs | 3 ++-
5 files changed, 69 insertions(+), 14 deletions(-)
diff --git a/rust/helpers/slab.c b/rust/helpers/slab.c
index a842bfbddcba..5db9450e9348 100644
--- a/rust/helpers/slab.c
+++ b/rust/helpers/slab.c
@@ -13,3 +13,16 @@ rust_helper_kvrealloc(const void *p, size_t size, gfp_t flags)
{
return kvrealloc(p, size, flags);
}
+
+void * __must_check __realloc_size(2)
+rust_helper_krealloc_node(const void *objp, size_t new_size, unsigned long align,
+ gfp_t flags, int node)
+{
+ return krealloc_node(objp, new_size, align, flags, node);
+}
+
+void * __must_check __realloc_size(2)
+rust_helper_kvrealloc_node(const void *p, size_t size, unsigned long align, gfp_t flags, int node)
+{
+ return kvrealloc_node(p, size, align, flags, node);
+}
diff --git a/rust/helpers/vmalloc.c b/rust/helpers/vmalloc.c
index 80d34501bbc0..7a71427a1e78 100644
--- a/rust/helpers/vmalloc.c
+++ b/rust/helpers/vmalloc.c
@@ -7,3 +7,9 @@ rust_helper_vrealloc(const void *p, size_t size, gfp_t flags)
{
return vrealloc(p, size, flags);
}
+
+void * __must_check __realloc_size(2)
+rust_helper_vrealloc_node(const void *p, size_t size, unsigned long align, gfp_t flags, int node)
+{
+ return vrealloc_node(p, size, align, flags, node);
+}
diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index a2c49e5494d3..1e26c2a7f47c 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -156,7 +156,30 @@ pub unsafe trait Allocator {
fn alloc(layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError> {
// SAFETY: Passing `None` to `realloc` is valid by its safety requirements and asks for a
// new memory allocation.
- unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags) }
+ unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags, None) }
+ }
+
+ /// 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` 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`] or [`Allocator::realloc`],
+ /// - 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: Option<i32>)
+ -> Result<NonNull<[u8]>, AllocError> {
+ // SAFETY: Passing `None` to `realloc` is valid by its safety requirements and asks for a
+ // new memory allocation.
+ unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags, nid) }
}
/// Re-allocate an existing memory allocation to satisfy the requested `layout`.
@@ -196,6 +219,7 @@ unsafe fn realloc(
layout: Layout,
old_layout: Layout,
flags: Flags,
+ nid: Option<i32>,
) -> Result<NonNull<[u8]>, AllocError>;
/// Free an existing memory allocation.
@@ -211,7 +235,7 @@ unsafe fn free(ptr: NonNull<u8>, layout: Layout) {
// SAFETY: The caller guarantees that `ptr` points at a valid allocation created by this
// allocator. We are passing a `Layout` with the smallest possible alignment, so it is
// smaller than or equal to the alignment previously used with this allocation.
- let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), layout, Flags(0)) };
+ let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), layout, Flags(0), None) };
}
}
diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
index aa2dfa9dca4c..4f0fe2b67593 100644
--- a/rust/kernel/alloc/allocator.rs
+++ b/rust/kernel/alloc/allocator.rs
@@ -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, crate::ffi::c_ulong, 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_NODE: 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_NODE: 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_NODE: Self = Self(bindings::kvrealloc_node);
/// # Safety
///
@@ -87,6 +89,7 @@ unsafe fn call(
layout: Layout,
old_layout: Layout,
flags: Flags,
+ nid: Option<i32>,
) -> Result<NonNull<[u8]>, AllocError> {
let size = aligned_size(layout);
let ptr = match ptr {
@@ -100,6 +103,11 @@ unsafe fn call(
None => ptr::null(),
};
+ let c_nid = match nid {
+ None => bindings::NUMA_NO_NODE,
+ Some(n) => n,
+ };
+
// SAFETY:
// - `self.0` is one of `krealloc`, `vrealloc`, `kvrealloc` and thus only requires that
// `ptr` is NULL or valid.
@@ -110,7 +118,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, 1, flags.0, c_nid).cast()
};
let ptr = if size == 0 {
@@ -134,9 +142,10 @@ unsafe fn realloc(
layout: Layout,
old_layout: Layout,
flags: Flags,
+ nid: Option<i32>,
) -> 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_NODE.call(ptr, layout, old_layout, flags, nid) }
}
}
@@ -151,6 +160,7 @@ unsafe fn realloc(
layout: Layout,
old_layout: Layout,
flags: Flags,
+ nid: Option<i32>,
) -> Result<NonNull<[u8]>, AllocError> {
// TODO: Support alignments larger than PAGE_SIZE.
if layout.align() > bindings::PAGE_SIZE {
@@ -160,7 +170,7 @@ 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_NODE.call(ptr, layout, old_layout, flags, nid) }
}
}
@@ -175,6 +185,7 @@ unsafe fn realloc(
layout: Layout,
old_layout: Layout,
flags: Flags,
+ nid: Option<i32>,
) -> Result<NonNull<[u8]>, AllocError> {
// TODO: Support alignments larger than PAGE_SIZE.
if layout.align() > bindings::PAGE_SIZE {
@@ -184,6 +195,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_NODE.call(ptr, layout, old_layout, flags, nid) }
}
}
diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 1a0dd852a468..ef4f977ba012 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -633,6 +633,7 @@ pub fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocEr
layout.into(),
self.layout.into(),
flags,
+ None,
)?
};
@@ -1058,7 +1059,7 @@ pub fn collect(self, flags: Flags) -> Vec<T, A> {
// the type invariant to be smaller than `cap`. Depending on `realloc` this operation
// may shrink the buffer or leave it as it is.
ptr = match unsafe {
- A::realloc(Some(buf.cast()), layout.into(), old_layout.into(), flags)
+ A::realloc(Some(buf.cast()), layout.into(), old_layout.into(), flags, None)
} {
// If we fail to shrink, which likely can't even happen, continue with the existing
// buffer.
--
2.39.2On Fri, Jun 27, 2025 at 11:38:58AM +0200, Vitaly Wool wrote:
> + /// Allocate memory based on `layout`, `flags` and `nid`.
The semantical meanting of layout and flags is rather obvious due to its type.
We should probably introduce a new type NumaNode with a constant for
NUMA_NODE_NONE.
> + ///
> + /// 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 `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`] or [`Allocator::realloc`],
> + /// - 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: Option<i32>)
> + -> Result<NonNull<[u8]>, AllocError> {
> + // SAFETY: Passing `None` to `realloc` is valid by its safety requirements and asks for a
> + // new memory allocation.
> + unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags, nid) }
> }
>
> /// Re-allocate an existing memory allocation to satisfy the requested `layout`.
> @@ -196,6 +219,7 @@ unsafe fn realloc(
> layout: Layout,
> old_layout: Layout,
> flags: Flags,
> + nid: Option<i32>,
> ) -> Result<NonNull<[u8]>, AllocError>;
Please rename to realloc_node() and expand the documentation explaining what
happens for invalid nid numbers and what happens if nid changes throughout
invocations.
Also, please introduce realloc() which calls realloc_node() with NUMA_NODE_NONE.
>
> /// Free an existing memory allocation.
> @@ -211,7 +235,7 @@ unsafe fn free(ptr: NonNull<u8>, layout: Layout) {
> // SAFETY: The caller guarantees that `ptr` points at a valid allocation created by this
> // allocator. We are passing a `Layout` with the smallest possible alignment, so it is
> // smaller than or equal to the alignment previously used with this allocation.
> - let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), layout, Flags(0)) };
> + let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), layout, Flags(0), None) };
> }
> }
>
> diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
> index aa2dfa9dca4c..4f0fe2b67593 100644
> --- a/rust/kernel/alloc/allocator.rs
> +++ b/rust/kernel/alloc/allocator.rs
> @@ -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, crate::ffi::c_ulong, 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_NODE: 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_NODE: 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_NODE: Self = Self(bindings::kvrealloc_node);
I think those constants can remain to be named just KREALLOC, etc.
© 2016 - 2025 Red Hat, Inc.