Add deadlock detection

Zoxc · Zoxc · commit 9be36f822fbe · 2023-03-11T15:36:34.000+01:00
diff --git a/rayon-core/src/lib.rs b/rayon-core/src/lib.rs
@@ -61,9 +61,6 @@
 //! conflicting requirements will need to be resolved before the build will
 //! succeed.
 
-#![deny(missing_debug_implementations)]
-#![deny(missing_docs)]
-#![deny(unreachable_pub)]
 #![warn(rust_2018_idioms)]
 
 use std::any::Any;
@@ -99,6 +96,7 @@ pub mod tlv;
 pub use self::broadcast::{broadcast, spawn_broadcast, BroadcastContext};
 pub use self::join::{join, join_context};
 pub use self::registry::ThreadBuilder;
+pub use self::registry::{mark_blocked, mark_unblocked, Registry};
 pub use self::scope::{in_place_scope, scope, Scope};
 pub use self::scope::{in_place_scope_fifo, scope_fifo, ScopeFifo};
 pub use self::spawn::{spawn, spawn_fifo};
@@ -189,6 +187,9 @@ pub struct ThreadPoolBuilder<S = DefaultSpawn> {
     /// 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>>,
 
@@ -217,6 +218,9 @@ pub struct Configuration {
 /// may be invoked multiple times in parallel.
 type PanicHandler = dyn Fn(Box<dyn Any + Send>) + Send + Sync;
 
+/// The type for a closure that gets invoked when the Rayon thread pool deadlocks
+type DeadlockHandler = dyn 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.
@@ -237,6 +241,7 @@ impl Default for ThreadPoolBuilder {
             stack_size: None,
             start_handler: None,
             exit_handler: None,
+            deadlock_handler: None,
             spawn_handler: DefaultSpawn,
             breadth_first: false,
         }
@@ -454,6 +459,7 @@ impl<S> ThreadPoolBuilder<S> {
             stack_size: self.stack_size,
             start_handler: self.start_handler,
             exit_handler: self.exit_handler,
+            deadlock_handler: self.deadlock_handler,
             breadth_first: self.breadth_first,
         }
     }
@@ -612,6 +618,20 @@ impl<S> ThreadPoolBuilder<S> {
         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) -> Self
+    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()
@@ -778,6 +798,7 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
             ref get_thread_name,
             ref panic_handler,
             ref stack_size,
+            ref deadlock_handler,
             ref start_handler,
             ref exit_handler,
             spawn_handler: _,
@@ -794,6 +815,7 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
         }
         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);
 
@@ -802,6 +824,7 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
             .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("breadth_first", &breadth_first)
diff --git a/rayon-core/src/registry.rs b/rayon-core/src/registry.rs
@@ -6,8 +6,8 @@ use crate::sleep::Sleep;
 use crate::tlv::Tlv;
 use crate::unwind;
 use crate::{
-    ErrorKind, ExitHandler, PanicHandler, StartHandler, ThreadPoolBuildError, ThreadPoolBuilder,
-    Yield,
+    DeadlockHandler, ErrorKind, ExitHandler, PanicHandler, StartHandler, ThreadPoolBuildError,
+    ThreadPoolBuilder, Yield,
 };
 use crossbeam_deque::{Injector, Steal, Stealer, Worker};
 use std::cell::Cell;
@@ -130,13 +130,14 @@ where
     }
 }
 
