tokio: avoid temporary references in linked_list::Link impls (#4841)

Author: Darksonn

tokio/build.rs

@@ -7,6 +7,10 @@ fn main() {
             if ac.probe_rustc_version(1, 59) {
                 autocfg::emit("tokio_const_thread_local")
             }
+
+            if !ac.probe_rustc_version(1, 51) {
+                autocfg::emit("tokio_no_addr_of")
+            }
         }
 
         Err(e) => {

tokio/src/io/driver/scheduled_io.rs

@@ -66,6 +66,14 @@ cfg_io_readiness! {
         _p: PhantomPinned,
     }
 
+    generate_addr_of_methods! {
+        impl<> Waiter {
+            unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> {
+                &self.pointers
+            }
+        }
+    }
+
     /// Future returned by `readiness()`.
     struct Readiness<'a> {
         scheduled_io: &'a ScheduledIo,
@@ -399,8 +407,8 @@ cfg_io_readiness! {
             ptr
         }
 
-        unsafe fn pointers(mut target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> {
-            NonNull::from(&mut target.as_mut().pointers)
+        unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> {
+            Waiter::addr_of_pointers(target)
         }
     }
 

tokio/src/macros/addr_of.rs

@@ -0,0 +1,53 @@
+//! This module defines a macro that lets you go from a raw pointer to a struct
+//! to a raw pointer to a field of the struct.
+
+#[cfg(not(tokio_no_addr_of))]
+macro_rules! generate_addr_of_methods {
+    (
+    impl<$($gen:ident)*> $struct_name:ty {$(
+        $(#[$attrs:meta])*
+        $vis:vis unsafe fn $fn_name:ident(self: NonNull<Self>) -> NonNull<$field_type:ty> {
+            &self$(.$field_name:tt)+
+        }
+    )*}
+    ) => {
+        impl<$($gen)*> $struct_name {$(
+            $(#[$attrs])*
+            $vis unsafe fn $fn_name(me: ::core::ptr::NonNull<Self>) -> ::core::ptr::NonNull<$field_type> {
+                let me = me.as_ptr();
+                let field = ::std::ptr::addr_of_mut!((*me) $(.$field_name)+ );
+                ::core::ptr::NonNull::new_unchecked(field)
+            }
+        )*}
+    };
+}
+
+// The `addr_of_mut!` macro is only available for MSRV at least 1.51.0. This
+// version of the macro uses a workaround for older versions of rustc.
+#[cfg(tokio_no_addr_of)]
+macro_rules! generate_addr_of_methods {
+    (
+    impl<$($gen:ident)*> $struct_name:ty {$(
+        $(#[$attrs:meta])*
+        $vis:vis unsafe fn $fn_name:ident(self: NonNull<Self>) -> NonNull<$field_type:ty> {
+            &self$(.$field_name:tt)+
+        }
+    )*}
+    ) => {
+        impl<$($gen)*> $struct_name {$(
+            $(#[$attrs])*
+            $vis unsafe fn $fn_name(me: ::core::ptr::NonNull<Self>) -> ::core::ptr::NonNull<$field_type> {
+                let me = me.as_ptr();
+                let me_u8 = me as *mut u8;
+
+                let field_offset = {
+                    let me_ref = &*me;
+                    let field_ref_u8 = (&me_ref $(.$field_name)+ ) as *const $field_type as *const u8;
+                    field_ref_u8.offset_from(me_u8)
+                };
+
+                ::core::ptr::NonNull::new_unchecked(me_u8.offset(field_offset).cast())
+            }
+        )*}
+    };
+}

tokio/src/macros/mod.rs

@@ -15,6 +15,9 @@ mod ready;
 #[macro_use]
 mod thread_local;
 
+#[macro_use]
+mod addr_of;
+
 cfg_trace! {
     #[macro_use]
     mod trace;

tokio/src/runtime/task/core.rs

@@ -60,7 +60,7 @@ pub(crate) struct Header {
     /// Task state.
     pub(super) state: State,
 
-    pub(super) owned: UnsafeCell<linked_list::Pointers<Header>>,
+    pub(super) owned: linked_list::Pointers<Header>,
 
     /// Pointer to next task, used with the injection queue.
     pub(super) queue_next: UnsafeCell<Option<NonNull<Header>>>,
@@ -86,6 +86,14 @@ pub(crate) struct Header {
     pub(super) id: Option<tracing::Id>,
 }
 
+generate_addr_of_methods! {
+    impl<> Header {
+        pub(super) unsafe fn addr_of_owned(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Header>> {
+            &self.owned
+        }
+    }
+}
+
 unsafe impl Send for Header {}
 unsafe impl Sync for Header {}
 
@@ -111,7 +119,7 @@ impl<T: Future, S: Schedule> Cell<T, S> {
         Box::new(Cell {
             header: Header {
                 state,
-                owned: UnsafeCell::new(linked_list::Pointers::new()),
+                owned: linked_list::Pointers::new(),
                 queue_next: UnsafeCell::new(None),
                 vtable: raw::vtable::<T, S>(),
                 owner_id: UnsafeCell::new(0),

tokio/src/runtime/task/mod.rs

@@ -473,8 +473,7 @@ unsafe impl<S> linked_list::Link for Task<S> {
     }
 
     unsafe fn pointers(target: NonNull<Header>) -> NonNull<linked_list::Pointers<Header>> {
-        // Not super great as it avoids some of looms checking...
-        NonNull::from(target.as_ref().owned.with_mut(|ptr| &mut *ptr))
+        Header::addr_of_owned(target)
     }
 }
 

tokio/src/sync/batch_semaphore.rs

@@ -112,6 +112,14 @@ struct Waiter {
     _p: PhantomPinned,
 }
 
+generate_addr_of_methods! {
+    impl<> Waiter {
+        unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> {
+            &self.pointers
+        }
+    }
+}
+
 impl Semaphore {
     /// The maximum number of permits which a semaphore can hold.
     ///
@@ -704,12 +712,6 @@ impl std::error::Error for TryAcquireError {}
 ///
 /// `Waiter` is forced to be !Unpin.
 unsafe impl linked_list::Link for Waiter {
-    // XXX: ideally, we would be able to use `Pin` here, to enforce the
-    // invariant that list entries may not move while in the list. However, we
-    // can't do this currently, as using `Pin<&'a mut Waiter>` as the `Handle`
-    // type would require `Semaphore` to be generic over a lifetime. We can't
-    // use `Pin<*mut Waiter>`, as raw pointers are `Unpin` regardless of whether
-    // or not they dereference to an `!Unpin` target.
     type Handle = NonNull<Waiter>;
     type Target = Waiter;
 
@@ -721,7 +723,7 @@ unsafe impl linked_list::Link for Waiter {
         ptr
     }
 
-    unsafe fn pointers(mut target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> {
-        NonNull::from(&mut target.as_mut().pointers)
+    unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> {
+        Waiter::addr_of_pointers(target)
     }
 }

tokio/src/sync/broadcast.rs

@@ -361,6 +361,14 @@ struct Waiter {
     _p: PhantomPinned,
 }
 
+generate_addr_of_methods! {
+    impl<> Waiter {
+        unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> {
+            &self.pointers
+        }
+    }
+}
+
 struct RecvGuard<'a, T> {
     slot: RwLockReadGuard<'a, Slot<T>>,
 }
@@ -1140,8 +1148,8 @@ unsafe impl linked_list::Link for Waiter {
         ptr
     }
 
-    unsafe fn pointers(mut target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> {
-        NonNull::from(&mut target.as_mut().pointers)
+    unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> {
+        Waiter::addr_of_pointers(target)
     }
 }
 

tokio/src/sync/notify.rs

@@ -214,7 +214,6 @@ enum NotificationType {
 }
 
 #[derive(Debug)]
-#[repr(C)] // required by `linked_list::Link` impl
 struct Waiter {
     /// Intrusive linked-list pointers.
     pointers: linked_list::Pointers<Waiter>,
@@ -229,6 +228,14 @@ struct Waiter {
     _p: PhantomPinned,
 }
 
+generate_addr_of_methods! {
+    impl<> Waiter {
+        unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> {
+            &self.pointers
+        }
+    }
+}
+
 /// Future returned from [`Notify::notified()`].
 ///
 /// This future is fused, so once it has completed, any future calls to poll
@@ -950,7 +957,7 @@ unsafe impl linked_list::Link for Waiter {
     }
 
     unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> {
-        target.cast()
+        Waiter::addr_of_pointers(target)
     }
 }
 

tokio/src/time/driver/entry.rs

@@ -326,7 +326,7 @@ pub(super) type EntryList = crate::util::linked_list::LinkedList<TimerShared, Ti
 ///
 /// Note that this structure is located inside the `TimerEntry` structure.
 #[derive(Debug)]
-#[repr(C)] // required by `link_list::Link` impl
+#[repr(C)]
 pub(crate) struct TimerShared {
     /// Data manipulated by the driver thread itself, only.
     driver_state: CachePadded<TimerSharedPadded>,
@@ -339,6 +339,14 @@ pub(crate) struct TimerShared {
     _p: PhantomPinned,
 }
 
+generate_addr_of_methods! {
+    impl<> TimerShared {
+        unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<TimerShared>> {
+            &self.driver_state.0.pointers
+        }
+    }
+}
+
 impl TimerShared {
     pub(super) fn new() -> Self {
         Self {
@@ -421,7 +429,6 @@ impl TimerShared {
 /// padded. This contains the information that the driver thread accesses most
 /// frequently to minimize contention. In particular, we move it away from the
 /// waker, as the waker is updated on every poll.
-#[repr(C)] // required by `link_list::Link` impl
 struct TimerSharedPadded {
     /// A link within the doubly-linked list of timers on a particular level and
     /// slot. Valid only if state is equal to Registered.
@@ -476,7 +483,7 @@ unsafe impl linked_list::Link for TimerShared {
     unsafe fn pointers(
         target: NonNull<Self::Target>,
     ) -> NonNull<linked_list::Pointers<Self::Target>> {
-        target.cast()
+        TimerShared::addr_of_pointers(target)
     }
 }
 

tokio/src/util/idle_notified_set.rs

@@ -88,10 +88,6 @@ enum List {
 /// move it out from this entry to prevent it from getting leaked. (Since the
 /// two linked lists are emptied in the destructor of `IdleNotifiedSet`, the
 /// value should not be leaked.)
-///
-/// This type is `#[repr(C)]` because its `linked_list::Link` implementation
-/// requires that `pointers` is the first field.
-#[repr(C)]
 struct ListEntry<T> {
     /// The linked list pointers of the list this entry is in.
     pointers: linked_list::Pointers<ListEntry<T>>,
@@ -105,6 +101,14 @@ struct ListEntry<T> {
     _pin: PhantomPinned,
 }
 
+generate_addr_of_methods! {
+    impl<T> ListEntry<T> {
+        unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<ListEntry<T>>> {
+            &self.pointers
+        }
+    }
+}
+
 // With mutable access to the `IdleNotifiedSet`, you can get mutable access to
 // the values.
 unsafe impl<T: Send> Send for IdleNotifiedSet<T> {}
@@ -453,11 +457,6 @@ unsafe impl<T> linked_list::Link for ListEntry<T> {
     unsafe fn pointers(
         target: NonNull<ListEntry<T>>,
     ) -> NonNull<linked_list::Pointers<ListEntry<T>>> {
-        // Safety: The pointers struct is the first field and ListEntry is
-        // `#[repr(C)]` so this cast is safe.
-        //
-        // We do this rather than doing a field access since `std::ptr::addr_of`
-        // is too new for our MSRV.
-        target.cast()
+        ListEntry::addr_of_pointers(target)
     }
 }