diff --git a/tokio-executor/src/park.rs b/tokio-executor/src/park.rs index 750281767..0770c50b7 100644 --- a/tokio-executor/src/park.rs +++ b/tokio-executor/src/park.rs @@ -44,11 +44,11 @@ //! [up]: trait.Unpark.html //! [mio]: https://docs.rs/mio/0.6/mio/struct.Poll.html -use crossbeam_utils::sync::{Parker, Unparker}; use std::marker::PhantomData; use std::mem; use std::rc::Rc; -use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Condvar, Mutex}; use std::task::{RawWaker, RawWakerVTable, Waker}; use std::time::Duration; @@ -166,16 +166,137 @@ pub struct ParkError { _p: (), } +struct Parker { + unparker: Arc, +} + /// Unblocks a thread that was blocked by `ParkThread`. #[derive(Clone, Debug)] pub struct UnparkThread { - inner: Unparker, + inner: Arc, } +#[derive(Debug)] +struct Inner { + state: AtomicUsize, + mutex: Mutex<()>, + condvar: Condvar, +} + +const IDLE: usize = 0; +const NOTIFY: usize = 1; +const SLEEP: usize = 2; + thread_local! { static CURRENT_PARKER: Parker = Parker::new(); } +// ==== impl Parker ==== + +impl Parker { + pub fn new() -> Self { + Self { + unparker: Arc::new(Inner { + state: AtomicUsize::new(IDLE), + mutex: Mutex::new(()), + condvar: Condvar::new(), + }), + } + } + + pub fn unparker(&self) -> &Arc { + &self.unparker + } + + pub fn park(&self) -> Result<(), ParkError> { + self.unparker.park(None) + } + + pub fn park_timeout(&self, timeout: Duration) -> Result<(), ParkError> { + self.unparker.park(Some(timeout)) + } +} + +// ==== impl Inner ==== + +impl Inner { + #[allow(clippy::wrong_self_convention)] + pub fn into_raw(this: Arc) -> *const () { + Arc::into_raw(this) as *const () + } + + pub unsafe fn from_raw(ptr: *const ()) -> Arc { + Arc::from_raw(ptr as *const Inner) + } + + /// Park the current thread for at most `dur`. + pub fn park(&self, timeout: Option) -> Result<(), ParkError> { + // If currently notified, then we skip sleeping. This is checked outside + // of the lock to avoid acquiring a mutex if not necessary. + match self.state.compare_and_swap(NOTIFY, IDLE, Ordering::SeqCst) { + NOTIFY => return Ok(()), + IDLE => {} + _ => unreachable!(), + } + + // The state is currently idle, so obtain the lock and then try to + // transition to a sleeping state. + let mut m = self.mutex.lock().unwrap(); + + // Transition to sleeping + match self.state.compare_and_swap(IDLE, SLEEP, Ordering::SeqCst) { + NOTIFY => { + // Notified before we could sleep, consume the notification and + // exit + self.state.store(IDLE, Ordering::SeqCst); + return Ok(()); + } + IDLE => {} + _ => unreachable!(), + } + + m = match timeout { + Some(timeout) => self.condvar.wait_timeout(m, timeout).unwrap().0, + None => self.condvar.wait(m).unwrap(), + }; + + // Transition back to idle. If the state has transitioned to `NOTIFY`, + // this will consume that notification + self.state.store(IDLE, Ordering::SeqCst); + + // Explicitly drop the mutex guard. There is no real point in doing it + // except that I find it helpful to make it explicit where we want the + // mutex to unlock. + drop(m); + + Ok(()) + } + + pub fn unpark(&self) { + // First, try transitioning from IDLE -> NOTIFY, this does not require a + // lock. + match self.state.compare_and_swap(IDLE, NOTIFY, Ordering::SeqCst) { + IDLE | NOTIFY => return, + SLEEP => {} + _ => unreachable!(), + } + + // The other half is sleeping, this requires a lock + let _m = self.mutex.lock().unwrap(); + + // Transition to NOTIFY + match self.state.swap(NOTIFY, Ordering::SeqCst) { + SLEEP => {} + NOTIFY => return, + IDLE => return, + _ => unreachable!(), + } + + // Wakeup the sleeper + self.condvar.notify_one(); + } +} + // ===== impl ParkThread ===== impl ParkThread { @@ -208,12 +329,12 @@ impl Park for ParkThread { } fn park(&mut self) -> Result<(), Self::Error> { - self.with_current(|inner| inner.park()); + self.with_current(|inner| inner.park())?; Ok(()) } fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { - self.with_current(|inner| inner.park_timeout(duration)); + self.with_current(|inner| inner.park_timeout(duration))?; Ok(()) } } @@ -243,12 +364,12 @@ impl UnparkThread { } } -unsafe fn unparker_to_raw_waker(unparker: Unparker) -> RawWaker { - RawWaker::new(Unparker::into_raw(unparker), &VTABLE) +unsafe fn unparker_to_raw_waker(unparker: Arc) -> RawWaker { + RawWaker::new(Inner::into_raw(unparker), &VTABLE) } unsafe fn clone(raw: *const ()) -> RawWaker { - let unparker = Unparker::from_raw(raw); + let unparker = Inner::from_raw(raw); // Increment the ref count mem::forget(unparker.clone()); @@ -257,18 +378,14 @@ unsafe fn clone(raw: *const ()) -> RawWaker { } unsafe fn wake(raw: *const ()) { - let unparker = Unparker::from_raw(raw); + let unparker = Inner::from_raw(raw); unparker.unpark(); } unsafe fn wake_by_ref(raw: *const ()) { - let unparker = Unparker::from_raw(raw); + let unparker = Inner::from_raw(raw); unparker.unpark(); // We don't actually own a reference to the unparker mem::forget(unparker); } - -unsafe fn drop(raw: *const ()) { - let _ = Unparker::from_raw(raw); -}