From 61cde0f542da575d20b853f5b71d28d2960b460b Mon Sep 17 00:00:00 2001 From: mox692 Date: Sun, 11 Feb 2024 21:18:38 +0900 Subject: [PATCH 1/7] add TaskLocalFuture::take_value --- tokio/src/task/task_local.rs | 70 ++++++++++++++++++++++++++++++++++++ tokio/tests/task_local.rs | 28 +++++++++++++++ 2 files changed, 98 insertions(+) diff --git a/tokio/src/task/task_local.rs b/tokio/src/task/task_local.rs index 4abeeb37e10..deff7199c0f 100644 --- a/tokio/src/task/task_local.rs +++ b/tokio/src/task/task_local.rs @@ -332,6 +332,76 @@ pin_project! { } } +impl TaskLocalFuture +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 { + /// // 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 { + 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; +} From 06c19a288bdf3573578a14b7d906ce7ccd2637c2 Mon Sep 17 00:00:00 2001 From: mox692 Date: Thu, 15 Feb 2024 23:50:32 +0900 Subject: [PATCH 2/7] Update docs --- tokio/src/task/task_local.rs | 65 +++++++++++++++++++----------------- 1 file changed, 34 insertions(+), 31 deletions(-) diff --git a/tokio/src/task/task_local.rs b/tokio/src/task/task_local.rs index deff7199c0f..2313df4295a 100644 --- a/tokio/src/task/task_local.rs +++ b/tokio/src/task/task_local.rs @@ -336,34 +336,13 @@ impl TaskLocalFuture 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`. + /// Returns the value stored in the task local by this `TaskLocalFuture`. /// - /// # 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); + /// The function returns: /// - /// // 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)); - /// # } - /// ``` + /// * `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 /// @@ -378,13 +357,13 @@ where /// ``` /// # async fn dox() { /// tokio::task_local! { - /// static KEY: u32; + /// 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()) + /// // 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); @@ -393,7 +372,31 @@ where /// assert_eq!(pinned.as_mut().take_value(), Some(42)); /// /// // Poll **after** invoking `take_value()` - /// let _ = (&mut pinned).as_mut().await; + /// let _ = pinned.as_mut().await; + /// # } + /// ``` + /// + /// # 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 { From 15decdbb353a0ca7fa428141eabc1a8c3f8b8e1b Mon Sep 17 00:00:00 2001 From: kim / Motoyuki Kimura Date: Tue, 20 Feb 2024 01:10:02 +0900 Subject: [PATCH 3/7] Update tokio/src/task/task_local.rs Co-authored-by: Alice Ryhl --- tokio/src/task/task_local.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tokio/src/task/task_local.rs b/tokio/src/task/task_local.rs index 2313df4295a..7f9713d8898 100644 --- a/tokio/src/task/task_local.rs +++ b/tokio/src/task/task_local.rs @@ -340,8 +340,7 @@ where /// /// The function returns: /// - /// * `Some(T)` if the task local value exists, and unset the task local value - /// in the Future. + /// * `Some(T)` if the task local value exists. /// * `None` if the task local value does not exist. /// /// # Note From daad89d1a552bc6dfa2d7ce6d9a3c45e3ccc332b Mon Sep 17 00:00:00 2001 From: kim / Motoyuki Kimura Date: Tue, 20 Feb 2024 01:10:16 +0900 Subject: [PATCH 4/7] Update tokio/src/task/task_local.rs Co-authored-by: Alice Ryhl --- tokio/src/task/task_local.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/tokio/src/task/task_local.rs b/tokio/src/task/task_local.rs index 7f9713d8898..79327f18dd4 100644 --- a/tokio/src/task/task_local.rs +++ b/tokio/src/task/task_local.rs @@ -343,7 +343,6 @@ where /// * `Some(T)` if the task local value exists. /// * `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 From 1159b0bbcf99fd7b107ac65b007f6a32b730363b Mon Sep 17 00:00:00 2001 From: kim / Motoyuki Kimura Date: Tue, 20 Feb 2024 01:10:28 +0900 Subject: [PATCH 5/7] Update tokio/src/task/task_local.rs Co-authored-by: Alice Ryhl --- tokio/src/task/task_local.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tokio/src/task/task_local.rs b/tokio/src/task/task_local.rs index 79327f18dd4..8027d80ba8f 100644 --- a/tokio/src/task/task_local.rs +++ b/tokio/src/task/task_local.rs @@ -341,7 +341,7 @@ where /// The function returns: /// /// * `Some(T)` if the task local value exists. - /// * `None` if the task local value does not exist. + /// * `None` if the task local value has already been taken. /// /// /// Note that this function attempts to take the task local value regardless of From 3a42b1cb2bd50553b6e2638f5e41a553a2b6b50c Mon Sep 17 00:00:00 2001 From: kim / Motoyuki Kimura Date: Tue, 20 Feb 2024 01:11:05 +0900 Subject: [PATCH 6/7] Update tokio/src/task/task_local.rs Co-authored-by: Alice Ryhl --- tokio/src/task/task_local.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tokio/src/task/task_local.rs b/tokio/src/task/task_local.rs index 8027d80ba8f..0a735ac88a1 100644 --- a/tokio/src/task/task_local.rs +++ b/tokio/src/task/task_local.rs @@ -344,10 +344,9 @@ where /// * `None` if the task local value has already been taken. /// /// - /// 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. + /// 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`. /// /// For example, the following code returns `Err` for accessing the `KEY` variable /// in the async block of the `scope` function. From 3888f77bafce82f417d60a86ffadb647e340fdc6 Mon Sep 17 00:00:00 2001 From: mox692 Date: Tue, 20 Feb 2024 01:16:20 +0900 Subject: [PATCH 7/7] Remove code snippet in note section --- tokio/src/task/task_local.rs | 26 -------------------------- 1 file changed, 26 deletions(-) diff --git a/tokio/src/task/task_local.rs b/tokio/src/task/task_local.rs index 0a735ac88a1..ba58ea6ae8b 100644 --- a/tokio/src/task/task_local.rs +++ b/tokio/src/task/task_local.rs @@ -343,36 +343,10 @@ where /// * `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`. /// - /// 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 _ = pinned.as_mut().await; - /// # } - /// ``` - /// /// # Examples /// /// ```