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

["Gary Guo" <gary@garyguo.net>]

On Mon Nov 9, 2026 at 10:37 AM GMT, 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 sets a bit in `current->flags` and returns
> the bits it added, which the restore call clears again. The pairing
> assumes strict LIFO nesting; restoring out of order corrupts the
> per-task state.
>
> Expose each pair as an implementation of a `ScopeKind` trait on the
> tag types `NoIo`, `NoFs`, `NoReclaim` and `MemallocPin`. The trait's
> provided `with` method saves the policy, runs a closure and restores
> the policy before returning.
>
> Signed-off-by: Andreas Hindborg <a.hindborg@kernel.org>
> ---
> Changes in v3:
> - Rebase on v7.4-rc1.
> - Replace the pinned `Scope` guard and `memalloc_scope!` macro with a
>   closure-based `ScopeKind::with` entry point (Alice, Gary).
> - Make the doctests compile and run under KUnit instead of `ignore`.
> - Document that `NoReclaim` has no effect in hard interrupt context.
> - Link to v2: https://msgid.link/20260605-gfp-noio-v2-1-cfa2e236c2a3@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
> ---
>  rust/bindings/bindings_helper.h |   1 +
>  rust/helpers/mm.c               |  40 ++++++++++
>  rust/kernel/alloc.rs            |   1 +
>  rust/kernel/alloc/scoped.rs     | 171 ++++++++++++++++++++++++++++++++++++++++
>  4 files changed, 213 insertions(+)
>
> diff --git a/rust/bindings/bindings_helper.h b/rust/bindings/bindings_helper.h
> index 3791af0aadc7..df21b5b129a4 100644
> --- a/rust/bindings/bindings_helper.h
> +++ b/rust/bindings/bindings_helper.h
> @@ -86,6 +86,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/serdev.h>
>  #include <linux/slab.h>
> diff --git a/rust/helpers/mm.c b/rust/helpers/mm.c
> index b5540997bd20..1e490b57057e 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);
>  }
> +
> +__rust_helper unsigned int rust_helper_memalloc_noio_save(void)
> +{
> +	return memalloc_noio_save();
> +}
> +
> +__rust_helper void rust_helper_memalloc_noio_restore(unsigned int flags)
> +{
> +	memalloc_noio_restore(flags);
> +}
> +
> +__rust_helper unsigned int rust_helper_memalloc_nofs_save(void)
> +{
> +	return memalloc_nofs_save();
> +}
> +
> +__rust_helper void rust_helper_memalloc_nofs_restore(unsigned int flags)
> +{
> +	memalloc_nofs_restore(flags);
> +}
> +
> +__rust_helper unsigned int rust_helper_memalloc_noreclaim_save(void)
> +{
> +	return memalloc_noreclaim_save();
> +}
> +
> +__rust_helper void rust_helper_memalloc_noreclaim_restore(unsigned int flags)
> +{
> +	memalloc_noreclaim_restore(flags);
> +}
> +
> +__rust_helper unsigned int rust_helper_memalloc_pin_save(void)
> +{
> +	return memalloc_pin_save();
> +}
> +
> +__rust_helper 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 68d80f3fa6c8..5e4c142b5682 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..6745c06d6725
> --- /dev/null
> +++ b/rust/kernel/alloc/scoped.rs
> @@ -0,0 +1,171 @@
> +// 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 bits it added, which a
> +//! later call clears again. The save/restore APIs assume strict LIFO nesting; restoring out of
> +//! order corrupts the per-task state.
> +//!
> +//! This module exposes each policy as a [`ScopeKind`] implementation. [`ScopeKind::with`] saves
> +//! the policy, runs a closure and restores the policy before returning.
> +//!
> +//! [`include/linux/sched/mm.h`]: srctree/include/linux/sched/mm.h
> +//!
> +//! # Examples
> +//!
> +//! ```
> +//! use kernel::alloc::scoped::{NoIo, ScopeKind};
> +//!
> +//! // Every allocation inside the closure behaves as if `GFP_NOIO` were set, even though the call
> +//! // site passes `GFP_KERNEL`.
> +//! let b = NoIo::with(|| KBox::new(42_u32, GFP_KERNEL))?;
> +//! assert_eq!(*b, 42);
> +//! # Ok::<(), Error>(())

I think this usage looks a bit strange, especially that you have to import the
`ScopeKind` trait which is a bit awkward API.

Any reason to not just have API like this?

    memalloc_flags_scope(MemallocScopeFlags::NO_IO, || KBox::new(42_u32, GFP_KERNEL))

Added benefit is that you can just define different scope kinds as bitflags:

    impl_flags!(
        #[derive(Clone, Copy)]
        pub struct MemallocScopeFlags(u32);

        #[derive(Clone, Copy)]
        pub enum MemallocScopeFlags {
            NO_IO = PF_MEMALLOC_NOIO,
            NO_FS = PF_MEMALLOC_NOFS,
            NO_RECLAIM = PF_MEMALLOC,
            PIN = PF_MEMALLOC_PIN,
        }
    );

Best,
Gary


> +//! ```
> +// # Implementation details
> +//
> +// Since the closure used in this API is synchronous it cannot contain a suspension point, so the
> +// save and restore always execute in the same call frame, on the same task, and nested scopes
> +// restore in LIFO order.
> +use crate::ffi::c_uint;
> +
> +/// A per-task allocation policy with a `memalloc_*` save/restore pair.
> +///
> +/// Implemented by the uninhabited types [`NoIo`], [`NoFs`], [`NoReclaim`] and [`MemallocPin`]. Use
> +/// [`ScopeKind::with`] to run code under the policy.
> +pub trait ScopeKind {
> +    /// Begin a scope on the current task and return the flags added to `current->flags`.
> +    ///
> +    /// Every call must be paired with a [`restore`] call on the same task, in LIFO order with
> +    /// respect to other scopes. Prefer [`with`], which arranges this automatically.
> +    ///
> +    /// [`restore`]: ScopeKind::restore
> +    /// [`with`]: ScopeKind::with
> +    fn save() -> c_uint;
> +
> +    /// End a scope on the current task.
> +    ///
> +    /// # Safety
> +    ///
> +    /// `prev` must be the value returned by the matching [`save`] call, the call must execute on
> +    /// the same task that ran [`save`], and no scope begun after the matching [`save`] may still be
> +    /// active.
> +    ///
> +    /// [`save`]: ScopeKind::save
> +    unsafe fn restore(prev: c_uint);
> +
> +    /// Run `f` with this policy applied.
> +    ///
> +    /// The current policy is saved before `f` runs and restored when `f` returns.
> +    fn with<R>(f: impl FnOnce() -> R) -> R {
> +        let prev = Self::save();
> +        let ret = f();
> +        // SAFETY: `prev` was produced by the `save` above on this task. `f` is a synchronous
> +        // closure, so it has already returned and any scope it began has already been restored.
> +        unsafe { Self::restore(prev) };
> +        ret
> +    }
> +}
> +