diff --git a/tokio/src/task/task_local.rs b/tokio/src/task/task_local.rs index 4abeeb37e10..ba58ea6ae8b 100644 --- a/tokio/src/task/task_local.rs +++ b/tokio/src/task/task_local.rs @@ -332,6 +332,50 @@ pin_project! { } } +impl TaskLocalFuture +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 { + let this = self.project(); + this.slot.take() + } +} + impl Future for TaskLocalFuture { type Output = F::Output; diff --git a/tokio/tests/task_local.rs b/tokio/tests/task_local.rs index fbc885c3599..a4718dc45bb 100644 --- a/tokio/tests/task_local.rs +++ b/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; +}