From 30ba405c97c00f88fa5a2d9a6ef2a194346404fd Mon Sep 17 00:00:00 2001
From: James Bornholt <bornholt@amazon.com>
Date: Wed, 6 Mar 2024 18:11:14 -0600
Subject: [PATCH] Fix dropping TLS for detached futures

When a detached future is being cleaned up after the end of an
execution, it won't be able to access the current task's storage any
more. We need to bail out early from the cleanup to avoid a panic.
---
 Cargo.toml            |  4 ++--
 src/future/mod.rs     |  7 +++++++
 tests/future/basic.rs | 26 ++++++++++++++++++++++++++
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index f7a102be..aee2d095 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,8 +18,8 @@ rand_core = "0.6.4"
 rand = "0.8.5"
 rand_pcg = "0.3.1"
 scoped-tls = "1.0.0"
-smallvec = { version = "1.10.0", features = ["const_new"] }
-tracing = { version = "0.1.21", default-features = false, features = ["std"] }
+smallvec = { version = "1.11.2", features = ["const_new"] }
+tracing = { version = "0.1.36", default-features = false, features = ["std"] }
 
 [dev-dependencies]
 criterion = { version = "0.4.0", features = ["html_reports"] }
diff --git a/src/future/mod.rs b/src/future/mod.rs
index 9d168b44..79fee118 100644
--- a/src/future/mod.rs
+++ b/src/future/mod.rs
@@ -145,6 +145,13 @@ where
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         match self.future.as_mut().poll(cx) {
             Poll::Ready(result) => {
+                // If we've finished execution already (this task was detached), don't clean up. We
+                // can't access the state any more to destroy thread locals, and don't want to run
+                // any more wakers (which will be no-ops anyway).
+                if ExecutionState::try_with(|state| state.is_finished()).unwrap_or(true) {
+                    return Poll::Ready(());
+                }
+
                 // Run thread-local destructors before publishing the result.
                 // See `pop_local` for details on why this loop looks this slightly funky way.
                 // TODO: thread locals and futures don't mix well right now. each task gets its own
diff --git a/tests/future/basic.rs b/tests/future/basic.rs
index 0255dbb4..c998406d 100644
--- a/tests/future/basic.rs
+++ b/tests/future/basic.rs
@@ -1,3 +1,4 @@
+use futures::task::{FutureObj, Spawn, SpawnError, SpawnExt as _};
 use futures::{try_join, Future};
 use shuttle::sync::{Barrier, Mutex};
 use shuttle::{check_dfs, check_random, future, scheduler::PctScheduler, thread, Runner};
@@ -385,3 +386,28 @@ fn is_finished_on_join_handle() {
         None,
     );
 }
+
+struct ShuttleSpawn;
+
+impl Spawn for ShuttleSpawn {
+    fn spawn_obj(&self, future: FutureObj<'static, ()>) -> Result<(), SpawnError> {
+        future::spawn(future);
+        Ok(())
+    }
+}
+
+// Make sure a spawned detached task gets cleaned up correctly after execution ends
+#[test]
+fn clean_up_detached_task() {
+    check_dfs(
+        || {
+            let atomic = shuttle::sync::atomic::AtomicUsize::new(0);
+            let _task_handle = ShuttleSpawn
+                .spawn_with_handle(async move {
+                    atomic.fetch_add(1, Ordering::SeqCst);
+                })
+                .unwrap();
+        },
+        None,
+    )
+}