Fix dropping TLS for detached futures (#139)

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

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