From e4cbc70279f8ee81c7ce7b57edf8a56844da1f39 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Tue, 13 Sep 2022 08:51:04 +0200 Subject: [PATCH] task: ignore failure to set TLS in LocalSet Drop (#4976) --- tokio/src/task/local.rs | 33 ++++++++++++++++++++++++++++++++- tokio/tests/task_local_set.rs | 21 +++++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index a5bd1bb8835..3786cf36669 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -633,6 +633,37 @@ impl LocalSet { f() }) } + + /// This method is like `with`, but it just calls `f` without setting the thread-local if that + /// fails. + fn with_if_possible(&self, f: impl FnOnce() -> T) -> T { + let mut f = Some(f); + + let res = CURRENT.try_with(|ctx| { + struct Reset<'a> { + ctx_ref: &'a Cell>>, + val: Option>, + } + impl<'a> Drop for Reset<'a> { + fn drop(&mut self) { + self.ctx_ref.replace(self.val.take()); + } + } + let old = ctx.replace(Some(self.context.clone())); + + let _reset = Reset { + ctx_ref: ctx, + val: old, + }; + + (f.take().unwrap())() + }); + + match res { + Ok(res) => res, + Err(_access_error) => (f.take().unwrap())(), + } + } } cfg_unstable! { @@ -744,7 +775,7 @@ impl Default for LocalSet { impl Drop for LocalSet { fn drop(&mut self) { - self.with(|| { + self.with_if_possible(|| { // Shut down all tasks in the LocalOwnedTasks and close it to // prevent new tasks from ever being added. self.context.owned.close_and_shutdown_all(); diff --git a/tokio/tests/task_local_set.rs b/tokio/tests/task_local_set.rs index b6245acf77a..e3c0c65f712 100644 --- a/tokio/tests/task_local_set.rs +++ b/tokio/tests/task_local_set.rs @@ -311,6 +311,27 @@ fn join_local_future_elsewhere() { }); } +// Tests for +#[cfg(not(tokio_wasi))] // Wasi doesn't support threads +#[tokio::test(flavor = "multi_thread")] +async fn localset_in_thread_local() { + thread_local! { + static LOCAL_SET: LocalSet = LocalSet::new(); + } + + // holds runtime thread until end of main fn. + let (_tx, rx) = oneshot::channel::<()>(); + let handle = tokio::runtime::Handle::current(); + + std::thread::spawn(move || { + LOCAL_SET.with(|local_set| { + handle.block_on(local_set.run_until(async move { + let _ = rx.await; + })) + }); + }); +} + #[test] fn drop_cancels_tasks() { use std::rc::Rc;