std: Reduce TLS access

2 files changed, 18 insertions, 15 deletions

diff --git a/src/libstd/rt/sched.rs b/src/libstd/rt/sched.rs
index 5ec2df32c48..a8f3d01351b 100644
--- a/src/libstd/rt/sched.rs
+++ b/src/libstd/rt/sched.rs
@@ -563,11 +563,10 @@ impl Scheduler {
         // run the cleanup job, as expected by the previously called
         // swap_contexts function.
         unsafe {
-            let sched = Local::unsafe_borrow::<Scheduler>();
-            (*sched).run_cleanup_job();
+            let task = Local::unsafe_borrow::<Task>();
+            (*task).sched.get_mut_ref().run_cleanup_job();
 
             // Must happen after running the cleanup job (of course).
-            let task = Local::unsafe_borrow::<Task>();
             (*task).death.check_killed((*task).unwinder.unwinding);
         }
     }
diff --git a/src/libstd/unstable/sync.rs b/src/libstd/unstable/sync.rs
index d7f9988edef..6fa0e0eb8c1 100644
--- a/src/libstd/unstable/sync.rs
+++ b/src/libstd/unstable/sync.rs
@@ -281,20 +281,24 @@ impl<T> Drop for UnsafeAtomicRcBox<T>{
  */
 // FIXME(#8140) should not be pub
 pub unsafe fn atomically<U>(f: &fn() -> U) -> U {
-    use rt::task::Task;
+    use rt::task::{Task, GreenTask, SchedTask};
     use rt::local::Local;
-    use rt::in_green_task_context;
-
-    if in_green_task_context() {
-        let t = Local::unsafe_borrow::<Task>();
-        do (|| {
-            (*t).death.inhibit_deschedule();
-            f()
-        }).finally {
-            (*t).death.allow_deschedule();
+
+    match Local::try_unsafe_borrow::<Task>() {
+        Some(t) => {
+            match (*t).task_type {
+                GreenTask(_) => {
+                    do (|| {
+                        (*t).death.inhibit_deschedule();
+                        f()
+                    }).finally {
+                        (*t).death.allow_deschedule();
+                    }
+                }
+                SchedTask => f()
+            }
         }
-    } else {
-        f()
+        None => f()
     }
 }