Fix case where cancelled future would be leaked and add a regression …

…test.

tokio/src/runtime/task/harness.rs

@@ -183,16 +183,22 @@ where
     // ====== internal ======
 
     fn complete(self, output: super::Result<T::Output>, is_join_interested: bool) {
+        let stage = &self.core().stage;
         if is_join_interested {
-            // Store the output. The future has already been dropped
+            // Store the output.
             //
             // Safety: Mutual exclusion is obtained by having transitioned the task
             // state -> Running
-            let stage = &self.core().stage;
             stage.store_output(output);
 
             // Transition to `Complete`, notifying the `JoinHandle` if necessary.
             transition_to_complete(self.header(), stage, &self.trailer());
+        } else {
+            // Drop the future.
+            //
+            // Safety: Mutual exclusion is obtained by having transitioned the task
+            // state -> Running
+            stage.drop_future_or_output();
         }
 
         // The task has completed execution and will no longer be scheduled.

tokio/tests/task_abort.rs

@@ -1,6 +1,7 @@
 #![warn(rust_2018_idioms)]
 #![cfg(feature = "full")]
 
+use std::sync::Arc;
 use std::thread::sleep;
 use std::time::Duration;
 
@@ -138,3 +139,38 @@ fn remote_abort_local_set_3929() {
     rt.block_on(local);
     jh2.join().unwrap();
 }
+
+/// Checks that a suspended task can be aborted even if the `JoinHandle` is immediately dropped.
+/// issue #3964: <https://github.com/tokio-rs/tokio/issues/3964>.
+#[test]
+fn test_abort_wakes_task_3964() {
+    let rt = tokio::runtime::Builder::new_current_thread()
+        .enable_time()
+        .build()
+        .unwrap();
+
+    rt.block_on(async move {
+        let notify_dropped = Arc::new(());
+        let weak_notify_dropped = Arc::downgrade(&notify_dropped);
+
+        let handle = tokio::spawn(async move {
+            // Make sure the Arc is moved into the task
+            let _notify_dropped = notify_dropped;
+            println!("task started");
+            tokio::time::sleep(std::time::Duration::new(100, 0)).await
+        });
+
+        // wait for task to sleep.
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+
+        handle.abort();
+        drop(handle);
+
+        // wait for task to abort.
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+
+        // Check that the Arc has been dropped.
+        assert!(weak_notify_dropped.upgrade().is_none());
+    });
+    //panic!();
+}