diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs index a512921c1..2b7050a21 100644 --- a/tokio/src/runtime/task/core.rs +++ b/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
, - /// Pointer to next task, used with the injection queue. pub(super) queue_next: UnsafeCell>>, @@ -86,23 +84,26 @@ pub(crate) struct Header { pub(super) id: Option, } +unsafe impl Send for Header {} +unsafe impl Sync for Header {} + +/// 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
, + /// Consumer task waiting on completion of this task. + pub(super) waker: UnsafeCell>, +} + generate_addr_of_methods! { - impl<> Header { + impl<> Trailer { pub(super) unsafe fn addr_of_owned(self: NonNull) -> NonNull> { &self.owned } } } -unsafe impl Send for Header {} -unsafe impl Sync for Header {} - -/// Cold data is stored after the future. -pub(super) struct Trailer { - /// Consumer task waiting on completion of this task. - pub(super) waker: UnsafeCell>, -} - /// Either the future or the output. pub(super) enum Stage { Running(T), @@ -116,10 +117,9 @@ impl Cell { pub(super) fn new(future: T, scheduler: S, state: State, task_id: Id) -> Box> { #[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::(), owner_id: UnsafeCell::new(0), @@ -135,8 +135,19 @@ impl Cell { }, 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
) -> NonNull { + let offset = me.as_ref().vtable.trailer_offset; + let trailer = me.as_ptr().cast::().add(offset).cast::(); + NonNull::new_unchecked(trailer) + } } impl Trailer { diff --git a/tokio/src/runtime/task/list.rs b/tokio/src/runtime/task/list.rs index 7a1dff0bb..ca06d459c 100644 --- a/tokio/src/runtime/task/list.rs +++ b/tokio/src/runtime/task/list.rs @@ -164,7 +164,7 @@ impl OwnedTasks { // 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 { diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index c2903f840..56911a731 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -334,6 +334,10 @@ impl Task { fn header(&self) -> &Header { self.raw.header() } + + fn header_ptr(&self) -> NonNull
{ + self.raw.header_ptr() + } } impl Notified { @@ -473,7 +477,7 @@ unsafe impl linked_list::Link for Task { } unsafe fn pointers(target: NonNull
) -> NonNull> { - Header::addr_of_owned(target) + self::core::Trailer::addr_of_owned(Header::get_trailer(target)) } } diff --git a/tokio/src/runtime/task/raw.rs b/tokio/src/runtime/task/raw.rs index 5555298a4..a24ac44bf 100644 --- a/tokio/src/runtime/task/raw.rs +++ b/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
), + + /// 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() -> &'static Vtable { drop_abort_handle: drop_abort_handle::, remote_abort: remote_abort::, shutdown: shutdown::, + trailer_offset: TrailerOffsetHelper::::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: +/// +struct TrailerOffsetHelper(T, S); +impl TrailerOffsetHelper { + // 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::
(), + std::mem::size_of::>(), + std::mem::align_of::>(), + std::mem::align_of::(), + ); +} + +/// Compute the offset of the `Trailer` field in `Cell` using the +/// `#[repr(C)]` algorithm. +/// +/// Pseudo-code for the `#[repr(C)]` algorithm can be found here: +/// +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(task: T, scheduler: S, id: Id) -> RawTask where