From 7797a377c36d7359630e3b5c267f55677ca747b7 Mon Sep 17 00:00:00 2001
From: Reto Kaiser <reto@retoonline.com>
Date: Mon, 8 Jul 2019 19:25:34 +0200
Subject: [PATCH] current-thread: make `tokio_current_thread::Handle` Sync
 (#1119)

---
 tokio-current-thread/Cargo.toml              |  1 +
 tokio-current-thread/src/lib.rs              | 27 +++++++++-----------
 tokio-current-thread/tests/current_thread.rs | 19 ++++++++++++++
 3 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/tokio-current-thread/Cargo.toml b/tokio-current-thread/Cargo.toml
index f25a59c81..9c00392e7 100644
--- a/tokio-current-thread/Cargo.toml
+++ b/tokio-current-thread/Cargo.toml
@@ -24,6 +24,7 @@ publish = false
 
 [dependencies]
 tokio-executor = { version = "0.2.0", path = "../tokio-executor" }
+crossbeam-channel = "0.3.8"
 
 [dev-dependencies]
 tokio-sync = { version = "0.2.0", path = "../tokio-sync" }
diff --git a/tokio-current-thread/src/lib.rs b/tokio-current-thread/src/lib.rs
index 7c21e78f1..bf10d59c2 100644
--- a/tokio-current-thread/src/lib.rs
+++ b/tokio-current-thread/src/lib.rs
@@ -36,7 +36,7 @@ use std::fmt;
 use std::future::Future;
 use std::pin::Pin;
 use std::rc::Rc;
-use std::sync::{atomic, mpsc, Arc};
+use std::sync::{atomic, Arc};
 use std::task::{Context, Poll, Waker};
 use std::thread;
 use std::time::{Duration, Instant};
@@ -61,7 +61,7 @@ pub struct CurrentThread<P: Park = ParkThread> {
     spawn_handle: Handle,
 
     /// Receiver for futures spawned from other threads
-    spawn_receiver: mpsc::Receiver<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
+    spawn_receiver: crossbeam_channel::Receiver<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
 
     /// The thread-local ID assigned to this executor.
     id: u64,
@@ -255,7 +255,7 @@ impl<P: Park> CurrentThread<P> {
     pub fn new_with_park(park: P) -> Self {
         let unpark = park.unpark();
 
-        let (spawn_sender, spawn_receiver) = mpsc::channel();
+        let (spawn_sender, spawn_receiver) = crossbeam_channel::unbounded();
         let thread = thread::current().id();
         let id = EXECUTOR_ID.with(|idc| {
             let id = idc.get();
@@ -277,7 +277,6 @@ impl<P: Park> CurrentThread<P> {
                 sender: spawn_sender,
                 num_futures,
                 waker,
-                shut_down: Cell::new(false),
                 thread,
                 id,
             },
@@ -616,9 +615,8 @@ impl<'a, P: Park> fmt::Debug for Entered<'a, P> {
 /// Handle to spawn a future on the corresponding `CurrentThread` instance
 #[derive(Clone)]
 pub struct Handle {
-    sender: mpsc::Sender<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
+    sender: crossbeam_channel::Sender<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
     num_futures: Arc<atomic::AtomicUsize>,
-    shut_down: Cell<bool>,
     /// Waker to the Scheduler
     waker: Waker,
     thread: thread::ThreadId,
@@ -631,7 +629,7 @@ pub struct Handle {
 impl fmt::Debug for Handle {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt.debug_struct("Handle")
-            .field("shut_down", &self.shut_down.get())
+            .field("shut_down", &self.is_shut_down())
             .finish()
     }
 }
@@ -654,19 +652,12 @@ impl Handle {
             }
         }
 
-        if self.shut_down.get() {
-            return Err(SpawnError::shutdown());
-        }
-
         // NOTE: += 2 since LSB is the shutdown bit
         let pending = self.num_futures.fetch_add(2, atomic::Ordering::SeqCst);
         if pending % 2 == 1 {
             // Bring the count back so we still know when the Runtime is idle.
             self.num_futures.fetch_sub(2, atomic::Ordering::SeqCst);
 
-            // Once the Runtime is shutting down, we know it won't come back.
-            self.shut_down.set(true);
-
             return Err(SpawnError::shutdown());
         }
 
@@ -687,12 +678,18 @@ impl Handle {
     /// This allows a caller to avoid creating the task if the call to `spawn`
     /// has a high likelihood of failing.
     pub fn status(&self) -> Result<(), SpawnError> {
-        if self.shut_down.get() {
+        if self.is_shut_down() {
             return Err(SpawnError::shutdown());
         }
 
         Ok(())
     }
+
+    fn is_shut_down(&self) -> bool {
+        // LSB of "num_futures" is the shutdown bit
+        let num_futures = self.num_futures.load(atomic::Ordering::SeqCst);
+        num_futures % 2 == 1
+    }
 }
 
 // ===== impl TaskExecutor =====
diff --git a/tokio-current-thread/tests/current_thread.rs b/tokio-current-thread/tests/current_thread.rs
index 794345cdb..49a407068 100644
--- a/tokio-current-thread/tests/current_thread.rs
+++ b/tokio-current-thread/tests/current_thread.rs
@@ -750,6 +750,25 @@ fn spawn_from_executor_with_handle() {
     current_thread.block_on(rx).unwrap();
 }
 
+#[test]
+fn handle_status() {
+    let current_thread = CurrentThread::new();
+    let handle = current_thread.handle();
+    assert!(handle.status().is_ok());
+
+    drop(current_thread);
+    assert!(handle.spawn(async { () }).is_err());
+    assert!(handle.status().is_err());
+}
+
+#[test]
+fn handle_is_sync() {
+    let current_thread = CurrentThread::new();
+    let handle = current_thread.handle();
+
+    let _box: Box<dyn Sync> = Box::new(handle);
+}
+
 async fn yield_once() {
     YieldOnce(false).await
 }