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

[Andreas Hindborg]

"Gary Guo" <gary@xxxxxxxxxxx> writes:

> 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@xxxxxxxxxx>
>> ---
>> 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@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
>> ---
>>  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,
>         }
>     );

Sure, we can do that. It's a bit long to write, but I guess that is OK.

Best regards,
Andreas Hindborg