executor: update park implementation (#1402)

Signed-off-by: Lucio Franco <[email protected]>

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<Inner>,
+}
+
 /// Unblocks a thread that was blocked by `ParkThread`.
 #[derive(Clone, Debug)]
 pub struct UnparkThread {
-    inner: Unparker,
+    inner: Arc<Inner>,
 }
 
+#[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<Inner> {
+        &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<Inner>) -> *const () {
+        Arc::into_raw(this) as *const ()
+    }
+
+    pub unsafe fn from_raw(ptr: *const ()) -> Arc<Inner> {
+        Arc::from_raw(ptr as *const Inner)
+    }
+
+    /// Park the current thread for at most `dur`.
+    pub fn park(&self, timeout: Option<Duration>) -> 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<Inner>) -> 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);
-}