core::rt Wire up task-local storage to the new scheduler

src/libcore/rt/local_services.rs

@@ -19,6 +19,7 @@
 //! (freestanding rust with local services?).
 
 use prelude::*;
+use libc::c_void;
 use super::sched::{Task, local_sched};
 use super::local_heap::LocalHeap;
 
@@ -27,11 +28,12 @@ pub struct LocalServices {
     gc: GarbageCollector,
     storage: LocalStorage,
     logger: Logger,
-    unwinder: Unwinder
+    unwinder: Unwinder,
+    destroyed: bool
 }
 
 pub struct GarbageCollector;
-pub struct LocalStorage;
+pub struct LocalStorage(*c_void, Option<~fn(*c_void)>);
 pub struct Logger;
 pub struct Unwinder;
 
@@ -40,11 +42,34 @@ impl LocalServices {
         LocalServices {
             heap: LocalHeap::new(),
             gc: GarbageCollector,
-            storage: LocalStorage,
+            storage: LocalStorage(ptr::null(), None),
             logger: Logger,
-            unwinder: Unwinder
+            unwinder: Unwinder,
+            destroyed: false
         }
     }
+
+    /// Must be called manually before finalization to clean up
+    /// thread-local resources. Some of the routines here expect
+    /// LocalServices to be available recursively so this must be
+    /// called unsafely, without removing LocalServices from
+    /// thread-local-storage.
+    pub fn destroy(&mut self) {
+        // This is just an assertion that `destroy` was called unsafely
+        // and this instance of LocalServices is still accessible.
+        do borrow_local_services |sched| {
+            assert!(ptr::ref_eq(sched, self));
+        }
+        match self.storage {
+            LocalStorage(ptr, Some(ref dtor)) => (*dtor)(ptr),
+            _ => ()
+        }
+        self.destroyed = true;
+    }
+}
+
+impl Drop for LocalServices {
+    fn finalize(&self) { assert!(self.destroyed) }
 }
 
 /// Borrow a pointer to the installed local services.
@@ -62,6 +87,19 @@ pub fn borrow_local_services(f: &fn(&mut LocalServices)) {
     }
 }
 
