Re: [PATCH v1 6/7] rust: workqueue: add safe API to workqueue

[Martin Rodriguez Reboredo <yakoyoku@gmail.com>]

On 5/17/23 17:31, Alice Ryhl wrote:
> This commit introduces `ArcWorkItem`, `BoxWorkItem`, and
> `define_work_adapter_newtype!` that make it possible to use the
> workqueue without any unsafe code whatsoever.
> 
> The `ArcWorkItem` and `BoxWorkItem` traits are used when a struct has a
> single `work_struct` field.
> 
> The `define_work_adapter_newtype!` macro is used when a struct has
> multiple `work_struct` fields. For each `work_struct` field, a newtype
> struct is defined that wraps `Arc<TheStruct>`, and pushing an instance
> of the newtype to a workqueue will enqueue it using the associated
> `work_struct` field. The newtypes are matched with `work_struct` fields
> by having the T in `Work<T>` be the newtype.
> 
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> ---
> [...]
> +
> +unsafe impl<T> WorkItem for Arc<T>
> +where
> +    T: ArcWorkItem + HasWork<Self> + ?Sized,
> +{
> +    type EnqueueOutput = Result<(), Self>;
> +
> +    unsafe fn __enqueue<F>(self, queue_work_on: F) -> Self::EnqueueOutput
> +    where
> +        F: FnOnce(*mut bindings::work_struct) -> bool,
> +    {
> +        let ptr = Arc::into_raw(self);
> +
> +        // Using `get_work_offset` here for object-safety.
> +        //
> +        // SAFETY: The pointer is valid since we just got it from `into_raw`.
> +        let off = unsafe { (&*ptr).get_work_offset() };
> +
> +        // SAFETY: The `HasWork` impl promises that this offset gives us a field of type
> +        // `Work<Self>` in the same allocation.
> +        let work_ptr = unsafe { (ptr as *const u8).add(off) as *const Work<Self> };
> +        // SAFETY: The pointer is not dangling.
> +        let work_ptr = unsafe { Work::raw_get(work_ptr) };
> +
> +        match (queue_work_on)(work_ptr) {

Match for boolean is not a good pattern in my eyes, if-else should be
used instead. Also aren't the parens around the closure unnecessary?

> +            true => Ok(()),
> +            // SAFETY: The work queue has not taken ownership of the pointer.
> +            false => Err(unsafe { Arc::from_raw(ptr) }),
> +        }
> +    }
> +}
> +
> [...]
> +
> +unsafe impl<T> WorkItem for Pin<Box<T>>
> +where
> +    T: BoxWorkItem + HasWork<Self> + ?Sized,
> +{
> +    // When a box is in a workqueue, the workqueue has exclusive ownership of the box. Therefore,
> +    // it's not possible to enqueue a box while it is in a workqueue.
> +    type EnqueueOutput = ();
> +
> +    unsafe fn __enqueue<F>(self, queue_work_on: F)
> +    where
> +        F: FnOnce(*mut bindings::work_struct) -> bool,
> +    {
> +        // SAFETY: We will not used the contents in an unpinned manner.
> +        let ptr = unsafe { Box::into_raw(Pin::into_inner_unchecked(self)) };
> +
> +        // Using `get_work_offset` here for object-safety.
> +        //
> +        // SAFETY: The pointer is valid since we just got it from `into_raw`.
> +        let off = unsafe { (&*ptr).get_work_offset() };
> +
> +        // SAFETY: The `HasWork` impl promises that this offset gives us a field of type
> +        // `Work<Self>` in the same allocation.
> +        let work_ptr = unsafe { (ptr as *mut u8).add(off) as *mut Work<Self> };
> +        // SAFETY: The pointer is not dangling.
> +        let work_ptr = unsafe { Work::raw_get(work_ptr) };
> +
> +        match (queue_work_on)(work_ptr) {

Same as above.

> +            true => {}
> +            // SAFETY: This method requires exclusive ownership of the box, so it cannot be in a
> +            // workqueue.
> +            false => unsafe { core::hint::unreachable_unchecked() },
> +        }
> +    }
> +}
> +
> [...]
> +
> +/// Helper macro for structs with several `Work` fields that can be in several queues at once.
> +///
> +/// For each `Work` field in your type `T`, a newtype struct that wraps an `Arc<T>` or
> +/// `Pin<Box<T>>` should be defined.
> +///
> +/// # Examples
> +///
> +/// ```

There must be those work macro and type imports here I think.

