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

[Andreas Hindborg <a.hindborg@kernel.org>]

Alice Ryhl <aliceryhl@google.com> 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@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.

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?

Best regards,
Andreas Hindborg