[PATCH v2] rust: alloc: add per-task memalloc scope abstractions

[Andreas Hindborg]

Add an abstraction for the per-task allocation policies exposed by
the kernel through paired save/restore helpers in `linux/sched/mm.h`:
`memalloc_noio`, `memalloc_nofs`, `memalloc_noreclaim` and
`memalloc_pin`. Each pair toggles a bit in `current->flags` and
returns the prior state for a later restore. The pairing assumes
strict LIFO nesting; restoring out of order corrupts the per-task
state.

Wrap the four pairs as a generic `Scope<K>` guard with a sealed
`ScopeKind` trait. Tag types `NoIo`, `NoFs`, `NoReclaim` and
`MemallocPin` select the underlying save/restore pair. `Scope` is
`!Unpin`, `!Send` and `!Sync`, and is only constructed through the
`memalloc_scope!` macro, which binds it via `core::pin::pin!` to a
hidden stack slot and hands out a `Pin<&Scope<K>>`. Safe code
therefore cannot move the guard across tasks, drop it ahead of its
lexical scope or otherwise violate the LIFO save/restore discipline.

Signed-off-by: Andreas Hindborg <a.hindborg@xxxxxxxxxx>
---
Changes in v2:
- Rewrite the patch to use scoped allocation flags instead of exposing
  a `GFP_NOIO` flag constant.
- Link to v1: https://lore.kernel.org/r/20260128-gfp-noio-v1-1-9a808fc49b44@xxxxxxxxxx

To: Miguel Ojeda <ojeda@xxxxxxxxxx>
To: Boqun Feng <boqun@xxxxxxxxxx>
To: Gary Guo <gary@xxxxxxxxxxx>
To: Björn Roy Baron <bjorn3_gh@xxxxxxxxxxxxxx>
To: Benno Lossin <lossin@xxxxxxxxxx>
To: Andreas Hindborg <a.hindborg@xxxxxxxxxx>
To: Alice Ryhl <aliceryhl@xxxxxxxxxx>
To: Trevor Gross <tmgross@xxxxxxxxx>
To: Danilo Krummrich <dakr@xxxxxxxxxx>
To: Lorenzo Stoakes <ljs@xxxxxxxxxx>
To: "Liam R. Howlett" <liam@xxxxxxxxxxxxx>
To: Vlastimil Babka <vbabka@xxxxxxxxxx>
To: Uladzislau Rezki <urezki@xxxxxxxxx>
Cc: linux-kernel@xxxxxxxxxxxxxxx
Cc: rust-for-linux@xxxxxxxxxxxxxxx
Cc: linux-mm@xxxxxxxxx
---
 rust/bindings/bindings_helper.h |   1 +
 rust/helpers/mm.c               |  40 +++++++
 rust/kernel/alloc.rs            |   1 +
 rust/kernel/alloc/scoped.rs     | 231 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 273 insertions(+)

diff --git a/rust/bindings/bindings_helper.h b/rust/bindings/bindings_helper.h
index 446dbeaf0866..1931b131345f 100644
--- a/rust/bindings/bindings_helper.h
+++ b/rust/bindings/bindings_helper.h
@@ -83,6 +83,7 @@
 #include <linux/refcount.h>
 #include <linux/regulator/consumer.h>
 #include <linux/sched.h>
+#include <linux/sched/mm.h>
 #include <linux/security.h>
 #include <linux/slab.h>
 #include <linux/sys_soc.h>
diff --git a/rust/helpers/mm.c b/rust/helpers/mm.c
index b5540997bd20..b8e7492512e8 100644
--- a/rust/helpers/mm.c
+++ b/rust/helpers/mm.c
@@ -48,3 +48,43 @@ __rust_helper void rust_helper_vma_end_read(struct vm_area_struct *vma)
 {
 	vma_end_read(vma);
 }
+
+unsigned int rust_helper_memalloc_noio_save(void)
+{
+	return memalloc_noio_save();
+}
+
+void rust_helper_memalloc_noio_restore(unsigned int flags)
+{
+	memalloc_noio_restore(flags);
+}
+
+unsigned int rust_helper_memalloc_nofs_save(void)
+{
+	return memalloc_nofs_save();
+}
+
+void rust_helper_memalloc_nofs_restore(unsigned int flags)
+{
+	memalloc_nofs_restore(flags);
+}
+
+unsigned int rust_helper_memalloc_noreclaim_save(void)
+{
+	return memalloc_noreclaim_save();
+}
+
+void rust_helper_memalloc_noreclaim_restore(unsigned int flags)
+{
+	memalloc_noreclaim_restore(flags);
+}
+
+unsigned int rust_helper_memalloc_pin_save(void)
+{
+	return memalloc_pin_save();
+}
+
+void rust_helper_memalloc_pin_restore(unsigned int flags)
+{
+	memalloc_pin_restore(flags);
+}
diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index e38720349dcf..8ebb8c9f3e67 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -6,6 +6,7 @@
 pub mod kbox;
 pub mod kvec;
 pub mod layout;
