From 6110b226c2c24dd6c5e0119f59f1f10588dbfc43 Mon Sep 17 00:00:00 2001
From: Diggory Blake <diggsey@googlemail.com>
Date: Sun, 18 Jul 2021 21:46:35 +0100
Subject: [PATCH] runtime: drop future when polling cancelled future (#3965)

---
 tokio/src/runtime/task/harness.rs |  2 +-
 tokio/tests/task_abort.rs         | 95 +++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+), 1 deletion(-)

diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs
index c9c99c76949..763310fe92d 100644
--- a/tokio/src/runtime/task/harness.rs
+++ b/tokio/src/runtime/task/harness.rs
@@ -415,7 +415,7 @@ fn poll_future<T: Future>(
     cx: Context<'_>,
 ) -> PollFuture<T::Output> {
     if snapshot.is_cancelled() {
-        PollFuture::Complete(Err(JoinError::cancelled()), snapshot.is_join_interested())
+        PollFuture::Complete(Err(cancel_task(core)), snapshot.is_join_interested())
     } else {
         let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
             struct Guard<'a, T: Future> {
diff --git a/tokio/tests/task_abort.rs b/tokio/tests/task_abort.rs
index e96dcf0c5b4..69f7a209035 100644
--- a/tokio/tests/task_abort.rs
+++ b/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,97 @@ 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(10)).await;
+
+        handle.abort();
+        drop(handle);
+
+        // wait for task to abort.
+        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+
+        // Check that the Arc has been dropped.
+        assert!(weak_notify_dropped.upgrade().is_none());
+    });
+}
+
+struct PanicOnDrop;
+
+impl Drop for PanicOnDrop {
+    fn drop(&mut self) {
+        panic!("Well what did you expect would happen...");
+    }
+}
+
+/// Checks that aborting a task whose destructor panics does not allow the
+/// panic to escape the task.
+#[test]
+fn test_abort_task_that_panics_on_drop_contained() {
+    let rt = tokio::runtime::Builder::new_current_thread()
+        .enable_time()
+        .build()
+        .unwrap();
+
+    rt.block_on(async move {
+        let handle = tokio::spawn(async move {
+            // Make sure the Arc is moved into the task
+            let _panic_dropped = PanicOnDrop;
+            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(10)).await;
+
+        handle.abort();
+        drop(handle);
+
+        // wait for task to abort.
+        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+    });
+}
+
+/// Checks that aborting a task whose destructor panics has the expected result.
+#[test]
+fn test_abort_task_that_panics_on_drop_returned() {
+    let rt = tokio::runtime::Builder::new_current_thread()
+        .enable_time()
+        .build()
+        .unwrap();
+
+    rt.block_on(async move {
+        let handle = tokio::spawn(async move {
+            // Make sure the Arc is moved into the task
+            let _panic_dropped = PanicOnDrop;
+            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(10)).await;
+
+        handle.abort();
+        assert!(handle.await.unwrap_err().is_panic());
+    });
+}