task: add TaskLocalFuture::take_value (#6340)

tokio/src/task/task_local.rs

@@ -332,6 +332,50 @@ pin_project! {
     }
 }
 
+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;
 

tokio/tests/task_local.rs

@@ -117,3 +117,31 @@ async fn task_local_available_on_completion_drop() {
     assert_eq!(rx.await.unwrap(), 42);
     h.await.unwrap();
 }
+
+#[tokio::test]
+async fn take_value() {
+    tokio::task_local! {
+        static KEY: u32
+    }
+    let fut = KEY.scope(1, async {});
+    let mut pinned = Box::pin(fut);
+    assert_eq!(pinned.as_mut().take_value(), Some(1));
+    assert_eq!(pinned.as_mut().take_value(), None);
+}
+
+#[tokio::test]
+async fn poll_after_take_value_should_fail() {
+    tokio::task_local! {
+        static KEY: u32
+    }
+    let fut = KEY.scope(1, async {
+        let result = KEY.try_with(|_| {});
+        // The task local value no longer exists.
+        assert!(result.is_err());
+    });
+    let mut fut = Box::pin(fut);
+    fut.as_mut().take_value();
+
+    // Poll the future after `take_value` has been called
+    fut.await;
+}