Add deadlock detection

Zoxc · Zoxc · commit 0f55a9956583 · 2019-01-31T19:26:13.000+01:00
diff --git a/rayon-core/src/lib.rs b/rayon-core/src/lib.rs
@@ -20,8 +20,6 @@
 //! succeed.
 
 #![doc(html_root_url = "https://docs.rs/rayon-core/1.4")]
-#![deny(missing_debug_implementations)]
-#![deny(missing_docs)]
 
 use std::any::Any;
 use std::env;
@@ -65,6 +63,7 @@ pub mod tlv;
 pub mod internal;
 pub use join::{join, join_context};
 pub use scope::{scope, Scope};
+pub use registry::{Registry, mark_blocked, mark_unblocked};
 pub use spawn::spawn;
 pub use worker_local::WorkerLocal;
 
@@ -141,6 +140,9 @@ pub struct ThreadPoolBuilder {
     /// The stack size for the created worker threads
     stack_size: Option<usize>,
 
+    /// Closure invoked on deadlock.
+    deadlock_handler: Option<Box<DeadlockHandler>>,
+
     /// Closure invoked on worker thread start.
     start_handler: Option<Box<StartHandler>>,
 
@@ -169,6 +171,9 @@ pub struct Configuration {
 /// may be invoked multiple times in parallel.
 type PanicHandler = Fn(Box<Any + Send>) + Send + Sync;
 
+/// The type for a closure that gets invoked when the Rayon thread pool deadlocks
+type DeadlockHandler = Fn() + Send + Sync;
+
 /// The type for a closure that gets invoked when a thread starts. The
 /// closure is passed the index of the thread on which it is invoked.
 /// Note that this same closure may be invoked multiple times in parallel.
@@ -358,6 +363,19 @@ impl ThreadPoolBuilder {
         self.breadth_first
     }
 
+    /// Takes the current deadlock callback, leaving `None`.
+    fn take_deadlock_handler(&mut self) -> Option<Box<DeadlockHandler>> {
+        self.deadlock_handler.take()
+    }
+
+    /// Set a callback to be invoked on current deadlock.
+    pub fn deadlock_handler<H>(mut self, deadlock_handler: H) -> ThreadPoolBuilder
+        where H: Fn() + Send + Sync + 'static
+    {
+        self.deadlock_handler = Some(Box::new(deadlock_handler));
+        self
+    }
+
     /// Takes the current thread start callback, leaving `None`.
     fn take_start_handler(&mut self) -> Option<Box<StartHandler>> {
         self.start_handler.take()
@@ -528,6 +546,7 @@ impl fmt::Debug for ThreadPoolBuilder {
             ref get_thread_name,
             ref panic_handler,
             ref stack_size,
+            ref deadlock_handler, 
             ref start_handler,
             ref main_handler,
             ref exit_handler,
@@ -544,6 +563,7 @@ impl fmt::Debug for ThreadPoolBuilder {
         }
         let get_thread_name = get_thread_name.as_ref().map(|_| ClosurePlaceholder);
         let panic_handler = panic_handler.as_ref().map(|_| ClosurePlaceholder);
+        let deadlock_handler = deadlock_handler.as_ref().map(|_| ClosurePlaceholder);
         let start_handler = start_handler.as_ref().map(|_| ClosurePlaceholder);
         let exit_handler = exit_handler.as_ref().map(|_| ClosurePlaceholder);
         let main_handler = main_handler.as_ref().map(|_| ClosurePlaceholder);
@@ -553,6 +573,7 @@ impl fmt::Debug for ThreadPoolBuilder {
             .field("get_thread_name", &get_thread_name)
             .field("panic_handler", &panic_handler)
             .field("stack_size", &stack_size)
+            .field("deadlock_handler", &deadlock_handler)
             .field("start_handler", &start_handler)
             .field("exit_handler", &exit_handler)
             .field("main_handler", &main_handler)
diff --git a/rayon-core/src/registry.rs b/rayon-core/src/registry.rs
@@ -18,7 +18,7 @@ use std::thread;
 use std::usize;
 use unwind;
 use util::leak;
-use {ErrorKind, ExitHandler, PanicHandler, StartHandler,
+use {ErrorKind, ExitHandler, PanicHandler, DeadlockHandler, StartHandler,
      MainHandler, ThreadPoolBuildError, ThreadPoolBuilder};
 
 pub struct Registry {
@@ -27,6 +27,7 @@ pub struct Registry {
     sleep: Sleep,
     job_uninjector: Stealer<JobRef>,
     panic_handler: Option<Box<PanicHandler>>,
+    deadlock_handler: Option<Box<DeadlockHandler>>,
     start_handler: Option<Box<StartHandler>>,
     exit_handler: Option<Box<ExitHandler>>,
     main_handler: Option<Box<MainHandler>>,
@@ -114,10 +115,11 @@ impl Registry {
         let registry = Arc::new(Registry {
             thread_infos: stealers.into_iter().map(|s| ThreadInfo::new(s)).collect(),
             state: Mutex::new(RegistryState::new(inj_worker)),
-            sleep: Sleep::new(),
+            sleep: Sleep::new(n_threads),
             job_uninjector: inj_stealer,
             terminate_latch: CountLatch::new(),
             panic_handler: builder.take_panic_handler(),
+            deadlock_handler: builder.take_deadlock_handler(),
             start_handler: builder.take_start_handler(),
             main_handler: builder.take_main_handler(),
             exit_handler: builder.take_exit_handler(),
@@ -367,14 +369,11 @@ impl Registry {
     {
         // This thread isn't a member of *any* thread pool, so just block.
         debug_assert!(WorkerThread::current().is_null());
-        let job = StackJob::new(
-            |injected| {
-                let worker_thread = WorkerThread::current();
-                assert!(injected && !worker_thread.is_null());
-                op(&*worker_thread, true)
-            },
-            LockLatch::new(),
-        );
+        let job = StackJob::new(|injected| {
+            let worker_thread = WorkerThread::current();
+            assert!(injected && !worker_thread.is_null());
+            op(&*worker_thread, true)
+        }, LockLatch::new());
         self.inject(&[job.as_job_ref()]);
         job.latch.wait();
         job.into_result()
@@ -436,6 +435,24 @@ impl Registry {
     }
 }
 
+/// Mark a Rayon worker thread as blocked. This triggers the deadlock handler
+/// if no other worker thread is active
+#[inline]
+pub fn mark_blocked() {
+    let worker_thread = WorkerThread::current();
+    assert!(!worker_thread.is_null());
+    unsafe {
+        let registry = &(*worker_thread).registry;
+        registry.sleep.mark_blocked(&registry.deadlock_handler)
+    }
+}
+
+/// Mark a previously blocked Rayon worker thread as unblocked
+#[inline]
+pub fn mark_unblocked(registry: &Registry) {
+    registry.sleep.mark_unblocked()
+}
+
 #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
 pub struct RegistryId {
     addr: usize,
@@ -594,7 +611,11 @@ impl WorkerThread {
                 yields = self.registry.sleep.work_found(self.index, yields);
                 self.execute(job);
             } else {
-                yields = self.registry.sleep.no_work_found(self.index, yields);
+                yields = self.registry.sleep.no_work_found(
+                    self.index,
+                    yields,
+                    &self.registry.deadlock_handler
+                );
             }
         }
 
diff --git a/rayon-core/src/sleep/README.md b/rayon-core/src/sleep/README.md
@@ -386,3 +386,36 @@ some of them were hit hard:
   - 8-10% overhead on nbody-parreduce
   - 35% overhead on increment-all
   - 245% overhead on join-recursively
+
+# Deadlock detection
+
+This module tracks a number of variables in order to detect deadlocks due to user code blocking.
+These variables are stored in the `SleepData` struct which itself is kept behind a mutex.
+It contains the following fields:
+- `worker_count` - The number of threads in the thread pool.
+- `active_threads` - The number of threads in the thread pool which are running
+                     and aren't blocked in user code or sleeping.
+- `blocked_threads` - The number of threads which are blocked in user code.
+                      This doesn't include threads blocked by Rayon.
+
+User code can indicate blocking by calling `mark_blocked` before blocking and
+calling `mark_unblocked` before unblocking a thread.
+This will adjust `active_threads` and `blocked_threads` accordingly.
+
+When we tickle the thread pool in `Sleep::tickle_cold`, we set `active_threads` to
+`worker_count` - `blocked_threads` since we wake up all Rayon threads, but not thread blocked
+by user code.
+
+A deadlock is detected by checking if `active_threads` is 0 and `blocked_threads` is above 0.
+If we ignored `blocked_threads` we would have a deadlock
+immediately when creating the thread pool.
+We would also deadlock once the thread pool ran out of work.
+It is not possible for Rayon itself to deadlock.
+Deadlocks can only be caused by user code blocking, so this condition doesn't miss any deadlocks.
+
+We check for the deadlock condition when
+threads fall asleep in `mark_unblocked` and in `Sleep::sleep`.
+If there's a deadlock detected we call the user provided deadlock handler while we hold the
+lock to `SleepData`. This means the deadlock handler cannot call `mark_blocked` and
+`mark_unblocked`. The user is expected to handle the deadlock in some non-Rayon thread.
+Once the deadlock handler returns, the thread which called the deadlock handler will go to sleep.
diff --git a/rayon-core/src/sleep/mod.rs b/rayon-core/src/sleep/mod.rs
@@ -1,15 +1,39 @@
 //! Code that decides when workers should go to sleep. See README.md
 //! for an overview.
 
+use DeadlockHandler;
 use log::Event::*;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Condvar, Mutex};
 use std::thread;
 use std::usize;
 
+struct SleepData {
+    /// The number of threads in the thread pool.
+    worker_count: usize,
+
+    /// The number of threads in the thread pool which are running and
+    /// aren't blocked in user code or sleeping.
+    active_threads: usize,
+
+    /// The number of threads which are blocked in user code.
+    /// This doesn't include threads blocked by this module.
+    blocked_threads: usize,
+}
+
+impl SleepData {
+    /// Checks if the conditions for a deadlock holds and if so calls the deadlock handler
+    #[inline]
+    pub fn deadlock_check(&self, deadlock_handler: &Option<Box<DeadlockHandler>>) {
+        if self.active_threads == 0 && self.blocked_threads > 0 {
+            (deadlock_handler.as_ref().unwrap())();
+        }
+    }
+}
+
 pub struct Sleep {
     state: AtomicUsize,
-    data: Mutex<()>,
+    data: Mutex<SleepData>,
     tickle: Condvar,
 }
 
@@ -20,14 +44,42 @@ const ROUNDS_UNTIL_SLEEPY: usize = 32;
 const ROUNDS_UNTIL_ASLEEP: usize = 64;
 
 impl Sleep {
-    pub fn new() -> Sleep {
+    pub fn new(worker_count: usize) -> Sleep {
         Sleep {
             state: AtomicUsize::new(AWAKE),
-            data: Mutex::new(()),
+            data: Mutex::new(SleepData {
+                worker_count,
+                active_threads: worker_count,
+                blocked_threads: 0,
+            }),
             tickle: Condvar::new(),
         }
     }
 
+    /// Mark a Rayon worker thread as blocked. This triggers the deadlock handler
+    /// if no other worker thread is active
+    #[inline]
+    pub fn mark_blocked(&self, deadlock_handler: &Option<Box<DeadlockHandler>>) {
+        let mut data = self.data.lock().unwrap();
+        debug_assert!(data.active_threads > 0);
+        debug_assert!(data.blocked_threads < data.worker_count);
+        debug_assert!(data.active_threads > 0);
+        data.active_threads -= 1;
+        data.blocked_threads += 1;
+
+        data.deadlock_check(deadlock_handler);
+    }
+
+    /// Mark a previously blocked Rayon worker thread as unblocked
+    #[inline]
+    pub fn mark_unblocked(&self) {
+        let mut data = self.data.lock().unwrap();
+        debug_assert!(data.active_threads < data.worker_count);
+        debug_assert!(data.blocked_threads > 0);
+        data.active_threads += 1;
+        data.blocked_threads -= 1;
+    }
+
     fn anyone_sleeping(&self, state: usize) -> bool {
         state & SLEEPING != 0
     }
@@ -61,7 +113,7 @@ impl Sleep {
     }
 
     #[inline]
-    pub fn no_work_found(&self, worker_index: usize, yields: usize) -> usize {
+    pub fn no_work_found(&self, worker_index: usize, yields: usize, deadlock_handler: &Option<Box<DeadlockHandler>>) -> usize {
         log!(DidNotFindWork {
             worker: worker_index,
             yields: yields,
@@ -88,7 +140,7 @@ impl Sleep {
             }
         } else {
             debug_assert_eq!(yields, ROUNDS_UNTIL_ASLEEP);
-            self.sleep(worker_index);
+            self.sleep(worker_index, deadlock_handler);
             0
         }
     }
@@ -122,7 +174,10 @@ impl Sleep {
             old_state: old_state,
         });
         if self.anyone_sleeping(old_state) {
-            let _data = self.data.lock().unwrap();
+            let mut data = self.data.lock().unwrap();
+            // Set the active threads to the number of workers,
+            // excluding threads blocked by the user since we won't wake those up
+            data.active_threads = data.worker_count - data.blocked_threads;
             self.tickle.notify_all();
         }
     }
@@ -188,7 +243,7 @@ impl Sleep {
         self.worker_is_sleepy(state, worker_index)
     }
 
-    fn sleep(&self, worker_index: usize) {
+    fn sleep(&self, worker_index: usize, deadlock_handler: &Option<Box<DeadlockHandler>>) {
         loop {
             // Acquire here suffices. If we observe that the current worker is still
             // sleepy, then in fact we know that no writes have occurred, and anyhow
@@ -235,7 +290,7 @@ impl Sleep {
                 // reason for the `compare_exchange` to fail is if an
                 // awaken comes, in which case the next cycle around
                 // the loop will just return.
-                let data = self.data.lock().unwrap();
+                let mut data = self.data.lock().unwrap();
 
                 // This must be SeqCst on success because we want to
                 // ensure:
@@ -264,6 +319,11 @@ impl Sleep {
                     log!(FellAsleep {
                         worker: worker_index
                     });
+
+                    // Decrement the number of active threads and check for a deadlock
+                    data.active_threads -= 1;
+                    data.deadlock_check(deadlock_handler);
+
                     let _ = self.tickle.wait(data).unwrap();
                     log!(GotAwoken {
                         worker: worker_index
