From 4152918a39214d22e4b9e8d2148641b38b4f1c6e Mon Sep 17 00:00:00 2001
From: Alice Ryhl <aliceryhl@google.com>
Date: Tue, 24 Aug 2021 11:58:30 +0000
Subject: [PATCH] runtime: fix leak in UnownedTask

---
 tokio/src/runtime/task/mod.rs   | 10 ++++++++--
 tokio/src/runtime/tests/task.rs |  6 ++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs
index cc3910d4b..5f0d5aba1 100644
--- a/tokio/src/runtime/task/mod.rs
+++ b/tokio/src/runtime/task/mod.rs
@@ -369,10 +369,16 @@ impl<S: Schedule> UnownedTask<S> {
         let raw = self.raw;
         mem::forget(self);
 
-        // Poll the task
+        // Transfer one ref-count to a Task object.
+        let task = Task::<S> {
+            raw,
+            _p: PhantomData,
+        };
+
+        // Use the other ref-count to poll the task.
         raw.poll();
         // Decrement our extra ref-count
-        raw.header().state.ref_dec();
+        drop(task);
     }
 
     pub(crate) fn shutdown(self) {
diff --git a/tokio/src/runtime/tests/task.rs b/tokio/src/runtime/tests/task.rs
index e93a1efe6..04e1b56e7 100644
--- a/tokio/src/runtime/tests/task.rs
+++ b/tokio/src/runtime/tests/task.rs
@@ -111,6 +111,12 @@ fn create_shutdown2() {
     drop(join);
 }
 
+#[test]
+fn unowned_poll() {
+    let (task, _) = unowned(async {}, NoopSchedule);
+    task.run();
+}
+
 #[test]
 fn schedule() {
     with(|rt| {