task: add LocalKey::try_get (#7666)

Signed-off-by: tison <[email protected]>

Author: tisonkun

tokio/src/task/task_local.rs

@@ -275,6 +275,16 @@ impl<T: Clone + 'static> LocalKey<T> {
     pub fn get(&'static self) -> T {
         self.with(|v| v.clone())
     }
+
+    /// Returns a copy of the task-local value
+    /// if the task-local value implements `Clone`.
+    ///
+    /// If the task-local with the associated key is not present, this
+    /// method will return an `AccessError`. For a panicking variant,
+    /// see `get`.
+    pub fn try_get(&'static self) -> Result<T, AccessError> {
+        self.try_with(|v| v.clone())
+    }
 }
 
 impl<T: 'static> fmt::Debug for LocalKey<T> {

tokio/tests/task_local.rs

@@ -145,3 +145,27 @@ async fn poll_after_take_value_should_fail() {
     // Poll the future after `take_value` has been called
     fut.await;
 }
+
+#[tokio::test]
+async fn get_value() {
+    tokio::task_local! {
+        static KEY: u32
+    }
+
+    KEY.scope(1, async {
+        assert_eq!(KEY.get(), 1);
+        assert_eq!(KEY.try_get().unwrap(), 1);
+    })
+    .await;
+
+    let fut = KEY.scope(1, async {
+        let result = KEY.try_get();
+        // 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;
+}