Restore TLV when needed

Author: Zoxc

rayon-core/src/job.rs

@@ -85,12 +85,12 @@ where
     F: FnOnce(bool) -> R + Send,
     R: Send,
 {
-    pub fn new(func: F, latch: L) -> StackJob<L, F, R> {
+    pub fn new(tlv: usize, func: F, latch: L) -> StackJob<L, F, R> {
         StackJob {
             latch: latch,
             func: UnsafeCell::new(Some(func)),
             result: UnsafeCell::new(JobResult::None),
-            tlv: tlv::get(),
+            tlv,
         }
     }
 

rayon-core/src/join/mod.rs

@@ -4,6 +4,7 @@ use log::Event::*;
 use registry::{self, WorkerThread};
 use std::any::Any;
 use unwind;
+use tlv;
 
 use FnContext;
 
@@ -120,10 +121,12 @@ 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(
+            tlv,
             |migrated| oper_b(FnContext::new(migrated)),
             SpinLatch::new(),
         );
@@ -134,7 +137,7 @@ where
         let status_a = unwind::halt_unwinding(move || oper_a(FnContext::new(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
@@ -151,7 +154,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 {
@@ -171,6 +178,9 @@ where
             }
         }
 
+        // Restore the TLV since we might have run some jobs overwriting it when waiting for job b.
+        tlv::set(tlv);
+
         return (result_a, job_b.into_result());
     })
 }
@@ -183,7 +193,12 @@ unsafe fn join_recover_from_panic(
     worker_thread: &WorkerThread,
     job_b_latch: &SpinLatch,
     err: Box<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)
 }

rayon-core/src/registry.rs

@@ -387,7 +387,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(|injected| {
+        let job = StackJob::new(0, |injected| {
             let worker_thread = WorkerThread::current();
             assert!(injected && !worker_thread.is_null());
             op(&*worker_thread, true)
@@ -410,6 +410,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());

rayon-core/src/scope/mod.rs

@@ -51,6 +51,8 @@ impl<'scope> ScopeBuilder<'scope> {
             in_worker(move |owner_thread, _| {
                 scope.steal_till_jobs_complete(owner_thread);
             });
+            // Restore the TLV if we ran some jobs while waiting
+            tlv::set(scope.tlv);
             result.unwrap() // only None if `op` panicked, and that would have been propagated
         }
     }
@@ -311,6 +313,8 @@ where
             };
             let result = scope.execute_job_closure(op);
             scope.steal_till_jobs_complete(owner_thread);
+            // Restore the TLV if we ran some jobs while waiting
+            tlv::set(scope.tlv);
             result.unwrap() // only None if `op` panicked, and that would have been propagated
         }
     })
@@ -460,6 +464,8 @@ impl<'scope> Scope<'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<Any + Send + 'static>> = mem::transmute(panic);
             unwind::resume_unwinding(*value);
         } else {