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::(); - (*sched).run_cleanup_job(); + let task = Local::unsafe_borrow::(); + (*task).sched.get_mut_ref().run_cleanup_job(); // Must happen after running the cleanup job (of course). - let task = Local::unsafe_borrow::(); (*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 Drop for UnsafeAtomicRcBox{ */ // FIXME(#8140) should not be pub pub unsafe fn atomically(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::(); - do (|| { - (*t).death.inhibit_deschedule(); - f() - }).finally { - (*t).death.allow_deschedule(); + match Local::try_unsafe_borrow::() { + Some(t) => { + match (*t).task_type { + GreenTask(_) => { + do (|| { + (*t).death.inhibit_deschedule(); + f() + }).finally { + (*t).death.allow_deschedule(); + } + } + SchedTask => f() + } } - } else { - f() + None => f() } }