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

[Alice Ryhl <aliceryhl@google.com>]

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