[10/13] rust: introduce `Task::current`

Message ID 20230330043954.562237-10-wedsonaf@gmail.com
State New
Headers
Series [01/13] rust: sync: introduce `LockClassKey` |

Commit Message

Wedson Almeida Filho March 30, 2023, 4:39 a.m. UTC
  From: Wedson Almeida Filho <walmeida@microsoft.com>

This allows Rust code to get a reference to the current task without
having to increment the refcount, but still guaranteeing memory safety.

Cc: Ingo Molnar <mingo@redhat.com>
Cc: Peter Zijlstra <peterz@infradead.org>
Signed-off-by: Wedson Almeida Filho <walmeida@microsoft.com>
---
 rust/helpers.c      |  6 ++++
 rust/kernel/task.rs | 83 ++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 88 insertions(+), 1 deletion(-)
  

Comments

Gary Guo March 31, 2023, 2:47 a.m. UTC | #1
On Thu, 30 Mar 2023 01:39:51 -0300
Wedson Almeida Filho <wedsonaf@gmail.com> wrote:

> From: Wedson Almeida Filho <walmeida@microsoft.com>
> 
> This allows Rust code to get a reference to the current task without
> having to increment the refcount, but still guaranteeing memory safety.
> 
> Cc: Ingo Molnar <mingo@redhat.com>
> Cc: Peter Zijlstra <peterz@infradead.org>
> Signed-off-by: Wedson Almeida Filho <walmeida@microsoft.com>
> ---
>  rust/helpers.c      |  6 ++++
>  rust/kernel/task.rs | 83 ++++++++++++++++++++++++++++++++++++++++++++-
>  2 files changed, 88 insertions(+), 1 deletion(-)
> 
> diff --git a/rust/helpers.c b/rust/helpers.c
> index 58a194042c86..96441744030e 100644
> --- a/rust/helpers.c
> +++ b/rust/helpers.c
> @@ -100,6 +100,12 @@ bool rust_helper_refcount_dec_and_test(refcount_t *r)
>  }
>  EXPORT_SYMBOL_GPL(rust_helper_refcount_dec_and_test);
>  
> +struct task_struct *rust_helper_get_current(void)
> +{
> +	return current;
> +}
> +EXPORT_SYMBOL_GPL(rust_helper_get_current);
> +
>  void rust_helper_get_task_struct(struct task_struct *t)
>  {
>  	get_task_struct(t);
> diff --git a/rust/kernel/task.rs b/rust/kernel/task.rs
> index 8d7a8222990f..8b2b56ba9c6d 100644
> --- a/rust/kernel/task.rs
> +++ b/rust/kernel/task.rs
> @@ -5,7 +5,7 @@
>  //! C header: [`include/linux/sched.h`](../../../../include/linux/sched.h).
>  
>  use crate::bindings;
> -use core::{cell::UnsafeCell, ptr};
> +use core::{cell::UnsafeCell, marker::PhantomData, ops::Deref, ptr};
>  
>  /// Wraps the kernel's `struct task_struct`.
>  ///
> @@ -13,6 +13,46 @@ use core::{cell::UnsafeCell, ptr};
>  ///
>  /// Instances of this type are always ref-counted, that is, a call to `get_task_struct` ensures
>  /// that the allocation remains valid at least until the matching call to `put_task_struct`.
> +///
> +/// # Examples
> +///
> +/// The following is an example of getting the PID of the current thread with zero additional cost
> +/// when compared to the C version:
> +///
> +/// ```
> +/// use kernel::task::Task;
> +///
> +/// let pid = Task::current().pid();
> +/// ```
> +///
> +/// Getting the PID of the current process, also zero additional cost:
> +///
> +/// ```
> +/// use kernel::task::Task;
> +///
> +/// let pid = Task::current().group_leader().pid();
> +/// ```
> +///
> +/// Getting the current task and storing it in some struct. The reference count is automatically
> +/// incremented when creating `State` and decremented when it is dropped:
> +///
> +/// ```
> +/// use kernel::{task::Task, ARef};
> +///
> +/// struct State {
> +///     creator: ARef<Task>,
> +///     index: u32,
> +/// }
> +///
> +/// impl State {
> +///     fn new() -> Self {
> +///         Self {
> +///             creator: Task::current().into(),
> +///             index: 0,
> +///         }
> +///     }
> +/// }
> +/// ```
>  #[repr(transparent)]
>  pub struct Task(pub(crate) UnsafeCell<bindings::task_struct>);
>  
> @@ -25,6 +65,20 @@ unsafe impl Sync for Task {}
>  type Pid = bindings::pid_t;
>  
>  impl Task {
> +    /// Returns a task reference for the currently executing task/thread.
> +    pub fn current<'a>() -> TaskRef<'a> {
> +        // SAFETY: Just an FFI call with no additional safety requirements.
> +        let ptr = unsafe { bindings::get_current() };
> +
> +        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 { &*ptr.cast() },
> +            _not_send: PhantomData,
> +        }
> +    }
> +

I don't think this API is sound, as you can do `&*Task::current()` and
get a `&'static Task`, which is very problematic.

A sound API would be

	pub fn with_current<R>(f: imp FnOnce(&Task) -> R) -> R { ... }

(which also is how thread local works in Rust)

You would have to write `Task::with_current(|cur| cur.pid())` though,
which unfortunately is a bit less ergonomic.

Best,
Gary
  
Alice Ryhl March 31, 2023, 7:32 a.m. UTC | #2
On 3/31/23 04:47, Gary Guo wrote:
> On Thu, 30 Mar 2023 01:39:51 -0300
> Wedson Almeida Filho <wedsonaf@gmail.com> wrote:
> 
>> From: Wedson Almeida Filho <walmeida@microsoft.com>
>>
>> This allows Rust code to get a reference to the current task without
>> having to increment the refcount, but still guaranteeing memory safety.
>>
>> Cc: Ingo Molnar <mingo@redhat.com>
>> Cc: Peter Zijlstra <peterz@infradead.org>
>> Signed-off-by: Wedson Almeida Filho <walmeida@microsoft.com>
>> ---
>>   rust/helpers.c      |  6 ++++
>>   rust/kernel/task.rs | 83 ++++++++++++++++++++++++++++++++++++++++++++-
>>   2 files changed, 88 insertions(+), 1 deletion(-)
>>
>> diff --git a/rust/helpers.c b/rust/helpers.c
>> index 58a194042c86..96441744030e 100644
>> --- a/rust/helpers.c
>> +++ b/rust/helpers.c
>> @@ -100,6 +100,12 @@ bool rust_helper_refcount_dec_and_test(refcount_t *r)
>>   }
>>   EXPORT_SYMBOL_GPL(rust_helper_refcount_dec_and_test);
>>   
>> +struct task_struct *rust_helper_get_current(void)
>> +{
>> +	return current;
>> +}
>> +EXPORT_SYMBOL_GPL(rust_helper_get_current);
>> +
>>   void rust_helper_get_task_struct(struct task_struct *t)
>>   {
>>   	get_task_struct(t);
>> diff --git a/rust/kernel/task.rs b/rust/kernel/task.rs
>> index 8d7a8222990f..8b2b56ba9c6d 100644
>> --- a/rust/kernel/task.rs
>> +++ b/rust/kernel/task.rs
>> @@ -5,7 +5,7 @@
>>   //! C header: [`include/linux/sched.h`](../../../../include/linux/sched.h).
>>   
>>   use crate::bindings;
>> -use core::{cell::UnsafeCell, ptr};
>> +use core::{cell::UnsafeCell, marker::PhantomData, ops::Deref, ptr};
>>   
>>   /// Wraps the kernel's `struct task_struct`.
>>   ///
>> @@ -13,6 +13,46 @@ use core::{cell::UnsafeCell, ptr};
>>   ///
>>   /// Instances of this type are always ref-counted, that is, a call to `get_task_struct` ensures
>>   /// that the allocation remains valid at least until the matching call to `put_task_struct`.
>> +///
>> +/// # Examples
>> +///
>> +/// The following is an example of getting the PID of the current thread with zero additional cost
>> +/// when compared to the C version:
>> +///
>> +/// ```
>> +/// use kernel::task::Task;
>> +///
>> +/// let pid = Task::current().pid();
>> +/// ```
>> +///
>> +/// Getting the PID of the current process, also zero additional cost:
>> +///
>> +/// ```
>> +/// use kernel::task::Task;
>> +///
>> +/// let pid = Task::current().group_leader().pid();
>> +/// ```
>> +///
>> +/// Getting the current task and storing it in some struct. The reference count is automatically
>> +/// incremented when creating `State` and decremented when it is dropped:
>> +///
>> +/// ```
>> +/// use kernel::{task::Task, ARef};
>> +///
>> +/// struct State {
>> +///     creator: ARef<Task>,
>> +///     index: u32,
>> +/// }
>> +///
>> +/// impl State {
>> +///     fn new() -> Self {
>> +///         Self {
>> +///             creator: Task::current().into(),
>> +///             index: 0,
>> +///         }
>> +///     }
>> +/// }
>> +/// ```
>>   #[repr(transparent)]
>>   pub struct Task(pub(crate) UnsafeCell<bindings::task_struct>);
>>   
>> @@ -25,6 +65,20 @@ unsafe impl Sync for Task {}
>>   type Pid = bindings::pid_t;
>>   
>>   impl Task {
>> +    /// Returns a task reference for the currently executing task/thread.
>> +    pub fn current<'a>() -> TaskRef<'a> {
>> +        // SAFETY: Just an FFI call with no additional safety requirements.
>> +        let ptr = unsafe { bindings::get_current() };
>> +
>> +        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 { &*ptr.cast() },
>> +            _not_send: PhantomData,
>> +        }
>> +    }
>> +
> 
> I don't think this API is sound, as you can do `&*Task::current()` and
> get a `&'static Task`, which is very problematic.
> 
> A sound API would be
> 
> 	pub fn with_current<R>(f: imp FnOnce(&Task) -> R) -> R { ... }
> 
> (which also is how thread local works in Rust)
> 
> You would have to write `Task::with_current(|cur| cur.pid())` though,
> which unfortunately is a bit less ergonomic.
> 
> Best,
> Gary

This is true, unfortunately. It would be possible to write a macro with 
a more similar API to the current implementation.
  
Wedson Almeida Filho April 1, 2023, 4:09 a.m. UTC | #3
Gary, thanks for reviewing!

On Fri, Mar 31, 2023 at 03:47:01AM +0100, Gary Guo wrote:
> 
> I don't think this API is sound, as you can do `&*Task::current()` and
> get a `&'static Task`, which is very problematic.

One thing that isn't clear to me is: how do you get a 'static lifetime in the
example above?

Altough `TaskRef` does have an arbitrary lifetime param, that's not the lifetime
that the returned `Task` reference gets. For illustration, I've explicitly added
a lifetime 'a in the impl below:

impl Deref for TaskRef<'_> {
    type Target = Task;
    fn deref(&'a self) -> &'a Self::Target {
        self.task
    }
}

Which means that the borrow of the `TaskRef` you use to call `deref` must
outlive the returned `Task`.

So how do you get a `TaskRef` with a static lifetime to begin with? Or is there
another trick to get the `&'static Task` that I can't see?

Thanks,
-Wedson
  
Gary Guo April 1, 2023, 7:01 a.m. UTC | #4
On Sat, 1 Apr 2023 01:09:18 -0300
Wedson Almeida Filho <wedsonaf@gmail.com> wrote:

> Gary, thanks for reviewing!
> 
> On Fri, Mar 31, 2023 at 03:47:01AM +0100, Gary Guo wrote:
> > 
> > I don't think this API is sound, as you can do `&*Task::current()` and
> > get a `&'static Task`, which is very problematic.  
> 
> One thing that isn't clear to me is: how do you get a 'static lifetime in the
> example above?
> 
> Altough `TaskRef` does have an arbitrary lifetime param, that's not the lifetime
> that the returned `Task` reference gets. For illustration, I've explicitly added
> a lifetime 'a in the impl below:
> 
> impl Deref for TaskRef<'_> {
>     type Target = Task;
>     fn deref(&'a self) -> &'a Self::Target {
>         self.task
>     }
> }
> 
> Which means that the borrow of the `TaskRef` you use to call `deref` must
> outlive the returned `Task`.
> 
> So how do you get a `TaskRef` with a static lifetime to begin with? Or is there
> another trick to get the `&'static Task` that I can't see?
> 
> Thanks,
> -Wedson

Hi Wedson,

My apologies for the confusion. `&*Task::current()` is not
sufficient. I typed too quick without double checking.

However it is still true that `TaskRef<'static>` is unsound, and it can
be retrieved from `current()`. The missing step is `&'static
TaskRef<'static>`.

So you can write `&*Box::leak(Box::try_new(Task::current()).unwrap())`
and get `&'static Task`.

Best,
Gary
  

Patch

diff --git a/rust/helpers.c b/rust/helpers.c
index 58a194042c86..96441744030e 100644
--- a/rust/helpers.c
+++ b/rust/helpers.c
@@ -100,6 +100,12 @@  bool rust_helper_refcount_dec_and_test(refcount_t *r)
 }
 EXPORT_SYMBOL_GPL(rust_helper_refcount_dec_and_test);
 
