From c659e4a757b5a2a63378ead33917f4e072dba5dc Mon Sep 17 00:00:00 2001
From: quininer <quininer@live.com>
Date: Tue, 23 Mar 2021 00:02:10 +0800
Subject: [PATCH] task: add sync_scope for LocalKey (#3612)

---
 tokio/src/task/task_local.rs | 45 ++++++++++++++++++++++++++++++++----
 1 file changed, 40 insertions(+), 5 deletions(-)

diff --git a/tokio/src/task/task_local.rs b/tokio/src/task/task_local.rs
index bc2e54a4e..e240cb6e5 100644
--- a/tokio/src/task/task_local.rs
+++ b/tokio/src/task/task_local.rs
@@ -127,6 +127,35 @@ impl<T: 'static> LocalKey<T> {
         .await
     }
 
+    /// Sets a value `T` as the task-local value for the closure `F`.
+    ///
+    /// On completion of `scope`, the task-local will be dropped.
+    ///
+    /// ### Examples
+    ///
+    /// ```
+    /// # async fn dox() {
+    /// tokio::task_local! {
+    ///     static NUMBER: u32;
+    /// }
+    ///
+    /// NUMBER.sync_scope(1, || {
+    ///     println!("task local value: {}", NUMBER.get());
+    /// });
+    /// # }
+    /// ```
+    pub fn sync_scope<F, R>(&'static self, value: T, f: F) -> R
+    where
+        F: FnOnce() -> R,
+    {
+        let mut scope = TaskLocalFuture {
+            local: &self,
+            slot: Some(value),
+            future: (),
+        };
+        Pin::new(&mut scope).with_task(|_| f())
+    }
+
     /// Accesses the current task-local and runs the provided closure.
     ///
     /// # Panics
@@ -185,10 +214,8 @@ pin_project! {
     }
 }
 
-impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> {
-    type Output = F::Output;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+impl<T: 'static, F> TaskLocalFuture<T, F> {
+    fn with_task<F2: FnOnce(Pin<&mut F>) -> R, R>(self: Pin<&mut Self>, f: F2) -> R {
         struct Guard<'a, T: 'static> {
             local: &'static LocalKey<T>,
             slot: &'a mut Option<T>,
@@ -213,7 +240,15 @@ impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> {
             local: *project.local,
         };
 
-        project.future.poll(cx)
+        f(project.future)
+    }
+}
+
+impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> {
+    type Output = F::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.with_task(|f| f.poll(cx))
     }
 }