rt: avoid early task shutdown (#3870)

Tokio 1.7.0 introduced a change intended to eagerly shutdown newly
spawned tasks if the runtime is in the process of shutting down.
However, it introduced a bug where already spawned tasks could be
shutdown too early, resulting in the potential introduction of deadlocks
if tasks acquired mutexes in drop handlers.

Fixes #3869

tokio/tests/rt_threaded.rs

@@ -12,8 +12,8 @@ use std::future::Future;
 use std::pin::Pin;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering::Relaxed;
-use std::sync::{mpsc, Arc};
-use std::task::{Context, Poll};
+use std::sync::{mpsc, Arc, Mutex};
+use std::task::{Context, Poll, Waker};
 
 #[test]
 fn single_thread() {
@@ -405,6 +405,74 @@ async fn hang_on_shutdown() {
     tokio::time::sleep(std::time::Duration::from_secs(1)).await;
 }
 
+/// Demonstrates tokio-rs/tokio#3869
+#[test]
+fn wake_during_shutdown() {
+    struct Shared {
+        waker: Option<Waker>,
+    }
+
+    struct MyFuture {
+        shared: Arc<Mutex<Shared>>,
+        put_waker: bool,
+    }
+
+    impl MyFuture {
+        fn new() -> (Self, Self) {
+            let shared = Arc::new(Mutex::new(Shared { waker: None }));
+            let f1 = MyFuture {
+                shared: shared.clone(),
+                put_waker: true,
+            };
+            let f2 = MyFuture {
+                shared,
+                put_waker: false,
+            };
+            (f1, f2)
+        }
+    }
+
+    impl Future for MyFuture {
+        type Output = ();
+
+        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+            let me = Pin::into_inner(self);
+            let mut lock = me.shared.lock().unwrap();
+            println!("poll {}", me.put_waker);
+            if me.put_waker {
+                println!("putting");
+                lock.waker = Some(cx.waker().clone());
+            }
+            Poll::Pending
+        }
+    }
+
+    impl Drop for MyFuture {
+        fn drop(&mut self) {
+            println!("drop {} start", self.put_waker);
+            let mut lock = self.shared.lock().unwrap();
+            if !self.put_waker {
+                lock.waker.take().unwrap().wake();
+            }
+            drop(lock);
+            println!("drop {} stop", self.put_waker);
+        }
+    }
+
+    let rt = tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(1)
+        .enable_all()
+        .build()
+        .unwrap();
+
+    let (f1, f2) = MyFuture::new();
+
+    rt.spawn(f1);
+    rt.spawn(f2);
+
+    rt.block_on(async { tokio::time::sleep(tokio::time::Duration::from_millis(20)).await });
+}
+
 fn rt() -> Runtime {
     Runtime::new().unwrap()
 }