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

[Gary Guo]

On Sat Jun 6, 2026 at 9:14 AM BST, Andreas Hindborg wrote:
> Alice Ryhl <aliceryhl@xxxxxxxxxx> writes:
>
>> 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@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.
>>
>> Your scope trick only works in normal fns, not in generators such as
>> async fn.
>
> Could you elaborate what would happen when this pattern is applied in a
> generator? I don't immediately see how the LIFO drop order can be broken
> when this is combined with a generator.
>
> Would a closure based approach have the same problem?

Generator/async allows you to do this:

fn bad_fn() {
    let mut p1 = Box::pin(async {
        memalloc_scope!(let _noio: NoIo);
        pending().await;
    });
    let mut p2 = Box::pin(async {
        memalloc_scope!(let _nofs: NoFs);
        pending().await;
    });

    let mut ctx = Context::from_waker(Waker::noop());
    p1.as_mut().poll(&mut ctx); // Save mask and set no IO
    p2.as_mut().poll(&mut ctx); // Save mask again and set no FS
    drop(p1); // Restore the original mask
    drop(p2); // Restore to the mask after no IO save, which is wrong
}

Best,
Gary

>
> Best regards,
> Andreas Hindborg