Allow customizing the threadpool's parker (#264)

* Allow customizing the threadpool's parker

This patch allows the user of threadpool to customize how the worker
threads park themselves. This allows custom parking logic to be
injected. For example, this allows embedding a timer on each worker
thread.

* Call `park` instance every so often.

Since the `park` is now customizable, it might have logic that must be
called every so often. For example, a timer might have timeouts that it
must expire.

Currently, if a worker is very busy, it won't call into the `park`
instance. This patch changes this so that after every 32 task
invocations, `park` is called with a duration of zero.

tokio-executor/CHANGELOG.md

@@ -1,3 +1,7 @@
+# 0.1.2 (unreleased)
+
+* Implement `Unpark` for `Box<Unpark>`.
+
 # 0.1.1 (March 22, 2018)
 
 * Optionally support futures 0.2.

tokio-executor/src/park.rs

@@ -127,6 +127,12 @@ pub trait Unpark: Sync + Send + 'static {
     fn unpark(&self);
 }
 
+impl Unpark for Box<Unpark> {
+    fn unpark(&self) {
+        (**self).unpark()
+    }
+}
+
 /// Blocks the current thread using a condition variable.
 ///
 /// Implements the [`Park`] functionality by using a condition variable. An

tokio-threadpool/CHANGELOG.md

@@ -1,3 +1,7 @@
+# 0.1.2 (unreleased)
+
+* Add the ability to specify a custom thread parker.
+
 # 0.1.1 (March 22, 2018)
 
 * Handle futures that panic on the threadpool.

tokio-threadpool/src/builder.rs

@@ -1,5 +1,6 @@
 use callback::Callback;
 use config::{Config, MAX_WORKERS};
+use park::{BoxPark, BoxedPark, DefaultPark};
 use sender::Sender;
 use shutdown_task::ShutdownTask;
 use sleep_stack::SleepStack;
@@ -9,12 +10,15 @@ use inner::Inner;
 use worker::Worker;
 use worker_entry::WorkerEntry;
 
+use std::error::Error;
+use std::fmt;
 use std::sync::Arc;
 use std::sync::atomic::AtomicUsize;
 use std::time::Duration;
 
 use num_cpus;
 use tokio_executor::Enter;
+use tokio_executor::park::Park;
 use futures::task::AtomicTask;
 
 #[cfg(feature = "unstable-futures")]
@@ -58,13 +62,15 @@ use futures2;
 /// thread_pool.shutdown().wait().unwrap();
 /// # }
 /// ```