+pub unsafe fn unsafe_borrow_local_services() -> &mut LocalServices {
+    use cast::transmute_mut_region;
+
+    match local_sched::unsafe_borrow().current_task {
+        Some(~ref mut task) => {
+            transmute_mut_region(&mut task.local_services)
+        }
+        None => {
+            fail!(~"no local services for schedulers yet")
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use rt::test::*;
@@ -75,4 +113,16 @@ mod test {
             assert!(*b == 5);
         }
     }
+
+    #[test]
+    fn tls() {
+        use task::local_data::*;
+        do run_in_newsched_task() {
+            unsafe {
+                fn key(_x: @~str) { }
+                local_data_set(key, @~"data");
+                assert!(*local_data_get(key).get() == ~"data");
+            }
+        }
+    }
 }

src/libcore/rt/sched/mod.rs

@@ -357,6 +357,13 @@ pub impl Task {
 
             start();
 
+            unsafe {
+                // Destroy the local heap, TLS, etc.
+                let sched = local_sched::unsafe_borrow();
+                let task = sched.current_task.get_mut_ref();
+                task.local_services.destroy();
+            }
+
             let sched = local_sched::take();
             sched.terminate_current_task();
         };

src/libcore/task/local_data.rs

@@ -27,7 +27,7 @@ magic.
 */
 
 use prelude::*;
-use task::local_data_priv::{local_get, local_pop, local_modify, local_set};
+use task::local_data_priv::{local_get, local_pop, local_modify, local_set, Handle};
 use task::rt;
 
 /**
@@ -53,7 +53,7 @@ pub type LocalDataKey<'self,T> = &'self fn(v: @T);
 pub unsafe fn local_data_pop<T:Durable>(
     key: LocalDataKey<T>) -> Option<@T> {
 
-    local_pop(rt::rust_get_task(), key)
+    local_pop(Handle::new(), key)
 }
 /**
  * Retrieve a task-local data value. It will also be kept alive in the
@@ -62,7 +62,7 @@ pub unsafe fn local_data_pop<T:Durable>(
 pub unsafe fn local_data_get<T:Durable>(
     key: LocalDataKey<T>) -> Option<@T> {
 
-    local_get(rt::rust_get_task(), key)
+    local_get(Handle::new(), key)
 }
 /**
  * Store a value in task-local data. If this key already has a value,
@@ -71,7 +71,7 @@ pub unsafe fn local_data_get<T:Durable>(
 pub unsafe fn local_data_set<T:Durable>(
     key: LocalDataKey<T>, data: @T) {
 
-    local_set(rt::rust_get_task(), key, data)
+    local_set(Handle::new(), key, data)
 }
 /**
  * Modify a task-local data value. If the function returns 'None', the
@@ -81,7 +81,7 @@ pub unsafe fn local_data_modify<T:Durable>(
     key: LocalDataKey<T>,
     modify_fn: &fn(Option<@T>) -> Option<@T>) {
 
-    local_modify(rt::rust_get_task(), key, modify_fn)
+    local_modify(Handle::new(), key, modify_fn)
 }
 
 #[test]

src/libcore/task/local_data_priv.rs

@@ -18,6 +18,30 @@ use task::rt;
 use task::local_data::LocalDataKey;
 
 use super::rt::rust_task;
+use rt::local_services::LocalStorage;
+
+pub enum Handle {
+    OldHandle(*rust_task),
+    NewHandle(*mut LocalStorage)
+}
+
+impl Handle {
+    pub fn new() -> Handle {
+        use rt::{context, OldTaskContext};
+        use rt::local_services::unsafe_borrow_local_services;
+        unsafe {
+            match context() {
+                OldTaskContext => {
+                    OldHandle(rt::rust_get_task())
+                }
+                _ => {
+                    let local_services = unsafe_borrow_local_services();
+                    NewHandle(&mut local_services.storage)
+                }
+            }
+        }
+    }
+}
 
 pub trait LocalData { }
 impl<T:Durable> LocalData for @T { }
@@ -39,7 +63,7 @@ type TaskLocalElement = (*libc::c_void, *libc::c_void, @LocalData);
 // Has to be a pointer at outermost layer; the foreign call returns void *.
 type TaskLocalMap = @mut ~[Option<TaskLocalElement>];
 
-extern fn cleanup_task_local_map(map_ptr: *libc::c_void) {
+fn cleanup_task_local_map(map_ptr: *libc::c_void) {
     unsafe {
         assert!(!map_ptr.is_null());
         // Get and keep the single reference that was created at the
@@ -50,8 +74,19 @@ extern fn cleanup_task_local_map(map_ptr: *libc::c_void) {
 }
 
 // Gets the map from the runtime. Lazily initialises if not done so already.
+unsafe fn get_local_map(handle: Handle) -> TaskLocalMap {
+    match handle {
+        OldHandle(task) => get_task_local_map(task),
+        NewHandle(local_storage) => get_newsched_local_map(local_storage)
+    }
+}
+
 unsafe fn get_task_local_map(task: *rust_task) -> TaskLocalMap {
 
+    extern fn cleanup_task_local_map_(map_ptr: *libc::c_void) {
+        cleanup_task_local_map(map_ptr);
+    }
+
     // Relies on the runtime initialising the pointer to null.
     // Note: The map's box lives in TLS invisibly referenced once. Each time
     // we retrieve it for get/set, we make another reference, which get/set
@@ -62,7 +97,7 @@ unsafe fn get_task_local_map(task: *rust_task) -> TaskLocalMap {
         // Use reinterpret_cast -- transmute would take map away from us also.
         rt::rust_set_task_local_data(
             task, cast::reinterpret_cast(&map));
-        rt::rust_task_local_data_atexit(task, cleanup_task_local_map);
+        rt::rust_task_local_data_atexit(task, cleanup_task_local_map_);
         // Also need to reference it an extra time to keep it for now.
         let nonmut = cast::transmute::<TaskLocalMap,
                                        @~[Option<TaskLocalElement>]>(map);
@@ -77,6 +112,32 @@ unsafe fn get_task_local_map(task: *rust_task) -> TaskLocalMap {
     }
 }
 
+unsafe fn get_newsched_local_map(local: *mut LocalStorage) -> TaskLocalMap {
+    match &mut *local {
+        &LocalStorage(map_ptr, Some(_)) => {
+            assert!(map_ptr.is_not_null());
+            let map = cast::transmute(map_ptr);
+            let nonmut = cast::transmute::<TaskLocalMap,
+            @~[Option<TaskLocalElement>]>(map);
+            cast::bump_box_refcount(nonmut);
+            return map;
+        }
+        &LocalStorage(ref mut map_ptr, ref mut at_exit) => {
+            assert!((*map_ptr).is_null());
+            let map: TaskLocalMap = @mut ~[];
+            // Use reinterpret_cast -- transmute would take map away from us also.
+            *map_ptr = cast::reinterpret_cast(&map);
+            let at_exit_fn: ~fn(*libc::c_void) = |p|cleanup_task_local_map(p);
+            *at_exit = Some(at_exit_fn);
+            // Also need to reference it an extra time to keep it for now.
+            let nonmut = cast::transmute::<TaskLocalMap,
+            @~[Option<TaskLocalElement>]>(map);
+            cast::bump_box_refcount(nonmut);
+            return map;
+        }
+    }
+}
+
 unsafe fn key_to_key_value<T:Durable>(
     key: LocalDataKey<T>) -> *libc::c_void {
 
@@ -106,10 +167,10 @@ unsafe fn local_data_lookup<T:Durable>(
 }
 
 unsafe fn local_get_helper<T:Durable>(
-    task: *rust_task, key: LocalDataKey<T>,
+    handle: Handle, key: LocalDataKey<T>,
     do_pop: bool) -> Option<@T> {
 
-    let map = get_task_local_map(task);
+    let map = get_local_map(handle);
     // Interpreturn our findings from the map
     do local_data_lookup(map, key).map |result| {
         // A reference count magically appears on 'data' out of thin air. It
@@ -128,23 +189,23 @@ unsafe fn local_get_helper<T:Durable>(
 
 
 pub unsafe fn local_pop<T:Durable>(
-    task: *rust_task,
+    handle: Handle,
     key: LocalDataKey<T>) -> Option<@T> {
 
-    local_get_helper(task, key, true)
+    local_get_helper(handle, key, true)
 }
 
 pub unsafe fn local_get<T:Durable>(
-    task: *rust_task,
+    handle: Handle,
     key: LocalDataKey<T>) -> Option<@T> {
 
-    local_get_helper(task, key, false)
+    local_get_helper(handle, key, false)
 }
 
 pub unsafe fn local_set<T:Durable>(
-    task: *rust_task, key: LocalDataKey<T>, data: @T) {
+    handle: Handle, key: LocalDataKey<T>, data: @T) {
 
-    let map = get_task_local_map(task);
+    let map = get_local_map(handle);
     // Store key+data as *voids. Data is invisibly referenced once; key isn't.
     let keyval = key_to_key_value(key);
     // We keep the data in two forms: one as an unsafe pointer, so we can get
@@ -174,12 +235,12 @@ pub unsafe fn local_set<T:Durable>(
 }
 
 pub unsafe fn local_modify<T:Durable>(
-    task: *rust_task, key: LocalDataKey<T>,
+    handle: Handle, key: LocalDataKey<T>,
     modify_fn: &fn(Option<@T>) -> Option<@T>) {
 
     // Could be more efficient by doing the lookup work, but this is easy.
-    let newdata = modify_fn(local_pop(task, key));
+    let newdata = modify_fn(local_pop(handle, key));
     if newdata.is_some() {
-        local_set(task, key, newdata.unwrap());
+        local_set(handle, key, newdata.unwrap());
     }
 }

src/libcore/task/spawn.rs

@@ -80,7 +80,7 @@ use prelude::*;
 use unstable;
 use ptr;
 use hashmap::HashSet;
-use task::local_data_priv::{local_get, local_set};
+use task::local_data_priv::{local_get, local_set, OldHandle};
 use task::rt::rust_task;
 use task::rt;
 use task::{Failure, ManualThreads, PlatformThread, SchedOpts, SingleThreaded};
@@ -451,7 +451,7 @@ fn gen_child_taskgroup(linked: bool, supervised: bool)
         /*##################################################################*
          * Step 1. Get spawner's taskgroup info.
          *##################################################################*/
-        let spawner_group = match local_get(spawner, taskgroup_key!()) {
+        let spawner_group = match local_get(OldHandle(spawner), taskgroup_key!()) {
             None => {
                 // Main task, doing first spawn ever. Lazily initialise here.
                 let mut members = new_taskset();
@@ -463,7 +463,7 @@ fn gen_child_taskgroup(linked: bool, supervised: bool)
                 // Main task/group has no ancestors, no notifier, etc.
                 let group =
                     @TCB(spawner, tasks, AncestorList(None), true, None);
-                local_set(spawner, taskgroup_key!(), group);
+                local_set(OldHandle(spawner), taskgroup_key!(), group);
                 group
             }
             Some(group) => group
@@ -627,7 +627,7 @@ fn spawn_raw_oldsched(opts: TaskOpts, f: ~fn()) {
                 let group = @TCB(child, child_arc, ancestors,
                                  is_main, notifier);
                 unsafe {
-                    local_set(child, taskgroup_key!(), group);
+                    local_set(OldHandle(child), taskgroup_key!(), group);
                 }
 
                 // Run the child's body.