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 EnterGuard<'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) -> EnterGuard<'_> {
+        CURRENT.inner.with(|c| {
+            c.set(&self.context as *const _ as *const ());
+            EnterGuard {_guard: &self}
+        })
+    }
+
     /// Spawns a `!Send` task onto the local task set.
     ///
     /// This task is guaranteed to be run on the current thread.

tokio/tests/task_local_set.rs

@@ -126,6 +126,23 @@ 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