Re: [PATCH v9 8/8] task: rust: rework how current is accessed

[Lorenzo Stoakes <lorenzo.stoakes@oracle.com>]

On Fri, Nov 22, 2024 at 03:40:33PM +0000, Alice Ryhl wrote:
> Introduce a new type called `CurrentTask` that lets you perform various
> operations that are only safe on the `current` task. Use the new type to
> provide a way to access the current mm without incrementing its
> refcount.

Nice!

>
> With this change, you can write stuff such as
>
> 	let vma = current!().mm().lock_vma_under_rcu(addr);
>
> without incrementing any refcounts.
>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>

On assumption that the problem you reference with the rust imports is
corrected in v10, and that what you are doing with current_raw() is
sensible, then:

Acked-by: Lorenzo Stoakes <lorenzo.stoakes@oracle.com>

Thanks!

> ---
> Reviewers: Does accessing task->mm on a non-current task require rcu
> protection?

Hm I am not actually sure, but it seems like you probably do, and I would say
you need the task lock right?

Looking at find_lock_task_mm() as used by the oomk for instance suggests as much.

>
> Christian: If you submit the PidNamespace abstractions this cycle, I'll
> update this to also apply to PidNamespace.
> ---
>  rust/kernel/mm.rs   | 22 ------------------
>  rust/kernel/task.rs | 64 ++++++++++++++++++++++++++++++++++++++++++-----------
>  2 files changed, 51 insertions(+), 35 deletions(-)
>
> diff --git a/rust/kernel/mm.rs b/rust/kernel/mm.rs
> index 50f4861ae4b9..f7d1079391ef 100644
> --- a/rust/kernel/mm.rs
> +++ b/rust/kernel/mm.rs
> @@ -142,28 +142,6 @@ fn deref(&self) -> &MmWithUser {
>
>  // These methods are safe to call even if `mm_users` is zero.
>  impl Mm {
> -    /// Call `mmgrab` on `current.mm`.
> -    #[inline]
> -    pub fn mmgrab_current() -> Option<ARef<Mm>> {
> -        // SAFETY: It's safe to get the `mm` field from current.
> -        let mm = unsafe {
> -            let current = bindings::get_current();
> -            (*current).mm
> -        };
> -
> -        if mm.is_null() {
> -            return None;
> -        }
> -
> -        // SAFETY: The value of `current->mm` is guaranteed to be null or a valid `mm_struct`. We
> -        // just checked that it's not null. Furthermore, the returned `&Mm` is valid only for the
> -        // duration of this function, and `current->mm` will stay valid for that long.
> -        let mm = unsafe { Mm::from_raw(mm) };
> -
> -        // This increments the refcount using `mmgrab`.
> -        Some(ARef::from(mm))
> -    }
> -

It's nice to drop this to discourage the unusual thing of grabbing current's mm
and incrementing reference count.

>      /// Returns a raw pointer to the inner `mm_struct`.
>      #[inline]
>      pub fn as_raw(&self) -> *mut bindings::mm_struct {
> diff --git a/rust/kernel/task.rs b/rust/kernel/task.rs
> index 9e59d86da42d..103d235eb844 100644
> --- a/rust/kernel/task.rs
> +++ b/rust/kernel/task.rs
> @@ -94,6 +94,26 @@ unsafe impl Send for Task {}
>  // synchronised by C code (e.g., `signal_pending`).
>  unsafe impl Sync for Task {}
>
> +/// Represents a [`Task`] obtained from the `current` global.
> +///
> +/// This type exists to provide more efficient operations that are only valid on the current task.
> +/// For example, to retrieve the pid-namespace of a task, you must use rcu protection unless it is
> +/// the current task.
> +///
> +/// # Invariants
> +///
> +/// Must be equal to `current` of some thread that is currently running somewhere.
> +pub struct CurrentTask(Task);

Nice, I do like the ability to express abstractions like this...

> +
> +// Make all `Task` methods available on `CurrentTask`.
> +impl Deref for CurrentTask {
> +    type Target = Task;
> +    #[inline]
> +    fn deref(&self) -> &Task {
> +        &self.0
> +    }
> +}
> +

It's nice to be able to 'alias' types like this too (ok I'm sure that's not
quite the write way of describing but you know what I mean), so you can abstract
something, then very simply create variants that have the same methods but
different attributes otherwise.

>  /// The type of process identifiers (PIDs).
>  type Pid = bindings::pid_t;
>
> @@ -121,27 +141,25 @@ pub fn current_raw() -> *mut bindings::task_struct {
>      /// # Safety
>      ///
>      /// Callers must ensure that the returned object doesn't outlive the current task/thread.
> -    pub unsafe fn current() -> impl Deref<Target = Task> {
> -        struct TaskRef<'a> {
> -            task: &'a Task,
> -            _not_send: NotThreadSafe,
> +    pub unsafe fn current() -> impl Deref<Target = CurrentTask> {
> +        struct TaskRef {
> +            task: *const CurrentTask,
>          }

Why do we drop the NotThreadSafe bit here? And it seems like the 'a lifetime
stuff has gone too?

I'm guessing the lifetime stuff is because of the SAFETY comment below about
assumptions about lifetime?

>
> -        impl Deref for TaskRef<'_> {
> -            type Target = Task;
> +        impl Deref for TaskRef {
> +            type Target = CurrentTask;
>
>              fn deref(&self) -> &Self::Target {
> -                self.task
> +                // SAFETY: The returned reference borrows from this `TaskRef`, so it cannot outlive
> +                // the `TaskRef`, which the caller of `Task::current()` has promised will not
> +                // outlive the task/thread for which `self.task` is the `current` pointer. Thus, it
> +                // is okay to return a `CurrentTask` reference here.
> +                unsafe { &*self.task }
>              }
>          }
>
> -        let current = Task::current_raw();
>          TaskRef {
> -            // SAFETY: If the current thread is still running, the current task is valid. Given
> -            // that `TaskRef` is not `Send`, we know it cannot be transferred to another thread
> -            // (where it could potentially outlive the caller).
> -            task: unsafe { &*current.cast() },
> -            _not_send: NotThreadSafe,
> +            task: Task::current_raw().cast(),
>          }

I guess these changes align with the changes above?

>      }
>
> @@ -203,6 +221,26 @@ pub fn wake_up(&self) {
>      }
>  }
>
> +impl CurrentTask {
> +    /// Access the address space of this task.
> +    ///
> +    /// To increment the refcount of the referenced `mm`, you can use `ARef::from`.
> +    #[inline]
> +    pub fn mm(&self) -> Option<&MmWithUser> {
> +        let mm = unsafe { (*self.as_ptr()).mm };
> +
> +        if mm.is_null() {
> +            None
> +        } else {
> +            // SAFETY: If `current->mm` is non-null, then it references a valid mm with a non-zero
> +            // value of `mm_users`. The returned `&MmWithUser` borrows from `CurrentTask`, so the
> +            // `&MmWithUser` cannot escape the current task, meaning `mm_users` can't reach zero
> +            // while the reference is still live.
> +            Some(unsafe { MmWithUser::from_raw(mm) })
> +        }
> +    }
> +}

Nice!

> +
>  // SAFETY: The type invariants guarantee that `Task` is always refcounted.
>  unsafe impl crate::types::AlwaysRefCounted for Task {
>      fn inc_ref(&self) {
>
> --
> 2.47.0.371.ga323438b13-goog
>