rt: reduce usage of ManuallyDrop (#2449)

2 files changed, 37 insertions, 22 deletions

diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs
index 573b9f3c..f4756c23 100644
--- a/tokio/src/runtime/task/core.rs
+++ b/tokio/src/runtime/task/core.rs
@@ -1,3 +1,14 @@
+//! Core task module.
+//!
+//! # Safety
+//!
+//! The functions in this module are private to the `task` module. All of them
+//! should be considered `unsafe` to use, but are not marked as such since it
+//! would be too noisy.
+//!
+//! Make sure to consult the relevant safety section of each function before
+//! use.
+
 use crate::loom::cell::UnsafeCell;
 use crate::runtime::task::raw::{self, Vtable};
 use crate::runtime::task::state::State;
@@ -95,15 +106,16 @@ impl<T: Future, S: Schedule> Cell<T, S> {
 }
 
 impl<T: Future, S: Schedule> Core<T, S> {
-    /// If needed, bind a scheduler to the task.
+    /// Bind a scheduler to the task.
+    ///
+    /// This only happens on the first poll and must be preceeded by a call to
+    /// `is_bound` to determine if binding is appropriate or not.
     ///
-    /// This only happens on the first poll.
+    /// # Safety
+    ///
+    /// Binding must not be done concurrently since it will mutate the task
+    /// core through a shared reference.
     pub(super) fn bind_scheduler(&self, task: Task<S>) {
-        use std::mem::ManuallyDrop;
-
-        // TODO: it would be nice to not have to wrap with a ManuallyDrop
-        let task = ManuallyDrop::new(task);
-
         // This function may be called concurrently, but the __first__ time it
         // is called, the caller has unique access to this field. All subsequent
         // concurrent calls will be via the `Waker`, which will "happens after"
@@ -111,12 +123,10 @@ impl<T: Future, S: Schedule> Core<T, S> {
         //
         // In other words, it is always safe to read the field and it is safe to
         // write to the field when it is `None`.
-        if self.is_bound() {
-            return;
-        }
+        debug_assert!(!self.is_bound());
 
         // Bind the task to the scheduler
-        let scheduler = S::bind(ManuallyDrop::into_inner(task));
+        let scheduler = S::bind(task);
 
         // Safety: As `scheduler` is not set, this is the first poll
         self.scheduler.with_mut(|ptr| unsafe {
diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs
index 29b231ea..a2488aa2 100644
--- a/tokio/src/runtime/task/harness.rs
+++ b/tokio/src/runtime/task/harness.rs
@@ -51,13 +51,13 @@ where
         // If this is the first time the task is polled, the task will be bound
         // to the scheduler, in which case the task ref count must be
         // incremented.
-        let ref_inc = !self.core().is_bound();
+        let is_not_bound = !self.core().is_bound();
 
         // Transition the task to the running state.
         //
         // A failure to transition here indicates the task has been cancelled
         // while in the run queue pending execution.
-        let snapshot = match self.header().state.transition_to_running(ref_inc) {
+        let snapshot = match self.header().state.transition_to_running(is_not_bound) {
             Ok(snapshot) => snapshot,
             Err(_) => {
                 // The task was shutdown while in the run queue. At this point,
@@ -67,15 +67,20 @@ where
             }
         };
 
-        // Ensure the task is bound to a scheduler instance. If this is the
-        // first time polling the task, a scheduler instance is pulled from the
-        // local context and assigned to the task.
-        //
-        // The scheduler maintains ownership of the task and responds to `wake`
-        // calls.
-        //
-        // The task reference count has been incremented.
-        self.core().bind_scheduler(self.to_task());
+        if is_not_bound {
+            // Ensure the task is bound to a scheduler instance. Since this is
+            // the first time polling the task, a scheduler instance is pulled
+            // from the local context and assigned to the task.
+            //
+            // The scheduler maintains ownership of the task and responds to
+            // `wake` calls.
+            //
+            // The task reference count has been incremented.
+            //
+            // Safety: Since we have unique access to the task so that we can
+            // safely call `bind_scheduler`.
+            self.core().bind_scheduler(self.to_task());
+        }
 
         // The transition to `Running` done above ensures that a lock on the
         // future has been obtained. This also ensures the `*mut T` pointer