> +/// struct MyStruct {
> +///     work1: Work<MyStructWork1>,
> +///     work2: Work<MyStructWork2>,
> +/// }
> +///
> +/// impl_has_work! {
> +///     impl HasWork<MyStructWork1> for MyStruct { self.work1 }
> +///     impl HasWork<MyStructWork2> for MyStruct { self.work2 }
> +/// }
> +///
> +/// define_work_adapter_newtype! {
> +///     struct MyStructWork1(Arc<MyStruct>);
> +///     struct MyStructWork2(Arc<MyStruct>);
> +/// }
> +///
> +/// impl MyStructWork1 {
> +///     fn run(self) {
> +///         // ...
> +///     }
> +/// }
> +///
> +/// impl MyStructWork2 {
> +///     fn run(self) {
> +///         // ...
> +///     }
> +/// }
> +/// ```
> +///
> +/// This will let you push a `MyStructWork1(arc)` or `MyStructWork2(arc)` to a work queue. The [`Arc`]
> +/// can be in two work queues at the same time, and the `run` method on the wrapper type is called
> +/// when the work item is called.
> +///
> +/// [`Arc`]: crate::sync::Arc
> +#[macro_export]
> +macro_rules! define_work_adapter_newtype {
> +    (
> +        $(#[$outer:meta])*
> +        $pub:vis struct $name:ident(
> +            $(#[$innermeta:meta])*
> +            $fpub:vis Arc<$inner:ty> $(,)?
> +        );
> +        $($rest:tt)*
> +    ) => {
> +        $(#[$outer])*
> +        $pub struct $name($(#[$innermeta])* $fpub $crate::sync::Arc<$inner>);
> +
> +        unsafe impl $crate::workqueue::WorkItem for $name {
> +            type EnqueueOutput = ::core::result::Result<(), $name>;
> +
> +            unsafe fn __enqueue<F>(self, queue_work_on: F) -> Self::EnqueueOutput
> +            where
> +                F: ::core::ops::FnOnce(*mut $crate::bindings::work_struct) -> bool,
> +            {
> +                let ptr = $crate::sync::Arc::into_raw(self.0);
> +
> +                // SAFETY: The pointer is not dangling since we just got it from Arc::into_raw.
> +                let work_ptr = unsafe { <$inner as $crate::workqueue::HasWork::<$name>>::raw_get_work(ptr.cast_mut()) };
> +
> +                // SAFETY: The pointer is not dangling.
> +                let work_ptr = unsafe { $crate::workqueue::Work::raw_get(work_ptr) };
> +
> +                match (queue_work_on)(work_ptr) {

Same as what I've said on those `WorkItem` impls.

> +                    true => Ok(()),
> +                    // SAFETY: The work queue has not taken ownership of the pointer.
> +                    false => Err($name(unsafe { $crate::sync::Arc::from_raw(ptr) })),
> +                }
> +            }
> +        }
> [...]
> +    };
> +
> +    (
> +        $(#[$outer:meta])*
> +        $pub:vis struct $name:ident(
> +            $(#[$innermeta:meta])*
> +            $fpub:vis Pin<Box<$inner:ty>> $(,)?
> +        );
> +        $($rest:tt)*
> +    ) => {
> +        $(#[$outer])*
> +        $pub struct $name($(#[$innermeta])* $fpub ::core::pin::Pin<::alloc::boxed::Box<$inner>>);
> +
> +        unsafe impl $crate::workqueue::WorkItem for $name {
> +            type EnqueueOutput = ();
> +
> +            unsafe fn __enqueue<F>(self, queue_work_on: F)
> +            where
> +                F: ::core::ops::FnOnce(*mut $crate::bindings::work_struct) -> bool,
> +            {
> +                // SAFETY: We will not used the contents in an unpinned manner.
> +                let boxed = unsafe { ::core::pin::Pin::into_inner_unchecked(self.0) };
> +                let ptr = ::alloc::boxed::Box::into_raw(boxed);
> +
> +                // SAFETY: The pointer is not dangling since we just got it from Box::into_raw.
> +                let work_ptr = unsafe { <$inner as $crate::workqueue::HasWork::<$name>>::raw_get_work(ptr) };
> +
> +                // SAFETY: The pointer is not dangling.
> +                let work_ptr = unsafe { $crate::workqueue::Work::raw_get(work_ptr) };
> +
> +                match (queue_work_on)(work_ptr) {

Same as above.

> +                    true => {},
> +                    // SAFETY: This method requires exclusive ownership of the box, so it cannot be in a
> +                    // workqueue.
> +                    false => unsafe { ::core::hint::unreachable_unchecked() },
> +                }
> +            }
> +        }
> +
> [...]