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

[Alice Ryhl <aliceryhl@google.com>]

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(),
+            )
+        };
+        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