Rust for Linux List
 help / color / mirror / Atom feed
* [PATCH v2] rust: alloc: add per-task memalloc scope abstractions
@ 2026-06-05 10:54 Andreas Hindborg
  2026-06-05 11:47 ` Alice Ryhl
  0 siblings, 1 reply; 2+ messages in thread
From: Andreas Hindborg @ 2026-06-05 10:54 UTC (permalink / raw)
  To: Miguel Ojeda, Gary Guo, Björn Roy Baron, Benno Lossin,
	Alice Ryhl, Trevor Gross, Danilo Krummrich, Uladzislau Rezki,
	Boqun Feng, Lorenzo Stoakes, Vlastimil Babka, Liam R. Howlett
  Cc: rust-for-linux, linux-kernel, Andreas Hindborg, Boqun Feng,
	Lorenzo Stoakes, Liam R. Howlett, Vlastimil Babka, linux-mm

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@kernel.org>
---
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@kernel.org

To: Miguel Ojeda <ojeda@kernel.org>
To: Boqun Feng <boqun@kernel.org>
To: Gary Guo <gary@garyguo.net>
To: Björn Roy Baron <bjorn3_gh@protonmail.com>
To: Benno Lossin <lossin@kernel.org>
To: Andreas Hindborg <a.hindborg@kernel.org>
To: Alice Ryhl <aliceryhl@google.com>
To: Trevor Gross <tmgross@umich.edu>
To: Danilo Krummrich <dakr@kernel.org>
To: Lorenzo Stoakes <ljs@kernel.org>
To: "Liam R. Howlett" <liam@infradead.org>
To: Vlastimil Babka <vbabka@kernel.org>
To: Uladzislau Rezki <urezki@gmail.com>
Cc: linux-kernel@vger.kernel.org
Cc: rust-for-linux@vger.kernel.org
Cc: linux-mm@kvack.org
---
 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@kernel.org>



^ permalink raw reply related	[flat|nested] 2+ messages in thread

* Re: [PATCH v2] rust: alloc: add per-task memalloc scope abstractions
  2026-06-05 10:54 [PATCH v2] rust: alloc: add per-task memalloc scope abstractions Andreas Hindborg
@ 2026-06-05 11:47 ` Alice Ryhl
  0 siblings, 0 replies; 2+ messages in thread
From: Alice Ryhl @ 2026-06-05 11:47 UTC (permalink / raw)
  To: Andreas Hindborg
  Cc: Miguel Ojeda, Gary Guo, Björn Roy Baron, Benno Lossin,
	Trevor Gross, Danilo Krummrich, Uladzislau Rezki, Boqun Feng,
	Lorenzo Stoakes, Vlastimil Babka, Liam R. Howlett, rust-for-linux,
	linux-kernel, linux-mm

On Fri, Jun 05, 2026 at 12:54:41PM +0200, Andreas Hindborg wrote:
> 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@kernel.org>
> ---
> 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@kernel.org
> 
> To: Miguel Ojeda <ojeda@kernel.org>
> To: Boqun Feng <boqun@kernel.org>
> To: Gary Guo <gary@garyguo.net>
> To: Björn Roy Baron <bjorn3_gh@protonmail.com>
> To: Benno Lossin <lossin@kernel.org>
> To: Andreas Hindborg <a.hindborg@kernel.org>
> To: Alice Ryhl <aliceryhl@google.com>
> To: Trevor Gross <tmgross@umich.edu>
> To: Danilo Krummrich <dakr@kernel.org>
> To: Lorenzo Stoakes <ljs@kernel.org>
> To: "Liam R. Howlett" <liam@infradead.org>
> To: Vlastimil Babka <vbabka@kernel.org>
> To: Uladzislau Rezki <urezki@gmail.com>
> Cc: linux-kernel@vger.kernel.org
> Cc: rust-for-linux@vger.kernel.org
> Cc: linux-mm@kvack.org
> ---
>  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.

Your scope trick only works in normal fns, not in generators such as
async fn.

> +//! [`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);

If we're not going to access this value, then I'd just do:

fn process_io_request() {
    memalloc_scope!(NoIo);
}

or

fn process_io_request() {
    memalloc_noio_scope!();
}

> +/// 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);
> +}

I think all this doc(hidden) + sealing + defining structs via macros is
unnecessary. Just make a normal trait. Or even just define four structs.

Alice

^ permalink raw reply	[flat|nested] 2+ messages in thread

end of thread, other threads:[~2026-06-05 11:47 UTC | newest]

Thread overview: 2+ messages (download: mbox.gz follow: Atom feed
-- links below jump to the message on this page --
2026-06-05 10:54 [PATCH v2] rust: alloc: add per-task memalloc scope abstractions Andreas Hindborg
2026-06-05 11:47 ` Alice Ryhl

This is a public inbox, see mirroring instructions
for how to clone and mirror all data and code used for this inbox