Fix dropping TLS for detached futures

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"] }

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

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,
+    )
+}