Make the ParallelGuard depth thread-local

ext/src/utils.rs

@@ -595,22 +595,34 @@ pub use logging::{LogWriter, ext_tracing_subscriber, init_logging};
 
 pub(crate) mod parallel {
 
-    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::cell::Cell;
 
-    static PARALLEL_DEPTH: AtomicUsize = AtomicUsize::new(0);
+    thread_local! {
+        static PARALLEL_DEPTH: Cell<usize> = const { Cell::new(0) };
+    }
 
     /// RAII guard that increments [`PARALLEL_DEPTH`] on creation and decrements on drop. Used to mark
     /// regions where `par_iter_mut` work is active, so that `step_resolution` jobs can detect priority
     /// inversion and retry.
+    ///
+    /// The depth is thread-local because the priority inversion we are guarding against only happens
+    /// when the guarded thread itself steals a long job: while it is blocked inside a `par_iter_mut`,
+    /// rayon may hand it another job from the pool, and if that job is a `step_resolution` the parent
+    /// cannot resume until the long job finishes. A stolen job always runs on the OS thread that
+    /// stole it, so checking this thread's depth is exactly the condition we need. Other threads are
+    /// free to pick up long jobs in the meantime.
     pub(crate) struct ParallelGuard {
         #[allow(dead_code)]
         span: tracing::span::EnteredSpan,
     }
 
     impl ParallelGuard {
         pub(crate) fn new() -> Self {
-            // We use Release to synchronize with `is_in_parallel`
-            let counter_start = PARALLEL_DEPTH.fetch_add(1, Ordering::Release);
+            let counter_start = PARALLEL_DEPTH.with(|depth| {
+                let old = depth.get();
+                depth.set(old + 1);
+                old
+            });
             Self {
                 span: tracing::info_span!(
                     "parallel_guard",
@@ -624,14 +636,17 @@ pub(crate) mod parallel {
 
     impl Drop for ParallelGuard {
         fn drop(&mut self) {
-            // We use Release to synchronize with `is_in_parallel`
-            let counter_end = PARALLEL_DEPTH.fetch_sub(1, Ordering::Release);
-            self.span.record("counter_end", counter_end - 1);
+            let counter_end = PARALLEL_DEPTH.with(|depth| {
+                let new = depth.get() - 1;
+                depth.set(new);
+                new
+            });
+            self.span.record("counter_end", counter_end);
         }
     }
 
     pub(crate) fn is_in_parallel() -> bool {
-        PARALLEL_DEPTH.load(Ordering::Acquire) > 0
+        PARALLEL_DEPTH.with(|depth| depth.get()) > 0
     }
 }