-pub(super) struct Registry {
+pub struct Registry {
     logger: Logger,
     thread_infos: Vec<ThreadInfo>,
     sleep: Sleep,
     injected_jobs: Injector<JobRef>,
     broadcasts: Mutex<Vec<Worker<JobRef>>>,
     panic_handler: Option<Box<PanicHandler>>,
+    deadlock_handler: Option<Box<DeadlockHandler>>,
     start_handler: Option<Box<StartHandler>>,
     exit_handler: Option<Box<ExitHandler>>,
 
@@ -290,6 +291,7 @@ impl Registry {
             broadcasts: Mutex::new(broadcasts),
             terminate_count: AtomicUsize::new(1),
             panic_handler: builder.take_panic_handler(),
+            deadlock_handler: builder.take_deadlock_handler(),
             start_handler: builder.take_start_handler(),
             exit_handler: builder.take_exit_handler(),
         });
@@ -317,7 +319,7 @@ impl Registry {
         Ok(registry)
     }
 
-    pub(super) fn current() -> Arc<Registry> {
+    pub fn current() -> Arc<Registry> {
         unsafe {
             let worker_thread = WorkerThread::current();
             let registry = if worker_thread.is_null() {
@@ -624,6 +626,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(super) struct RegistryId {
     addr: usize,
@@ -823,9 +843,12 @@ impl WorkerThread {
                 self.execute(job);
                 idle_state = self.registry.sleep.start_looking(self.index, latch);
             } else {
-                self.registry
-                    .sleep
-                    .no_work_found(&mut idle_state, latch, || self.has_injected_job())
+                self.registry.sleep.no_work_found(
+                    &mut idle_state,
+                    latch,
+                    || self.has_injected_job(),
+                    &self.registry.deadlock_handler,
+                )
             }
         }
 
diff --git a/rayon-core/src/sleep/README.md b/rayon-core/src/sleep/README.md
@@ -216,4 +216,37 @@ Meanwhile, the sleepy thread does the following:
 Either PushFence or SleepFence must come first:
 
 * If PushFence comes first, then PushJob must be visible to ReadJob.
-* If SleepFence comes first, then IncSleepers is visible to ReadSleepers.
+* If SleepFence comes first, then IncSleepers is visible to ReadSleepers.
+
+# 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
@@ -4,6 +4,7 @@
 use crate::latch::CoreLatch;
 use crate::log::Event::*;
 use crate::log::Logger;
+use crate::DeadlockHandler;
 use crossbeam_utils::CachePadded;
 use std::sync::atomic::Ordering;
 use std::sync::{Condvar, Mutex};
@@ -14,6 +15,29 @@ mod counters;
 pub(crate) use self::counters::THREADS_MAX;
 use self::counters::{AtomicCounters, JobsEventCounter};
 
+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())();
+        }
+    }
+}
+
 /// The `Sleep` struct is embedded into each registry. It governs the waking and sleeping
 /// of workers. It has callbacks that are invoked periodically at significant events,
 /// such as when workers are looping and looking for work, when latches are set, or when
@@ -29,6 +53,8 @@ pub(super) struct Sleep {
     worker_sleep_states: Vec<CachePadded<WorkerSleepState>>,
 
     counters: AtomicCounters,
+
+    data: Mutex<SleepData>,
 }
 
 /// An instance of this struct is created when a thread becomes idle.
@@ -68,9 +94,38 @@ impl Sleep {
             logger,
             worker_sleep_states: (0..n_threads).map(|_| Default::default()).collect(),
             counters: AtomicCounters::new(),
+            data: Mutex::new(SleepData {
+                worker_count: n_threads,
+                active_threads: n_threads,
+                blocked_threads: 0,
+            }),
         }
     }
 
+    /// 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;
+    }
+
     #[inline]
     pub(super) fn start_looking(&self, worker_index: usize, latch: &CoreLatch) -> IdleState {
         self.logger.log(|| ThreadIdle {
@@ -106,6 +161,7 @@ impl Sleep {
         idle_state: &mut IdleState,
         latch: &CoreLatch,
         has_injected_jobs: impl FnOnce() -> bool,
+        deadlock_handler: &Option<Box<DeadlockHandler>>,
     ) {
         if idle_state.rounds < ROUNDS_UNTIL_SLEEPY {
             thread::yield_now();
@@ -119,7 +175,7 @@ impl Sleep {
             thread::yield_now();
         } else {
             debug_assert_eq!(idle_state.rounds, ROUNDS_UNTIL_SLEEPING);
-            self.sleep(idle_state, latch, has_injected_jobs);
+            self.sleep(idle_state, latch, has_injected_jobs, deadlock_handler);
         }
     }
 
@@ -142,6 +198,7 @@ impl Sleep {
         idle_state: &mut IdleState,
         latch: &CoreLatch,
         has_injected_jobs: impl FnOnce() -> bool,
+        deadlock_handler: &Option<Box<DeadlockHandler>>,
     ) {
         let worker_index = idle_state.worker_index;
 
@@ -215,6 +272,13 @@ impl Sleep {
             // the one that wakes us.)
             self.counters.sub_sleeping_thread();
         } else {
+            {
+                // Decrement the number of active threads and check for a deadlock
+                let mut data = self.data.lock().unwrap();
+                data.active_threads -= 1;
+                data.deadlock_check(deadlock_handler);
+            }
+
             // If we don't see an injected job (the normal case), then flag
             // ourselves as asleep and wait till we are notified.
             //
@@ -372,6 +436,9 @@ impl Sleep {
             // do.
             self.counters.sub_sleeping_thread();
 
+            // Increment the number of active threads
+            self.data.lock().unwrap().active_threads += 1;
+
             self.logger.log(|| ThreadNotify { worker: index });
 
             true
