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

Code before:
    impl<T: 'static, F> PinnedDrop for TaskLocalFuture<T, F> {
        fn drop(this: Pin<&mut Self>) {
            let this = this.project();
            if mem::needs_drop::<F>() && this.future.is_some() {
                // Drop the future while the task-local is set, if possible. Otherwise
                // the future is dropped normally when the `Option<F>` field drops.
                let mut future = this.future;
                let _ = this.local.scope_inner(this.slot, || {
                    future.set(None);
                });
            }
        }
    }
}

impl<T, F> TaskLocalFuture<T, F>
where
    T: 'static,
{
    /// Returns the value stored in the task local by this `TaskLocalFuture`.
    ///
    /// The function returns:
    ///
    /// * `Some(T)` if the task local value exists, and unset the task local value
    ///   in the Future.
    /// * `None` if the task local value does not exist.
    ///
    /// # Note
    ///
    /// Note that this function attempts to take the task local value regardless of
    /// whether the `TaskLocalFuture` is completed or not. This means that if you
    /// call this function before the `TaskLocalFuture` is completed, you need to
    /// make sure that the access to the task local value in the `TaskLocalFuture` is safe.
    ///
    /// For example, the following code returns `Err` for accessing the `KEY` variable
    /// in the async block of the `scope` function.
    ///
    /// ```
    /// # async fn dox() {
    /// tokio::task_local! {
    ///     static KEY: u32;
    /// }
    ///
    /// let fut = KEY.scope(42, async {
    ///     // Since `take_value()` has already been called at this point,
    ///     // `try_with` here will fail.
    ///     assert!(KEY.try_with(|_| {}).is_err())
    /// });
    ///
    /// let mut pinned = Box::pin(fut);
    ///

Review comment (suggestion):
Are there other cases than calling `take_value` multiple times that can trigger this?
```suggestion
    /// * `None` if the task local value has already been taken.
```

Code after:
    impl<T: 'static, F> PinnedDrop for TaskLocalFuture<T, F> {
        fn drop(this: Pin<&mut Self>) {
            let this = this.project();
            if mem::needs_drop::<F>() && this.future.is_some() {
                // Drop the future while the task-local is set, if possible. Otherwise
                // the future is dropped normally when the `Option<F>` field drops.
                let mut future = this.future;
                let _ = this.local.scope_inner(this.slot, || {
                    future.set(None);
                });
            }
        }
    }
}

impl<T, F> TaskLocalFuture<T, F>
where
    T: 'static,
{
    /// Returns the value stored in the task local by this `TaskLocalFuture`.
    ///
    /// The function returns:
    ///
    /// * `Some(T)` if the task local value exists.
    /// * `None` if the task local value has already been taken.
    ///
    /// Note that this function attempts to take the task local value even if
    /// the future has not yet completed. In that case, the value will no longer
    /// be available via the task local after the call to `take_value`.
    ///
    /// # Examples
    ///
    /// ```
    /// # async fn dox() {
    /// tokio::task_local! {
    ///     static KEY: u32;
    /// }
    ///
    /// let fut = KEY.scope(42, async {
    ///     // Do some async work
    /// });
    ///
    /// let mut pinned = Box::pin(fut);
    ///
    /// // Complete the TaskLocalFuture
    /// let _ = pinned.as_mut().await;
    ///
    /// // And here, we can take task local value
    /// let value = pinned.as_mut().take_value();
    ///
    /// assert_eq!(value, Some(42));