-// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
-// file at the top-level directory of this distribution and at
-// http://rust-lang.org/COPYRIGHT.
-//
-// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
-// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
-// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
-// option. This file may not be copied, modified, or distributed
-// except according to those terms.
-
-/*!
-
-Task local data management
-
-Allows storing arbitrary types inside task-local-data (TLD), to be accessed
-anywhere within a task, keyed by a global pointer parameterized over the type of
-the TLD slot. Useful for dynamic variables, singletons, and interfacing with
-foreign code with bad callback interfaces.
-
-To declare a new key for storing local data of a particular type, use the
-`local_data_key!` macro. This macro will expand to a `static` item appropriately
-named and annotated. This name is then passed to the functions in this module to
-modify/read the slot specified by the key.
-
-```rust
-local_data_key!(key_int: int)
-local_data_key!(key_vector: Vec<int>)
-
-key_int.replace(Some(3));
-assert_eq!(*key_int.get().unwrap(), 3);
-
-key_vector.replace(Some(vec![4]));
-assert_eq!(*key_vector.get().unwrap(), vec![4]);
-```
-
-*/
-
-// Casting 'Arcane Sight' reveals an overwhelming aura of Transmutation
-// magic.
-
-pub use self::KeyValue::*;
-
-use core::prelude::*;
-
-use alloc::heap;
-use collections::TreeMap;
-use core::cmp;
-use core::kinds::marker;
-use core::mem;
-use core::ptr;
-use core::fmt;
-use core::cell::UnsafeCell;
-
-use local::Local;
-use task::{Task, LocalStorage};
-
-/**
- * Indexes a task-local data slot. This pointer is used for comparison to
- * differentiate keys from one another. The actual type `T` is not used anywhere
- * as a member of this type, except that it is parameterized with it to define
- * the type of each key's value.
- *
- * The value of each Key is of the singleton enum KeyValue. These also have the
- * same name as `Key` and their purpose is to take up space in the programs data
- * sections to ensure that each value of the `Key` type points to a unique
- * location.
- */
-pub type Key<T> = &'static KeyValue<T>;
-
-#[allow(missing_docs)]
-pub enum KeyValue<T> { KeyValueKey }
-
-// The task-local-map stores all TLD information for the currently running
-// task. It is stored as an owned pointer into the runtime, and it's only
-// allocated when TLD is used for the first time.
-//
-// TLD values are boxed up, with a loan count stored in the box. The box is
-// necessary given how TLD maps are constructed, but theoretically in the
-// future this could be rewritten to statically construct TLD offsets at
-// compile-time to get O(1) lookup. At that time, the box can be removed.
-//
-// A very common usage pattern for TLD is to use replace(None) to extract a
-// value from TLD, work with it, and then store it (or a derived/new value)
-// back with replace(v). We take special care to reuse the allocation in this
-// case for performance reasons.
-//
-// However, that does mean that if a value is replaced with None, the
-// allocation will stay alive and the entry will stay in the TLD map until the
-// task deallocates. This makes the assumption that every key inserted into a
-// given task's TLD is going to be present for a majority of the rest of the
-// task's lifetime, but that's a fairly safe assumption, and there's very
-// little downside as long as it holds true for most keys.
-//
-// The Map type must be public in order to allow rustrt to see it.
-//
-// We'd like to use HashMap here, but it uses TLD in its construction (it uses
-// the task-local rng). We could try to provide our own source of randomness,
-// except it also lives in libstd (which is a client of us) so we can't even
-// reference it. Instead, use TreeMap, which provides reasonable performance.
-#[doc(hidden)]
-pub type Map = TreeMap<uint, TLDValue>;
-#[unsafe_no_drop_flag]
-struct TLDValue {
- // box_ptr is a pointer to TLDValueBox<T>. It can never be null.
- box_ptr: *mut (),
- // drop_fn is the function that knows how to drop the box_ptr.
- drop_fn: unsafe fn(p: *mut ())
-}
-
-struct TLDValueBox<T> {
- // value is only initialized when refcount >= 1.
- value: T,
- // refcount of 0 means uninitialized value, 1 means initialized, 2+ means
- // borrowed.
- // NB: we use UnsafeCell instead of Cell because Ref should be allowed to
- // be Sync. The only mutation occurs when a Ref is created or destroyed,
- // so there's no issue with &Ref being thread-safe.
- refcount: UnsafeCell<uint>
-}
-
-// Gets the map from the runtime. Lazily initialises if not done so already.
-unsafe fn get_local_map<'a>() -> Option<&'a mut Map> {
- if !Local::exists(None::<Task>) { return None }
-
- let task: *mut Task = Local::unsafe_borrow();
- match &mut (*task).storage {
- // If the at_exit function is already set, then we just need to take
- // a loan out on the TLD map stored inside
- &LocalStorage(Some(ref mut map_ptr)) => {
- return Some(map_ptr);
- }
- // If this is the first time we've accessed TLD, perform similar
- // actions to the oldsched way of doing things.
- &LocalStorage(ref mut slot) => {
- *slot = Some(TreeMap::new());
- match *slot {
- Some(ref mut map_ptr) => { return Some(map_ptr) }
- None => panic!("unreachable code"),
- }
- }
- }
-}
-
-/// A RAII immutable reference to a task-local value.
-///
-/// The task-local data can be accessed through this value, and when this
-/// structure is dropped it will return the borrow on the data.
-pub struct Ref<T:'static> {
- // FIXME #12808: strange names to try to avoid interfering with
- // field accesses of the contained type via Deref
- _inner: &'static TLDValueBox<T>,
- _marker: marker::NoSend
-}
-
-fn key_to_key_value<T: 'static>(key: Key<T>) -> uint {
- key as *const _ as uint
-}
-
-impl<T: 'static> KeyValue<T> {
- /// Replaces a value in task local data.
- ///
- /// If this key is already present in TLD, then the previous value is
- /// replaced with the provided data, and then returned.
- ///
- /// # Panics
- ///
- /// This function will panic if the key is present in TLD and currently on
- /// loan with the `get` method.
- ///
- /// It will also panic if there is no local task (because the current thread
- /// is not owned by the runtime).
- ///
- /// # Example
- ///
- /// ```
- /// local_data_key!(foo: int)
- ///
- /// assert_eq!(foo.replace(Some(10)), None);
- /// assert_eq!(foo.replace(Some(4)), Some(10));
- /// assert_eq!(foo.replace(None), Some(4));
- /// ```
- pub fn replace(&'static self, data: Option<T>) -> Option<T> {
- let map = match unsafe { get_local_map() } {
- Some(map) => map,
- None => panic!("must have a local task to insert into TLD"),
- };
- let keyval = key_to_key_value(self);
-
- // The following match takes a mutable borrow on the map. In order to insert
- // our data if the key isn't present, we need to let the match end first.
- let data = match (map.get_mut(&keyval), data) {
- (None, Some(data)) => {
- // The key doesn't exist and we need to insert it. To make borrowck
- // happy, return it up a scope and insert it there.
- data
- }
- (None, None) => {
- // The key doesn't exist and we're trying to replace it with nothing.
- // Do nothing.
- return None
- }
- (Some(slot), data) => {
- // We have a slot with a box.
- let value_box = slot.box_ptr as *mut TLDValueBox<T>;
- let refcount = unsafe { *(*value_box).refcount.get() };
- return match (refcount, data) {
- (0, None) => {
- // The current value is uninitialized and we have no new value.
- // Do nothing.
- None
- }
- (0, Some(new_value)) => {
- // The current value is uninitialized and we're storing a new value.
- unsafe {
- ptr::write(&mut (*value_box).value, new_value);
- *(*value_box).refcount.get() = 1;
- None
- }
- }
- (1, None) => {
- // We have an initialized value and we're removing it.
- unsafe {
- let ret = ptr::read(&(*value_box).value);
- *(*value_box).refcount.get() = 0;
- Some(ret)
- }
- }
- (1, Some(new_value)) => {
- // We have an initialized value and we're replacing it.
- let value_ref = unsafe { &mut (*value_box).value };
- let ret = mem::replace(value_ref, new_value);
- // Refcount is already 1, leave it as that.
- Some(ret)
- }
- _ => {
- // Refcount is 2+, which means we have a live borrow.
- panic!("TLD value cannot be replaced because it is already borrowed");
- }
- }
- }
- };
- // If we've reached this point, we need to insert into the map.
- map.insert(keyval, TLDValue::new(data));
- None
- }
-
- /// Borrows a value from TLD.
- ///
- /// If `None` is returned, then this key is not present in TLD. If `Some`
- /// is returned, then the returned data is a smart pointer representing a
- /// new loan on this TLD key. While on loan, this key cannot be altered via
- /// the `replace` method.
- ///
- /// # Example
- ///
- /// ```
- /// local_data_key!(key: int)
- ///
- /// assert!(key.get().is_none());
- ///
- /// key.replace(Some(3));
- /// assert_eq!(*key.get().unwrap(), 3);
- /// ```
- pub fn get(&'static self) -> Option<Ref<T>> {
- let map = match unsafe { get_local_map() } {
- Some(map) => map,
- None => return None,
- };
- let keyval = key_to_key_value(self);
-
- match map.get(&keyval) {
- Some(slot) => {
- let value_box = slot.box_ptr as *mut TLDValueBox<T>;
- if unsafe { *(*value_box).refcount.get() } >= 1 {
- unsafe {
- *(*value_box).refcount.get() += 1;
- Some(Ref {
- _inner: &*value_box,
- _marker: marker::NoSend
- })
- }
- } else {
- None
- }
- }
- None => None
- }
- }
-
- // it's not clear if this is the right design for a public API, or if
- // there's even a need for this as a public API, but our benchmarks need
- // this to ensure consistent behavior on each run.
- #[cfg(test)]
- fn clear(&'static self) {
- let map = match unsafe { get_local_map() } {
- Some(map) => map,
- None => return
- };
- let keyval = key_to_key_value(self);
- self.replace(None); // ensure we have no outstanding borrows
- map.remove(&keyval);
- }
-}
-
-impl<T: 'static> Deref<T> for Ref<T> {
- #[inline(always)]
- fn deref<'a>(&'a self) -> &'a T {
- &self._inner.value
- }
-}
-
-impl<T: 'static + fmt::Show> fmt::Show for Ref<T> {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- (**self).fmt(f)
- }
-}
-
-impl<T: cmp::PartialEq + 'static> cmp::PartialEq for Ref<T> {
- fn eq(&self, other: &Ref<T>) -> bool {
- (**self).eq(&**other)
- }
- fn ne(&self, other: &Ref<T>) -> bool {
- (**self).ne(&**other)
- }
-}
-
-impl<T: cmp::Eq + 'static> cmp::Eq for Ref<T> {}
-
-impl<T: cmp::PartialOrd + 'static> cmp::PartialOrd for Ref<T> {
- fn partial_cmp(&self, other: &Ref<T>) -> Option<cmp::Ordering> {
- (**self).partial_cmp(&**other)
- }
- fn lt(&self, other: &Ref<T>) -> bool { (**self).lt(&**other) }
- fn le(&self, other: &Ref<T>) -> bool { (**self).le(&**other) }
- fn gt(&self, other: &Ref<T>) -> bool { (**self).gt(&**other) }
- fn ge(&self, other: &Ref<T>) -> bool { (**self).ge(&**other) }
-}
-
-impl<T: cmp::Ord + 'static> cmp::Ord for Ref<T> {
- fn cmp(&self, other: &Ref<T>) -> cmp::Ordering {
- (**self).cmp(&**other)
- }
-}
-
-#[unsafe_destructor]
-impl<T: 'static> Drop for Ref<T> {
- fn drop(&mut self) {
- unsafe {
- *self._inner.refcount.get() -= 1;
- }
- }
-}
-
-impl TLDValue {
- fn new<T>(value: T) -> TLDValue {
- let box_ptr = unsafe {
- let allocation = heap::allocate(mem::size_of::<TLDValueBox<T>>(),
- mem::min_align_of::<TLDValueBox<T>>());
- if allocation.is_null() { ::alloc::oom() }
- let value_box = allocation as *mut TLDValueBox<T>;
- ptr::write(value_box, TLDValueBox {
- value: value,
- refcount: UnsafeCell::new(1)
- });
- value_box as *mut ()
- };
- // Destruction of TLDValue needs to know how to properly deallocate the TLDValueBox,
- // so we need our own custom destructor function.
- unsafe fn d<T>(p: *mut ()) {
- let value_box = p as *mut TLDValueBox<T>;
- debug_assert!(*(*value_box).refcount.get() < 2, "TLDValue destructed while borrowed");
- // use a RAII type here to ensure we always deallocate even if we panic while
- // running the destructor for the value.
- struct Guard<T> {
- p: *mut TLDValueBox<T>
- }
- #[unsafe_destructor]
- impl<T> Drop for Guard<T> {
- fn drop(&mut self) {
- let size = mem::size_of::<TLDValueBox<T>>();
- let align = mem::align_of::<TLDValueBox<T>>();
- unsafe { heap::deallocate(self.p as *mut u8, size, align); }
- }
- }
- let _guard = Guard::<T> { p: value_box };
- if *(*value_box).refcount.get() != 0 {
- // the contained value is valid; drop it
- ptr::read(&(*value_box).value);
- }
- // the box will be deallocated by the guard
- }
- TLDValue {
- box_ptr: box_ptr,
- drop_fn: d::<T>
- }
- }
-}
-
-
-impl Drop for TLDValue {
- fn drop(&mut self) {
- // box_ptr should always be non-null. Check it anyway just to be thorough
- if !self.box_ptr.is_null() {
- unsafe { (self.drop_fn)(self.box_ptr) }
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- extern crate test;
-
- use std::prelude::*;
- use super::*;
- use std::task;
-
- #[test]
- fn test_tls_multitask() {
- static MY_KEY: Key<String> = &KeyValueKey;
- MY_KEY.replace(Some("parent data".to_string()));
- task::spawn(proc() {
- // TLD shouldn't carry over.
- assert!(MY_KEY.get().is_none());
- MY_KEY.replace(Some("child data".to_string()));
- assert!(MY_KEY.get().as_ref().unwrap().as_slice() == "child data");
- // should be cleaned up for us
- });
-
- // Must work multiple times
- assert!(MY_KEY.get().unwrap().as_slice() == "parent data");
- assert!(MY_KEY.get().unwrap().as_slice() == "parent data");
- assert!(MY_KEY.get().unwrap().as_slice() == "parent data");
- }
-
- #[test]
- fn test_tls_overwrite() {
- static MY_KEY: Key<String> = &KeyValueKey;
- MY_KEY.replace(Some("first data".to_string()));
- MY_KEY.replace(Some("next data".to_string())); // Shouldn't leak.
- assert!(MY_KEY.get().unwrap().as_slice() == "next data");
- }
-
- #[test]
- fn test_tls_pop() {
- static MY_KEY: Key<String> = &KeyValueKey;
- MY_KEY.replace(Some("weasel".to_string()));
- assert!(MY_KEY.replace(None).unwrap() == "weasel".to_string());
- // Pop must remove the data from the map.
- assert!(MY_KEY.replace(None).is_none());
- }
-
- #[test]
- fn test_tls_crust_automorestack_memorial_bug() {
- // This might result in a stack-canary clobber if the runtime fails to
- // set sp_limit to 0 when calling the cleanup extern - it might
- // automatically jump over to the rust stack, which causes next_c_sp
- // to get recorded as something within a rust stack segment. Then a
- // subsequent upcall (esp. for logging, think vsnprintf) would run on
- // a stack smaller than 1 MB.
- static MY_KEY: Key<String> = &KeyValueKey;
- task::spawn(proc() {
- MY_KEY.replace(Some("hax".to_string()));
- });
- }
-
- #[test]
- fn test_tls_multiple_types() {
- static STR_KEY: Key<String> = &KeyValueKey;
- static BOX_KEY: Key<Box<int>> = &KeyValueKey;
- static INT_KEY: Key<int> = &KeyValueKey;
- task::spawn(proc() {
- STR_KEY.replace(Some("string data".to_string()));
- BOX_KEY.replace(Some(box 0));
- INT_KEY.replace(Some(42));
- });
- }
-
- #[test]
- fn test_tls_overwrite_multiple_types() {
- static STR_KEY: Key<String> = &KeyValueKey;
- static BOX_KEY: Key<Box<int>> = &KeyValueKey;
- static INT_KEY: Key<int> = &KeyValueKey;
- task::spawn(proc() {
- STR_KEY.replace(Some("string data".to_string()));
- STR_KEY.replace(Some("string data 2".to_string()));
- BOX_KEY.replace(Some(box 0));
- BOX_KEY.replace(Some(box 1));
- INT_KEY.replace(Some(42));
- // This could cause a segfault if overwriting-destruction is done
- // with the crazy polymorphic transmute rather than the provided
- // finaliser.
- INT_KEY.replace(Some(31337));
- });
- }
-
- #[test]
- #[should_fail]
- fn test_tls_cleanup_on_panic() {
- static STR_KEY: Key<String> = &KeyValueKey;
- static BOX_KEY: Key<Box<int>> = &KeyValueKey;
- static INT_KEY: Key<int> = &KeyValueKey;
- STR_KEY.replace(Some("parent data".to_string()));
- BOX_KEY.replace(Some(box 0));
- task::spawn(proc() {
- STR_KEY.replace(Some("string data".to_string()));
- BOX_KEY.replace(Some(box 2));
- INT_KEY.replace(Some(42));
- panic!();
- });
- // Not quite nondeterministic.
- INT_KEY.replace(Some(31337));
- panic!();
- }
-
- #[test]
- fn test_cleanup_drops_values() {
- let (tx, rx) = channel::<()>();
- struct Dropper {
- tx: Sender<()>
- };
- impl Drop for Dropper {
- fn drop(&mut self) {
- self.tx.send(());
- }
- }
- static KEY: Key<Dropper> = &KeyValueKey;
- let _ = task::try(proc() {
- KEY.replace(Some(Dropper{ tx: tx }));
- });
- // At this point the task has been cleaned up and the TLD dropped.
- // If the channel doesn't have a value now, then the Sender was leaked.
- assert_eq!(rx.try_recv(), Ok(()));
- }
-
- #[test]
- fn test_static_pointer() {
- static KEY: Key<&'static int> = &KeyValueKey;
- static VALUE: int = 0;
- KEY.replace(Some(&VALUE));
- }
-
- #[test]
- fn test_owned() {
- static KEY: Key<Box<int>> = &KeyValueKey;
- KEY.replace(Some(box 1));
-
- {
- let k1 = KEY.get().unwrap();
- let k2 = KEY.get().unwrap();
- let k3 = KEY.get().unwrap();
- assert_eq!(**k1, 1);
- assert_eq!(**k2, 1);
- assert_eq!(**k3, 1);
- }
- KEY.replace(Some(box 2));
- assert_eq!(**KEY.get().unwrap(), 2);
- }
-
- #[test]
- fn test_same_key_type() {
- static KEY1: Key<int> = &KeyValueKey;
- static KEY2: Key<int> = &KeyValueKey;
- static KEY3: Key<int> = &KeyValueKey;
- static KEY4: Key<int> = &KeyValueKey;
- static KEY5: Key<int> = &KeyValueKey;
- KEY1.replace(Some(1));
- KEY2.replace(Some(2));
- KEY3.replace(Some(3));
- KEY4.replace(Some(4));
- KEY5.replace(Some(5));
-
- assert_eq!(*KEY1.get().unwrap(), 1);
- assert_eq!(*KEY2.get().unwrap(), 2);
- assert_eq!(*KEY3.get().unwrap(), 3);
- assert_eq!(*KEY4.get().unwrap(), 4);
- assert_eq!(*KEY5.get().unwrap(), 5);
- }
-
- #[test]
- #[should_fail]
- fn test_nested_get_set1() {
- static KEY: Key<int> = &KeyValueKey;
- assert_eq!(KEY.replace(Some(4)), None);
-
- let _k = KEY.get();
- KEY.replace(Some(4));
- }
-
- // ClearKey is a RAII class that ensures the keys are cleared from the map.
- // This is so repeated runs of a benchmark don't bloat the map with extra
- // keys and distort the measurements.
- // It's not used on the tests because the tests run in separate tasks.
- struct ClearKey<T>(Key<T>);
- #[unsafe_destructor]
- impl<T: 'static> Drop for ClearKey<T> {
- fn drop(&mut self) {
- let ClearKey(ref key) = *self;
- key.clear();
- }
- }
-
- #[bench]
- fn bench_replace_none(b: &mut test::Bencher) {
- static KEY: Key<uint> = &KeyValueKey;
- let _clear = ClearKey(KEY);
- KEY.replace(None);
- b.iter(|| {
- KEY.replace(None)
- });
- }
-
- #[bench]
- fn bench_replace_some(b: &mut test::Bencher) {
- static KEY: Key<uint> = &KeyValueKey;
- let _clear = ClearKey(KEY);
- KEY.replace(Some(1u));
- b.iter(|| {
- KEY.replace(Some(2))
- });
- }
-
- #[bench]
- fn bench_replace_none_some(b: &mut test::Bencher) {
- static KEY: Key<uint> = &KeyValueKey;
- let _clear = ClearKey(KEY);
- KEY.replace(Some(0u));
- b.iter(|| {
- let old = KEY.replace(None).unwrap();
- let new = old + 1;
- KEY.replace(Some(new))
- });
- }
-
- #[bench]
- fn bench_100_keys_replace_last(b: &mut test::Bencher) {
- static KEYS: [KeyValue<uint>, ..100] = [KeyValueKey, ..100];
- let _clear = KEYS.iter().map(ClearKey).collect::<Vec<ClearKey<uint>>>();
- for (i, key) in KEYS.iter().enumerate() {
- key.replace(Some(i));
- }
- b.iter(|| {
- let key: Key<uint> = &KEYS[99];
- key.replace(Some(42))
- });
- }
-
- #[bench]
- fn bench_1000_keys_replace_last(b: &mut test::Bencher) {
- static KEYS: [KeyValue<uint>, ..1000] = [KeyValueKey, ..1000];
- let _clear = KEYS.iter().map(ClearKey).collect::<Vec<ClearKey<uint>>>();
- for (i, key) in KEYS.iter().enumerate() {
- key.replace(Some(i));
- }
- b.iter(|| {
- let key: Key<uint> = &KEYS[999];
- key.replace(Some(42))
- });
- for key in KEYS.iter() { key.clear(); }
- }
-
- #[bench]
- fn bench_get(b: &mut test::Bencher) {
- static KEY: Key<uint> = &KeyValueKey;
- let _clear = ClearKey(KEY);
- KEY.replace(Some(42));
- b.iter(|| {
- KEY.get()
- });
- }
-
- #[bench]
- fn bench_100_keys_get_last(b: &mut test::Bencher) {
- static KEYS: [KeyValue<uint>, ..100] = [KeyValueKey, ..100];
- let _clear = KEYS.iter().map(ClearKey).collect::<Vec<ClearKey<uint>>>();
- for (i, key) in KEYS.iter().enumerate() {
- key.replace(Some(i));
- }
- b.iter(|| {
- let key: Key<uint> = &KEYS[99];
- key.get()
- });
- }
-
- #[bench]
- fn bench_1000_keys_get_last(b: &mut test::Bencher) {
- static KEYS: [KeyValue<uint>, ..1000] = [KeyValueKey, ..1000];
- let _clear = KEYS.iter().map(ClearKey).collect::<Vec<ClearKey<uint>>>();
- for (i, key) in KEYS.iter().enumerate() {
- key.replace(Some(i));
- }
- b.iter(|| {
- let key: Key<uint> = &KEYS[999];
- key.get()
- });
- }
-}