Skip to content

Commit

Permalink
Fix dropping TLS for detached futures
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
jamesbornholt committed Mar 7, 2024
1 parent 29c782f commit ab1b48b
Show file tree
Hide file tree
Showing 4 changed files with 35 additions and 2 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ 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"] }
smallvec = { version = "1.11.2", features = ["const_new"] }
tracing = { version = "0.1.21", default-features = false, features = ["std"] }

[dev-dependencies]
Expand Down
7 changes: 7 additions & 0 deletions src/future/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/runtime/execution.rs
Original file line number Diff line number Diff line change
Expand Up @@ -165,7 +165,7 @@ impl Execution {
}

if state.config.record_steps_in_span {
state.current().step_span.record("i", state.current_schedule.len());
state.current().step_span.record("i", &state.current_schedule.len());
}
});
});
Expand Down
26 changes: 26 additions & 0 deletions tests/future/basic.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -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,
)
}

0 comments on commit ab1b48b

Please sign in to comment.