Skip to content

Commit

Permalink
task: ignore failure to set TLS in LocalSet Drop
Browse files Browse the repository at this point in the history
  • Loading branch information
Darksonn committed Sep 5, 2022
1 parent 1fd7f82 commit 2c25b73
Showing 1 changed file with 32 additions and 1 deletion.
33 changes: 32 additions & 1 deletion tokio/src/task/local.rs
Expand Up @@ -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! {
Expand Down Expand Up @@ -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();
Expand Down

0 comments on commit 2c25b73

Please sign in to comment.