From ed3ece266b5ab2ae3795a91d8dac25bc0fdfeec2 Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Tue, 20 Nov 2018 09:17:07 -0800
Subject: [PATCH] current-thread: fix shutdown on idle (#763)

When spawning using `Handle` while on the executor, tasks were being
double counted. This prevented the number of active tasks to reach zero,
thus preventing the executor from shutting down.

This changes `spawn` to check if being called from the executor
**before** incrementing the number of active tasks.

Fixes #760
---
 tokio-current-thread/src/lib.rs              | 14 +++++++-------
 tokio-current-thread/tests/current_thread.rs | 19 +++++++++++++++++++
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/tokio-current-thread/src/lib.rs b/tokio-current-thread/src/lib.rs
index 9af62244f3a..34f033a55be 100644
--- a/tokio-current-thread/src/lib.rs
+++ b/tokio-current-thread/src/lib.rs
@@ -653,6 +653,13 @@ impl Handle {
     where
         F: Future<Item = (), Error = ()> + Send + 'static,
     {
+        if thread::current().id() == self.thread {
+            let mut e = TaskExecutor::current();
+            if e.id() == Some(self.id) {
+                return e.spawn_local(Box::new(future));
+            }
+        }
+
         if self.shut_down.get() {
             return Err(SpawnError::shutdown());
         }
@@ -669,13 +676,6 @@ impl Handle {
             return Err(SpawnError::shutdown());
         }
 
-        if thread::current().id() == self.thread {
-            let mut e = TaskExecutor::current();
-            if e.id() == Some(self.id) {
-                return e.spawn_local(Box::new(future));
-            }
-        }
-
         self.sender.send(Box::new(future))
             .expect("CurrentThread does not exist anymore");
         // use 0 for the id, CurrentThread does not make use of it
diff --git a/tokio-current-thread/tests/current_thread.rs b/tokio-current-thread/tests/current_thread.rs
index a98ed792976..d2b6a6908fa 100644
--- a/tokio-current-thread/tests/current_thread.rs
+++ b/tokio-current-thread/tests/current_thread.rs
@@ -778,6 +778,25 @@ fn spawn_from_other_thread_unpark() {
     ).unwrap();
 }
 
+#[test]
+fn spawn_from_executor_with_handle() {
+    let mut current_thread = CurrentThread::new();
+    let handle = current_thread.handle();
+    let (tx, rx) = oneshot::channel();
+
+    current_thread.spawn(lazy(move || {
+        handle.spawn(lazy(move || {
+            tx.send(()).unwrap();
+            Ok(())
+        })).unwrap();
+        Ok::<_, ()>(())
+    }));
+
+    current_thread.run();
+
+    rx.wait().unwrap();
+}
+
 fn ok() -> future::FutureResult<(), ()> {
     future::ok(())
 }