src/librustrt/local_data.rs

   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.
   4 //
   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.
  10
  11 /*!
  12
  13 Task local data management
  14
  15 Allows storing arbitrary types inside task-local-storage (TLS), to be accessed
  16 anywhere within a task, keyed by a global pointer parameterized over the type of
  17 the TLS slot.  Useful for dynamic variables, singletons, and interfacing with
  18 foreign code with bad callback interfaces.
  19
  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.
  24
  25 ```rust
  26 local_data_key!(key_int: int)
  27 local_data_key!(key_vector: Vec<int>)
  28
  29 key_int.replace(Some(3));
  30 assert_eq!(*key_int.get().unwrap(), 3);
  31
  32 key_vector.replace(Some(vec![4]));
  33 assert_eq!(*key_vector.get().unwrap(), vec![4]);
  34 ```
  35
  36 */
  37
  38 // Casting 'Arcane Sight' reveals an overwhelming aura of Transmutation
  39 // magic.
  40
  41 use core::prelude::*;
  42
  43 use alloc::boxed::Box;
  44 use collections::vec::Vec;
  45 use core::kinds::marker;
  46 use core::mem;
  47 use core::raw;
  48
  49 use local::Local;
  50 use task::{Task, LocalStorage};
  51
  52 /**
  53  * Indexes a task-local data slot. This pointer is used for comparison to
  54  * differentiate keys from one another. The actual type `T` is not used anywhere
  55  * as a member of this type, except that it is parameterized with it to define
  56  * the type of each key's value.
  57  *
  58  * The value of each Key is of the singleton enum KeyValue. These also have the
  59  * same name as `Key` and their purpose is to take up space in the programs data
  60  * sections to ensure that each value of the `Key` type points to a unique
  61  * location.
  62  */
  63 pub type Key<T> = &'static KeyValue<T>;
  64
  65 #[allow(missing_doc)]
  66 pub enum KeyValue<T> { Key }
  67
  68 #[doc(hidden)]
  69 trait LocalData {}
  70 impl<T: 'static> LocalData for T {}
  71
  72 // The task-local-map stores all TLS information for the currently running task.
  73 // It is stored as an owned pointer into the runtime, and it's only allocated
  74 // when TLS is used for the first time. This map must be very carefully
  75 // constructed because it has many mutable loans unsoundly handed out on it to
  76 // the various invocations of TLS requests.
  77 //
  78 // One of the most important operations is loaning a value via `get` to a
  79 // caller. In doing so, the slot that the TLS entry is occupying cannot be
  80 // invalidated because upon returning its loan state must be updated. Currently
  81 // the TLS map is a vector, but this is possibly dangerous because the vector
  82 // can be reallocated/moved when new values are pushed onto it.
  83 //
  84 // This problem currently isn't solved in a very elegant way. Inside the `get`
  85 // function, it internally "invalidates" all references after the loan is
  86 // finished and looks up into the vector again. In theory this will prevent
  87 // pointers from being moved under our feet so long as LLVM doesn't go too crazy
  88 // with the optimizations.
  89 //
  90 // n.b. If TLS is used heavily in future, this could be made more efficient with
  91 //      a proper map.
  92 #[doc(hidden)]
  93 pub type Map = Vec<Option<(*const u8, TLSValue, uint)>>;
  94 type TLSValue = Box<LocalData + Send>;
  95
  96 // Gets the map from the runtime. Lazily initialises if not done so already.
  97 unsafe fn get_local_map() -> Option<&mut Map> {
  98     if !Local::exists(None::<Task>) { return None }
  99
 100     let task: *mut Task = Local::unsafe_borrow();
 101     match &mut (*task).storage {
 102         // If the at_exit function is already set, then we just need to take
 103         // a loan out on the TLS map stored inside
 104         &LocalStorage(Some(ref mut map_ptr)) => {
 105             return Some(map_ptr);
 106         }
 107         // If this is the first time we've accessed TLS, perform similar
 108         // actions to the oldsched way of doing things.
 109         &LocalStorage(ref mut slot) => {
 110             *slot = Some(Vec::new());
 111             match *slot {
 112                 Some(ref mut map_ptr) => { return Some(map_ptr) }
 113                 None => fail!("unreachable code"),
 114             }
 115         }
 116     }
 117 }
 118
 119 fn key_to_key_value<T: 'static>(key: Key<T>) -> *const u8 {
 120     key as *const KeyValue<T> as *const u8
 121 }
 122
 123 /// An RAII immutable reference to a task-local value.
 124 ///
 125 /// The task-local data can be accessed through this value, and when this
 126 /// structure is dropped it will return the borrow on the data.
 127 pub struct Ref<T> {
 128     // FIXME #12808: strange names to try to avoid interfering with
 129     // field accesses of the contained type via Deref
 130     _ptr: &'static T,
 131     _key: Key<T>,
 132     _index: uint,
 133     _nosend: marker::NoSend,
 134 }
 135
 136 impl<T: 'static> KeyValue<T> {
 137     /// Replaces a value in task local storage.
 138     ///
 139     /// If this key is already present in TLS, then the previous value is
 140     /// replaced with the provided data, and then returned.
 141     ///
 142     /// # Failure
 143     ///
 144     /// This function will fail if this key is present in TLS and currently on
 145     /// loan with the `get` method.
 146     ///
 147     /// # Example
 148     ///
 149     /// ```
 150     /// local_data_key!(foo: int)
 151     ///
 152     /// assert_eq!(foo.replace(Some(10)), None);
 153     /// assert_eq!(foo.replace(Some(4)), Some(10));
 154     /// assert_eq!(foo.replace(None), Some(4));
 155     /// ```
 156     pub fn replace(&'static self, data: Option<T>) -> Option<T> {
 157         let map = match unsafe { get_local_map() } {
 158             Some(map) => map,
 159             None => fail!("must have a local task to insert into TLD"),
 160         };
 161         let keyval = key_to_key_value(self);
 162
 163         // When the task-local map is destroyed, all the data needs to be
 164         // cleaned up. For this reason we can't do some clever tricks to store
 165         // '~T' as a '*c_void' or something like that. To solve the problem, we
 166         // cast everything to a trait (LocalData) which is then stored inside
 167         // the map.  Upon destruction of the map, all the objects will be
 168         // destroyed and the traits have enough information about them to
 169         // destroy themselves.
 170         //
 171         // Additionally, the type of the local data map must ascribe to Send, so
 172         // we do the transmute here to add the Send bound back on. This doesn't
 173         // actually matter because TLS will always own the data (until its moved
 174         // out) and we're not actually sending it to other schedulers or
 175         // anything.
 176         let newval = data.map(|d| {
 177             let d = box d as Box<LocalData>;
 178             let d: Box<LocalData + Send> = unsafe { mem::transmute(d) };
 179             (keyval, d, 0)
 180         });
 181
 182         let pos = match self.find(map) {
 183             Some((i, _, &0)) => Some(i),
 184             Some((_, _, _)) => fail!("TLS value cannot be replaced because it \
 185                                       is already borrowed"),
 186             None => map.iter().position(|entry| entry.is_none()),
 187         };
 188
 189         match pos {
 190             Some(i) => {
 191                 mem::replace(map.get_mut(i), newval).map(|(_, data, _)| {
 192                     // Move `data` into transmute to get out the memory that it
 193                     // owns, we must free it manually later.
 194                     let t: raw::TraitObject = unsafe { mem::transmute(data) };
 195                     let alloc: Box<T> = unsafe { mem::transmute(t.data) };
 196
 197                     // Now that we own `alloc`, we can just move out of it as we
 198                     // would with any other data.
 199                     *alloc
 200                 })
 201             }
 202             None => {
 203                 map.push(newval);
 204                 None
 205             }
 206         }
 207     }
 208
 209     /// Borrows a value from TLS.
 210     ///
 211     /// If `None` is returned, then this key is not present in TLS. If `Some` is
 212     /// returned, then the returned data is a smart pointer representing a new
 213     /// loan on this TLS key. While on loan, this key cannot be altered via the
 214     /// `replace` method.
 215     ///
 216     /// # Example
 217     ///
 218     /// ```
 219     /// local_data_key!(key: int)
 220     ///
 221     /// assert!(key.get().is_none());
 222     ///
 223     /// key.replace(Some(3));
 224     /// assert_eq!(*key.get().unwrap(), 3);
 225     /// ```
 226     pub fn get(&'static self) -> Option<Ref<T>> {
 227         let map = match unsafe { get_local_map() } {
 228             Some(map) => map,
 229             None => return None,
 230         };
 231
 232         self.find(map).map(|(pos, data, loan)| {
 233             *loan += 1;
 234
 235             // data was created with `~T as ~LocalData`, so we extract
 236             // pointer part of the trait, (as ~T), and then use
 237             // compiler coercions to achieve a '&' pointer.
 238             let ptr = unsafe {
 239                 let data = data as *const Box<LocalData + Send>
 240                                 as *const raw::TraitObject;
 241                 &mut *((*data).data as *mut T)
 242             };
 243             Ref { _ptr: ptr, _index: pos, _nosend: marker::NoSend, _key: self }
 244         })
 245     }
 246
 247     fn find<'a>(&'static self,
 248                 map: &'a mut Map) -> Option<(uint, &'a TLSValue, &'a mut uint)>{
 249         let key_value = key_to_key_value(self);
 250         map.mut_iter().enumerate().filter_map(|(i, entry)| {
 251             match *entry {
 252                 Some((k, ref data, ref mut loan)) if k == key_value => {
 253                     Some((i, data, loan))
 254                 }
 255                 _ => None
 256             }
 257         }).next()
 258     }
 259 }
 260
 261 impl<T: 'static> Deref<T> for Ref<T> {
 262     fn deref<'a>(&'a self) -> &'a T { self._ptr }
 263 }
 264
 265 #[unsafe_destructor]
 266 impl<T: 'static> Drop for Ref<T> {
 267     fn drop(&mut self) {
 268         let map = unsafe { get_local_map().unwrap() };
 269
 270         let (_, _, ref mut loan) = *map.get_mut(self._index).get_mut_ref();
 271         *loan -= 1;
 272     }
 273 }
 274
 275 #[cfg(test)]
 276 mod tests {
 277     use std::prelude::*;
 278     use std::gc::{Gc, GC};
 279     use super::*;
 280     use std::task;
 281
 282     #[test]
 283     fn test_tls_multitask() {
 284         static my_key: Key<String> = &Key;
 285         my_key.replace(Some("parent data".to_string()));
 286         task::spawn(proc() {
 287             // TLS shouldn't carry over.
 288             assert!(my_key.get().is_none());
 289             my_key.replace(Some("child data".to_string()));
 290             assert!(my_key.get().get_ref().as_slice() == "child data");
 291             // should be cleaned up for us
 292         });
 293
 294         // Must work multiple times
 295         assert!(my_key.get().unwrap().as_slice() == "parent data");
 296         assert!(my_key.get().unwrap().as_slice() == "parent data");
 297         assert!(my_key.get().unwrap().as_slice() == "parent data");
 298     }
 299
 300     #[test]
 301     fn test_tls_overwrite() {
 302         static my_key: Key<String> = &Key;
 303         my_key.replace(Some("first data".to_string()));
 304         my_key.replace(Some("next data".to_string())); // Shouldn't leak.
 305         assert!(my_key.get().unwrap().as_slice() == "next data");
 306     }
 307
 308     #[test]
 309     fn test_tls_pop() {
 310         static my_key: Key<String> = &Key;
 311         my_key.replace(Some("weasel".to_string()));
 312         assert!(my_key.replace(None).unwrap() == "weasel".to_string());
 313         // Pop must remove the data from the map.
 314         assert!(my_key.replace(None).is_none());
 315     }
 316
 317     #[test]
 318     fn test_tls_crust_automorestack_memorial_bug() {
 319         // This might result in a stack-canary clobber if the runtime fails to
 320         // set sp_limit to 0 when calling the cleanup extern - it might
 321         // automatically jump over to the rust stack, which causes next_c_sp
 322         // to get recorded as something within a rust stack segment. Then a
 323         // subsequent upcall (esp. for logging, think vsnprintf) would run on
 324         // a stack smaller than 1 MB.
 325         static my_key: Key<String> = &Key;
 326         task::spawn(proc() {
 327             my_key.replace(Some("hax".to_string()));
 328         });
 329     }
 330
 331     #[test]
 332     fn test_tls_multiple_types() {
 333         static str_key: Key<String> = &Key;
 334         static box_key: Key<Gc<()>> = &Key;
 335         static int_key: Key<int> = &Key;
 336         task::spawn(proc() {
 337             str_key.replace(Some("string data".to_string()));
 338             box_key.replace(Some(box(GC) ()));
 339             int_key.replace(Some(42));
 340         });
 341     }
 342
 343     #[test]
 344     fn test_tls_overwrite_multiple_types() {
 345         static str_key: Key<String> = &Key;
 346         static box_key: Key<Gc<()>> = &Key;
 347         static int_key: Key<int> = &Key;
 348         task::spawn(proc() {
 349             str_key.replace(Some("string data".to_string()));
 350             str_key.replace(Some("string data 2".to_string()));
 351             box_key.replace(Some(box(GC) ()));
 352             box_key.replace(Some(box(GC) ()));
 353             int_key.replace(Some(42));
 354             // This could cause a segfault if overwriting-destruction is done
 355             // with the crazy polymorphic transmute rather than the provided
 356             // finaliser.
 357             int_key.replace(Some(31337));
 358         });
 359     }
 360
 361     #[test]
 362     #[should_fail]
 363     fn test_tls_cleanup_on_failure() {
 364         static str_key: Key<String> = &Key;
 365         static box_key: Key<Gc<()>> = &Key;
 366         static int_key: Key<int> = &Key;
 367         str_key.replace(Some("parent data".to_string()));
 368         box_key.replace(Some(box(GC) ()));
 369         task::spawn(proc() {
 370             str_key.replace(Some("string data".to_string()));
 371             box_key.replace(Some(box(GC) ()));
 372             int_key.replace(Some(42));
 373             fail!();
 374         });
 375         // Not quite nondeterministic.
 376         int_key.replace(Some(31337));
 377         fail!();
 378     }
 379
 380     #[test]
 381     fn test_static_pointer() {
 382         static key: Key<&'static int> = &Key;
 383         static VALUE: int = 0;
 384         key.replace(Some(&VALUE));
 385     }
 386
 387     #[test]
 388     fn test_owned() {
 389         static key: Key<Box<int>> = &Key;
 390         key.replace(Some(box 1));
 391
 392         {
 393             let k1 = key.get().unwrap();
 394             let k2 = key.get().unwrap();
 395             let k3 = key.get().unwrap();
 396             assert_eq!(**k1, 1);
 397             assert_eq!(**k2, 1);
 398             assert_eq!(**k3, 1);
 399         }
 400         key.replace(Some(box 2));
 401         assert_eq!(**key.get().unwrap(), 2);
 402     }
 403
 404     #[test]
 405     fn test_same_key_type() {
 406         static key1: Key<int> = &Key;
 407         static key2: Key<int> = &Key;
 408         static key3: Key<int> = &Key;
 409         static key4: Key<int> = &Key;
 410         static key5: Key<int> = &Key;
 411         key1.replace(Some(1));
 412         key2.replace(Some(2));
 413         key3.replace(Some(3));
 414         key4.replace(Some(4));
 415         key5.replace(Some(5));
 416
 417         assert_eq!(*key1.get().unwrap(), 1);
 418         assert_eq!(*key2.get().unwrap(), 2);
 419         assert_eq!(*key3.get().unwrap(), 3);
 420         assert_eq!(*key4.get().unwrap(), 4);
 421         assert_eq!(*key5.get().unwrap(), 5);
 422     }
 423
 424     #[test]
 425     #[should_fail]
 426     fn test_nested_get_set1() {
 427         static key: Key<int> = &Key;
 428         key.replace(Some(4));
 429
 430         let _k = key.get();
 431         key.replace(Some(4));
 432     }
 433 }