-#[derive(Debug)]
 pub struct Builder {
     /// Thread pool specific configuration values
     config: Config,
 
     /// Number of workers to spawn
     pool_size: usize,
+
+    /// Generates the `Park` instances
+    new_park: Box<Fn() -> BoxPark>,
 }
 
 impl Builder {
@@ -92,6 +98,11 @@ impl Builder {
     pub fn new() -> Builder {
         let num_cpus = num_cpus::get();
 
+        let new_park = Box::new(|| {
+            Box::new(BoxedPark::new(DefaultPark::new()))
+                as BoxPark
+        });
+
         Builder {
             pool_size: num_cpus,
             config: Config {
@@ -100,6 +111,7 @@ impl Builder {
                 stack_size: None,
                 around_worker: None,
             },
+            new_park,
         }
     }
 
@@ -249,6 +261,45 @@ impl Builder {
         self
     }
 
+    /// Customize the `park` instance used by each worker thread.
+    ///
+    /// The provided closure `f` is called once per worker and returns a `Park`
+    /// instance that is used by the worker to put itself to sleep.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # extern crate tokio_threadpool;
+    /// # extern crate futures;
+    /// # use tokio_threadpool::Builder;
+    /// # fn decorate<F>(f: F) -> F { f }
+    ///
+    /// # pub fn main() {
+    /// // Create a thread pool with default configuration values
+    /// let thread_pool = Builder::new()
+    ///     .custom_park(|| {
+    ///         use tokio_threadpool::park::DefaultPark;
+    ///
+    ///         // This is the default park type that the worker would use if we
+    ///         // did not customize it.
+    ///         let park = DefaultPark::new();
+    ///
+    ///         // Decorate the `park` instance, allowing us to customize work
+    ///         // that happens when a worker therad goes to sleep.
+    ///         decorate(park)
+    ///     })
+    ///     .build();
+    /// # }
+    /// ```
+    pub fn custom_park<F, P>(&mut self, f: F) -> &mut Self
+    where F: Fn() -> P + 'static,
+          P: Park + Send + 'static,
+          P::Error: Error,
+    {
+        self.new_park = Box::new(move || Box::new(BoxedPark::new(f())));
+        self
+    }
+
     /// Create the configured `ThreadPool`.
     ///
     /// The returned `ThreadPool` instance is ready to spawn tasks.
@@ -272,7 +323,10 @@ impl Builder {
         trace!("build; num-workers={}", self.pool_size);
 
         for _ in 0..self.pool_size {
-            workers.push(WorkerEntry::new());
+            let park = (self.new_park)();
+            let unpark = park.unpark();
+
+            workers.push(WorkerEntry::new(park, unpark));
         }
 
         let inner = Arc::new(Inner {
@@ -299,3 +353,13 @@ impl Builder {
         ThreadPool { inner }
     }
 }
+
+impl fmt::Debug for Builder {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.debug_struct("Builder")
+            .field("config", &self.config)
+            .field("pool_size", &self.pool_size)
+            .field("new_park", &"Box<Fn() -> BoxPark>")
+            .finish()
+    }
+}

tokio-threadpool/src/lib.rs

@@ -15,6 +15,8 @@ extern crate log;
 #[cfg(feature = "unstable-futures")]
 extern crate futures2;
 
+pub mod park;
+
 mod builder;
 mod callback;
 mod config;

tokio-threadpool/src/park/boxed.rs

@@ -0,0 +1,40 @@
+use tokio_executor::park::{Park, Unpark};
+
+use std::error::Error;
+use std::time::Duration;
+
+pub(crate) type BoxPark = Box<Park<Unpark = BoxUnpark, Error = ()> + Send>;
+pub(crate) type BoxUnpark = Box<Unpark>;
+
+pub(crate) struct BoxedPark<T>(T);
+
+impl<T> BoxedPark<T> {
+    pub fn new(inner: T) -> Self {
+        BoxedPark(inner)
+    }
+}
+
+impl<T: Park + Send> Park for BoxedPark<T>
+where T::Error: Error,
+{
+    type Unpark = BoxUnpark;
+    type Error = ();
+
+    fn unpark(&self) -> Self::Unpark {
+        Box::new(self.0.unpark())
+    }
+
+    fn park(&mut self) -> Result<(), Self::Error> {
+        self.0.park()
+            .map_err(|e| {
+                warn!("calling `park` on worker thread errored -- shutting down thread: {}", e);
+            })
+    }
+
+    fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
+        self.0.park_timeout(duration)
+            .map_err(|e| {
+                warn!("calling `park` on worker thread errored -- shutting down thread: {}", e);
+            })
+    }
+}

tokio-threadpool/src/park/default_park.rs

@@ -0,0 +1,170 @@
+use tokio_executor::park::{Park, Unpark};
+
+use std::error::Error;
+use std::fmt;
+use std::sync::{Arc, Mutex, Condvar};
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering::SeqCst;
+use std::time::Duration;
+
+/// Parks the thread.
+#[derive(Debug)]
+pub struct DefaultPark {
+    inner: Arc<Inner>,
+}
+
+/// Unparks threads that were parked by `DefaultPark`.
+#[derive(Debug)]
+pub struct DefaultUnpark {
+    inner: Arc<Inner>,
+}
+
+/// Error returned by [`ParkThread`]
+///
+/// This currently is never returned, but might at some point in the future.
+///
+/// [`ParkThread`]: struct.ParkThread.html
+#[derive(Debug)]
+pub struct ParkError {
+    _p: (),
+}
+
+#[derive(Debug)]
+struct Inner {
+    state: AtomicUsize,
+    mutex: Mutex<()>,
+    condvar: Condvar,
+}
+
+const IDLE: usize = 0;
+const NOTIFY: usize = 1;
+const SLEEP: usize = 2;
+
+// ===== impl DefaultPark =====
+
+impl DefaultPark {
+    /// Creates a new `DefaultPark` instance.
+    pub fn new() -> DefaultPark {
+        let inner = Arc::new(Inner {
+            state: AtomicUsize::new(IDLE),
+            mutex: Mutex::new(()),
+            condvar: Condvar::new(),
+        });
+
+        DefaultPark { inner }
+    }
+}
+
+impl Park for DefaultPark {
+    type Unpark = DefaultUnpark;
+    type Error = ParkError;
+
+    fn unpark(&self) -> Self::Unpark {
+        let inner = self.inner.clone();
+        DefaultUnpark { inner }
+    }
+
+    fn park(&mut self) -> Result<(), Self::Error> {
+        self.inner.park(None)
+    }
+
+    fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
+        self.inner.park(Some(duration))
+    }
+}
+
+// ===== impl DefaultUnpark =====
+
+impl Unpark for DefaultUnpark {
+    fn unpark(&self) {
+        self.inner.unpark();
+    }
+}
+
+impl Inner {
+    /// Park the current thread for at most `dur`.
+    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, SeqCst) {
+            NOTIFY => return Ok(()),
+            IDLE => {},
+            _ => unreachable!(),
+        }
+
+        // If the duration is zero, then there is no need to actually block
+        if let Some(ref dur) = timeout {
+            if *dur == Duration::from_millis(0) {
+                return Ok(());
+            }
+        }
+
+        // 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, SeqCst) {
+            NOTIFY => {
+                // Notified before we could sleep, consume the notification and
+                // exit
+                self.state.store(IDLE, 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 transitione dto `NOTIFY`,
+        // this will consume that notification
+        self.state.store(IDLE, 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(())
+    }
+
+    fn unpark(&self) {
+        // First, try transitioning from IDLE -> NOTIFY, this does not require a
+        // lock.
+        match self.state.compare_and_swap(IDLE, NOTIFY, SeqCst) {
+            IDLE | NOTIFY => return,
+            SLEEP => {}
+            _ => unreachable!(),
+        }
+
+        // The other half is sleeping, this requires a lock
+        let _m = self.mutex.lock().unwrap();
+
+        // Transition from SLEEP -> NOTIFY
+        match self.state.compare_and_swap(SLEEP, NOTIFY, SeqCst) {
+            SLEEP => {}
+            _ => return,
+        }
+
+        // Wakeup the sleeper
+        self.condvar.notify_one();
+    }
+}
+
+// ===== impl ParkError =====
+
+impl fmt::Display for ParkError {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        self.description().fmt(fmt)
+    }
+}
+
+impl Error for ParkError {
+    fn description(&self) -> &str {
+        "unknown park error"
+    }
+}

tokio-threadpool/src/park/mod.rs

@@ -0,0 +1,8 @@
+//! Thread parking utilities.
+
+mod boxed;
+mod default_park;
+
+pub use self::default_park::{DefaultPark, DefaultUnpark, ParkError};
+
+pub(crate) use self::boxed::{BoxPark, BoxUnpark, BoxedPark};

tokio-threadpool/src/worker.rs

@@ -13,15 +13,15 @@ use worker_state::{
     WORKER_SIGNALED,
 };
 
+use tokio_executor;
+
 use std::cell::Cell;
 use std::marker::PhantomData;
 use std::rc::Rc;
-use std::thread;
-use std::time::Instant;
 use std::sync::atomic::Ordering::{AcqRel, Acquire};
 use std::sync::Arc;
-
-use tokio_executor;
+use std::thread;
+use std::time::{Duration, Instant};
 
 /// Thread worker
 ///
@@ -42,6 +42,9 @@ pub struct Worker {
     _p: PhantomData<Rc<()>>,
 }
 
+// Pointer to the current worker info
+thread_local!(static CURRENT_WORKER: Cell<*const Worker> = Cell::new(0 as *const _));
+
 impl Worker {
     pub(crate) fn spawn(idx: usize, inner: &Arc<Inner>) {
         trace!("spawning new worker thread; idx={}", idx);
@@ -107,6 +110,8 @@ impl Worker {
     ///
     /// This function blocks until the worker is shutting down.
     pub fn run(&self) {
+        const LIGHT_SLEEP_INTERVAL: usize = 32;
+
         // Get the notifier.
         let notify = Arc::new(Notifier {
             inner: Arc::downgrade(&self.inner),
@@ -115,6 +120,7 @@ impl Worker {
 
         let mut first = true;
         let mut spin_cnt = 0;
+        let mut tick = 0;
 
         while self.check_run_state(first) {
             first = false;
@@ -125,13 +131,24 @@ impl Worker {
 
             // Run the next available task
             if self.try_run_task(&notify, &mut sender) {
+                if tick % LIGHT_SLEEP_INTERVAL == 0 {
+                    self.sleep_light();
+                }
+
+                tick = tick.wrapping_add(1);
                 spin_cnt = 0;
+
                 // As long as there is work, keep looping.
                 continue;
             }
 
             // No work in this worker's queue, it is time to try stealing.
             if self.try_steal_task(&notify, &mut sender) {
+                if tick % LIGHT_SLEEP_INTERVAL == 0 {
+                    self.sleep_light();
+                }
+
+                tick = tick.wrapping_add(1);
                 spin_cnt = 0;
                 continue;
             }
@@ -142,16 +159,11 @@ impl Worker {
             }
 
             // Starting to get sleeeeepy
-            if spin_cnt < 32 {
+            if spin_cnt < 61 {
                 spin_cnt += 1;
-
-                // Don't do anything further
-            } else if spin_cnt < 256 {
-                spin_cnt += 1;
-
-                // Yield the thread
-                thread::yield_now();
             } else {
+                tick = 0;
+
                 if !self.sleep() {
                     return;
                 }
@@ -357,7 +369,6 @@ impl Worker {
     /// Put the worker to sleep
     ///
     /// Returns `true` if woken up due to new work arriving.
-    #[inline]
     fn sleep(&self) -> bool {
         trace!("Worker::sleep; idx={}", self.idx);
 
@@ -365,9 +376,7 @@ impl Worker {
 
         // The first part of the sleep process is to transition the worker state
         // to "pushed". Now, it may be that the worker is already pushed on the
-        // sleeper stack, in which case, we don't push again. However, part of
-        // this process is also to do some final state checks to avoid entering
-        // the mutex if at all possible.
+        // sleeper stack, in which case, we don't push again.
 
         loop {
             let mut next = state;
@@ -376,6 +385,9 @@ impl Worker {
                 WORKER_RUNNING => {
                     // Try setting the pushed state
                     next.set_pushed();
+
+                    // Transition the worker state to sleeping
+                    next.set_lifecycle(WORKER_SLEEPING);
                 }
                 WORKER_NOTIFIED | WORKER_SIGNALED => {
                     // No need to sleep, transition back to running and move on.
@@ -417,66 +429,18 @@ impl Worker {
             state = actual;
         }
 
-        // Acquire the sleep mutex, the state is transitioned to sleeping within
-        // the mutex in order to avoid losing wakeup notifications.
-        let mut lock = self.entry().park_mutex.lock().unwrap();
-
-        // Transition the state to sleeping, a CAS is still needed as other
-        // state transitions could happen unrelated to the sleep / wakeup
-        // process. We also have to redo the lifecycle check done above as
-        // the state could have been transitioned before entering the mutex.
-        loop {
-            let mut next = state;
-
-            match state.lifecycle() {
-                WORKER_RUNNING => {}
-                WORKER_NOTIFIED | WORKER_SIGNALED => {
-                    // Release the lock, sleep will not happen this call.
-                    drop(lock);
-
-                    // Transition back to running
-                    loop {
-                        let mut next = state;
-                        next.set_lifecycle(WORKER_RUNNING);
-
-                        let actual = self.entry().state.compare_and_swap(
-                            state.into(), next.into(), AcqRel).into();
-
-                        if actual == state {
-                            return true;
-                        }
-
-                        state = actual;
-                    }
-                }
-                _ => unreachable!(),
-            }
-
-            trace!(" sleeping -- set WORKER_SLEEPING; idx={}", self.idx);
-
-            next.set_lifecycle(WORKER_SLEEPING);
-
-            let actual = self.entry().state.compare_and_swap(
-                state.into(), next.into(), AcqRel).into();
-
-            if actual == state {
-                break;
-            }
-
-            state = actual;
-        }
-
         trace!("    -> starting to sleep; idx={}", self.idx);
 
         let sleep_until = self.inner.config.keep_alive
             .map(|dur| Instant::now() + dur);
 
-        // The state has been transitioned to sleeping, we can now wait on the
-        // condvar. This is done in a loop as condvars can wakeup spuriously.
+        // The state has been transitioned to sleeping, we can now wait by
+        // calling the parker. This is done in a loop as condvars can wakeup
+        // spuriously.
         loop {
             let mut drop_thread = false;
 
-            lock = match sleep_until {
+            match sleep_until {
                 Some(when) => {
                     let now = Instant::now();
 
@@ -486,14 +450,20 @@ impl Worker {
 
                     let dur = when - now;
 
-                    self.entry().park_condvar
-                        .wait_timeout(lock, dur)
-                        .unwrap().0
+                    unsafe {
+                        (*self.entry().park.get())
+                            .park_timeout(dur)
+                            .unwrap();
+                    }
                 }
                 None => {
-                    self.entry().park_condvar.wait(lock).unwrap()
+                    unsafe {
+                        (*self.entry().park.get())
+                            .park()
+                            .unwrap();
+                    }
                 }
-            };
+            }
 
             trace!("    -> wakeup; idx={}", self.idx);
 
@@ -504,9 +474,6 @@ impl Worker {
                 match state.lifecycle() {
                     WORKER_SLEEPING => {}
                     WORKER_NOTIFIED | WORKER_SIGNALED => {
-                        // Release the lock, done sleeping
-                        drop(lock);
-
                         // Transition back to running
                         loop {
                             let mut next = state;
@@ -526,6 +493,7 @@ impl Worker {
                 }
 
                 if !drop_thread {
+                    // This goees back to the outer loop.
                     break;
                 }
 
@@ -547,6 +515,17 @@ impl Worker {
         }
     }
 
+    /// This doesn't actually put the thread to sleep. It calls
+    /// `park.park_timeout` with a duration of 0. This allows the park
+    /// implementation to perform any work that might be done on an interval.
+    fn sleep_light(&self) {
+        unsafe {
+            (*self.entry().park.get())
+                .park_timeout(Duration::from_millis(0))
+                .unwrap();
+        }
+    }
+
     fn entry(&self) -> &WorkerEntry {
         &self.inner.workers[self.idx]
     }
@@ -568,6 +547,3 @@ impl Drop for Worker {
         }
     }
 }
-
-// Pointer to the current worker info
-thread_local!(static CURRENT_WORKER: Cell<*const Worker> = Cell::new(0 as *const _));

tokio-threadpool/src/worker_entry.rs

@@ -1,3 +1,4 @@
+use park::{BoxPark, BoxUnpark};
 use task::{Task, Queue};
 use worker_state::{
     WorkerState,
@@ -6,13 +7,12 @@ use worker_state::{
 };
 
 use std::cell::UnsafeCell;
-use std::sync::atomic::Ordering::{AcqRel};
+use std::fmt;
+use std::sync::atomic::Ordering::{AcqRel, Relaxed};
 use std::sync::atomic::AtomicUsize;
-use std::sync::{Mutex, Condvar};
 
 use deque;
 
-#[derive(Debug)]
 pub(crate) struct WorkerEntry {
     // Worker state. This is mutated when notifying the worker.
     pub state: AtomicUsize,
@@ -26,18 +26,18 @@ pub(crate) struct WorkerEntry {
     // Stealer half of deque
     pub steal: deque::Stealer<Task>,
 
-    // Park mutex
-    pub park_mutex: Mutex<()>,
+    // Thread parker
+    pub park: UnsafeCell<BoxPark>,
 
-    // Park condvar
-    pub park_condvar: Condvar,
+    // Thread unparker
+    pub unpark: BoxUnpark,
 
     // MPSC queue of jobs submitted to the worker from an external source.
     pub inbound: Queue,
 }
 
 impl WorkerEntry {
-    pub fn new() -> Self {
+    pub fn new(park: BoxPark, unpark: BoxUnpark) -> Self {
         let w = deque::Deque::new();
         let s = w.stealer();
 
@@ -47,8 +47,8 @@ impl WorkerEntry {
             deque: w,
             steal: s,
             inbound: Queue::new(),
-            park_mutex: Mutex::new(()),
-            park_condvar: Condvar::new(),
+            park: UnsafeCell::new(park),
+            unpark,
         }
     }
 
@@ -104,8 +104,7 @@ impl WorkerEntry {
 
     #[inline]
     pub fn wakeup(&self) {
-        let _lock = self.park_mutex.lock().unwrap();
-        self.park_condvar.notify_one();
+        self.unpark.unpark();
     }
 
     #[inline]
@@ -118,3 +117,17 @@ impl WorkerEntry {
         unsafe { *self.next_sleeper.get() = val; }
     }
 }
+
+impl fmt::Debug for WorkerEntry {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.debug_struct("WorkerEntry")
+            .field("state", &self.state.load(Relaxed))
+            .field("next_sleeper", &"UnsafeCell<usize>")
+            .field("deque", &self.deque)
+            .field("steal", &self.steal)
+            .field("park", &"UnsafeCell<BoxPark>")
+            .field("unpark", &"BoxUnpark")
+            .field("inbound", &self.inbound)
+            .finish()
+    }
+}