rt: cleanup and simplify scheduler (scheduler v2.5) (#2273)

A refactor of the scheduler internals focusing on simplifying and
reducing unsafety. There are no fundamental logic changes.

* The state transitions of the core task component are refined and
reduced.
* `basic_scheduler` has most unsafety removed.
* `local_set` has most unsafety removed.
* `threaded_scheduler` limits most unsafety to its queue implementation.

3 files changed, 84 insertions, 22 deletions

diff --git a/tokio/src/macros/assert.rs b/tokio/src/macros/assert.rs
deleted file mode 100644
index 4b1cf272..00000000
--- a/tokio/src/macros/assert.rs
+++ /dev/null
@@ -1,18 +0,0 @@
-/// Asserts option is some
-macro_rules! assert_some {
-    ($e:expr) => {{
-        match $e {
-            Some(v) => v,
-            _ => panic!("expected some, was none"),
-        }
-    }};
-}
-
-/// Asserts option is none
-macro_rules! assert_none {
-    ($e:expr) => {{
-        if let Some(v) = $e {
-            panic!("expected none, was {:?}", v);
-        }
-    }};
-}
diff --git a/tokio/src/macros/mod.rs b/tokio/src/macros/mod.rs
index a37b3e49..2643c360 100644
--- a/tokio/src/macros/mod.rs
+++ b/tokio/src/macros/mod.rs
@@ -1,10 +1,6 @@
 #![cfg_attr(not(feature = "full"), allow(unused_macros))]
 
 #[macro_use]
-#[cfg(test)]
-mod assert;
-
-#[macro_use]
 mod cfg;
 
 #[macro_use]
@@ -19,6 +15,10 @@ mod ready;
 #[macro_use]
 mod thread_local;
 
+#[macro_use]
+#[cfg(feature = "rt-core")]
+pub(crate) mod scoped_tls;
+
 cfg_macros! {
     #[macro_use]
     mod select;
diff --git a/tokio/src/macros/scoped_tls.rs b/tokio/src/macros/scoped_tls.rs
new file mode 100644
index 00000000..666f382b
--- /dev/null
+++ b/tokio/src/macros/scoped_tls.rs
@@ -0,0 +1,80 @@
+use crate::loom::thread::LocalKey;
+
+use std::cell::Cell;
+use std::marker;
+
+/// Set a reference as a thread-local
+#[macro_export]
+macro_rules! scoped_thread_local {
+    ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty) => (
+        $(#[$attrs])*
+        $vis static $name: $crate::macros::scoped_tls::ScopedKey<$ty>
+            = $crate::macros::scoped_tls::ScopedKey {
+                inner: {
+                    thread_local!(static FOO: ::std::cell::Cell<*const ()> = {
+                        std::cell::Cell::new(::std::ptr::null())
+                    });
+                    &FOO
+                },
+                _marker: ::std::marker::PhantomData,
+            };
+    )
+}
+
+/// Type representing a thread local storage key corresponding to a reference
+/// to the type parameter `T`.
+pub(crate) struct ScopedKey<T> {
+    #[doc(hidden)]
+    pub(crate) inner: &'static LocalKey<Cell<*const ()>>,
+    #[doc(hidden)]
+    pub(crate) _marker: marker::PhantomData<T>,
+}
+
+unsafe impl<T> Sync for ScopedKey<T> {}
+
+impl<T> ScopedKey<T> {
+    /// Inserts a value into this scoped thread local storage slot for a
+    /// duration of a closure.
+    pub(crate) fn set<F, R>(&'static self, t: &T, f: F) -> R
+    where
+        F: FnOnce() -> R,
+    {
+        struct Reset {
+            key: &'static LocalKey<Cell<*const ()>>,
+            val: *const (),
+        }
+
+        impl Drop for Reset {
+            fn drop(&mut self) {
+                self.key.with(|c| c.set(self.val));
+            }
+        }
+
+        let prev = self.inner.with(|c| {
+            let prev = c.get();
+            c.set(t as *const _ as *const ());
+            prev
+        });
+
+        let _reset = Reset {
+            key: self.inner,
+            val: prev,
+        };
+
+        f()
+    }
+
+    /// Gets a value out of this scoped variable.
+    pub(crate) fn with<F, R>(&'static self, f: F) -> R
+    where
+        F: FnOnce(Option<&T>) -> R,
+    {
+        let val = self.inner.with(|c| c.get());
+
+        if val.is_null() {
+            f(None)
+        } else {
+            unsafe { f(Some(&*(val as *const T))) }
+        }
+    }
+}