From e4cbc70279f8ee81c7ce7b57edf8a56844da1f39 Mon Sep 17 00:00:00 2001
From: Alice Ryhl <alice@ryhl.io>
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<T>(&self, f: impl FnOnce() -> T) -> T {
+        let mut f = Some(f);
+
+        let res = CURRENT.try_with(|ctx| {
+            struct Reset<'a> {
+                ctx_ref: &'a Cell<Option<Rc<Context>>>,
+                val: Option<Rc<Context>>,
+            }
+            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 <https://github.com/tokio-rs/tokio/issues/4973>
+#[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;