diff --git a/tokio-current-thread/src/lib.rs b/tokio-current-thread/src/lib.rs index 08f614e7a..fb2519f93 100644 --- a/tokio-current-thread/src/lib.rs +++ b/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 { @@ -64,6 +65,9 @@ pub struct CurrentThread { /// Receiver for futures spawned from other threads spawn_receiver: mpsc::Receiver + Send + 'static>>, + + /// The thread-local ID assigned to this executor. + id: u64, } /// Executes futures on the current thread. @@ -176,6 +180,7 @@ impl Error for BlockError { /// This is mostly split out to make the borrow checker happy. struct Borrow<'a, U: 'a> { + id: u64, scheduler: &'a mut Scheduler, num_futures: &'a atomic::AtomicUsize, } @@ -186,13 +191,21 @@ trait SpawnLocal { struct CurrentRunner { spawn: Cell>, + id: Cell>, } /// 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 = 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 CurrentThread

{ 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 CurrentThread

{ 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 CurrentThread

{ fn borrow(&mut self) -> Borrow { 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, shut_down: Cell, 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 { + CURRENT.with(|current| { + current.id.get() + }) + } + /// Spawn a future onto the current `CurrentThread` instance. pub fn spawn_local(&mut self, future: Box>) -> 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); } } diff --git a/tokio-current-thread/src/scheduler.rs b/tokio-current-thread/src/scheduler.rs index de65cc03b..f7a8f3b3b 100644 --- a/tokio-current-thread/src/scheduler.rs +++ b/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, };