Fix #2119 and failing state assertion (#2212)

Add a test for #2119 and failing state assertion,
and a fix to go with it.

Co-authored-by: Carl Lerche <[email protected]>

tokio/src/runtime/blocking/mod.rs

@@ -5,7 +5,7 @@
 
 cfg_blocking_impl! {
     mod pool;
-    pub(crate) use pool::{spawn_blocking, BlockingPool, Spawner};
+    pub(crate) use pool::{spawn_blocking, try_spawn_blocking, BlockingPool, Spawner};
 
     mod schedule;
     mod shutdown;

tokio/src/runtime/blocking/pool.rs

@@ -65,10 +65,21 @@ where
     let rt = Handle::current();
 
     let (task, handle) = task::joinable(BlockingTask::new(func));
-    rt.blocking_spawner.spawn(task, &rt);
+    let _ = rt.blocking_spawner.spawn(task, &rt);
     handle
 }
 
+#[allow(dead_code)]
+pub(crate) fn try_spawn_blocking<F, R>(func: F) -> Result<(), ()>
+where
+    F: FnOnce() -> R + Send + 'static,
+{
+    let rt = Handle::current();
+
+    let (task, _handle) = task::joinable(BlockingTask::new(func));
+    rt.blocking_spawner.spawn(task, &rt)
+}
+
 // ===== impl BlockingPool =====
 
 impl BlockingPool {
@@ -137,7 +148,7 @@ impl fmt::Debug for BlockingPool {
 // ===== impl Spawner =====
 
 impl Spawner {
-    fn spawn(&self, task: Task, rt: &Handle) {
+    fn spawn(&self, task: Task, rt: &Handle) -> Result<(), ()> {
         let shutdown_tx = {
             let mut shared = self.inner.shared.lock().unwrap();
 
@@ -146,7 +157,7 @@ impl Spawner {
                 task.shutdown();
 
                 // no need to even push this task; it would never get picked up
-                return;
+                return Err(());
             }
 
             shared.queue.push_back(task);
@@ -178,6 +189,8 @@ impl Spawner {
         if let Some(shutdown_tx) = shutdown_tx {
             self.spawn_thread(shutdown_tx, rt);
         }
+
+        Ok(())
     }
 
     fn spawn_thread(&self, shutdown_tx: shutdown::Sender, rt: &Handle) {
@@ -217,9 +230,6 @@ impl Inner {
                 run_task(task);
 
                 shared = self.shared.lock().unwrap();
-                if shared.shutdown {
-                    break; // Need to increment idle before we exit
-                }
             }
 
             // IDLE

tokio/src/runtime/mod.rs

@@ -198,7 +198,8 @@ mod blocking;
 use blocking::BlockingPool;
 
 cfg_blocking_impl! {
-    pub(crate) use blocking::spawn_blocking;
+    #[allow(unused_imports)]
+    pub(crate) use blocking::{spawn_blocking, try_spawn_blocking};
 }
 
 mod builder;

tokio/src/runtime/thread_pool/tests/loom_pool.rs

@@ -8,6 +8,28 @@ use loom::sync::{Arc, Mutex};
 use std::future::Future;
 use std::sync::atomic::Ordering::{Acquire, Relaxed, Release};
 
+#[test]
+fn racy_shutdown() {
+    loom::model(|| {
+        let pool = mk_pool(1);
+
+        // here's the case we want to exercise:
+        //
+        // a worker that still has tasks in its local queue gets sent to the blocking pool (due to
+        // block_in_place). the blocking pool is shut down, so drops the worker. the worker's
+        // shutdown method never gets run.
+        //
+        // we do this by spawning two tasks on one worker, the first of which does block_in_place,
+        // and then immediately drop the pool.
+
+        pool.spawn(async {
+            crate::task::block_in_place(|| {});
+        });
+        pool.spawn(async {});
+        drop(pool);
+    });
+}
+
 #[test]
 fn pool_multi_spawn() {
     loom::model(|| {

tokio/src/runtime/thread_pool/worker.rs

@@ -212,12 +212,6 @@ impl Worker {
             return;
         }
 
-        // make sure no subsequent code thinks that it is on a worker
-        current::clear();
-
-        // Track that the worker is gone
-        self.gone.set(true);
-
         // If this method is called, we need to move the entire worker onto a
         // separate (blocking) thread before returning. Once we return, the
         // caller is going to execute some blocking code which would otherwise
@@ -259,7 +253,20 @@ impl Worker {
         };
 
         // Give away the worker
-        runtime::spawn_blocking(move || worker.run());
+        //
+        // Returns `Err` if the spawn failed due to the runtime shutting down
+        let res = runtime::try_spawn_blocking(move || worker.run());
+
+        // If the worker hand-off was successful, clear the local state.
+        // Otherwise, the runtime is in the process of shutting down, so we will
+        // just block on the worker.
+        if res.is_ok() {
+            // make sure no subsequent code thinks that it is on a worker
+            current::clear();
+
+            // Track that the worker is gone
+            self.gone.set(true);
+        }
     }
 }