Use spawn_local to spawn from local Handles (#565)

Previously, every call to `current_thread::Handle::spawn` would go
through a `mpsc` channel. This is unnecessary when the `Handle` is still
on the same thread as the current thread executor. This patch fixes that
by storing the `ThreadId` of the executor when it is created, and then
comparing against that when `Handle::spawn` is called. If the call is
made from the same thread, `spawn_local` is used directly.

Fixes #562.

tokio-current-thread/src/lib.rs

@@ -44,6 +44,7 @@ use std::error::Error;
 use std::rc::Rc;
 use std::sync::{atomic, mpsc, Arc};
 use std::time::{Duration, Instant};
+use std::thread;
 
 /// Executes tasks on the current thread
 pub struct CurrentThread<P: Park = ParkThread> {
@@ -64,6 +65,9 @@ pub struct CurrentThread<P: Park = ParkThread> {
 
     /// Receiver for futures spawned from other threads
     spawn_receiver: mpsc::Receiver<Box<Future<Item = (), Error = ()> + Send + 'static>>,
+
+    /// The thread-local ID assigned to this executor.
+    id: u64,
 }
 
 /// Executes futures on the current thread.
@@ -176,6 +180,7 @@ impl<T: fmt::Debug> Error for BlockError<T> {
 
 /// This is mostly split out to make the borrow checker happy.
 struct Borrow<'a, U: 'a> {
+    id: u64,
     scheduler: &'a mut Scheduler<U>,
     num_futures: &'a atomic::AtomicUsize,
 }
@@ -186,13 +191,21 @@ trait SpawnLocal {
 
 struct CurrentRunner {
     spawn: Cell<Option<*mut SpawnLocal>>,
+    id: Cell<Option<u64>>,
 }
 
 /// Current thread's task runner. This is set in `TaskRunner::with`
 thread_local!(static CURRENT: CurrentRunner = CurrentRunner {
     spawn: Cell::new(None),
+    id: Cell::new(None),
 });
 
+/// Unique ID to assign to each new executor launched on this thread.
+///
+/// The unique ID is used to determine if the currently running executor matches the one referred
+/// to by a `Handle` so that direct task dispatch can be used.
+thread_local!(static EXECUTOR_ID: Cell<u64> = Cell::new(0));
+
 /// Run the executor bootstrapping the execution with the provided future.
 ///
 /// This creates a new [`CurrentThread`] executor, spawns the provided future,
@@ -256,6 +269,12 @@ impl<P: Park> CurrentThread<P> {
         let unpark = park.unpark();
 
         let (spawn_sender, spawn_receiver) = mpsc::channel();
+        let thread = thread::current().id();
+        let id = EXECUTOR_ID.with(|idc| {
+            let id = idc.get();
+            idc.set(id + 1);
+            id
+        });
 
         let scheduler = Scheduler::new(unpark);
         let notify = scheduler.notify();
@@ -266,11 +285,14 @@ impl<P: Park> CurrentThread<P> {
             scheduler: scheduler,
             num_futures: num_futures.clone(),
             park,
+            id,
             spawn_handle: Handle {
                 sender: spawn_sender,
                 num_futures: num_futures,
                 notify: notify,
                 shut_down: Cell::new(false),
+                thread: thread,
+                id,
             },
             spawn_receiver: spawn_receiver,
         }
@@ -367,6 +389,7 @@ impl<P: Park> CurrentThread<P> {
 
     fn borrow(&mut self) -> Borrow<P::Unpark> {
         Borrow {
+            id: self.id,
             scheduler: &mut self.scheduler,
             num_futures: &*self.num_futures,
         }
@@ -567,6 +590,7 @@ impl<'a, P: Park> Entered<'a, P> {
         // FIXME: Slightly ugly but needed to make the borrow checker happy
         let (mut borrow, spawn_receiver) = (
             Borrow {
+                id: self.executor.id,
                 scheduler: &mut self.executor.scheduler,
                 num_futures: &*self.executor.num_futures,
             },
@@ -579,6 +603,7 @@ impl<'a, P: Park> Entered<'a, P> {
 
         // After any pending futures were scheduled, do the actual tick
         borrow.scheduler.tick(
+            borrow.id,
             &mut *self.enter,
             borrow.num_futures)
     }
@@ -602,6 +627,10 @@ pub struct Handle {
     num_futures: Arc<atomic::AtomicUsize>,
     shut_down: Cell<bool>,
     notify: executor::NotifyHandle,
+    thread: thread::ThreadId,
+
+    /// The thread-local ID assigned to this Handle's executor.
+    id: u64,
 }
 
 // Manual implementation because the Sender does not implement Debug
@@ -640,12 +669,17 @@ impl Handle {
             return Err(SpawnError::shutdown());
         }
 
-        self.sender
-            .send(Box::new(future))
+        if thread::current().id() == self.thread {
+            let mut e = TaskExecutor::current();
+            if e.id() == Some(self.id) {
+                return e.spawn_local(Box::new(future));
+            }
+        }
+
+        self.sender.send(Box::new(future))
             .expect("CurrentThread does not exist anymore");
         // use 0 for the id, CurrentThread does not make use of it
         self.notify.notify(0);
-
         Ok(())
     }
 }
@@ -665,6 +699,13 @@ impl TaskExecutor {
         }
     }
 
+    /// Get the current executor's thread-local ID.
+    fn id(&self) -> Option<u64> {
+        CURRENT.with(|current| {
+            current.id.get()
+        })
+    }
+
     /// Spawn a future onto the current `CurrentThread` instance.
     pub fn spawn_local(&mut self, future: Box<Future<Item = (), Error = ()>>)
         -> Result<(), SpawnError>
@@ -716,6 +757,7 @@ impl<'a, U: Unpark> Borrow<'a, U> {
     where F: FnOnce() -> R,
     {
         CURRENT.with(|current| {
+            current.id.set(Some(self.id));
             current.set_spawn(self, || {
                 f()
             })
@@ -745,6 +787,7 @@ impl CurrentRunner {
         impl<'a> Drop for Reset<'a> {
             fn drop(&mut self) {
                 self.0.spawn.set(None);
+                self.0.id.set(None);
             }
         }
 

tokio-current-thread/src/scheduler.rs

@@ -210,7 +210,7 @@ where U: Unpark,
     ///
     /// This function should be called whenever the caller is notified via a
     /// wakeup.
-    pub fn tick(&mut self, enter: &mut Enter, num_futures: &AtomicUsize) -> bool
+    pub fn tick(&mut self, eid: u64, enter: &mut Enter, num_futures: &AtomicUsize) -> bool
     {
         let mut ret = false;
         let tick = self.inner.tick_num.fetch_add(1, SeqCst)
@@ -279,6 +279,7 @@ where U: Unpark,
                 let node = self.nodes.remove(node);
 
                 let mut borrow = Borrow {
+                    id: eid,
                     scheduler: self,
                     num_futures,
                 };