Re: [PATCH 1/2] rust: page: add Page::alloc_page_memcg()

["T.J. Mercier" <tjmercier@google.com>]

On Fri, Oct 16, 2026 at 6:30 AM Alice Ryhl <aliceryhl@google.com> wrote:
>
> Pages allocated on behalf of a userspace process and mapped into its
> address space should be accounted to the process's memory cgroup. Thus,
> add `Page::alloc_page_memcg()` to charge newly allocated pages using
> `mem_cgroup_charge()`, and switch `Page::drop()` from `__free_pages()`
> to `put_page()` so that charged folios are uncharged when freed.
>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> ---
>  rust/helpers/page.c | 17 +++++++++++++++++
>  rust/kernel/page.rs | 50 +++++++++++++++++++++++++++++++++++++++++++++++++-
>  2 files changed, 66 insertions(+), 1 deletion(-)
>
> diff --git a/rust/helpers/page.c b/rust/helpers/page.c
> index f8463fbed2a2..145758a8a954 100644
> --- a/rust/helpers/page.c
> +++ b/rust/helpers/page.c
> @@ -2,6 +2,7 @@
>
>  #include <linux/gfp.h>
>  #include <linux/highmem.h>
> +#include <linux/memcontrol.h>
>  #include <linux/mm.h>
>
>  __rust_helper struct page *rust_helper_alloc_pages(gfp_t gfp_mask,
> @@ -10,6 +11,22 @@ __rust_helper struct page *rust_helper_alloc_pages(gfp_t gfp_mask,
>         return alloc_pages(gfp_mask, order);
>  }
>
> +__rust_helper void rust_helper_put_page(struct page *page)
> +{
> +       put_page(page);
> +}
> +
> +__rust_helper struct folio *rust_helper_page_folio(struct page *page)
> +{
> +       return page_folio(page);
> +}
> +
> +__rust_helper int rust_helper_mem_cgroup_charge(struct folio *folio,
> +                                               struct mm_struct *mm, gfp_t gfp)
> +{
> +       return mem_cgroup_charge(folio, mm, gfp);
> +}
> +
>  __rust_helper void *rust_helper_kmap_local_page(struct page *page)
>  {
>         return kmap_local_page(page);
> diff --git a/rust/kernel/page.rs b/rust/kernel/page.rs
> index 1c0796ea229f..7a36b7dd6c9c 100644
> --- a/rust/kernel/page.rs
> +++ b/rust/kernel/page.rs
> @@ -10,8 +10,10 @@
>      bindings,
>      error::{
>          code::*,
> +        Error,
>          Result, //
>      },
> +    mm::Mm,
>      uaccess::UserSliceReader, //
>  };
>  use core::{
> @@ -195,6 +197,28 @@ pub fn alloc_page(flags: Flags) -> Result<Self, AllocError> {
>          Ok(Self { page })
>      }
>
> +    /// Allocates a new page and charges it to the memory cgroup of `mm`.
> +    #[inline]
> +    pub fn alloc_page_memcg(flags: Flags, mm: &Mm) -> Result<Self, MemcgAllocError> {
> +        let page = Self::alloc_page(flags)?;
> +        // Note: Charging the folio with `mem_cgroup_charge()` only works because `Page::drop` uses
> +        // `put_page()` (which uncharges the folio before freeing it) rather than `__free_pages()`.
> +        //
> +        // SAFETY: `page` was just allocated, so we hold the only reference to it and it is not
> +        // yet charged to any memcg. `mm.as_raw()` is a valid pointer to an `mm_struct`.
> +        let ret = unsafe {
> +            bindings::mem_cgroup_charge(
> +                bindings::page_folio(page.as_ptr()),
> +                mm.as_raw(),
> +                flags.as_raw(),

I wonder if we should we mask out __GFP_ACCOUNT from flags? It's not
provided by the new caller in the next patch, but it could be included
inadvertently in a future call elsewhere. The VM_BUG_ON_FOLIO that
checks if the page would be double charged only happens when
CONFIG_DEBUG_VM=y. Without it on production builds, folio->memcg_data
just gets overwritten. We'd probably catch that eventually in a
CONFIG_DEBUG_VM build, but not sure if we should rely on that. I don't
feel strongly, just


> +            )
> +        };
> +        if ret != 0 {
> +            return Err(MemcgAllocError::Memcg);
> +        }
> +        Ok(page)
> +    }
> +
>      /// Returns a raw pointer to the page.
>      pub fn as_ptr(&self) -> *mut bindings::page {
>          self.page.as_ptr()
> @@ -376,7 +400,31 @@ pub unsafe fn copy_from_user_slice_raw(
>  impl Drop for Page {
>      #[inline]
>      fn drop(&mut self) {
> +        // We need to use `put_page()` instead of `__free_pages()` to properly handle memcg
> +        // accounting.
> +        //
>          // SAFETY: By the type invariants, we have ownership of the page and can free it.
> -        unsafe { bindings::__free_pages(self.page.as_ptr(), 0) };
> +        unsafe { bindings::put_page(self.page.as_ptr()) };
> +    }
> +}
> +
> +/// Error returned by [`Page::alloc_page_memcg`].
> +#[derive(Copy, Clone, PartialEq, Eq, Debug)]
> +pub enum MemcgAllocError {
> +    /// Allocating the page failed.
> +    Alloc,
> +    /// Charging the page to the memory cgroup failed.
> +    Memcg,
> +}
> +
> +impl From<AllocError> for MemcgAllocError {
> +    fn from(_: AllocError) -> Self {
> +        Self::Alloc
> +    }
> +}
> +
> +impl From<MemcgAllocError> for Error {
> +    fn from(_: MemcgAllocError) -> Error {
> +        ENOMEM
>      }
>  }
>
> --
> 2.56.0.212.g129ede20d9-goog
>