[PATCH v3 1/2] rust: page: use the page's reference count to decide when to free the allocation

Abdiel Janulgue posted 2 patches 4 days, 3 hours ago
[PATCH v3 1/2] rust: page: use the page's reference count to decide when to free the allocation
Posted by Abdiel Janulgue 4 days, 3 hours ago
Ensure pages returned by the constructor are always reference counted.
This requires that we replace the page pointer wrapper with Opaque instead
of NonNull to make it possible to cast to a Page pointer from a raw struct
page pointer which is needed to create an ARef instance.

Signed-off-by: Abdiel Janulgue <abdiel.janulgue@gmail.com>
---
 rust/bindings/bindings_helper.h |  1 +
 rust/helpers/page.c             | 10 +++++++++
 rust/kernel/page.rs             | 38 ++++++++++++++++++++++-----------
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/rust/bindings/bindings_helper.h b/rust/bindings/bindings_helper.h
index a80783fcbe04..daa3225a185f 100644
--- a/rust/bindings/bindings_helper.h
+++ b/rust/bindings/bindings_helper.h
@@ -15,6 +15,7 @@
 #include <linux/firmware.h>
 #include <linux/jiffies.h>
 #include <linux/mdio.h>
+#include <linux/mm.h>
 #include <linux/phy.h>
 #include <linux/refcount.h>
 #include <linux/sched.h>
diff --git a/rust/helpers/page.c b/rust/helpers/page.c
index b3f2b8fbf87f..48d4481c1e33 100644
--- a/rust/helpers/page.c
+++ b/rust/helpers/page.c
@@ -17,3 +17,13 @@ void rust_helper_kunmap_local(const void *addr)
 {
 	kunmap_local(addr);
 }
+
+void rust_helper_put_page(struct page *page)
+{
+	put_page(page);
+}
+
+void rust_helper_get_page(struct page *page)
+{
+	get_page(page);
+}
diff --git a/rust/kernel/page.rs b/rust/kernel/page.rs
index fdac6c375fe4..fdf7ee203597 100644
--- a/rust/kernel/page.rs
+++ b/rust/kernel/page.rs
@@ -8,6 +8,7 @@
     error::code::*,
     error::Result,
     uaccess::UserSliceReader,
+    types::{Opaque, ARef},
 };
 use core::ptr::{self, NonNull};
 
@@ -30,13 +31,14 @@ pub const fn page_align(addr: usize) -> usize {
     (addr + (PAGE_SIZE - 1)) & PAGE_MASK
 }
 
-/// A pointer to a page that owns the page allocation.
+/// A pointer to a reference-counted page.
 ///
 /// # Invariants
 ///
-/// The pointer is valid, and has ownership over the page.
+/// The pointer is valid.
+#[repr(transparent)]
 pub struct Page {
-    page: NonNull<bindings::page>,
+    page: Opaque<bindings::page>,
 }
 
 // SAFETY: Pages have no logic that relies on them staying on a given thread, so moving them across
