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

Code before:

impl<T, F> TaskLocalFuture<T, F>
where
    T: 'static,
{
    /// Takes the task local value `T` owned by the `TaskLocalFuture`. If the
    /// task local value exists, then returns `Some(T)` and the task local value
    /// inside the `TaskLocalFuture` becomes unset. If it does not exist,
    /// it returns `None`.
    ///
    /// # 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 _ = (&mut pinned).as_mut().await;
    ///
    ///  // And here, we can take task local value
    ///  let value = pinned.as_mut().take_value();
    ///
    ///  assert_eq!(value, Some(42));
    /// # }
    /// ```
    ///
    /// # 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 {

Review comment (suggestion):
I would think that the `.as_mut()` is enough to make this compile.
```suggestion
    ///  // Complete the TaskLocalFuture
    ///  let _ = pinned.as_mut().await;
```

Code after:

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));
    /// # }
    /// ```
    pub fn take_value(self: Pin<&mut Self>) -> Option<T> {
        let this = self.project();
        this.slot.take()
    }
}

impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> {
    type Output = F::Output;

    #[track_caller]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();