Make StackJob take a &mut latch.

Author: ralfbiedert

rayon-core/src/job.rs

@@ -64,24 +64,24 @@ impl JobRef {
 /// executes it need not free any heap data, the cleanup occurs when
 /// the stack frame is later popped.  The function parameter indicates
 /// `true` if the job was stolen -- executed on a different thread.
-pub(super) struct StackJob<L, F, R>
+pub(super) struct StackJob<'a, L, F, R>
 where
     L: Latch + Sync,
     F: FnOnce(bool) -> R + Send,
     R: Send,
 {
-    pub(super) latch: L,
+    pub(super) latch: &'a mut L,
     func: UnsafeCell<Option<F>>,
     result: UnsafeCell<JobResult<R>>,
 }
 
-impl<L, F, R> StackJob<L, F, R>
+impl<'a, L, F, R> StackJob<'a, L, F, R>
 where
     L: Latch + Sync,
     F: FnOnce(bool) -> R + Send,
     R: Send,
 {
-    pub(super) fn new(func: F, latch: L) -> StackJob<L, F, R> {
+    pub(super) fn new(func: F, latch: &'a mut L) -> StackJob<'a, L, F, R> {
         StackJob {
             latch,
             func: UnsafeCell::new(Some(func)),
@@ -102,7 +102,7 @@ where
     }
 }
 
-impl<L, F, R> Job for StackJob<L, F, R>
+impl<'a, L, F, R> Job for StackJob<'a, L, F, R>
 where
     L: Latch + Sync,
     F: FnOnce(bool) -> R + Send,

rayon-core/src/join/mod.rs

@@ -120,13 +120,12 @@ where
             worker: worker_thread.index()
         });
 
+        let mut 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)), &mut latch);
         let job_b_ref = job_b.as_job_ref();
         worker_thread.push(job_b_ref);
 
@@ -165,7 +164,7 @@ where
                 log!(LostJob {
                     worker: worker_thread.index()
                 });
-                worker_thread.wait_until(&job_b.latch);
+                worker_thread.wait_until(job_b.latch);
                 debug_assert!(job_b.latch.probe());
                 break;
             }

rayon-core/src/registry.rs

@@ -9,7 +9,7 @@ use latch::{CountLatch, Latch, LatchProbe, LockLatch, SpinLatch, TickleLatch};
 use log::Event::*;
 use sleep::Sleep;
 use std::any::Any;
-use std::cell::Cell;
+use std::cell::{Cell, RefCell};
 use std::collections::hash_map::DefaultHasher;
 use std::fmt;
 use std::hash::Hasher;
@@ -486,6 +486,8 @@ impl Registry {
         OP: FnOnce(&WorkerThread, bool) -> R + Send,
         R: Send,
     {
+        let mut latch = LockLatch::new();
+
         // This thread isn't a member of *any* thread pool, so just block.
         debug_assert!(WorkerThread::current().is_null());
         let job = StackJob::new(
@@ -494,7 +496,7 @@ impl Registry {
                 assert!(injected && !worker_thread.is_null());
                 op(&*worker_thread, true)
             },
-            LockLatch::new(),
+            &mut latch,
         );
         self.inject(&[job.as_job_ref()]);
         job.latch.wait();
@@ -510,17 +512,17 @@ impl Registry {
         // This thread is a member of a different pool, so let it process
         // other work while waiting for this `op` to complete.
         debug_assert!(current_thread.registry().id() != self.id());
-        let latch = TickleLatch::new(SpinLatch::new(), &current_thread.registry().sleep);
+        let mut latch = TickleLatch::new(SpinLatch::new(), &current_thread.registry().sleep);
         let job = StackJob::new(
             |injected| {
                 let worker_thread = WorkerThread::current();
                 assert!(injected && !worker_thread.is_null());
                 op(&*worker_thread, true)
             },
-            latch,
+            &mut latch,
         );
         self.inject(&[job.as_job_ref()]);
-        current_thread.wait_until(&job.latch);
+        current_thread.wait_until(job.latch);
         job.into_result()
     }