runtime: add owner id for tasks in OwnedTasks (#3979)

tokio/src/loom/std/atomic_u64.rs

@@ -2,19 +2,15 @@
 //! re-export of `AtomicU64`. On 32 bit platforms, this is implemented using a
 //! `Mutex`.
 
-pub(crate) use self::imp::AtomicU64;
-
 // `AtomicU64` can only be used on targets with `target_has_atomic` is 64 or greater.
 // Once `cfg_target_has_atomic` feature is stable, we can replace it with
 // `#[cfg(target_has_atomic = "64")]`.
 // Refs: https://github.com/rust-lang/rust/tree/master/src/librustc_target
-#[cfg(not(any(target_arch = "arm", target_arch = "mips", target_arch = "powerpc")))]
+cfg_has_atomic_u64! {
-mod imp {
     pub(crate) use std::sync::atomic::AtomicU64;
 }
 
-#[cfg(any(target_arch = "arm", target_arch = "mips", target_arch = "powerpc"))]
+cfg_not_has_atomic_u64! {
-mod imp {
     use crate::loom::sync::Mutex;
     use std::sync::atomic::Ordering;
 

tokio/src/macros/cfg.rs

@@ -384,3 +384,29 @@ macro_rules! cfg_not_coop {
         )*
     }
 }
+
+macro_rules! cfg_has_atomic_u64 {
+    ($($item:item)*) => {
+        $(
+            #[cfg(not(any(
+                    target_arch = "arm",
+                    target_arch = "mips",
+                    target_arch = "powerpc"
+                    )))]
+            $item
+        )*
+    }
+}
+
+macro_rules! cfg_not_has_atomic_u64 {
+    ($($item:item)*) => {
+        $(
+            #[cfg(any(
+                    target_arch = "arm",
+                    target_arch = "mips",
+                    target_arch = "powerpc"
+                    ))]
+            $item
+        )*
+    }
+}

tokio/src/runtime/basic_scheduler.rs

@@ -246,7 +246,10 @@ impl<P: Park> Inner<P> {
                     };
 
                     match entry {
-                        RemoteMsg::Schedule(task) => crate::coop::budget(|| task.run()),
+                        RemoteMsg::Schedule(task) => {
+                            let task = context.shared.owned.assert_owner(task);
+                            crate::coop::budget(|| task.run())
+                        }
                     }
                 }
 
@@ -319,29 +322,25 @@ impl<P: Park> Drop for BasicScheduler<P> {
             }
 
             // Drain local queue
+            // We already shut down every task, so we just need to drop the task.
             for task in context.tasks.borrow_mut().queue.drain(..) {
-                task.shutdown();
+                drop(task);
             }
 
             // Drain remote queue and set it to None
-            let mut remote_queue = scheduler.spawner.shared.queue.lock();
+            let remote_queue = scheduler.spawner.shared.queue.lock().take();
 
             // Using `Option::take` to replace the shared queue with `None`.
-            if let Some(remote_queue) = remote_queue.take() {
+            // We already shut down every task, so we just need to drop the task.
+            if let Some(remote_queue) = remote_queue {
                 for entry in remote_queue {
                     match entry {
                         RemoteMsg::Schedule(task) => {
-                            task.shutdown();
+                            drop(task);
                         }
                     }
                 }
             }
-            // By dropping the mutex lock after the full duration of the above loop,
-            // any thread that sees the queue in the `None` state is guaranteed that
-            // the runtime has fully shut down.
-            //
-            // The assert below is unrelated to this mutex.
-            drop(remote_queue);
 
             assert!(context.shared.owned.is_empty());
         });
