Repository: tokio-rs/tokio
File: tokio/src/task/local.rs (Rust)

Code before:
        })
    }

    fn pop_local(&self) -> Option<task::Notified<Arc<Shared>>> {
        unsafe {
            // Safety: because the `LocalSet` itself is `!Send`, we know we are
            // on the same thread if we have access to the `LocalSet`, and can
            // therefore access the local run queue.
            self.context.shared.local_state.task_pop_front()
        }
    }

    fn with<T>(&self, f: impl FnOnce() -> T) -> T {
        CURRENT.with(|local_data| {
            let _guard = local_data.enter(self.context.clone());
            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(|local_data| {
            local_data.enter(self.context.clone());
            (f.take().unwrap())()
        });

        match res {
            Ok(res) => res,
            Err(_access_error) => (f.take().unwrap())(),
        }
    }
}

cfg_unstable! {
    impl LocalSet {
        /// Configure how the `LocalSet` responds to an unhandled panic on a
        /// spawned task.
        ///
        /// By default, an unhandled panic (i.e. a panic not caught by
        /// [`std::panic::catch_unwind`]) has no impact on the `LocalSet`'s
        /// execution. The panic is error value is forwarded to the task's
        /// [`JoinHandle`] and all other spawned tasks continue running.
        ///
        /// The `unhandled_panic` option enables configuring this behavior.
        ///
        /// * `UnhandledPanic::Ignore` is the default behavior. Panics on
        ///   spawned tasks have no impact on the `LocalSet`'s execution.
        /// * `UnhandledPanic::ShutdownRuntime` will force the `LocalSet` to

Review comment (suggestion):
it looks like this is dropping the guard immediately, resulting in the reset happening _before_ we call `f()`. I don't think that's right. This needs to be:
```suggestion
            let _guard = local_data.enter(self.context.clone());
```

Code after:
                .queue
                .lock()
                .as_mut()
                .and_then(|queue| queue.pop_front())
                .or_else(|| self.pop_local())
        } else {
            self.pop_local().or_else(|| {
                self.context
                    .shared
                    .queue
                    .lock()
                    .as_mut()
                    .and_then(|queue| queue.pop_front())
            })
        };

        task.map(|task| unsafe {
            // Safety: because the `LocalSet` itself is `!Send`, we know we are
            // on the same thread if we have access to the `LocalSet`, and can
            // therefore access the local run queue.
            self.context.shared.local_state.assert_owner(task)
        })
    }

    fn pop_local(&self) -> Option<task::Notified<Arc<Shared>>> {
        unsafe {
            // Safety: because the `LocalSet` itself is `!Send`, we know we are
            // on the same thread if we have access to the `LocalSet`, and can
            // therefore access the local run queue.
            self.context.shared.local_state.task_pop_front()
        }
    }

    fn with<T>(&self, f: impl FnOnce() -> T) -> T {
        CURRENT.with(|local_data| {
            let _guard = local_data.enter(self.context.clone());
            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(|local_data| {
            let _guard = local_data.enter(self.context.clone());
            (f.take().unwrap())()
        });

        match res {