@@ -71,19 +73,23 @@ impl Page {
     /// let page = Page::alloc_page(GFP_KERNEL | __GFP_ZERO)?;
     /// # Ok(()) }
     /// ```
-    pub fn alloc_page(flags: Flags) -> Result<Self, AllocError> {
+    pub fn alloc_page(flags: Flags) -> Result<ARef<Self>, AllocError> {
         // SAFETY: Depending on the value of `gfp_flags`, this call may sleep. Other than that, it
         // is always safe to call this method.
         let page = unsafe { bindings::alloc_pages(flags.as_raw(), 0) };
-        let page = NonNull::new(page).ok_or(AllocError)?;
-        // INVARIANT: We just successfully allocated a page, so we now have ownership of the newly
-        // allocated page. We transfer that ownership to the new `Page` object.
-        Ok(Self { page })
+        if page.is_null() {
+            return Err(AllocError);
+        }
+        // CAST: Self` is a `repr(transparent)` wrapper around `bindings::page`.
+        let ptr = page.cast::<Self>();
+        // INVARIANT: We just successfully allocated a page, ptr points to the new `Page` object.
+        // SAFETY: According to invariant above ptr is valid.
+        Ok(unsafe { ARef::from_raw(NonNull::new_unchecked(ptr)) })
     }
 
     /// Returns a raw pointer to the page.
     pub fn as_ptr(&self) -> *mut bindings::page {
-        self.page.as_ptr()
+        self.page.get()
     }
 
     /// Runs a piece of code with this page mapped to an address.
@@ -252,9 +258,15 @@ pub unsafe fn copy_from_user_slice_raw(
     }
 }
 
-impl Drop for Page {
-    fn drop(&mut self) {
-        // SAFETY: By the type invariants, we have ownership of the page and can free it.
-        unsafe { bindings::__free_pages(self.page.as_ptr(), 0) };
+// SAFETY: Instances of `Page` are always reference-counted.
+unsafe impl crate::types::AlwaysRefCounted for Page {
+    fn inc_ref(&self) {
+        // SAFETY: The existence of a shared reference guarantees that the refcount is non-zero.
+        unsafe { bindings::get_page(self.as_ptr()) };
+    }
+
+    unsafe fn dec_ref(obj: ptr::NonNull<Self>) {
+        // SAFETY: The safety requirements guarantee that the refcount is non-zero.
+        unsafe { bindings::put_page(obj.cast().as_ptr()) }
     }
 }
-- 
2.43.0
Re: [PATCH v3 1/2] rust: page: use the page's reference count to decide when to free the allocation
Posted by Alice Ryhl 4 days, 2 hours ago
On Tue, Nov 19, 2024 at 12:24 PM Abdiel Janulgue
<abdiel.janulgue@gmail.com> wrote:
>
> Ensure pages returned by the constructor are always reference counted.
> This requires that we replace the page pointer wrapper with Opaque instead
> of NonNull to make it possible to cast to a Page pointer from a raw struct
> page pointer which is needed to create an ARef instance.
>
> Signed-off-by: Abdiel Janulgue <abdiel.janulgue@gmail.com>

> -/// A pointer to a page that owns the page allocation.
> +/// A pointer to a reference-counted page.
>  ///
>  /// # Invariants
>  ///
> -/// The pointer is valid, and has ownership over the page.
> +/// The pointer is valid.
> +#[repr(transparent)]
>  pub struct Page {
> -    page: NonNull<bindings::page>,
> +    page: Opaque<bindings::page>,

With this change, Page is no longer a pointer, nor does it contain a
pointer. The documentation should be updated to reflect this.

>  // SAFETY: Pages have no logic that relies on them staying on a given thread, so moving them across
> @@ -71,19 +73,23 @@ impl Page {
>      /// let page = Page::alloc_page(GFP_KERNEL | __GFP_ZERO)?;
>      /// # Ok(()) }
>      /// ```
> -    pub fn alloc_page(flags: Flags) -> Result<Self, AllocError> {
> +    pub fn alloc_page(flags: Flags) -> Result<ARef<Self>, AllocError> {
>          // SAFETY: Depending on the value of `gfp_flags`, this call may sleep. Other than that, it
>          // is always safe to call this method.
>          let page = unsafe { bindings::alloc_pages(flags.as_raw(), 0) };
> -        let page = NonNull::new(page).ok_or(AllocError)?;
> -        // INVARIANT: We just successfully allocated a page, so we now have ownership of the newly
> -        // allocated page. We transfer that ownership to the new `Page` object.
> -        Ok(Self { page })
> +        if page.is_null() {
> +            return Err(AllocError);
> +        }
> +        // CAST: Self` is a `repr(transparent)` wrapper around `bindings::page`.
> +        let ptr = page.cast::<Self>();
> +        // INVARIANT: We just successfully allocated a page, ptr points to the new `Page` object.
> +        // SAFETY: According to invariant above ptr is valid.
> +        Ok(unsafe { ARef::from_raw(NonNull::new_unchecked(ptr)) })

Why did you change the null check? You should be able to avoid
changing anything but the last line.

Alice
Re: [PATCH v3 1/2] rust: page: use the page's reference count to decide when to free the allocation
Posted by Abdiel Janulgue 4 days, 2 hours ago
On 19/11/2024 13:45, Alice Ryhl wrote:
>> +    pub fn alloc_page(flags: Flags) -> Result<ARef<Self>, AllocError> {
>>           // SAFETY: Depending on the value of `gfp_flags`, this call may sleep. Other than that, it
>>           // is always safe to call this method.
>>           let page = unsafe { bindings::alloc_pages(flags.as_raw(), 0) };
>> -        let page = NonNull::new(page).ok_or(AllocError)?;
>> -        // INVARIANT: We just successfully allocated a page, so we now have ownership of the newly
>> -        // allocated page. We transfer that ownership to the new `Page` object.
>> -        Ok(Self { page })
>> +        if page.is_null() {
>> +            return Err(AllocError);
>> +        }
>> +        // CAST: Self` is a `repr(transparent)` wrapper around `bindings::page`.
>> +        let ptr = page.cast::<Self>();
>> +        // INVARIANT: We just successfully allocated a page, ptr points to the new `Page` object.
>> +        // SAFETY: According to invariant above ptr is valid.
>> +        Ok(unsafe { ARef::from_raw(NonNull::new_unchecked(ptr)) })
> 
> Why did you change the null check? You should be able to avoid
> changing anything but the last line.

Changing only the line, it complains:

86  |         Ok(unsafe { ARef::from_raw(page) })
     |                     -------------- ^^^^ expected `NonNull<Page>`, 
found `NonNull<page>`

Unless this is what you mean?

         let page = unsafe { bindings::alloc_pages(flags.as_raw(), 0) };
         let page = page.cast::<Self>();
         let page = NonNull::new(page).ok_or(AllocError)?;
         Ok(unsafe { ARef::from_raw(page) })

But what if alloc_pages returns null in the place? Would that be a valid 
cast still?

Regards,
Abdiel
Re: [PATCH v3 1/2] rust: page: use the page's reference count to decide when to free the allocation
Posted by Alice Ryhl 4 days, 2 hours ago
On Tue, Nov 19, 2024 at 1:06 PM Abdiel Janulgue
<abdiel.janulgue@gmail.com> wrote:
>
>
> On 19/11/2024 13:45, Alice Ryhl wrote:
> >> +    pub fn alloc_page(flags: Flags) -> Result<ARef<Self>, AllocError> {
> >>           // SAFETY: Depending on the value of `gfp_flags`, this call may sleep. Other than that, it
> >>           // is always safe to call this method.
> >>           let page = unsafe { bindings::alloc_pages(flags.as_raw(), 0) };
> >> -        let page = NonNull::new(page).ok_or(AllocError)?;
> >> -        // INVARIANT: We just successfully allocated a page, so we now have ownership of the newly
> >> -        // allocated page. We transfer that ownership to the new `Page` object.
> >> -        Ok(Self { page })
> >> +        if page.is_null() {
> >> +            return Err(AllocError);
> >> +        }
> >> +        // CAST: Self` is a `repr(transparent)` wrapper around `bindings::page`.
> >> +        let ptr = page.cast::<Self>();
> >> +        // INVARIANT: We just successfully allocated a page, ptr points to the new `Page` object.
> >> +        // SAFETY: According to invariant above ptr is valid.
> >> +        Ok(unsafe { ARef::from_raw(NonNull::new_unchecked(ptr)) })
> >
> > Why did you change the null check? You should be able to avoid
> > changing anything but the last line.
>
> Changing only the line, it complains:
>
> 86  |         Ok(unsafe { ARef::from_raw(page) })
>      |                     -------------- ^^^^ expected `NonNull<Page>`,
> found `NonNull<page>`
>
> Unless this is what you mean?
>
>          let page = unsafe { bindings::alloc_pages(flags.as_raw(), 0) };
>          let page = page.cast::<Self>();
>          let page = NonNull::new(page).ok_or(AllocError)?;
>          Ok(unsafe { ARef::from_raw(page) })

You can put the cast here:

let page = unsafe { bindings::alloc_pages(flags.as_raw(), 0) };
let page = NonNull::new(page).ok_or(AllocError)?;
Ok(unsafe { ARef::from_raw(page.cast()) })

> But what if alloc_pages returns null in the place? Would that be a valid
> cast still?

The cast is correct both before and after the null check. The above
code returns Err(AllocError) when the pointer is null.

Alice