Add LocalSet::enter

tokio/src/task/local.rs

@@ -319,6 +319,11 @@ const MAX_TASKS_PER_TICK: usize = 61;
 /// How often it check the remote queue first.
 const REMOTE_FIRST_INTERVAL: u8 = 31;
 
+#[derive(Debug)]
+pub struct LocalEnterGuard<'a> {
+    _guard: &'a LocalSet,
+}
+
 impl LocalSet {
     /// Returns a new local task set.
     pub fn new() -> LocalSet {
@@ -336,6 +341,14 @@ impl LocalSet {
         }
     }
 
+    /// Enter current LocalSet context
+    pub fn enter(&self) -> LocalEnterGuard<'_> {
+        CURRENT.inner.with(|c| {
+            c.set(&self.context as *const _ as *const ());
+            LocalEnterGuard { _guard: &self }
+        })
+    }
+
     /// Spawns a `!Send` task onto the local task set.
     ///
     /// This task is guaranteed to be run on the current thread.

tokio/src/task/mod.rs

@@ -297,7 +297,7 @@ cfg_rt! {
     }
 
     mod local;
-    pub use local::{spawn_local, LocalSet};
+    pub use local::{spawn_local, LocalSet, LocalEnterGuard};
 
     mod task_local;
     pub use task_local::LocalKey;

tokio/tests/task_local_set.rs

@@ -126,6 +126,21 @@ async fn local_threadpool_timer() {
         })
         .await;
 }
+#[test]
+fn enter_guard_spawn() {
+    let local = LocalSet::new();
+    let _guard = local.enter();
+    // Run the local task set.
+
+    let join = task::spawn_local(async { true });
+    let rt = runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap();
+    local.block_on(&rt, async move {
+        assert!(join.await.unwrap());
+    });
+}
 
 #[test]
 // This will panic, since the thread that calls `block_on` cannot use