Add thread-local values which are preserved with jobs

Zoxc · cuviper · commit c7ba2822b7ee · 2019-09-09T16:52:42.000-07:00
diff --git a/rayon-core/src/job.rs b/rayon-core/src/job.rs
@@ -3,6 +3,7 @@ use latch::Latch;
 use std::any::Any;
 use std::cell::UnsafeCell;
 use std::mem;
+use tlv;
 use unwind;
 
 pub(super) enum JobResult<T> {
@@ -73,6 +74,7 @@ where
     pub(super) latch: L,
     func: UnsafeCell<Option<F>>,
     result: UnsafeCell<JobResult<R>>,
+    tlv: usize,
 }
 
 impl<L, F, R> StackJob<L, F, R>
@@ -81,11 +83,12 @@ where
     F: FnOnce(bool) -> R + Send,
     R: Send,
 {
-    pub(super) fn new(func: F, latch: L) -> StackJob<L, F, R> {
+    pub(super) fn new(tlv: usize, func: F, latch: L) -> StackJob<L, F, R> {
         StackJob {
             latch,
             func: UnsafeCell::new(Some(func)),
             result: UnsafeCell::new(JobResult::None),
+            tlv,
         }
     }
 
@@ -114,6 +117,7 @@ where
         }
 
         let this = &*this;
+        tlv::set(this.tlv);
         let abort = unwind::AbortIfPanic;
         let func = (*this.func.get()).take().unwrap();
         (*this.result.get()) = match unwind::halt_unwinding(call(func)) {
@@ -136,15 +140,17 @@ where
     BODY: FnOnce() + Send,
 {
     job: UnsafeCell<Option<BODY>>,
+    tlv: usize,
 }
 
 impl<BODY> HeapJob<BODY>
 where
     BODY: FnOnce() + Send,
 {
-    pub(super) fn new(func: BODY) -> Self {
+    pub(super) fn new(tlv: usize, func: BODY) -> Self {
         HeapJob {
             job: UnsafeCell::new(Some(func)),
+            tlv,
         }
     }
 
@@ -163,6 +169,7 @@ where
 {
     unsafe fn execute(this: *const Self) {
         let this: Box<Self> = mem::transmute(this);
+        tlv::set(this.tlv);
         let job = (*this.job.get()).take().unwrap();
         job();
     }
diff --git a/rayon-core/src/join/mod.rs b/rayon-core/src/join/mod.rs
@@ -3,6 +3,7 @@ use latch::{LatchProbe, SpinLatch};
 use log::Event::*;
 use registry::{self, WorkerThread};
 use std::any::Any;
+use tlv;
 use unwind;
 
 use FnContext;
@@ -135,18 +136,19 @@ where
             worker: worker_thread.index()
         });
 
+        let tlv = tlv::get();
         // 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(call_b(oper_b), SpinLatch::new());
+        let job_b = StackJob::new(tlv, call_b(oper_b), SpinLatch::new());
         let job_b_ref = job_b.as_job_ref();
         worker_thread.push(job_b_ref);
 
         // Execute task a; hopefully b gets stolen in the meantime.
         let status_a = unwind::halt_unwinding(call_a(oper_a, injected));
         let result_a = match status_a {
             Ok(v) => v,
-            Err(err) => join_recover_from_panic(worker_thread, &job_b.latch, err),
+            Err(err) => join_recover_from_panic(worker_thread, &job_b.latch, err, tlv),
         };
 
         // Now that task A has finished, try to pop job B from the
@@ -163,7 +165,11 @@ where
                     log!(PoppedRhs {
                         worker: worker_thread.index()
                     });
+                    // Restore the TLV since we might have run some jobs overwriting it when waiting for job b.
+                    tlv::set(tlv);
+
                     let result_b = job_b.run_inline(injected);
+
                     return (result_a, result_b);
                 } else {
                     log!(PoppedJob {
@@ -183,6 +189,9 @@ where
             }
         }
 
+        // Restore the TLV since we might have run some jobs overwriting it when waiting for job b.
+        tlv::set(tlv);
+
         (result_a, job_b.into_result())
     })
 }
@@ -195,7 +204,12 @@ unsafe fn join_recover_from_panic(
     worker_thread: &WorkerThread,
     job_b_latch: &SpinLatch,
     err: Box<dyn Any + Send>,
+    tlv: usize,
 ) -> ! {
     worker_thread.wait_until(job_b_latch);
+
+    // Restore the TLV since we might have run some jobs overwriting it when waiting for job b.
+    tlv::set(tlv);
+
     unwind::resume_unwinding(err)
 }
