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

Code before:
    /// # async fn dox() {
    /// tokio::task_local! {
    ///     static NUMBER: u32;
    /// }
    ///
    /// NUMBER.sync_scope(1, || {
    ///     println!("task local value: {}", NUMBER.get());
    /// });
    /// # }
    /// ```
    ///
    /// [`with`]: fn@Self::with
    /// [`try_with`]: fn@Self::try_with
    #[track_caller]
    pub fn sync_scope<F, R>(&'static self, value: T, f: F) -> R
    where
        F: FnOnce() -> R,
    {
        let mut value = Some(value);
        match self.scope_inner(&mut value, f) {
            Ok(res) => res,
            Err(ScopeInnerErr::BorrowError) => {
                panic!("sync_scope called while Task Local Storage is borrowed")
            }
            Err(ScopeInnerErr::AccessError) => {
                panic!("cannot access a Task Local Storage value during or after destruction")
            }
        }
    }

    fn scope_inner<F, R>(&'static self, slot: &mut Option<T>, f: F) -> Result<R, ScopeInnerErr>
    where
        F: FnOnce() -> R,
    {
        struct Guard<'a, T: 'static> {
            local: &'static LocalKey<T>,
            slot: &'a mut Option<T>,
        }

        impl<'a, T: 'static> Drop for Guard<'a, T> {
            fn drop(&mut self) {
                // This should not panic.
                //
                // We know that the RefCell was not borrowed before the call to
                // `scope_inner`, so the only way for this to panic is if the
                // closure has created but not destroyed a RefCell guard.
                // However, we never give user-code access to the guards, so
                // there's no way for user-code to forget to destroy a guard.
                //
                // The call to `with` also should not panic, since the
                // thread-local wasn't destroyed when we first called

Review comment (suggestion):
very tiny nit: in the docs, we write "task-local storage", but in the panics, we write "Task Local Storage". i have a slight preference for keeping this consistent:
```suggestion
                panic!("sync_scope called while task-local storage is borrowed")
            }
            Err(ScopeInnerErr::AccessError) => {
                panic!("cannot access a task-local storage value during or after destruction")
```

Code after:
    /// # async fn dox() {
    /// tokio::task_local! {
    ///     static NUMBER: u32;
    /// }
    ///
    /// NUMBER.sync_scope(1, || {
    ///     println!("task local value: {}", NUMBER.get());
    /// });
    /// # }
    /// ```
    ///
    /// [`with`]: fn@Self::with
    /// [`try_with`]: fn@Self::try_with
    #[track_caller]
    pub fn sync_scope<F, R>(&'static self, value: T, f: F) -> R
    where
        F: FnOnce() -> R,
    {
        let mut value = Some(value);
        match self.scope_inner(&mut value, f) {
            Ok(res) => res,
            Err(err) => err.panic(),
        }
    }

    fn scope_inner<F, R>(&'static self, slot: &mut Option<T>, f: F) -> Result<R, ScopeInnerErr>
    where
        F: FnOnce() -> R,
    {
        struct Guard<'a, T: 'static> {
            local: &'static LocalKey<T>,
            slot: &'a mut Option<T>,
        }

        impl<'a, T: 'static> Drop for Guard<'a, T> {
            fn drop(&mut self) {
                // This should not panic.
                //
                // We know that the RefCell was not borrowed before the call to
                // `scope_inner`, so the only way for this to panic is if the
                // closure has created but not destroyed a RefCell guard.
                // However, we never give user-code access to the guards, so
                // there's no way for user-code to forget to destroy a guard.
                //
                // The call to `with` also should not panic, since the
                // thread-local wasn't destroyed when we first called
                // `scope_inner`, and it shouldn't have gotten destroyed since
                // then.
                self.local.inner.with(|inner| {
                    let mut ref_mut = inner.borrow_mut();
                    mem::swap(self.slot, &mut *ref_mut);