WIP

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<F>(&self, future: F, id: super::task::Id) -> JoinHandle<F::Output>
+    pub(crate) fn spawn<F>(&self, future: OwningPtr<'_, F>, id: super::task::Id) -> JoinHandle<F::Output>
     where
         F: crate::future::Future + Send + 'static,
         F::Output: Send + 'static,

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<F>(&self, future: F) -> JoinHandle<F::Output>
     where
         F: Future + Send + 'static,
@@ -199,6 +202,7 @@ impl Handle {
     /// });
     /// # }
     #[track_caller]
+    #[inline(always)]
     pub fn spawn_blocking<F, R>(&self, func: F) -> JoinHandle<R>
     where
         F: FnOnce() -> R + Send + 'static,
@@ -300,6 +304,7 @@ impl Handle {
     }
 
     #[track_caller]
+    #[inline(always)]
     pub(crate) fn spawn_named<F>(&self, future: F, _name: Option<&str>) -> JoinHandle<F::Output>
     where
         F: Future + Send + 'static,
@@ -336,20 +341,13 @@ cfg_metrics! {
 
 impl HandleInner {
     #[track_caller]
+    #[inline(always)]
     pub(crate) fn spawn_blocking<F, R>(&self, rt: &dyn ToHandle, func: F) -> JoinHandle<R>
     where
         F: FnOnce() -> R + Send + 'static,
         R: Send + 'static,
     {
-        let (join_handle, _was_spawned) = if cfg!(debug_assertions)
-            && std::mem::size_of::<F>() > 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::<F>() > 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<F, R>(
         &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;
 

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<T: Future> {
 impl<T: Future, S: Schedule> Cell<T, S> {
     /// 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<Cell<T, S>> {
+    pub(super) fn new(future: OwningPtr<'_, 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 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<MaybeUninit<Cell<T, S>>> = {
+            let layout = std::alloc::Layout::new::<Cell<T, S>>();
+            // 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<Cell<T, S>>)
+            }
+        };
+
+        *uninit_box = MaybeUninit::new(Cell {
             header: Header {
                 state,
                 owned: UnsafeCell::new(linked_list::Pointers::new()),
@@ -116,19 +146,26 @@ impl<T: Future, S: Schedule> Cell<T, S> {
                 vtable: raw::vtable::<T, S>(),
                 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<Cell<T, S>> = unsafe {
+            Box::from_raw(Box::into_raw(uninit_box).cast())
+        };
+
+        init_box
     }
 }
 

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<S: 'static> OwnedTasks<S> {
     /// OwnedTasks has been closed.
     pub(crate) fn bind<T>(
         &self,
-        task: T,
+        task: OwningPtr<'_, T>,
         scheduler: S,
         id: super::Id,
     ) -> (JoinHandle<T::Output>, Option<Notified<S>>)
@@ -186,7 +187,7 @@ impl<S: 'static> LocalOwnedTasks<S> {
 
     pub(crate) fn bind<T>(
         &self,
-        task: T,
+        task: OwningPtr<'_, T>,
         scheduler: S,
         id: super::Id,
     ) -> (JoinHandle<T::Output>, Option<Notified<S>>)

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<T, S>(
-        task: T,
+        task: OwningPtr<'_, T>,
         scheduler: S,
         id: Id,
     ) -> (Task<S>, Notified<S>, JoinHandle<T::Output>)
@@ -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<T, S>(task: T, scheduler: S, id: Id) -> (UnownedTask<S>, JoinHandle<T::Output>)
+    pub(crate) fn unowned<T, S>(task: OwningPtr<'_, T>, scheduler: S, id: Id) -> (UnownedTask<S>, JoinHandle<T::Output>)
     where
         S: Schedule,
         T: Send + Future + 'static,

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<T: Future, S: Schedule>() -> &'static Vtable {
 }
 
 impl RawTask {
-    pub(super) fn new<T, S>(task: T, scheduler: S, id: Id) -> RawTask
+    pub(super) fn new<T, S>(task: OwningPtr<'_, T>, scheduler: S, id: Id) -> RawTask
     where
         T: Future,
         S: Schedule,

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! {

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<T>,
+}
+
+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<T>) -> 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)
+        }
+    }
+}