1 // Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
2 // file at the top-level directory of this distribution and at
3 // http://rust-lang.org/COPYRIGHT.
5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
8 // option. This file may not be copied, modified, or distributed
9 // except according to those terms.
13 Task local data management
15 Allows storing arbitrary types inside task-local-data (TLD), to be accessed
16 anywhere within a task, keyed by a global pointer parameterized over the type of
17 the TLD slot. Useful for dynamic variables, singletons, and interfacing with
18 foreign code with bad callback interfaces.
20 To declare a new key for storing local data of a particular type, use the
21 `local_data_key!` macro. This macro will expand to a `static` item appropriately
22 named and annotated. This name is then passed to the functions in this module to
23 modify/read the slot specified by the key.
26 local_data_key!(key_int: int)
27 local_data_key!(key_vector: Vec<int>)
29 key_int.replace(Some(3));
30 assert_eq!(*key_int.get().unwrap(), 3);
32 key_vector.replace(Some(vec![4]));
33 assert_eq!(*key_vector.get().unwrap(), vec![4]);
38 // Casting 'Arcane Sight' reveals an overwhelming aura of Transmutation
41 pub use self::KeyValue::*;
46 use collections::TreeMap;
48 use core::kinds::marker;
52 use core::cell::UnsafeCell;
55 use task::{Task, LocalStorage};
58 * Indexes a task-local data slot. This pointer is used for comparison to
59 * differentiate keys from one another. The actual type `T` is not used anywhere
60 * as a member of this type, except that it is parameterized with it to define
61 * the type of each key's value.
63 * The value of each Key is of the singleton enum KeyValue. These also have the
64 * same name as `Key` and their purpose is to take up space in the programs data
65 * sections to ensure that each value of the `Key` type points to a unique
68 pub type Key<T> = &'static KeyValue<T>;
70 #[allow(missing_docs)]
71 pub enum KeyValue<T> { KeyValueKey }
73 // The task-local-map stores all TLD information for the currently running
74 // task. It is stored as an owned pointer into the runtime, and it's only
75 // allocated when TLD is used for the first time.
77 // TLD values are boxed up, with a loan count stored in the box. The box is
78 // necessary given how TLD maps are constructed, but theoretically in the
79 // future this could be rewritten to statically construct TLD offsets at
80 // compile-time to get O(1) lookup. At that time, the box can be removed.
82 // A very common usage pattern for TLD is to use replace(None) to extract a
83 // value from TLD, work with it, and then store it (or a derived/new value)
84 // back with replace(v). We take special care to reuse the allocation in this
85 // case for performance reasons.
87 // However, that does mean that if a value is replaced with None, the
88 // allocation will stay alive and the entry will stay in the TLD map until the
89 // task deallocates. This makes the assumption that every key inserted into a
90 // given task's TLD is going to be present for a majority of the rest of the
91 // task's lifetime, but that's a fairly safe assumption, and there's very
92 // little downside as long as it holds true for most keys.
94 // The Map type must be public in order to allow rustrt to see it.
96 // We'd like to use HashMap here, but it uses TLD in its construction (it uses
97 // the task-local rng). We could try to provide our own source of randomness,
98 // except it also lives in libstd (which is a client of us) so we can't even
99 // reference it. Instead, use TreeMap, which provides reasonable performance.
101 pub type Map = TreeMap<uint, TLDValue>;
102 #[unsafe_no_drop_flag]
104 // box_ptr is a pointer to TLDValueBox<T>. It can never be null.
106 // drop_fn is the function that knows how to drop the box_ptr.
107 drop_fn: unsafe fn(p: *mut ())
110 struct TLDValueBox<T> {
111 // value is only initialized when refcount >= 1.
113 // refcount of 0 means uninitialized value, 1 means initialized, 2+ means
115 // NB: we use UnsafeCell instead of Cell because Ref should be allowed to
116 // be Sync. The only mutation occurs when a Ref is created or destroyed,
117 // so there's no issue with &Ref being thread-safe.
118 refcount: UnsafeCell<uint>
121 // Gets the map from the runtime. Lazily initialises if not done so already.
122 unsafe fn get_local_map<'a>() -> Option<&'a mut Map> {
123 if !Local::exists(None::<Task>) { return None }
125 let task: *mut Task = Local::unsafe_borrow();
126 match &mut (*task).storage {
127 // If the at_exit function is already set, then we just need to take
128 // a loan out on the TLD map stored inside
129 &LocalStorage(Some(ref mut map_ptr)) => {
130 return Some(map_ptr);
132 // If this is the first time we've accessed TLD, perform similar
133 // actions to the oldsched way of doing things.
134 &LocalStorage(ref mut slot) => {
135 *slot = Some(TreeMap::new());
137 Some(ref mut map_ptr) => { return Some(map_ptr) }
138 None => panic!("unreachable code"),
144 /// A RAII immutable reference to a task-local value.
146 /// The task-local data can be accessed through this value, and when this
147 /// structure is dropped it will return the borrow on the data.
148 pub struct Ref<T:'static> {
149 // FIXME #12808: strange names to try to avoid interfering with
150 // field accesses of the contained type via Deref
151 _inner: &'static TLDValueBox<T>,
152 _marker: marker::NoSend
155 fn key_to_key_value<T: 'static>(key: Key<T>) -> uint {
156 key as *const _ as uint
159 impl<T: 'static> KeyValue<T> {
160 /// Replaces a value in task local data.
162 /// If this key is already present in TLD, then the previous value is
163 /// replaced with the provided data, and then returned.
167 /// This function will panic if the key is present in TLD and currently on
168 /// loan with the `get` method.
170 /// It will also panic if there is no local task (because the current thread
171 /// is not owned by the runtime).
176 /// local_data_key!(foo: int)
178 /// assert_eq!(foo.replace(Some(10)), None);
179 /// assert_eq!(foo.replace(Some(4)), Some(10));
180 /// assert_eq!(foo.replace(None), Some(4));
182 pub fn replace(&'static self, data: Option<T>) -> Option<T> {
183 let map = match unsafe { get_local_map() } {
185 None => panic!("must have a local task to insert into TLD"),
187 let keyval = key_to_key_value(self);
189 // The following match takes a mutable borrow on the map. In order to insert
190 // our data if the key isn't present, we need to let the match end first.
191 let data = match (map.get_mut(&keyval), data) {
192 (None, Some(data)) => {
193 // The key doesn't exist and we need to insert it. To make borrowck
194 // happy, return it up a scope and insert it there.
198 // The key doesn't exist and we're trying to replace it with nothing.
202 (Some(slot), data) => {
203 // We have a slot with a box.
204 let value_box = slot.box_ptr as *mut TLDValueBox<T>;
205 let refcount = unsafe { *(*value_box).refcount.get() };
206 return match (refcount, data) {
208 // The current value is uninitialized and we have no new value.
212 (0, Some(new_value)) => {
213 // The current value is uninitialized and we're storing a new value.
215 ptr::write(&mut (*value_box).value, new_value);
216 *(*value_box).refcount.get() = 1;
221 // We have an initialized value and we're removing it.
223 let ret = ptr::read(&(*value_box).value);
224 *(*value_box).refcount.get() = 0;
228 (1, Some(new_value)) => {
229 // We have an initialized value and we're replacing it.
230 let value_ref = unsafe { &mut (*value_box).value };
231 let ret = mem::replace(value_ref, new_value);
232 // Refcount is already 1, leave it as that.
236 // Refcount is 2+, which means we have a live borrow.
237 panic!("TLD value cannot be replaced because it is already borrowed");
242 // If we've reached this point, we need to insert into the map.
243 map.insert(keyval, TLDValue::new(data));
247 /// Borrows a value from TLD.
249 /// If `None` is returned, then this key is not present in TLD. If `Some`
250 /// is returned, then the returned data is a smart pointer representing a
251 /// new loan on this TLD key. While on loan, this key cannot be altered via
252 /// the `replace` method.
257 /// local_data_key!(key: int)
259 /// assert!(key.get().is_none());
261 /// key.replace(Some(3));
262 /// assert_eq!(*key.get().unwrap(), 3);
264 pub fn get(&'static self) -> Option<Ref<T>> {
265 let map = match unsafe { get_local_map() } {
269 let keyval = key_to_key_value(self);
271 match map.get(&keyval) {
273 let value_box = slot.box_ptr as *mut TLDValueBox<T>;
274 if unsafe { *(*value_box).refcount.get() } >= 1 {
276 *(*value_box).refcount.get() += 1;
279 _marker: marker::NoSend
290 // it's not clear if this is the right design for a public API, or if
291 // there's even a need for this as a public API, but our benchmarks need
292 // this to ensure consistent behavior on each run.
294 fn clear(&'static self) {
295 let map = match unsafe { get_local_map() } {
299 let keyval = key_to_key_value(self);
300 self.replace(None); // ensure we have no outstanding borrows
305 impl<T: 'static> Deref<T> for Ref<T> {
307 fn deref<'a>(&'a self) -> &'a T {
312 impl<T: 'static + fmt::Show> fmt::Show for Ref<T> {
313 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
318 impl<T: cmp::PartialEq + 'static> cmp::PartialEq for Ref<T> {
319 fn eq(&self, other: &Ref<T>) -> bool {
320 (**self).eq(&**other)
322 fn ne(&self, other: &Ref<T>) -> bool {
323 (**self).ne(&**other)
327 impl<T: cmp::Eq + 'static> cmp::Eq for Ref<T> {}
329 impl<T: cmp::PartialOrd + 'static> cmp::PartialOrd for Ref<T> {
330 fn partial_cmp(&self, other: &Ref<T>) -> Option<cmp::Ordering> {
331 (**self).partial_cmp(&**other)
333 fn lt(&self, other: &Ref<T>) -> bool { (**self).lt(&**other) }
334 fn le(&self, other: &Ref<T>) -> bool { (**self).le(&**other) }
335 fn gt(&self, other: &Ref<T>) -> bool { (**self).gt(&**other) }
336 fn ge(&self, other: &Ref<T>) -> bool { (**self).ge(&**other) }
339 impl<T: cmp::Ord + 'static> cmp::Ord for Ref<T> {
340 fn cmp(&self, other: &Ref<T>) -> cmp::Ordering {
341 (**self).cmp(&**other)
346 impl<T: 'static> Drop for Ref<T> {
349 *self._inner.refcount.get() -= 1;
355 fn new<T>(value: T) -> TLDValue {
356 let box_ptr = unsafe {
357 let allocation = heap::allocate(mem::size_of::<TLDValueBox<T>>(),
358 mem::min_align_of::<TLDValueBox<T>>());
359 if allocation.is_null() { ::alloc::oom() }
360 let value_box = allocation as *mut TLDValueBox<T>;
361 ptr::write(value_box, TLDValueBox {
363 refcount: UnsafeCell::new(1)
367 // Destruction of TLDValue needs to know how to properly deallocate the TLDValueBox,
368 // so we need our own custom destructor function.
369 unsafe fn d<T>(p: *mut ()) {
370 let value_box = p as *mut TLDValueBox<T>;
371 debug_assert!(*(*value_box).refcount.get() < 2, "TLDValue destructed while borrowed");
372 // use a RAII type here to ensure we always deallocate even if we panic while
373 // running the destructor for the value.
375 p: *mut TLDValueBox<T>
378 impl<T> Drop for Guard<T> {
380 let size = mem::size_of::<TLDValueBox<T>>();
381 let align = mem::align_of::<TLDValueBox<T>>();
382 unsafe { heap::deallocate(self.p as *mut u8, size, align); }
385 let _guard = Guard::<T> { p: value_box };
386 if *(*value_box).refcount.get() != 0 {
387 // the contained value is valid; drop it
388 ptr::read(&(*value_box).value);
390 // the box will be deallocated by the guard
400 impl Drop for TLDValue {
402 // box_ptr should always be non-null. Check it anyway just to be thorough
403 if !self.box_ptr.is_null() {
404 unsafe { (self.drop_fn)(self.box_ptr) }
418 fn test_tls_multitask() {
419 static MY_KEY: Key<String> = &KeyValueKey;
420 MY_KEY.replace(Some("parent data".to_string()));
422 // TLD shouldn't carry over.
423 assert!(MY_KEY.get().is_none());
424 MY_KEY.replace(Some("child data".to_string()));
425 assert!(MY_KEY.get().as_ref().unwrap().as_slice() == "child data");
426 // should be cleaned up for us
429 // Must work multiple times
430 assert!(MY_KEY.get().unwrap().as_slice() == "parent data");
431 assert!(MY_KEY.get().unwrap().as_slice() == "parent data");
432 assert!(MY_KEY.get().unwrap().as_slice() == "parent data");
436 fn test_tls_overwrite() {
437 static MY_KEY: Key<String> = &KeyValueKey;
438 MY_KEY.replace(Some("first data".to_string()));
439 MY_KEY.replace(Some("next data".to_string())); // Shouldn't leak.
440 assert!(MY_KEY.get().unwrap().as_slice() == "next data");
445 static MY_KEY: Key<String> = &KeyValueKey;
446 MY_KEY.replace(Some("weasel".to_string()));
447 assert!(MY_KEY.replace(None).unwrap() == "weasel".to_string());
448 // Pop must remove the data from the map.
449 assert!(MY_KEY.replace(None).is_none());
453 fn test_tls_crust_automorestack_memorial_bug() {
454 // This might result in a stack-canary clobber if the runtime fails to
455 // set sp_limit to 0 when calling the cleanup extern - it might
456 // automatically jump over to the rust stack, which causes next_c_sp
457 // to get recorded as something within a rust stack segment. Then a
458 // subsequent upcall (esp. for logging, think vsnprintf) would run on
459 // a stack smaller than 1 MB.
460 static MY_KEY: Key<String> = &KeyValueKey;
462 MY_KEY.replace(Some("hax".to_string()));
467 fn test_tls_multiple_types() {
468 static STR_KEY: Key<String> = &KeyValueKey;
469 static BOX_KEY: Key<Box<int>> = &KeyValueKey;
470 static INT_KEY: Key<int> = &KeyValueKey;
472 STR_KEY.replace(Some("string data".to_string()));
473 BOX_KEY.replace(Some(box 0));
474 INT_KEY.replace(Some(42));
479 fn test_tls_overwrite_multiple_types() {
480 static STR_KEY: Key<String> = &KeyValueKey;
481 static BOX_KEY: Key<Box<int>> = &KeyValueKey;
482 static INT_KEY: Key<int> = &KeyValueKey;
484 STR_KEY.replace(Some("string data".to_string()));
485 STR_KEY.replace(Some("string data 2".to_string()));
486 BOX_KEY.replace(Some(box 0));
487 BOX_KEY.replace(Some(box 1));
488 INT_KEY.replace(Some(42));
489 // This could cause a segfault if overwriting-destruction is done
490 // with the crazy polymorphic transmute rather than the provided
492 INT_KEY.replace(Some(31337));
498 fn test_tls_cleanup_on_panic() {
499 static STR_KEY: Key<String> = &KeyValueKey;
500 static BOX_KEY: Key<Box<int>> = &KeyValueKey;
501 static INT_KEY: Key<int> = &KeyValueKey;
502 STR_KEY.replace(Some("parent data".to_string()));
503 BOX_KEY.replace(Some(box 0));
505 STR_KEY.replace(Some("string data".to_string()));
506 BOX_KEY.replace(Some(box 2));
507 INT_KEY.replace(Some(42));
510 // Not quite nondeterministic.
511 INT_KEY.replace(Some(31337));
516 fn test_cleanup_drops_values() {
517 let (tx, rx) = channel::<()>();
521 impl Drop for Dropper {
526 static KEY: Key<Dropper> = &KeyValueKey;
527 let _ = task::try(proc() {
528 KEY.replace(Some(Dropper{ tx: tx }));
530 // At this point the task has been cleaned up and the TLD dropped.
531 // If the channel doesn't have a value now, then the Sender was leaked.
532 assert_eq!(rx.try_recv(), Ok(()));
536 fn test_static_pointer() {
537 static KEY: Key<&'static int> = &KeyValueKey;
538 static VALUE: int = 0;
539 KEY.replace(Some(&VALUE));
544 static KEY: Key<Box<int>> = &KeyValueKey;
545 KEY.replace(Some(box 1));
548 let k1 = KEY.get().unwrap();
549 let k2 = KEY.get().unwrap();
550 let k3 = KEY.get().unwrap();
555 KEY.replace(Some(box 2));
556 assert_eq!(**KEY.get().unwrap(), 2);
560 fn test_same_key_type() {
561 static KEY1: Key<int> = &KeyValueKey;
562 static KEY2: Key<int> = &KeyValueKey;
563 static KEY3: Key<int> = &KeyValueKey;
564 static KEY4: Key<int> = &KeyValueKey;
565 static KEY5: Key<int> = &KeyValueKey;
566 KEY1.replace(Some(1));
567 KEY2.replace(Some(2));
568 KEY3.replace(Some(3));
569 KEY4.replace(Some(4));
570 KEY5.replace(Some(5));
572 assert_eq!(*KEY1.get().unwrap(), 1);
573 assert_eq!(*KEY2.get().unwrap(), 2);
574 assert_eq!(*KEY3.get().unwrap(), 3);
575 assert_eq!(*KEY4.get().unwrap(), 4);
576 assert_eq!(*KEY5.get().unwrap(), 5);
581 fn test_nested_get_set1() {
582 static KEY: Key<int> = &KeyValueKey;
583 assert_eq!(KEY.replace(Some(4)), None);
586 KEY.replace(Some(4));
589 // ClearKey is a RAII class that ensures the keys are cleared from the map.
590 // This is so repeated runs of a benchmark don't bloat the map with extra
591 // keys and distort the measurements.
592 // It's not used on the tests because the tests run in separate tasks.
593 struct ClearKey<T>(Key<T>);
595 impl<T: 'static> Drop for ClearKey<T> {
597 let ClearKey(ref key) = *self;
603 fn bench_replace_none(b: &mut test::Bencher) {
604 static KEY: Key<uint> = &KeyValueKey;
605 let _clear = ClearKey(KEY);
613 fn bench_replace_some(b: &mut test::Bencher) {
614 static KEY: Key<uint> = &KeyValueKey;
615 let _clear = ClearKey(KEY);
616 KEY.replace(Some(1u));
623 fn bench_replace_none_some(b: &mut test::Bencher) {
624 static KEY: Key<uint> = &KeyValueKey;
625 let _clear = ClearKey(KEY);
626 KEY.replace(Some(0u));
628 let old = KEY.replace(None).unwrap();
630 KEY.replace(Some(new))
635 fn bench_100_keys_replace_last(b: &mut test::Bencher) {
636 static KEYS: [KeyValue<uint>, ..100] = [KeyValueKey, ..100];
637 let _clear = KEYS.iter().map(ClearKey).collect::<Vec<ClearKey<uint>>>();
638 for (i, key) in KEYS.iter().enumerate() {
639 key.replace(Some(i));
642 let key: Key<uint> = &KEYS[99];
643 key.replace(Some(42))
648 fn bench_1000_keys_replace_last(b: &mut test::Bencher) {
649 static KEYS: [KeyValue<uint>, ..1000] = [KeyValueKey, ..1000];
650 let _clear = KEYS.iter().map(ClearKey).collect::<Vec<ClearKey<uint>>>();
651 for (i, key) in KEYS.iter().enumerate() {
652 key.replace(Some(i));
655 let key: Key<uint> = &KEYS[999];
656 key.replace(Some(42))
658 for key in KEYS.iter() { key.clear(); }
662 fn bench_get(b: &mut test::Bencher) {
663 static KEY: Key<uint> = &KeyValueKey;
664 let _clear = ClearKey(KEY);
665 KEY.replace(Some(42));
672 fn bench_100_keys_get_last(b: &mut test::Bencher) {
673 static KEYS: [KeyValue<uint>, ..100] = [KeyValueKey, ..100];
674 let _clear = KEYS.iter().map(ClearKey).collect::<Vec<ClearKey<uint>>>();
675 for (i, key) in KEYS.iter().enumerate() {
676 key.replace(Some(i));
679 let key: Key<uint> = &KEYS[99];
685 fn bench_1000_keys_get_last(b: &mut test::Bencher) {
686 static KEYS: [KeyValue<uint>, ..1000] = [KeyValueKey, ..1000];
687 let _clear = KEYS.iter().map(ClearKey).collect::<Vec<ClearKey<uint>>>();
688 for (i, key) in KEYS.iter().enumerate() {
689 key.replace(Some(i));
692 let key: Key<uint> = &KEYS[999];