From ca087319251da0740a26ab2afee3700a91af739f Mon Sep 17 00:00:00 2001
From: Alice Ryhl <aliceryhl@google.com>
Date: Mon, 5 Sep 2022 11:20:26 +0200
Subject: [PATCH] task: ignore failure to set TLS in LocalSet Drop

---
 tokio/src/task/local.rs | 33 ++++++++++++++++++++++++++++++++-
 1 file changed, 32 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();