diff --git a/rayon-core/src/lib.rs b/rayon-core/src/lib.rs
@@ -64,6 +64,8 @@ mod util;
 mod compile_fail;
 mod test;
 
+pub mod tlv;
+
 #[cfg(rayon_unstable)]
 pub mod internal;
 pub use join::{join, join_context};
diff --git a/rayon-core/src/registry.rs b/rayon-core/src/registry.rs
@@ -495,6 +495,7 @@ 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(
+                0,
                 |injected| {
                     let worker_thread = WorkerThread::current();
                     assert!(injected && !worker_thread.is_null());
@@ -519,6 +520,7 @@ impl Registry {
         debug_assert!(current_thread.registry().id() != self.id());
         let latch = TickleLatch::new(SpinLatch::new(), &current_thread.registry().sleep);
         let job = StackJob::new(
+            0,
             |injected| {
                 let worker_thread = WorkerThread::current();
                 assert!(injected && !worker_thread.is_null());
diff --git a/rayon-core/src/scope/mod.rs b/rayon-core/src/scope/mod.rs
@@ -15,6 +15,7 @@ use std::mem;
 use std::ptr;
 use std::sync::atomic::{AtomicPtr, Ordering};
 use std::sync::Arc;
+use tlv;
 use unwind;
 
 mod internal;
@@ -60,6 +61,9 @@ struct ScopeBase<'scope> {
     /// `Sync`, but it's still safe to let the `Scope` implement `Sync` because
     /// the closures are only *moved* across threads to be executed.
     marker: PhantomData<Box<dyn FnOnce(&Scope<'scope>) + Send + Sync + 'scope>>,
+
+    /// The TLV at the scope's creation. Used to set the TLV for spawned jobs.
+    tlv: usize,
 }
 
 /// Create a "fork-join" scope `s` and invokes the closure with a
@@ -452,7 +456,7 @@ impl<'scope> Scope<'scope> {
     {
         self.base.increment();
         unsafe {
-            let job_ref = Box::new(HeapJob::new(move || {
+            let job_ref = Box::new(HeapJob::new(self.base.tlv, move || {
                 self.base.execute_job(move || body(self))
             }))
             .as_job_ref();
@@ -493,7 +497,7 @@ impl<'scope> ScopeFifo<'scope> {
     {
         self.base.increment();
         unsafe {
-            let job_ref = Box::new(HeapJob::new(move || {
+            let job_ref = Box::new(HeapJob::new(self.base.tlv, move || {
                 self.base.execute_job(move || body(self))
             }))
             .as_job_ref();
@@ -520,6 +524,7 @@ impl<'scope> ScopeBase<'scope> {
             panic: AtomicPtr::new(ptr::null_mut()),
             job_completed_latch: CountLatch::new(),
             marker: PhantomData,
+            tlv: tlv::get(),
         }
     }
 
@@ -537,6 +542,8 @@ impl<'scope> ScopeBase<'scope> {
     {
         let result = self.execute_job_closure(func);
         self.steal_till_jobs_complete(owner_thread);
+        // Restore the TLV if we ran some jobs while waiting
+        tlv::set(self.tlv);
         result.unwrap() // only None if `op` panicked, and that would have been propagated
     }
 
@@ -613,6 +620,8 @@ impl<'scope> ScopeBase<'scope> {
             log!(ScopeCompletePanicked {
                 owner_thread: owner_thread.index()
             });
+            // Restore the TLV if we ran some jobs while waiting
+            tlv::set(self.tlv);
             let value: Box<Box<dyn Any + Send + 'static>> = mem::transmute(panic);
             unwind::resume_unwinding(*value);
         } else {
diff --git a/rayon-core/src/spawn/mod.rs b/rayon-core/src/spawn/mod.rs
@@ -92,7 +92,7 @@ where
     // executed. This ref is decremented at the (*) below.
     registry.increment_terminate_count();
 
-    Box::new(HeapJob::new({
+    Box::new(HeapJob::new(0, {
         let registry = registry.clone();
         move || {
             match unwind::halt_unwinding(func) {
diff --git a/rayon-core/src/tlv.rs b/rayon-core/src/tlv.rs
@@ -0,0 +1,30 @@
+//! Allows access to the Rayon's thread local value
+//! which is preserved when moving jobs across threads
+
+use std::cell::Cell;
+
+thread_local!(pub(crate) static TLV: Cell<usize> = Cell::new(0));
+
+/// Sets the current thread-local value to `value` inside the closure.
+/// The old value is restored when the closure ends
+pub fn with<F: FnOnce() -> R, R>(value: usize, f: F) -> R {
+    struct Reset(usize);
+    impl Drop for Reset {
+        fn drop(&mut self) {
+            TLV.with(|tlv| tlv.set(self.0));
+        }
+    }
+    let _reset = Reset(get());
+    TLV.with(|tlv| tlv.set(value));
+    f()
+}
+
+/// Sets the current thread-local value
+pub fn set(value: usize) {
+    TLV.with(|tlv| tlv.set(value));
+}
+
+/// Returns the current thread-local value
+pub fn get() -> usize {
+    TLV.with(|tlv| tlv.get())
+}

Original file line number	Diff line number	Diff line change
`@@ -3,6 +3,7 @@ use latch::Latch;`
`3`	`3`	`use std::any::Any;`
`4`	`4`	`use std::cell::UnsafeCell;`
`5`	`5`	`use std::mem;`
	`6`	`+use tlv;`
`6`	`7`	`use unwind;`
`7`	`8`
`8`	`9`	`pub(super) enum JobResult<T> {`
`@@ -73,6 +74,7 @@ where`
`73`	`74`	`pub(super) latch: L,`
`74`	`75`	`func: UnsafeCell<Option<F>>,`
`75`	`76`	`result: UnsafeCell<JobResult<R>>,`
	`77`	`+ tlv: usize,`
`76`	`78`	`}`
`77`	`79`
`78`	`80`	`impl<L, F, R> StackJob<L, F, R>`
`@@ -81,11 +83,12 @@ where`
`81`	`83`	`F: FnOnce(bool) -> R + Send,`
`82`	`84`	`R: Send,`
`83`	`85`	`{`
`84`		`- pub(super) fn new(func: F, latch: L) -> StackJob<L, F, R> {`
	`86`	`+ pub(super) fn new(tlv: usize, func: F, latch: L) -> StackJob<L, F, R> {`
`85`	`87`	`StackJob {`
`86`	`88`	`latch,`
`87`	`89`	`func: UnsafeCell::new(Some(func)),`
`88`	`90`	`result: UnsafeCell::new(JobResult::None),`
	`91`	`+ tlv,`
`89`	`92`	`}`
`90`	`93`	`}`
`91`	`94`
`@@ -114,6 +117,7 @@ where`
`114`	`117`	`}`
`115`	`118`
`116`	`119`	`let this = &*this;`
	`120`	`+ tlv::set(this.tlv);`
`117`	`121`	`let abort = unwind::AbortIfPanic;`
`118`	`122`	`let func = (*this.func.get()).take().unwrap();`
`119`	`123`	`(*this.result.get()) = match unwind::halt_unwinding(call(func)) {`
`@@ -136,15 +140,17 @@ where`
`136`	`140`	`BODY: FnOnce() + Send,`
`137`	`141`	`{`
`138`	`142`	`job: UnsafeCell<Option<BODY>>,`
	`143`	`+ tlv: usize,`
`139`	`144`	`}`
`140`	`145`
`141`	`146`	`impl<BODY> HeapJob<BODY>`
`142`	`147`	`where`
`143`	`148`	`BODY: FnOnce() + Send,`
`144`	`149`	`{`
`145`		`- pub(super) fn new(func: BODY) -> Self {`
	`150`	`+ pub(super) fn new(tlv: usize, func: BODY) -> Self {`
`146`	`151`	`HeapJob {`
`147`	`152`	`job: UnsafeCell::new(Some(func)),`
	`153`	`+ tlv,`
`148`	`154`	`}`
`149`	`155`	`}`
`150`	`156`
`@@ -163,6 +169,7 @@ where`
`163`	`169`	`{`
`164`	`170`	`unsafe fn execute(this: *const Self) {`
`165`	`171`	`let this: Box<Self> = mem::transmute(this);`
	`172`	`+ tlv::set(this.tlv);`
`166`	`173`	`let job = (*this.job.get()).take().unwrap();`
`167`	`174`	`job();`
`168`	`175`	`}`