core::rt: A little bit of cleanup to thread-local scheduler

Author: brson

src/libcore/rt/sched/local.rs

@@ -39,18 +39,11 @@ pub fn take() -> ~Scheduler {
     }
 }
 
-/// Give the Scheduler to thread-local storage for the duration of the block
-pub fn install(sched: ~Scheduler, f: &fn()) -> ~Scheduler {
-    put(sched);
-    f();
-    return take();
-}
-
 /// Borrow a mutable reference to the thread-local Scheduler
 /// # Safety Note
 /// Because this leaves the Scheduler in thread-local storage it is possible
 /// For the Scheduler pointer to be aliased
-pub fn borrow(f: &fn(&mut Scheduler)) {
+pub unsafe fn borrow(f: &fn(&mut Scheduler)) {
     unsafe {
         let key = tls_key();
         let mut void_sched: *mut c_void = tls::get(key);
@@ -96,11 +89,13 @@ fn thread_local_scheduler_two_instances() {
 }
 
 #[test]
-fn install_borrow_smoke_test() {
+fn borrow_smoke_test() {
     let scheduler = ~UvEventLoop::new_scheduler();
-    let _scheduler = do install(scheduler) {
+    put(scheduler);
+    unsafe {
         do borrow |_sched| {
         }
-    };
+    }
+    let _scheduler = take();
 }
 

src/libcore/rt/sched/mod.rs

@@ -93,7 +93,9 @@ pub impl Scheduler {
         assert!(!self.in_task_context());
 
         // Give ownership of the scheduler (self) to the thread
-        do self.install |scheduler| {
+        local::put(self);
+
+        do Scheduler::local |scheduler| {
             fn run_scheduler_once() {
                 do Scheduler::local |scheduler| {
                     if scheduler.resume_task_from_queue() {
@@ -106,16 +108,12 @@ pub impl Scheduler {
             scheduler.event_loop.callback(run_scheduler_once);
             scheduler.event_loop.run();
         }
-    }
 
-    fn install(~self, f: &fn(&mut Scheduler)) -> ~Scheduler {
-        do local::install(self) {
-            local::borrow(f)
-        }
+        return local::take();
     }
 
     fn local(f: &fn(&mut Scheduler)) {
-        local::borrow(f)
+        unsafe { local::borrow(f) }
     }
 
     // * Scheduler-context operations
@@ -206,7 +204,7 @@ pub impl Scheduler {
         }
 
         // We could be executing in a different thread now
-        do local::borrow |sched| {
+        do Scheduler::local |sched| {
             sched.run_cleanup_job();
         }
     }
@@ -230,7 +228,7 @@ pub impl Scheduler {
         }
 
         // We could be executing in a different thread now
-        do local::borrow |sched| {
+        do Scheduler::local |sched| {
             sched.run_cleanup_job();
         }
     }