diff --git a/tokio/src/runtime/basic_scheduler.rs b/tokio/src/runtime/basic_scheduler.rs index 9792ef57b..4f6de85e5 100644 --- a/tokio/src/runtime/basic_scheduler.rs +++ b/tokio/src/runtime/basic_scheduler.rs @@ -9,7 +9,7 @@ use crate::runtime::{Callback, HandleInner}; use crate::runtime::{MetricsBatch, SchedulerMetrics, WorkerMetrics}; use crate::sync::notify::Notify; use crate::util::atomic_cell::AtomicCell; -use crate::util::{waker_ref, Wake, WakerRef}; +use crate::util::{waker_ref, OwningPtr, Wake, WakerRef}; use std::cell::RefCell; use std::collections::VecDeque; @@ -376,7 +376,7 @@ impl Context { impl Spawner { /// Spawns a future onto the basic scheduler - pub(crate) fn spawn(&self, future: F, id: super::task::Id) -> JoinHandle + pub(crate) fn spawn(&self, future: OwningPtr<'_, F>, id: super::task::Id) -> JoinHandle where F: crate::future::Future + Send + 'static, F::Output: Send + 'static, diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index 14101070c..fae0505ef 100644 --- a/tokio/src/runtime/handle.rs +++ b/tokio/src/runtime/handle.rs @@ -2,9 +2,11 @@ use crate::runtime::blocking::{BlockingTask, NoopSchedule}; use crate::runtime::task::{self, JoinHandle}; use crate::runtime::{blocking, context, driver, Spawner}; use crate::util::error::{CONTEXT_MISSING_ERROR, THREAD_LOCAL_DESTROYED_ERROR}; +use crate::util::OwningPtr; use std::future::Future; use std::marker::PhantomData; +use std::mem::ManuallyDrop; use std::{error, fmt}; /// Handle to the runtime. @@ -171,6 +173,7 @@ impl Handle { /// # } /// ``` #[track_caller] + #[inline(always)] pub fn spawn(&self, future: F) -> JoinHandle where F: Future + Send + 'static, @@ -199,6 +202,7 @@ impl Handle { /// }); /// # } #[track_caller] + #[inline(always)] pub fn spawn_blocking(&self, func: F) -> JoinHandle where F: FnOnce() -> R + Send + 'static, @@ -300,6 +304,7 @@ impl Handle { } #[track_caller] + #[inline(always)] pub(crate) fn spawn_named(&self, future: F, _name: Option<&str>) -> JoinHandle where F: Future + Send + 'static, @@ -336,20 +341,13 @@ cfg_metrics! { impl HandleInner { #[track_caller] + #[inline(always)] pub(crate) fn spawn_blocking(&self, rt: &dyn ToHandle, func: F) -> JoinHandle where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { - let (join_handle, _was_spawned) = if cfg!(debug_assertions) - && std::mem::size_of::() > 2048 - { - self.spawn_blocking_inner(Box::new(func), blocking::Mandatory::NonMandatory, None, rt) - } else { - self.spawn_blocking_inner(func, blocking::Mandatory::NonMandatory, None, rt) - }; - - join_handle + self.spawn_blocking_inner(func, blocking::Mandatory::NonMandatory, None, rt) } cfg_fs! { @@ -363,21 +361,7 @@ impl HandleInner { F: FnOnce() -> R + Send + 'static, R: Send + 'static, { - let (join_handle, was_spawned) = if cfg!(debug_assertions) && std::mem::size_of::() > 2048 { - self.spawn_blocking_inner( - Box::new(func), - blocking::Mandatory::Mandatory, - None, - rt, - ) - } else { - self.spawn_blocking_inner( - func, - blocking::Mandatory::Mandatory, - None, - rt, - ) - }; + let (join_handle, was_spawned) = self.spawn_blocking_inner( func, blocking::Mandatory::Mandatory, None, rt); if was_spawned { Some(join_handle) @@ -388,6 +372,7 @@ impl HandleInner { } #[track_caller] + #[inline(always)] pub(crate) fn spawn_blocking_inner( &self, func: F, @@ -417,6 +402,10 @@ impl HandleInner { fut.instrument(span) }; + // safety: We don't touch `fut_storage` after passing it to OwningPtr. + let mut fut_storage = ManuallyDrop::new(fut); + let fut = unsafe { OwningPtr::new(&mut fut_storage) }; + #[cfg(not(all(tokio_unstable, feature = "tracing")))] let _ = name; diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs index 548c56da3..65b6231ec 100644 --- a/tokio/src/runtime/task/core.rs +++ b/tokio/src/runtime/task/core.rs @@ -14,8 +14,9 @@ use crate::loom::cell::UnsafeCell; use crate::runtime::task::raw::{self, Vtable}; use crate::runtime::task::state::State; use crate::runtime::task::{Id, Schedule}; -use crate::util::linked_list; +use crate::util::{linked_list, OwningPtr}; +use std::mem::MaybeUninit; use std::pin::Pin; use std::ptr::NonNull; use std::task::{Context, Poll, Waker}; @@ -105,10 +106,39 @@ pub(super) enum Stage { impl Cell { /// Allocates a new task cell, containing the header, trailer, and core /// structures. - pub(super) fn new(future: T, scheduler: S, state: State, task_id: Id) -> Box> { + pub(super) fn new(future: OwningPtr<'_, T>, scheduler: S, state: State, task_id: Id) -> Box> { #[cfg(all(tokio_unstable, feature = "tracing"))] - let id = future.id(); - Box::new(Cell { + let tracing_id = future.id(); + + // Using `Box::new` to allocate memory executes the operations in the following order: + // + // 1. Create value on the stack. + // 2. Allocate memory for the Box. + // 3. Move value into allocation. + // + // Unfortunately, this ordering optimizes poorly. If the allocation panics, then the value + // must run its destructor. By allocating the memory before creating the value, there's a + // good chance that the compiler can optimize away the value on the stack. + let uninit_box: Box>> = { + let layout = std::alloc::Layout::new::>(); + // safety: This is safe because the layout has a non-zero size. This is true even if T + // and S are zero-sized because of the header. + let alloc = unsafe { std::alloc::alloc(layout) }; + + if alloc.is_null() { + std::alloc::handle_alloc_error(layout); + } + + // safety: The pointer was allocated using the correct memory layout, and we checked + // that the allocation did not fail. + // + // We do not need to initialize the memory because the type is MaybeUninit. + unsafe { + Box::from_raw(alloc as *mut MaybeUninit>) + } + }; + + *uninit_box = MaybeUninit::new(Cell { header: Header { state, owned: UnsafeCell::new(linked_list::Pointers::new()), @@ -116,19 +146,26 @@ impl Cell { vtable: raw::vtable::(), owner_id: UnsafeCell::new(0), #[cfg(all(tokio_unstable, feature = "tracing"))] - id, + id: tracing_id, }, core: Core { scheduler, stage: CoreStage { - stage: UnsafeCell::new(Stage::Running(future)), + stage: UnsafeCell::new(Stage::Running(future.into_inner())), }, task_id, }, trailer: Trailer { waker: UnsafeCell::new(None), - }, - }) + } + }); + + // safety: We just initialized the value. + let init_box: Box> = unsafe { + Box::from_raw(Box::into_raw(uninit_box).cast()) + }; + + init_box } } diff --git a/tokio/src/runtime/task/list.rs b/tokio/src/runtime/task/list.rs index 7a1dff0bb..251aca792 100644 --- a/tokio/src/runtime/task/list.rs +++ b/tokio/src/runtime/task/list.rs @@ -11,6 +11,7 @@ use crate::loom::cell::UnsafeCell; use crate::loom::sync::Mutex; use crate::runtime::task::{JoinHandle, LocalNotified, Notified, Schedule, Task}; use crate::util::linked_list::{Link, LinkedList}; +use crate::util::OwningPtr; use std::marker::PhantomData; @@ -82,7 +83,7 @@ impl OwnedTasks { /// OwnedTasks has been closed. pub(crate) fn bind( &self, - task: T, + task: OwningPtr<'_, T>, scheduler: S, id: super::Id, ) -> (JoinHandle, Option>) @@ -186,7 +187,7 @@ impl LocalOwnedTasks { pub(crate) fn bind( &self, - task: T, + task: OwningPtr<'_, T>, scheduler: S, id: super::Id, ) -> (JoinHandle, Option>) diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index e73b3f35a..1fa730bab 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -178,7 +178,7 @@ use self::state::State; mod waker; use crate::future::Future; -use crate::util::linked_list; +use crate::util::{linked_list, OwningPtr}; use std::marker::PhantomData; use std::ptr::NonNull; @@ -275,7 +275,7 @@ cfg_rt! { /// immediately. The Notified is sent to the scheduler as an ordinary /// notification. fn new_task( - task: T, + task: OwningPtr<'_, T>, scheduler: S, id: Id, ) -> (Task, Notified, JoinHandle) @@ -302,7 +302,7 @@ cfg_rt! { /// only when the task is not going to be stored in an `OwnedTasks` list. /// /// Currently only blocking tasks use this method. - pub(crate) fn unowned(task: T, scheduler: S, id: Id) -> (UnownedTask, JoinHandle) + pub(crate) fn unowned(task: OwningPtr<'_, T>, scheduler: S, id: Id) -> (UnownedTask, JoinHandle) where S: Schedule, T: Send + Future + 'static, diff --git a/tokio/src/runtime/task/raw.rs b/tokio/src/runtime/task/raw.rs index 5555298a4..c337d36e8 100644 --- a/tokio/src/runtime/task/raw.rs +++ b/tokio/src/runtime/task/raw.rs @@ -1,5 +1,6 @@ use crate::future::Future; use crate::runtime::task::{Cell, Harness, Header, Id, Schedule, State}; +use crate::util::OwningPtr; use std::ptr::NonNull; use std::task::{Poll, Waker}; @@ -52,7 +53,7 @@ pub(super) fn vtable() -> &'static Vtable { } impl RawTask { - pub(super) fn new(task: T, scheduler: S, id: Id) -> RawTask + pub(super) fn new(task: OwningPtr<'_, T>, scheduler: S, id: Id) -> RawTask where T: Future, S: Schedule, diff --git a/tokio/src/util/mod.rs b/tokio/src/util/mod.rs index 618f55438..d6228b4f9 100644 --- a/tokio/src/util/mod.rs +++ b/tokio/src/util/mod.rs @@ -58,6 +58,9 @@ cfg_rt! { mod vec_deque_cell; pub(crate) use vec_deque_cell::VecDequeCell; + + mod owning_ptr; + pub(crate) use owning_ptr::OwningPtr; } cfg_rt_multi_thread! { diff --git a/tokio/src/util/owning_ptr.rs b/tokio/src/util/owning_ptr.rs new file mode 100644 index 000000000..0cda5a082 --- /dev/null +++ b/tokio/src/util/owning_ptr.rs @@ -0,0 +1,58 @@ +use std::mem::ManuallyDrop; +use std::ops::{Deref, DerefMut}; + +/// A pointer type similar to `Box` that owns a value, but doesn't own the +/// storage location containing the value. +pub(crate) struct OwningPtr<'a, T> { + storage: &'a mut ManuallyDrop, +} + +impl<'a, T> OwningPtr<'a, T> { + /// Take ownership of the value in the `ManuallyDrop`. + /// + /// # Safety + /// + /// The storage must contain a valid value before this call, and it should be + /// considered uninitialized after the `OwningPtr` is destroyed. + pub(crate) unsafe fn new(storage: &'a mut ManuallyDrop) -> Self { + Self { storage } + } + + /// Take ownership of the value. + #[inline] + pub(crate) fn into_inner(self) -> T { + // Don't run destructor of `self`. + let this = ManuallyDrop::new(self); + + // safety: The creator of the OwningPtr guarantees that the storage + // contains a valid value, and promises to treat it as uninitialized + // afterwards. + unsafe { + ManuallyDrop::take(this.storage) + } + } +} + +impl<'a, T> Deref for OwningPtr<'a, T> { + type Target = T; + fn deref(&self) -> &T { + &**self.storage + } +} + +impl<'a, T> DerefMut for OwningPtr<'a, T> { + fn deref_mut(&mut self) -> &mut T { + &mut **self.storage + } +} + +impl<'a, T> Drop for OwningPtr<'a, T> { + fn drop(&mut self) { + // safety: The creator of the OwningPtr guarantees that the storage + // contains a valid value, and promises to treat it as uninitialized + // afterwards. + unsafe { + ManuallyDrop::drop(self.storage) + } + } +}