From f03b3ba785a6d336884bfc525046906f8c2a9904 Mon Sep 17 00:00:00 2001
From: Joel Galenson <jgalenson@google.com>
Date: Mon, 5 Oct 2020 08:16:15 -0700
Subject: Import tokio-0.2.22

Test: None
Change-Id: Iea7ee5e62819c9b16dbfad05a6146775df72506a
---
 src/runtime/context.rs | 73 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 73 insertions(+)
 create mode 100644 src/runtime/context.rs

(limited to 'src/runtime/context.rs')

diff --git a/src/runtime/context.rs b/src/runtime/context.rs
new file mode 100644
index 0000000..1b267f4
--- /dev/null
+++ b/src/runtime/context.rs
@@ -0,0 +1,73 @@
+//! Thread local runtime context
+use crate::runtime::Handle;
+
+use std::cell::RefCell;
+
+thread_local! {
+    static CONTEXT: RefCell<Option<Handle>> = RefCell::new(None)
+}
+
+pub(crate) fn current() -> Option<Handle> {
+    CONTEXT.with(|ctx| ctx.borrow().clone())
+}
+
+cfg_io_driver! {
+    pub(crate) fn io_handle() -> crate::runtime::io::Handle {
+        CONTEXT.with(|ctx| match *ctx.borrow() {
+            Some(ref ctx) => ctx.io_handle.clone(),
+            None => Default::default(),
+        })
+    }
+}
+
+cfg_time! {
+    pub(crate) fn time_handle() -> crate::runtime::time::Handle {
+        CONTEXT.with(|ctx| match *ctx.borrow() {
+            Some(ref ctx) => ctx.time_handle.clone(),
+            None => Default::default(),
+        })
+    }
+
+    cfg_test_util! {
+        pub(crate) fn clock() -> Option<crate::runtime::time::Clock> {
+            CONTEXT.with(|ctx| match *ctx.borrow() {
+                Some(ref ctx) => Some(ctx.clock.clone()),
+                None => None,
+            })
+        }
+    }
+}
+
+cfg_rt_core! {
+    pub(crate) fn spawn_handle() -> Option<crate::runtime::Spawner> {
+        CONTEXT.with(|ctx| match *ctx.borrow() {
+            Some(ref ctx) => Some(ctx.spawner.clone()),
+            None => None,
+        })
+    }
+}
+
+/// Set this [`Handle`] as the current active [`Handle`].
+///
+/// [`Handle`]: Handle
+pub(crate) fn enter<F, R>(new: Handle, f: F) -> R
+where
+    F: FnOnce() -> R,
+{
+    struct DropGuard(Option<Handle>);
+
+    impl Drop for DropGuard {
+        fn drop(&mut self) {
+            CONTEXT.with(|ctx| {
+                *ctx.borrow_mut() = self.0.take();
+            });
+        }
+    }
+
+    let _guard = CONTEXT.with(|ctx| {
+        let old = ctx.borrow_mut().replace(new);
+        DropGuard(old)
+    });
+
+    f()
+}
-- 
cgit v1.2.3