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

Code before:
    ///
    /// 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);
    ///
    /// // With this call, the task local value of fut is unset.
    /// assert_eq!(pinned.as_mut().take_value(), Some(42));
    ///
    /// // Poll **after** invoking `take_value()`
    /// let _ = (&mut pinned).as_mut().await;
    /// # }
    /// ```
    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();
        let mut future_opt = this.future;

        let res = this
            .local
            .scope_inner(this.slot, || match future_opt.as_mut().as_pin_mut() {
                Some(fut) => {
                    let res = fut.poll(cx);
                    if res.is_ready() {
                        future_opt.set(None);
                    }
                    Some(res)
                }
                None => None,
            });

Review comment (suggestion):
The change itself looks good, but I'd like the documentation to be a bit improved.

A few thoughts:

* The first line of documentation should be a short one-line summary of what it does. The body of the documentation then follows after an empty line. This could be e.g. "Returns the value stored in the task local by this future."
* We generally don't have sections after the examples, so your "Note" section should probably be moved. It can be combined with the body you already have at the top.
* For a function like this, a bulleted list often works well for explaining the behavior in each case.
* I would probably put the note about what happens if you poll the future after calling this method after the bulleted list.
* Your examples have one space too many. Currently, they will appear in the documentation indented by one space.

Code after:
        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();
        let mut future_opt = this.future;

        let res = this
            .local
            .scope_inner(this.slot, || match future_opt.as_mut().as_pin_mut() {
                Some(fut) => {
                    let res = fut.poll(cx);
                    if res.is_ready() {
                        future_opt.set(None);
                    }
                    Some(res)
                }
                None => None,
            });

        match res {
            Ok(Some(res)) => res,
            Ok(None) => panic!("`TaskLocalFuture` polled after completion"),
            Err(err) => err.panic(),
        }
    }
}

impl<T: 'static, F> fmt::Debug for TaskLocalFuture<T, F>
where
    T: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        /// Format the Option without Some.
        struct TransparentOption<'a, T> {
            value: &'a Option<T>,
        }
        impl<'a, T: fmt::Debug> fmt::Debug for TransparentOption<'a, T> {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                match self.value.as_ref() {
                    Some(value) => value.fmt(f),
                    // Hitting the None branch should not be possible.
                    None => f.pad("<missing>"),
                }
            }