runtime: move owned field to Trailer (#4843)

tokio/src/runtime/task/core.rs

@@ -60,8 +60,6 @@ pub(crate) struct Header {
     /// Task state.
     pub(super) state: State,
 
-    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,23 +84,26 @@ 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 {}
 
-/// Cold data is stored after the future.
+/// Cold data is stored after the future. Data is considered cold if it is only
+/// used during creation or shutdown of the task.
 pub(super) struct Trailer {
+    /// Pointers for the linked list in the `OwnedTasks` that owns this task.
+    pub(super) owned: linked_list::Pointers<Header>,
     /// Consumer task waiting on completion of this task.
     pub(super) waker: UnsafeCell<Option<Waker>>,
 }
 
+generate_addr_of_methods! {
+    impl<> Trailer {
+        pub(super) unsafe fn addr_of_owned(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Header>> {
+            &self.owned
+        }
+    }
+}
+
 /// Either the future or the output.
 pub(super) enum Stage<T: Future> {
     Running(T),
@@ -116,10 +117,9 @@ impl<T: Future, S: Schedule> Cell<T, S> {
     pub(super) fn new(future: T, scheduler: S, state: State, task_id: Id) -> Box<Cell<T, S>> {
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         let id = future.id();
-        Box::new(Cell {
+        let result = Box::new(Cell {
             header: Header {
                 state,
-                owned: linked_list::Pointers::new(),
                 queue_next: UnsafeCell::new(None),
                 vtable: raw::vtable::<T, S>(),
                 owner_id: UnsafeCell::new(0),
@@ -135,8 +135,19 @@ impl<T: Future, S: Schedule> Cell<T, S> {
             },
             trailer: Trailer {
                 waker: UnsafeCell::new(None),
+                owned: linked_list::Pointers::new(),
             },
-        })
+        });
+
+        #[cfg(debug_assertions)]
+        {
+            let trailer_addr = (&result.trailer) as *const Trailer as usize;
+            let trailer_ptr = unsafe { Header::get_trailer(NonNull::from(&result.header)) };
+
+            assert_eq!(trailer_addr, trailer_ptr.as_ptr() as usize);
+        }
+
+        result
     }
 }
 
@@ -248,6 +259,17 @@ impl Header {
         // the safety requirements on `set_owner_id`.
         unsafe { self.owner_id.with(|ptr| *ptr) }
     }
+
+    /// Gets a pointer to the `Trailer` of the task containing this `Header`.
+    ///
+    /// # Safety
+    ///
+    /// The provided raw pointer must point at the header of a task.
+    pub(super) unsafe fn get_trailer(me: NonNull<Header>) -> NonNull<Trailer> {
+        let offset = me.as_ref().vtable.trailer_offset;
+        let trailer = me.as_ptr().cast::<u8>().add(offset).cast::<Trailer>();
+        NonNull::new_unchecked(trailer)
+    }
 }
 
 impl Trailer {

tokio/src/runtime/task/list.rs

@@ -164,7 +164,7 @@ impl<S: 'static> OwnedTasks<S> {
 
         // safety: We just checked that the provided task is not in some other
         // linked list.
-        unsafe { self.inner.lock().list.remove(task.header().into()) }
+        unsafe { self.inner.lock().list.remove(task.header_ptr()) }
     }
 
     pub(crate) fn is_empty(&self) -> bool {

tokio/src/runtime/task/mod.rs

@@ -334,6 +334,10 @@ impl<S: 'static> Task<S> {
     fn header(&self) -> &Header {
         self.raw.header()
     }
+
+    fn header_ptr(&self) -> NonNull<Header> {
+        self.raw.header_ptr()
+    }
 }
 
 impl<S: 'static> Notified<S> {
@@ -473,7 +477,7 @@ unsafe impl<S> linked_list::Link for Task<S> {
     }
 
     unsafe fn pointers(target: NonNull<Header>) -> NonNull<linked_list::Pointers<Header>> {
-        Header::addr_of_owned(target)
+        self::core::Trailer::addr_of_owned(Header::get_trailer(target))
     }
 }
 

tokio/src/runtime/task/raw.rs

@@ -1,4 +1,5 @@
 use crate::future::Future;
+use crate::runtime::task::core::{Core, Trailer};
 use crate::runtime::task::{Cell, Harness, Header, Id, Schedule, State};
 
 use std::ptr::NonNull;
@@ -35,6 +36,9 @@ pub(super) struct Vtable {
 
     /// Scheduler is being shutdown.
     pub(super) shutdown: unsafe fn(NonNull<Header>),
+
+    /// The number of bytes that the `trailer` field is offset from the header.
+    pub(super) trailer_offset: usize,
 }
 
 /// Get the vtable for the requested `T` and `S` generics.
@@ -48,9 +52,55 @@ pub(super) fn vtable<T: Future, S: Schedule>() -> &'static Vtable {
         drop_abort_handle: drop_abort_handle::<T, S>,
         remote_abort: remote_abort::<T, S>,
         shutdown: shutdown::<T, S>,
+        trailer_offset: TrailerOffsetHelper::<T, S>::OFFSET,
     }
 }
 
+/// Calling `get_trailer_offset` directly in vtable doesn't work because it
+/// prevents the vtable from being promoted to a static reference.
+///
+/// See this thread for more info:
+/// <https://users.rust-lang.org/t/custom-vtables-with-integers/78508>
+struct TrailerOffsetHelper<T, S>(T, S);
+impl<T: Future, S: Schedule> TrailerOffsetHelper<T, S> {
+    // Pass `size_of`/`align_of` as arguments rather than calling them directly
+    // inside `get_trailer_offset` because trait bounds on generic parameters
+    // of const fn are unstable on our MSRV.
+    const OFFSET: usize = get_trailer_offset(
+        std::mem::size_of::<Header>(),
+        std::mem::size_of::<Core<T, S>>(),
+        std::mem::align_of::<Core<T, S>>(),
+        std::mem::align_of::<Trailer>(),
+    );
+}
+
+/// Compute the offset of the `Trailer` field in `Cell<T, S>` using the
+/// `#[repr(C)]` algorithm.
+///
+/// Pseudo-code for the `#[repr(C)]` algorithm can be found here:
+/// <https://doc.rust-lang.org/reference/type-layout.html#reprc-structs>
+const fn get_trailer_offset(
+    header_size: usize,
+    core_size: usize,
+    core_align: usize,
+    trailer_align: usize,
+) -> usize {
+    let mut offset = header_size;
+
+    let core_misalign = offset % core_align;
+    if core_misalign > 0 {
+        offset += core_align - core_misalign;
+    }
+    offset += core_size;
+
+    let trailer_misalign = offset % trailer_align;
+    if trailer_misalign > 0 {
+        offset += trailer_align - trailer_misalign;
+    }
+
+    offset
+}
+
 impl RawTask {
     pub(super) fn new<T, S>(task: T, scheduler: S, id: Id) -> RawTask
     where