+struct task_struct *rust_helper_get_current(void)
+{
+	return current;
+}
+EXPORT_SYMBOL_GPL(rust_helper_get_current);
+
 void rust_helper_get_task_struct(struct task_struct *t)
 {
 	get_task_struct(t);
diff --git a/rust/kernel/task.rs b/rust/kernel/task.rs
index 8d7a8222990f..8b2b56ba9c6d 100644
--- a/rust/kernel/task.rs
+++ b/rust/kernel/task.rs
@@ -5,7 +5,7 @@ 
 //! C header: [`include/linux/sched.h`](../../../../include/linux/sched.h).
 
 use crate::bindings;
-use core::{cell::UnsafeCell, ptr};
+use core::{cell::UnsafeCell, marker::PhantomData, ops::Deref, ptr};
 
 /// Wraps the kernel's `struct task_struct`.
 ///
@@ -13,6 +13,46 @@  use core::{cell::UnsafeCell, ptr};
 ///
 /// Instances of this type are always ref-counted, that is, a call to `get_task_struct` ensures
 /// that the allocation remains valid at least until the matching call to `put_task_struct`.
+///
+/// # Examples
+///
+/// The following is an example of getting the PID of the current thread with zero additional cost
+/// when compared to the C version:
+///
+/// ```
+/// use kernel::task::Task;
+///
+/// let pid = Task::current().pid();
+/// ```
+///
+/// Getting the PID of the current process, also zero additional cost:
+///
+/// ```
+/// use kernel::task::Task;
+///
+/// let pid = Task::current().group_leader().pid();
+/// ```
+///
+/// Getting the current task and storing it in some struct. The reference count is automatically
+/// incremented when creating `State` and decremented when it is dropped:
+///
+/// ```
+/// use kernel::{task::Task, ARef};
+///
+/// struct State {
+///     creator: ARef<Task>,
+///     index: u32,
+/// }
+///
+/// impl State {
+///     fn new() -> Self {
+///         Self {
+///             creator: Task::current().into(),
+///             index: 0,
+///         }
+///     }
+/// }
+/// ```
 #[repr(transparent)]
 pub struct Task(pub(crate) UnsafeCell<bindings::task_struct>);
 
@@ -25,6 +65,20 @@  unsafe impl Sync for Task {}
 type Pid = bindings::pid_t;
 
 impl Task {
+    /// Returns a task reference for the currently executing task/thread.
+    pub fn current<'a>() -> TaskRef<'a> {
+        // SAFETY: Just an FFI call with no additional safety requirements.
+        let ptr = unsafe { bindings::get_current() };
+
+        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 { &*ptr.cast() },
+            _not_send: PhantomData,
+        }
+    }
+
     /// Returns the group leader of the given task.
     pub fn group_leader(&self) -> &Task {
         // SAFETY: By the type invariant, we know that `self.0` is valid.
@@ -69,3 +123,30 @@  unsafe impl crate::types::AlwaysRefCounted for Task {
         unsafe { bindings::put_task_struct(obj.cast().as_ptr()) }
     }
 }
+
+/// A wrapper for a shared reference to [`Task`] that isn't [`Send`].
+///
+/// We make this explicitly not [`Send`] so that we can use it to represent the current thread
+/// without having to increment/decrement the task's reference count.
+///
+/// # Invariants
+///
+/// The wrapped [`Task`] remains valid for the lifetime of the object.
+pub struct TaskRef<'a> {
+    task: &'a Task,
+    _not_send: PhantomData<*mut ()>,
+}
+
+impl Deref for TaskRef<'_> {
+    type Target = Task;
+
+    fn deref(&self) -> &Self::Target {
+        self.task
+    }
+}
+
+impl From<TaskRef<'_>> for crate::types::ARef<Task> {
+    fn from(t: TaskRef<'_>) -> Self {
+        t.deref().into()
+    }
+}