Merge #668

668: Reusing LockLatch to reduce allocation  r=cuviper a=ralfbiedert

This addresses the first half of #666.

I kept the `LockLatch` and added a `reset()` method because that had overall fewer lines of code. If you want I can still change it to a separate `ReusableLockLatch` type.


Co-authored-by: Ralf Biedert <[email protected]>
Co-authored-by: Josh Stone <[email protected]>

.gitignore

@@ -3,3 +3,4 @@ target
 *~
 TAGS
 *.bk
+.idea

rayon-core/src/join/mod.rs

@@ -120,13 +120,12 @@ where
             worker: worker_thread.index()
         });
 
+        let latch = SpinLatch::new();
+
         // Create virtual wrapper for task b; this all has to be
         // done here so that the stack frame can keep it all live
         // long enough.
-        let job_b = StackJob::new(
-            |migrated| oper_b(FnContext::new(migrated)),
-            SpinLatch::new(),
-        );
+        let job_b = StackJob::new(|migrated| oper_b(FnContext::new(migrated)), latch);
         let job_b_ref = job_b.as_job_ref();
         worker_thread.push(job_b_ref);
 

rayon-core/src/latch.rs

@@ -86,6 +86,15 @@ impl LockLatch {
         }
     }
 
+    /// Block until latch is set, then resets this lock latch so it can be reused again.
+    pub(super) fn wait_and_reset(&self) {
+        let mut guard = self.m.lock().unwrap();
+        while !*guard {
+            guard = self.v.wait(guard).unwrap();
+        }
+        *guard = false;
+    }
+
     /// Block until latch is set.
     pub(super) fn wait(&self) {
         let mut guard = self.m.lock().unwrap();
@@ -186,3 +195,21 @@ impl<'a, L: Latch> Latch for TickleLatch<'a, L> {
         self.sleep.tickle(usize::MAX);
     }
 }
+
+impl<'a, L> LatchProbe for &'a L
+where
+    L: LatchProbe,
+{
+    fn probe(&self) -> bool {
+        L::probe(self)
+    }
+}
+
+impl<'a, L> Latch for &'a L
+where
+    L: Latch,
+{
+    fn set(&self) {
+        L::set(self);
+    }
+}

rayon-core/src/registry.rs

@@ -486,19 +486,23 @@ impl Registry {
         OP: FnOnce(&WorkerThread, bool) -> R + Send,
         R: Send,
     {
-        // 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(),
-        );
-        self.inject(&[job.as_job_ref()]);
-        job.latch.wait();
-        job.into_result()
+        thread_local!(static LOCK_LATCH: LockLatch = LockLatch::new());
+
+        LOCK_LATCH.with(|l| {
+            // 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)
+                },
+                l,
+            );
+            self.inject(&[job.as_job_ref()]);
+            job.latch.wait_and_reset(); // Make sure we can use the same latch again next time.
+            job.into_result()
+        })
     }
 
     #[cold]