+pub mod scoped;
 
 pub use self::kbox::Box;
 pub use self::kbox::KBox;
diff --git a/rust/kernel/alloc/scoped.rs b/rust/kernel/alloc/scoped.rs
new file mode 100644
index 000000000000..0251792c9f3c
--- /dev/null
+++ b/rust/kernel/alloc/scoped.rs
@@ -0,0 +1,231 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Scoped allocation policies for the current task.
+//!
+//! The kernel exposes several per-task allocation policies through
+//! save/restore pairs in [`include/linux/sched/mm.h`]: `memalloc_noio`,
+//! `memalloc_nofs`, `memalloc_noreclaim` and `memalloc_pin`. Each pair
+//! sets a bit in `current->flags` and returns the prior state, which a
+//! later call restores. The save/restore APIs assume strict LIFO
+//! nesting; restoring out of order corrupts the per-task state.
+//!
+//! This module exposes the policies as a generic [`Scope<K>`] guard,
+//! parameterized over a [`ScopeKind`] tag. The type is `!Unpin` and
+//! constructed only through the [`memalloc_scope!`] macro, which binds
+//! it to a hidden stack slot via [`core::pin::pin!`] and rebinds the
+//! handle as a shared pinned reference. Safe code therefore has no path
+//! to either move the guard or drop it ahead of its lexical scope, so
+//! nested scopes always restore in LIFO order.
+//!
+//! [`include/linux/sched/mm.h`]: srctree/include/linux/sched/mm.h
+//!
+//! # Examples
+//!
+//! ```ignore
+//! use kernel::memalloc_scope;
+//! use kernel::alloc::scoped::NoIo;
+//!
+//! fn process_io_request() {
+//!     memalloc_scope!(let _noio: NoIo);
+//!     // Every allocation in this scope behaves as if `GFP_NOIO` were
+//!     // set, even when the call site passes `GFP_KERNEL`.
+//! }
+//! ```
+
+use core::{
+    ffi::c_uint,
+    marker::{
+        PhantomData,
+        PhantomPinned, //
+    },
+};
+
+use crate::types::NotThreadSafe;
+
+pub use crate::memalloc_scope;
+
+mod private {
+    pub trait Sealed {}
+}
+
+/// Selects which `memalloc_*` save/restore pair a [`Scope`] wraps.
+///
+/// Implemented only by the zero-sized tag types in this module
+/// ([`NoIo`], [`NoFs`], [`NoReclaim`], [`MemallocPin`]). The trait is
+/// sealed.
+pub trait ScopeKind: private::Sealed {
+    /// Begin a scope on the current task and return the prior state.
+    #[doc(hidden)]
+    fn save() -> c_uint;
+
+    /// End a scope on the current task.
+    ///
+    /// # Safety
+    ///
+    /// `prev` must be the value returned by the matching [`save`] call,
+    /// and the call must execute on the same task that ran [`save`].
+    ///
+    /// [`save`]: ScopeKind::save
+    #[doc(hidden)]
+    unsafe fn restore(prev: c_uint);
+}
+
+/// A scope that imposes an allocation policy on the current task while
+/// it is live.
+///
+/// Construct one with [`memalloc_scope!`]. `Scope` is `!Unpin` and its
+/// constructor is hidden, so a `Scope` only ever exists pinned to a
+/// stack slot owned by the construction macro; safe code cannot drop
+/// it out of order or send it across tasks. The C-side state is
+/// restored in [`Drop`], which runs when the stack slot goes out of
+/// scope.
+pub struct Scope<K: ScopeKind> {
+    prev: c_uint,
+    _kind: PhantomData<K>,
+    _pin: PhantomPinned,
+    _not_thread_safe: NotThreadSafe,
+}
+
+impl<K: ScopeKind> Scope<K> {
+    /// Begin a scope of kind `K` on the current task.
+    ///
+    /// # Safety
+    ///
+    /// The returned value must be pinned to the stack frame that calls
+    /// this function and dropped on the same task. In practice, only
+    /// [`memalloc_scope!`] should call this — the macro arranges both.
+    #[doc(hidden)]
+    pub unsafe fn new() -> Self {
+        Self {
+            prev: K::save(),
+            _kind: PhantomData,
+            _pin: PhantomPinned,
+            _not_thread_safe: NotThreadSafe,
+        }
+    }
+}
+
+impl<K: ScopeKind> Drop for Scope<K> {
+    fn drop(&mut self) {
+        // SAFETY: `self.prev` was produced by `K::save` in `Self::new`.
+        // The caller of `new` upheld the contract that the value
+        // remains pinned to its construction stack frame, so this drop
+        // runs on the same task as the matching save.
+        unsafe { K::restore(self.prev) };
+    }
+}
+
+macro_rules! define_kind {
+    (
+        $(#[$meta:meta])*
+        $name:ident, $save:ident, $restore:ident $(,)?
+    ) => {
+        $(#[$meta])*
+        pub struct $name;
+
+        impl private::Sealed for $name {}
+
+        impl ScopeKind for $name {
+            fn save() -> c_uint {
+                // SAFETY: Updates a per-task flag and is documented as
+                // safe from any context.
+                unsafe { bindings::$save() }
+            }
+
+            unsafe fn restore(prev: c_uint) {
+                // SAFETY: Per the trait contract, `prev` is the value
+                // returned by the matching `save`, on the same task.
+                unsafe { bindings::$restore(prev) };
+            }
+        }
+    };
+}
+
+define_kind!(
+    /// `GFP_NOIO` scope.
+    ///
+    /// While a `Scope<NoIo>` is live, allocations on the current task
+    /// behave as if `GFP_NOIO` were set, making them safe to issue from
+    /// the IO completion path.
+    ///
+    /// Corresponds to `memalloc_noio_save` / `memalloc_noio_restore` in
+    /// `include/linux/sched/mm.h`.
+    NoIo,
+    memalloc_noio_save,
+    memalloc_noio_restore,
+);
+
+define_kind!(
+    /// `GFP_NOFS` scope.
+    ///
+    /// While a `Scope<NoFs>` is live, allocations on the current task
+    /// behave as if `GFP_NOFS` were set, making them safe to issue from
+    /// a filesystem critical section.
+    ///
+    /// Corresponds to `memalloc_nofs_save` / `memalloc_nofs_restore` in
+    /// `include/linux/sched/mm.h`.
+    NoFs,
+    memalloc_nofs_save,
+    memalloc_nofs_restore,
+);
+
+define_kind!(
+    /// No-reclaim scope.
+    ///
+    /// While a `Scope<NoReclaim>` is live, allocations on the current
+    /// task may dip into the memory reserves. Callers must be sure their
+    /// allocations will help free more memory shortly; see the kernel C
+    /// documentation for the full contract.
+    ///
+    /// Corresponds to `memalloc_noreclaim_save` /
+    /// `memalloc_noreclaim_restore` in `include/linux/sched/mm.h`.
+    NoReclaim,
+    memalloc_noreclaim_save,
+    memalloc_noreclaim_restore,
+);
+
+define_kind!(
+    /// Long-term pin scope.
+    ///
+    /// While a `Scope<MemallocPin>` is live, allocations on the current
+    /// task are restricted to zones that allow long-term pinning.
+    ///
+    /// Corresponds to `memalloc_pin_save` / `memalloc_pin_restore` in
+    /// `include/linux/sched/mm.h`.
+    MemallocPin,
+    memalloc_pin_save,
+    memalloc_pin_restore,
+);
+
+/// Bind a [`Scope`] of the given kind to the current stack frame.
+///
+/// `$kind` must name one of the zero-sized tag types defined in this
+/// module: [`NoIo`], [`NoFs`], [`NoReclaim`], [`MemallocPin`]. The
+/// macro shadows `$name` first with the owning pinned slot and then
+/// with a shared pinned reference, so the value lives until the end of
+/// the enclosing block and cannot be dropped early by safe code.
+///
+/// # Examples
+///
+/// ```ignore
+/// use kernel::memalloc_scope;
+/// use kernel::alloc::scoped::NoIo;
+///
+/// memalloc_scope!(let _scope: NoIo);
+/// // ... allocations here behave as if `GFP_NOIO` were set.
+/// ```
+#[macro_export]
+macro_rules! memalloc_scope {
+    (let $name:ident : $kind:ident) => {
+        // SAFETY: `pin!` places the value in a hidden stack slot and
+        // returns a `Pin<&mut _>`; combined with `Scope: !Unpin`, safe
+        // code can neither extract ownership nor reorder its drop
+        // relative to nested scopes, so the save/restore discipline of
+        // the underlying C API is preserved.
+        let $name = ::core::pin::pin!(unsafe {
+            $crate::alloc::scoped::Scope::<$crate::alloc::scoped::$kind>::new()
+        });
+        let $name: ::core::pin::Pin<&$crate::alloc::scoped::Scope<$crate::alloc::scoped::$kind>> =
+            $name.as_ref();
+    };
+}

---
base-commit: 7fd2df204f342fc17d1a0bfcd474b24232fb0f32
change-id: 20260128-gfp-noio-fbd41e135088

Best regards,
--  
Andreas Hindborg <a.hindborg@xxxxxxxxxx>