std::rt: Optimize TLS use in change_task_context

3 files changed, 25 insertions, 1 deletions

diff --git a/src/libstd/rt/local.rs b/src/libstd/rt/local.rs
index 80beb5a2835..18b7394700f 100644
--- a/src/libstd/rt/local.rs
+++ b/src/libstd/rt/local.rs
@@ -21,6 +21,7 @@ pub trait Local {
     fn take() -> ~Self;
     fn exists() -> bool;
     fn borrow<T>(f: &fn(&mut Self) -> T) -> T;
+    unsafe fn unsafe_take() -> ~Self;
     unsafe fn unsafe_borrow() -> *mut Self;
     unsafe fn try_unsafe_borrow() -> Option<*mut Self>;
 }
@@ -46,6 +47,8 @@ impl Local for Task {
         }
     }
     #[inline]
+    unsafe fn unsafe_take() -> ~Task { local_ptr::unsafe_take() }
+    #[inline]
     unsafe fn unsafe_borrow() -> *mut Task { local_ptr::unsafe_borrow() }
     #[inline]
     unsafe fn try_unsafe_borrow() -> Option<*mut Task> {
@@ -89,6 +92,7 @@ impl Local for Scheduler {
             }
         }
     }
+    unsafe fn unsafe_take() -> ~Scheduler { rtabort!("unimpl") }
     unsafe fn unsafe_borrow() -> *mut Scheduler {
         match (*Local::unsafe_borrow::<Task>()).sched {
             Some(~ref mut sched) => {
@@ -122,6 +126,7 @@ impl Local for IoFactoryObject {
     fn take() -> ~IoFactoryObject { rtabort!("unimpl") }
     fn exists() -> bool { rtabort!("unimpl") }
     fn borrow<T>(_f: &fn(&mut IoFactoryObject) -> T) -> T { rtabort!("unimpl") }
+    unsafe fn unsafe_take() -> ~IoFactoryObject { rtabort!("unimpl") }
     unsafe fn unsafe_borrow() -> *mut IoFactoryObject {
         let sched = Local::unsafe_borrow::<Scheduler>();
         let io: *mut IoFactoryObject = (*sched).event_loop.io().unwrap();
diff --git a/src/libstd/rt/local_ptr.rs b/src/libstd/rt/local_ptr.rs
index 491a864ebfe..3125a1da937 100644
--- a/src/libstd/rt/local_ptr.rs
+++ b/src/libstd/rt/local_ptr.rs
@@ -64,6 +64,23 @@ pub unsafe fn take<T>() -> ~T {
     return ptr;
 }
 
+/// Take ownership of a pointer from thread-local storage.
+///
+/// # Safety note
+///
+/// Does not validate the pointer type.
+/// Leaves the old pointer in TLS for speed.
+#[inline]
+pub unsafe fn unsafe_take<T>() -> ~T {
+    let key = tls_key();
+    let void_ptr: *mut c_void = tls::get(key);
+    if void_ptr.is_null() {
+        rtabort!("thread-local pointer is null. bogus!");
+    }
+    let ptr: ~T = cast::transmute(void_ptr);
+    return ptr;
+}
+
 /// Check whether there is a thread-local pointer installed.
 pub fn exists() -> bool {
     unsafe {
diff --git a/src/libstd/rt/sched.rs b/src/libstd/rt/sched.rs
index 111e60ccb2f..b161864a74f 100644
--- a/src/libstd/rt/sched.rs
+++ b/src/libstd/rt/sched.rs
@@ -505,7 +505,9 @@ impl Scheduler {
         let mut this = self;
 
         // The current task is grabbed from TLS, not taken as an input.
-        let current_task: ~Task = Local::take::<Task>();
+        // Doing an unsafe_take to avoid writing back a null pointer -
+        // We're going to call `put` later to do that.
+        let current_task: ~Task = unsafe { Local::unsafe_take::<Task>() };
 
         // Check that the task is not in an atomically() section (e.g.,
         // holding a pthread mutex, which could deadlock the scheduler).