rt: coalesce thread-locals used by the runtime (#1925)

Previously, thread-locals used by the various drivers were situated
with the driver code. This resulted in state being spread out and many
thread-locals being required to run a runtime.

This PR coalesces the thread-locals into a single struct.

1 files changed, 146 insertions, 0 deletions

diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs
new file mode 100644
index 00000000..6ac60b07
--- /dev/null
+++ b/tokio/src/runtime/context.rs
@@ -0,0 +1,146 @@
+//! Thread local runtime context
+use crate::runtime::Spawner;
+use std::cell::RefCell;
+
+thread_local! {
+    static CONTEXT: RefCell<Option<ThreadContext>> = RefCell::new(None)
+}
+
+/// ThreadContext makes Runtime context accessible to each Runtime thread.
+#[derive(Debug, Clone)]
+pub(crate) struct ThreadContext {
+    /// Handles to the executor.
+    spawner: Spawner,
+
+    /// Handles to the I/O drivers
+    io_handle: crate::runtime::io::Handle,
+
+    /// Handles to the time drivers
+    time_handle: crate::runtime::time::Handle,
+
+    /// Source of `Instant::now()`
+    clock: Option<crate::runtime::time::Clock>,
+}
+
+impl Default for ThreadContext {
+    fn default() -> Self {
+        ThreadContext {
+            spawner: Spawner::Shell,
+            #[cfg(all(feature = "io-driver", not(loom)))]
+            io_handle: None,
+            #[cfg(any(not(feature = "io-driver"), loom))]
+            io_handle: (),
+            #[cfg(all(feature = "time", not(loom)))]
+            time_handle: None,
+            #[cfg(any(not(feature = "time"), loom))]
+            time_handle: (),
+            clock: None,
+        }
+    }
+}
+
+impl ThreadContext {
+    /// Construct a new [`ThreadContext`]
+    ///
+    /// [`ThreadContext`]: struct.ThreadContext.html
+    pub(crate) fn new(
+        spawner: Spawner,
+        io_handle: crate::runtime::io::Handle,
+        time_handle: crate::runtime::time::Handle,
+        clock: Option<crate::runtime::time::Clock>,
+    ) -> Self {
+        ThreadContext {
+            spawner,
+            #[cfg(all(feature = "io-driver", not(loom)))]
+            io_handle,
+            #[cfg(any(not(feature = "io-driver"), loom))]
+            io_handle,
+            #[cfg(all(feature = "time", not(loom)))]
+            time_handle,
+            #[cfg(any(not(feature = "time"), loom))]
+            time_handle,
+            clock,
+        }
+    }
+
+    /// Clone the current [`ThreadContext`] if one is set, otherwise construct a new [`ThreadContext`].
+    ///
+    /// [`ThreadContext`]: struct.ThreadContext.html
+    #[allow(dead_code)]
+    pub(crate) fn clone_current() -> Self {
+        CONTEXT.with(|ctx| ctx.borrow().clone().unwrap_or_else(Default::default))
+    }
+
+    /// Set this [`ThreadContext`] as the current active [`ThreadContext`].
+    ///
+    /// [`ThreadContext`]: struct.ThreadContext.html
+    pub(crate) fn enter(self) -> ThreadContextDropGuard {
+        CONTEXT.with(|ctx| {
+            let previous = ctx.borrow_mut().replace(self);
+            ThreadContextDropGuard { previous }
+        })
+    }
+
+    #[cfg(all(feature = "test-util", feature = "time", test))]
+    pub(crate) fn with_time_handle(mut self, handle: crate::runtime::time::Handle) -> Self {
+        self.time_handle = handle;
+        self
+    }
+
+    #[cfg(all(feature = "test-util", feature = "time", test))]
+    pub(crate) fn with_clock(mut self, clock: crate::runtime::time::Clock) -> Self {
+        self.clock.replace(clock);
+        self
+    }
+
+    #[cfg(all(feature = "io-driver", not(loom)))]
+    pub(crate) fn io_handle() -> crate::runtime::io::Handle {
+        CONTEXT.with(|ctx| match *ctx.borrow() {
+            Some(ref ctx) => ctx.io_handle.clone(),
+            None => None,
+        })
+    }
+
+    #[cfg(all(feature = "time", not(loom)))]
+    pub(crate) fn time_handle() -> crate::runtime::time::Handle {
+        CONTEXT.with(|ctx| match *ctx.borrow() {
+            Some(ref ctx) => ctx.time_handle.clone(),
+            None => None,
+        })
+    }
+
+    #[cfg(feature = "rt-core")]
+    pub(crate) fn spawn_handle() -> Option<Spawner> {
+        CONTEXT.with(|ctx| match *ctx.borrow() {
+            Some(ref ctx) => Some(ctx.spawner.clone()),
+            None => None,
+        })
+    }
+
+    #[cfg(all(feature = "test-util", feature = "time"))]
+    pub(crate) fn clock() -> Option<crate::runtime::time::Clock> {
+        CONTEXT.with(
+            |ctx| match ctx.borrow().as_ref().map(|ctx| ctx.clock.clone()) {
+                Some(Some(clock)) => Some(clock),
+                _ => None,
+            },
+        )
+    }
+}
+
+/// [`ThreadContextDropGuard`] will replace the `previous` thread context on drop.
+///
+/// [`ThreadContextDropGuard`]: struct.ThreadContextDropGuard.html
+#[derive(Debug)]
+pub(crate) struct ThreadContextDropGuard {
+    previous: Option<ThreadContext>,
+}
+
+impl Drop for ThreadContextDropGuard {
+    fn drop(&mut self) {
+        CONTEXT.with(|ctx| match self.previous.clone() {
+            Some(prev) => ctx.borrow_mut().replace(prev),
+            None => ctx.borrow_mut().take(),
+        });
+    }
+}