@@ -400,8 +399,7 @@ impl fmt::Debug for Spawner {
 
 impl Schedule for Arc<Shared> {
     fn release(&self, task: &Task<Self>) -> Option<Task<Self>> {
-        // SAFETY: Inserted into the list in bind above.
+        self.owned.remove(task)
-        unsafe { self.owned.remove(task) }
     }
 
     fn schedule(&self, task: task::Notified<Self>) {

tokio/src/runtime/blocking/pool.rs

@@ -71,7 +71,7 @@ struct Shared {
     worker_thread_index: usize,
 }
 
-type Task = task::Notified<NoopSchedule>;
+type Task = task::UnownedTask<NoopSchedule>;
 
 const KEEP_ALIVE: Duration = Duration::from_secs(10);
 

tokio/src/runtime/task/core.rs

@@ -65,6 +65,19 @@ pub(crate) struct Header {
     /// Table of function pointers for executing actions on the task.
     pub(super) vtable: &'static Vtable,
 
+    /// This integer contains the id of the OwnedTasks or LocalOwnedTasks that
+    /// this task is stored in. If the task is not in any list, should be the
+    /// id of the list that it was previously in, or zero if it has never been
+    /// in any list.
+    ///
+    /// Once a task has been bound to a list, it can never be bound to another
+    /// list, even if removed from the first list.
+    ///
+    /// The id is not unset when removed from a list because we want to be able
+    /// to read the id without synchronization, even if it is concurrently being
+    /// removed from the list.
+    pub(super) owner_id: UnsafeCell<u64>,
+
     /// The tracing ID for this instrumented task.
     #[cfg(all(tokio_unstable, feature = "tracing"))]
     pub(super) id: Option<tracing::Id>,
@@ -98,6 +111,7 @@ impl<T: Future, S: Schedule> Cell<T, S> {
                 owned: UnsafeCell::new(linked_list::Pointers::new()),
                 queue_next: UnsafeCell::new(None),
                 vtable: raw::vtable::<T, S>(),
+                owner_id: UnsafeCell::new(0),
                 #[cfg(all(tokio_unstable, feature = "tracing"))]
                 id,
             },
@@ -203,12 +217,27 @@ impl<T: Future> CoreStage<T> {
 
 cfg_rt_multi_thread! {
     impl Header {
-        pub(crate) unsafe fn set_next(&self, next: Option<NonNull<Header>>) {
+        pub(super) unsafe fn set_next(&self, next: Option<NonNull<Header>>) {
             self.queue_next.with_mut(|ptr| *ptr = next);
         }
     }
 }
 
+impl Header {
+    // safety: The caller must guarantee exclusive access to this field, and
+    // must ensure that the id is either 0 or the id of the OwnedTasks
+    // containing this task.
+    pub(super) unsafe fn set_owner_id(&self, owner: u64) {
+        self.owner_id.with_mut(|ptr| *ptr = owner);
+    }
+
+    pub(super) fn get_owner_id(&self) -> u64 {
+        // safety: If there are concurrent writes, then that write has violated
+        // the safety requirements on `set_owner_id`.
+        unsafe { self.owner_id.with(|ptr| *ptr) }
+    }
+}
+
 impl Trailer {
     pub(crate) unsafe fn set_waker(&self, waker: Option<Waker>) {
         self.waker.with_mut(|ptr| {

tokio/src/runtime/task/list.rs

@@ -8,13 +8,53 @@
 
 use crate::future::Future;
 use crate::loom::sync::Mutex;
-use crate::runtime::task::{JoinHandle, Notified, Schedule, Task};
+use crate::runtime::task::{JoinHandle, LocalNotified, Notified, Schedule, Task};
 use crate::util::linked_list::{Link, LinkedList};
 
 use std::marker::PhantomData;
 
+// The id from the module below is used to verify whether a given task is stored
+// in this OwnedTasks, or some other task. The counter starts at one so we can
+// use zero for tasks not owned by any list.
+//
+// The safety checks in this file can technically be violated if the counter is
+// overflown, but the checks are not supposed to ever fail unless there is a
+// bug in Tokio, so we accept that certain bugs would not be caught if the two
+// mixed up runtimes happen to have the same id.
+
+cfg_has_atomic_u64! {
+    use std::sync::atomic::{AtomicU64, Ordering};
+
+    static NEXT_OWNED_TASKS_ID: AtomicU64 = AtomicU64::new(1);
+
+    fn get_next_id() -> u64 {
+        loop {
+            let id = NEXT_OWNED_TASKS_ID.fetch_add(1, Ordering::Relaxed);
+            if id != 0 {
+                return id;
+            }
+        }
+    }
+}
+
+cfg_not_has_atomic_u64! {
+    use std::sync::atomic::{AtomicU32, Ordering};
+
+    static NEXT_OWNED_TASKS_ID: AtomicU32 = AtomicU32::new(1);
+
+    fn get_next_id() -> u64 {
+        loop {
+            let id = NEXT_OWNED_TASKS_ID.fetch_add(1, Ordering::Relaxed);
+            if id != 0 {
+                return u64::from(id);
+            }
+        }
+    }
+}
+
 pub(crate) struct OwnedTasks<S: 'static> {
     inner: Mutex<OwnedTasksInner<S>>,
+    id: u64,
 }
 struct OwnedTasksInner<S: 'static> {
     list: LinkedList<Task<S>, <Task<S> as Link>::Target>,
@@ -24,7 +64,8 @@ struct OwnedTasksInner<S: 'static> {
 pub(crate) struct LocalOwnedTasks<S: 'static> {
     list: LinkedList<Task<S>, <Task<S> as Link>::Target>,
     closed: bool,
-    _not_send: PhantomData<*const ()>,
+    id: u64,
+    _not_send_or_sync: PhantomData<*const ()>,
 }
 
 impl<S: 'static> OwnedTasks<S> {
@@ -34,6 +75,7 @@ impl<S: 'static> OwnedTasks<S> {
                 list: LinkedList::new(),
                 closed: false,
             }),
+            id: get_next_id(),
         }
     }
 
@@ -51,11 +93,17 @@ impl<S: 'static> OwnedTasks<S> {
     {
         let (task, notified, join) = super::new_task(task, scheduler);
 
+        unsafe {
+            // safety: We just created the task, so we have exclusive access
+            // to the field.
+            task.header().set_owner_id(self.id);
+        }
+
         let mut lock = self.inner.lock();
         if lock.closed {
             drop(lock);
-            drop(task);
+            drop(notified);
-            notified.shutdown();
+            task.shutdown();
             (join, None)
         } else {
             lock.list.push_front(task);
@@ -63,14 +111,36 @@ impl<S: 'static> OwnedTasks<S> {
         }
     }
 
+    /// Assert that the given task is owned by this OwnedTasks and convert it to
+    /// a LocalNotified, giving the thread permission to poll this task.
+    #[inline]
+    pub(crate) fn assert_owner(&self, task: Notified<S>) -> LocalNotified<S> {
+        assert_eq!(task.0.header().get_owner_id(), self.id);
+
+        // safety: All tasks bound to this OwnedTasks are Send, so it is safe
+        // to poll it on this thread no matter what thread we are on.
+        LocalNotified {
+            task: task.0,
+            _not_send: PhantomData,
+        }
+    }
+
     pub(crate) fn pop_back(&self) -> Option<Task<S>> {
         self.inner.lock().list.pop_back()
     }
 
-    /// The caller must ensure that if the provided task is stored in a
+    pub(crate) fn remove(&self, task: &Task<S>) -> Option<Task<S>> {
-    /// linked list, then it is in this linked list.
+        let task_id = task.header().get_owner_id();
-    pub(crate) unsafe fn remove(&self, task: &Task<S>) -> Option<Task<S>> {
+        if task_id == 0 {
-        self.inner.lock().list.remove(task.header().into())
+            // The task is unowned.
+            return None;
+        }
+
+        assert_eq!(task_id, self.id);
+
+        // safety: We just checked that the provided task is not in some other
+        // linked list.
+        unsafe { self.inner.lock().list.remove(task.header().into()) }
     }
 
     pub(crate) fn is_empty(&self) -> bool {
@@ -93,7 +163,8 @@ impl<S: 'static> LocalOwnedTasks<S> {
         Self {
             list: LinkedList::new(),
             closed: false,
-            _not_send: PhantomData,
+            id: get_next_id(),
+            _not_send_or_sync: PhantomData,
         }
     }
 
@@ -109,9 +180,15 @@ impl<S: 'static> LocalOwnedTasks<S> {
     {
         let (task, notified, join) = super::new_task(task, scheduler);
 
+        unsafe {
+            // safety: We just created the task, so we have exclusive access
+            // to the field.
+            task.header().set_owner_id(self.id);
+        }
+
         if self.closed {
-            drop(task);
+            drop(notified);
-            notified.shutdown();
+            task.shutdown();
             (join, None)
         } else {
             self.list.push_front(task);
@@ -123,10 +200,33 @@ impl<S: 'static> LocalOwnedTasks<S> {
         self.list.pop_back()
     }
 
-    /// The caller must ensure that if the provided task is stored in a
+    pub(crate) fn remove(&mut self, task: &Task<S>) -> Option<Task<S>> {
-    /// linked list, then it is in this linked list.
+        let task_id = task.header().get_owner_id();
-    pub(crate) unsafe fn remove(&mut self, task: &Task<S>) -> Option<Task<S>> {
+        if task_id == 0 {
-        self.list.remove(task.header().into())
+            // The task is unowned.
+            return None;
+        }
+
+        assert_eq!(task_id, self.id);
+
+        // safety: We just checked that the provided task is not in some other
+        // linked list.
+        unsafe { self.list.remove(task.header().into()) }
+    }
+
+    /// Assert that the given task is owned by this LocalOwnedTasks and convert
+    /// it to a LocalNotified, giving the thread permission to poll this task.
+    #[inline]
+    pub(crate) fn assert_owner(&self, task: Notified<S>) -> LocalNotified<S> {
+        assert_eq!(task.0.header().get_owner_id(), self.id);
+
+        // safety: The task was bound to this LocalOwnedTasks, and the
+        // LocalOwnedTasks is not Send or Sync, so we are on the right thread
+        // for polling this task.
+        LocalNotified {
+            task: task.0,
+            _not_send: PhantomData,
+        }
     }
 
     pub(crate) fn is_empty(&self) -> bool {
@@ -139,3 +239,23 @@ impl<S: 'static> LocalOwnedTasks<S> {
         self.closed = true;
     }
 }
+
+#[cfg(all(test))]
+mod tests {
+    use super::*;
+
+    // This test may run in parallel with other tests, so we only test that ids
+    // come in increasing order.
+    #[test]
+    fn test_id_not_broken() {
+        let mut last_id = get_next_id();
+        assert_ne!(last_id, 0);
+
+        for _ in 0..1000 {
+            let next_id = get_next_id();
+            assert_ne!(next_id, 0);
+            assert!(last_id < next_id);
+            last_id = next_id;
+        }
+    }
+}

tokio/src/runtime/task/mod.rs

@@ -1,6 +1,6 @@
 mod core;
 use self::core::Cell;
-pub(crate) use self::core::Header;
+use self::core::Header;
 
 mod error;
 #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
@@ -46,13 +46,34 @@ pub(crate) struct Task<S: 'static> {
 unsafe impl<S> Send for Task<S> {}
 unsafe impl<S> Sync for Task<S> {}
 
-/// A task was notified
+/// A task was notified.
 #[repr(transparent)]
 pub(crate) struct Notified<S: 'static>(Task<S>);
 
+// safety: This type cannot be used to touch the task without first verifying
+// that the value is on a thread where it is safe to poll the task.
 unsafe impl<S: Schedule> Send for Notified<S> {}
 unsafe impl<S: Schedule> Sync for Notified<S> {}
 
+/// A non-Send variant of Notified with the invariant that it is on a thread
+/// where it is safe to poll it.
+#[repr(transparent)]
+pub(crate) struct LocalNotified<S: 'static> {
+    task: Task<S>,
+    _not_send: PhantomData<*const ()>,
+}
+
+/// A task that is not owned by any OwnedTasks. Used for blocking tasks.
+/// This type holds two ref-counts.
+pub(crate) struct UnownedTask<S: 'static> {
+    raw: RawTask,
+    _p: PhantomData<S>,
+}
+
+// safety: This type can only be created given a Send task.
+unsafe impl<S> Send for UnownedTask<S> {}
+unsafe impl<S> Sync for UnownedTask<S> {}
+
 /// Task result sent back
 pub(crate) type Result<T> = std::result::Result<T, JoinError>;
 
@@ -105,41 +126,50 @@ cfg_rt! {
     /// Create a new task with an associated join handle. This method is used
     /// only when the task is not going to be stored in an `OwnedTasks` list.
     ///
-    /// Currently only blocking tasks and tests use this method.
+    /// Currently only blocking tasks use this method.
-    pub(crate) fn unowned<T, S>(task: T, scheduler: S) -> (Notified<S>, JoinHandle<T::Output>)
+    pub(crate) fn unowned<T, S>(task: T, scheduler: S) -> (UnownedTask<S>, JoinHandle<T::Output>)
     where
         S: Schedule,
         T: Send + Future + 'static,
         T::Output: Send + 'static,
     {
         let (task, notified, join) = new_task(task, scheduler);
-        drop(task);
+
-        (notified, join)
+        // This transfers the ref-count of task and notified into an UnownedTask.
+        // This is valid because an UnownedTask holds two ref-counts.
+        let unowned = UnownedTask {
+            raw: task.raw,
+            _p: PhantomData,
+        };
+        std::mem::forget(task);
+        std::mem::forget(notified);
+
+        (unowned, join)
     }
 }
 
 impl<S: 'static> Task<S> {
-    pub(crate) unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
+    unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
         Task {
             raw: RawTask::from_raw(ptr),
             _p: PhantomData,
         }
     }
 
-    pub(crate) fn header(&self) -> &Header {
+    fn header(&self) -> &Header {
         self.raw.header()
     }
 }
 
 cfg_rt_multi_thread! {
     impl<S: 'static> Notified<S> {
-        pub(crate) unsafe fn from_raw(ptr: NonNull<Header>) -> Notified<S> {
+        unsafe fn from_raw(ptr: NonNull<Header>) -> Notified<S> {
             Notified(Task::from_raw(ptr))
         }
     }
 
     impl<S: 'static> Task<S> {
-        pub(crate) fn into_raw(self) -> NonNull<Header> {
+        fn into_raw(self) -> NonNull<Header> {
             let ret = self.header().into();
             mem::forget(self);
             ret
@@ -147,7 +177,7 @@ cfg_rt_multi_thread! {
     }
 
     impl<S: 'static> Notified<S> {
-        pub(crate) fn into_raw(self) -> NonNull<Header> {
+        fn into_raw(self) -> NonNull<Header> {
             self.0.into_raw()
         }
     }
@@ -160,16 +190,45 @@ impl<S: Schedule> Task<S> {
     }
 }
 
-impl<S: Schedule> Notified<S> {
+impl<S: Schedule> LocalNotified<S> {
     /// Run the task
     pub(crate) fn run(self) {
-        self.0.raw.poll();
+        self.task.raw.poll();
+        mem::forget(self);
+    }
+}
+
+impl<S: Schedule> UnownedTask<S> {
+    // Used in test of the inject queue.
+    #[cfg(test)]
+    pub(super) fn into_notified(self) -> Notified<S> {
+        Notified(self.into_task())
+    }
+
+    fn into_task(self) -> Task<S> {
+        // Convert into a task.
+        let task = Task {
+            raw: self.raw,
+            _p: PhantomData,
+        };
+        mem::forget(self);
+
+        // Drop a ref-count since an UnownedTask holds two.
+        task.header().state.ref_dec();
+
+        task
+    }
+
+    pub(crate) fn run(self) {
+        // Decrement the ref-count
+        self.raw.header().state.ref_dec();
+        // Poll the task
+        self.raw.poll();
         mem::forget(self);
     }
 
-    /// Pre-emptively cancel the task as part of the shutdown process.
     pub(crate) fn shutdown(self) {
-        self.0.shutdown();
+        self.into_task().shutdown()
     }
 }
 
@@ -183,6 +242,16 @@ impl<S: 'static> Drop for Task<S> {
     }
 }
 
+impl<S: 'static> Drop for UnownedTask<S> {
+    fn drop(&mut self) {
+        // Decrement the ref count
+        if self.raw.header().state.ref_dec_twice() {
+            // Deallocate if this is the final ref count
+            self.raw.dealloc();
+        }
+    }
+}
+
 impl<S> fmt::Debug for Task<S> {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(fmt, "Task({:p})", self.header())

tokio/src/runtime/task/state.rs

@@ -324,6 +324,12 @@ impl State {
         prev.ref_count() == 1
     }
 
+    /// Returns `true` if the task should be released.
+    pub(super) fn ref_dec_twice(&self) -> bool {
+        let prev = Snapshot(self.val.fetch_sub(2 * REF_ONE, AcqRel));
+        prev.ref_count() == 2
+    }
+
     fn fetch_update<F>(&self, mut f: F) -> Result<Snapshot, Snapshot>
     where
         F: FnMut(Snapshot) -> Option<Snapshot>,

tokio/src/runtime/tests/mod.rs

@@ -14,7 +14,7 @@ mod unowned_wrapper {
         let span = tracing::trace_span!("test_span");
         let task = task.instrument(span);
         let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule);
-        (task, handle)
+        (task.into_notified(), handle)
     }
 
     #[cfg(not(all(tokio_unstable, feature = "tracing")))]
@@ -24,7 +24,7 @@ mod unowned_wrapper {
         T::Output: Send + 'static,
     {
         let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule);
-        (task, handle)
+        (task.into_notified(), handle)
     }
 }
 

tokio/src/runtime/tests/task.rs

@@ -241,6 +241,7 @@ impl Runtime {
         while !self.is_empty() && n < max {
             let task = self.next_task();
             n += 1;
+            let task = self.0.owned.assert_owner(task);
             task.run();
         }
 
@@ -264,7 +265,7 @@ impl Runtime {
         }
 
         while let Some(task) = core.queue.pop_back() {
-            task.shutdown();
+            drop(task);
         }
 
         drop(core);
@@ -275,8 +276,7 @@ impl Runtime {
 
 impl Schedule for Runtime {
     fn release(&self, task: &Task<Self>) -> Option<Task<Self>> {
-        // safety: copying worker.rs
+        self.0.owned.remove(task)
-        unsafe { self.0.owned.remove(task) }
     }
 
     fn schedule(&self, task: task::Notified<Self>) {

tokio/src/runtime/thread_pool/worker.rs

@@ -384,6 +384,8 @@ impl Context {
     }
 
     fn run_task(&self, task: Notified, mut core: Box<Core>) -> RunResult {
+        let task = self.worker.shared.owned.assert_owner(task);
+
         // Make sure the worker is not in the **searching** state. This enables
         // another idle worker to try to steal work.
         core.transition_from_searching(&self.worker);
@@ -414,6 +416,7 @@ impl Context {
                 if coop::has_budget_remaining() {
                     // Run the LIFO task, then loop
                     *self.core.borrow_mut() = Some(core);
+                    let task = self.worker.shared.owned.assert_owner(task);
                     task.run();
                 } else {
                     // Not enough budget left to run the LIFO task, push it to
@@ -626,8 +629,7 @@ impl Worker {
 
 impl task::Schedule for Arc<Shared> {
     fn release(&self, task: &Task) -> Option<Task> {
-        // SAFETY: Inserted into owned in bind.
+        self.owned.remove(task)
-        unsafe { self.owned.remove(task) }
     }
 
     fn schedule(&self, task: Notified) {
@@ -762,8 +764,10 @@ impl Shared {
         }
 
         // Drain the injection queue
+        //
+        // We already shut down every task, so we can simply drop the tasks.
         while let Some(task) = self.inject.pop() {
-            task.shutdown();
+            drop(task);
         }
     }
 

tokio/src/task/local.rs

@@ -540,11 +540,11 @@ impl LocalSet {
         true
     }
 
-    fn next_task(&self) -> Option<task::Notified<Arc<Shared>>> {
+    fn next_task(&self) -> Option<task::LocalNotified<Arc<Shared>>> {
         let tick = self.tick.get();
         self.tick.set(tick.wrapping_add(1));
 
-        if tick % REMOTE_FIRST_INTERVAL == 0 {
+        let task = if tick % REMOTE_FIRST_INTERVAL == 0 {
             self.context
                 .shared
                 .queue
@@ -566,7 +566,9 @@ impl LocalSet {
                         .as_mut()
                         .and_then(|queue| queue.pop_front())
                 })
-        }
+        };
+
+        task.map(|task| self.context.tasks.borrow_mut().owned.assert_owner(task))
     }
 
     fn with<T>(&self, f: impl FnOnce() -> T) -> T {
@@ -631,15 +633,17 @@ impl Drop for LocalSet {
                 task.shutdown();
             }
 
+            // We already called shutdown on all tasks above, so there is no
+            // need to call shutdown.
             for task in self.context.tasks.borrow_mut().queue.drain(..) {
-                task.shutdown();
+                drop(task);
             }
 
             // Take the queue from the Shared object to prevent pushing
             // notifications to it in the future.
             let queue = self.context.shared.queue.lock().take().unwrap();
             for task in queue {
-                task.shutdown();
+                drop(task);
             }
 
             assert!(self.context.tasks.borrow().owned.is_empty());
@@ -711,12 +715,8 @@ impl task::Schedule for Arc<Shared> {
     fn release(&self, task: &Task<Self>) -> Option<Task<Self>> {
         CURRENT.with(|maybe_cx| {
             let cx = maybe_cx.expect("scheduler context missing");
-
             assert!(cx.shared.ptr_eq(self));
-
+            cx.tasks.borrow_mut().owned.remove(&task)
-            // safety: task must be contained by list. It is inserted into the
-            // list when spawning.
-            unsafe { cx.tasks.borrow_mut().owned.remove(&task) }
         })
     }