From 6610ba9bd6b6a48d59f80573b5fa307972ace55a Mon Sep 17 00:00:00 2001
From: Alice Ryhl <alice@ryhl.io>
Date: Mon, 12 Jul 2021 14:42:00 +0200
Subject: [PATCH] runtime: use OwnedTasks in LocalSet (#3950)

---
 tokio/src/runtime/task/mod.rs |  2 +-
 tokio/src/task/local.rs       | 12 ++++--------
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs
index a4d4146fe25..5e2477906c6 100644
--- a/tokio/src/runtime/task/mod.rs
+++ b/tokio/src/runtime/task/mod.rs
@@ -19,7 +19,7 @@ mod join;
 pub use self::join::JoinHandle;
 
 mod list;
-pub(super) use self::list::OwnedTasks;
+pub(crate) use self::list::OwnedTasks;
 
 mod raw;
 use self::raw::RawTask;
diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs
index 49b0ec6c4d4..01af63fc9c0 100644
--- a/tokio/src/task/local.rs
+++ b/tokio/src/task/local.rs
@@ -1,7 +1,6 @@
 //! Runs `!Send` futures on the current thread.
-use crate::runtime::task::{self, JoinHandle, Task};
+use crate::runtime::task::{self, JoinHandle, OwnedTasks, Task};
 use crate::sync::AtomicWaker;
-use crate::util::linked_list::{Link, LinkedList};
 
 use std::cell::{Cell, RefCell};
 use std::collections::VecDeque;
@@ -233,7 +232,7 @@ struct Context {
 
 struct Tasks {
     /// Collection of all active tasks spawned onto this executor.
-    owned: LinkedList<Task<Arc<Shared>>, <Task<Arc<Shared>> as Link>::Target>,
+    owned: OwnedTasks<Arc<Shared>>,
 
     /// Local run queue sender and receiver.
     queue: VecDeque<task::Notified<Arc<Shared>>>,
@@ -334,7 +333,7 @@ impl LocalSet {
             tick: Cell::new(0),
             context: Context {
                 tasks: RefCell::new(Tasks {
-                    owned: LinkedList::new(),
+                    owned: OwnedTasks::new(),
                     queue: VecDeque::with_capacity(INITIAL_CAPACITY),
                 }),
                 shared: Arc::new(Shared {
@@ -682,17 +681,14 @@ impl task::Schedule for Arc<Shared> {
     }
 
     fn release(&self, task: &Task<Self>) -> Option<Task<Self>> {
-        use std::ptr::NonNull;
-
         CURRENT.with(|maybe_cx| {
             let cx = maybe_cx.expect("scheduler context missing");
 
             assert!(cx.shared.ptr_eq(self));
 
-            let ptr = NonNull::from(task.header());
             // safety: task must be contained by list. It is inserted into the
             // list in `bind`.
-            unsafe { cx.tasks.borrow_mut().owned.remove(ptr) }
+            unsafe { cx.tasks.borrow_mut().owned.remove(&task) }
         })
     }