src/libstd/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 use, declare a static variable of the type you wish to store. The
  21 initialization should be `&local_data::Key`. This is then the key to what you
  22 wish to store.
  23
  24 ~~~{.rust}
  25 use std::local_data;
  26
  27 static key_int: local_data::Key<int> = &local_data::Key;
  28 static key_vector: local_data::Key<~[int]> = &local_data::Key;
  29
  30 local_data::set(key_int, 3);
  31 local_data::get(key_int, |opt| assert_eq!(opt, Some(&3)));
  32
  33 local_data::set(key_vector, ~[4]);
  34 local_data::get(key_int, |opt| assert_eq!(opt, Some(&~[4])));
  35 ~~~
  36
  37 Casting 'Arcane Sight' reveals an overwhelming aura of Transmutation
  38 magic.
  39
  40 */
  41
  42 use prelude::*;
  43
  44 use task::local_data_priv::*;
  45
  46 #[cfg(test)] use task;
  47
  48 /**
  49  * Indexes a task-local data slot. This pointer is used for comparison to
  50  * differentiate keys from one another. The actual type `T` is not used anywhere
  51  * as a member of this type, except that it is parameterized with it to define
  52  * the type of each key's value.
  53  *
  54  * The value of each Key is of the singleton enum KeyValue. These also have the
  55  * same name as `Key` and their purpose is to take up space in the programs data
  56  * sections to ensure that each value of the `Key` type points to a unique
  57  * location.
  58  */
  59 pub type Key<T> = &'static KeyValue<T>;
  60
  61 pub enum KeyValue<T> { Key }
  62
  63 /**
  64  * Remove a task-local data value from the table, returning the
  65  * reference that was originally created to insert it.
  66  */
  67 pub fn pop<T: 'static>(key: Key<T>) -> Option<T> {
  68     unsafe { local_pop(Handle::new(), key) }
  69 }
  70
  71 /**
  72  * Retrieve a task-local data value. It will also be kept alive in the
  73  * table until explicitly removed.
  74  */
  75 pub fn get<T: 'static, U>(key: Key<T>, f: &fn(Option<&T>) -> U) -> U {
  76     unsafe { local_get(Handle::new(), key, f) }
  77 }
  78
  79 /**
  80  * Retrieve a mutable borrowed pointer to a task-local data value.
  81  */
  82 pub fn get_mut<T: 'static, U>(key: Key<T>, f: &fn(Option<&mut T>) -> U) -> U {
  83     unsafe { local_get_mut(Handle::new(), key, f) }
  84 }
  85
  86 /**
  87  * Store a value in task-local data. If this key already has a value,
  88  * that value is overwritten (and its destructor is run).
  89  */
  90 pub fn set<T: 'static>(key: Key<T>, data: T) {
  91     unsafe { local_set(Handle::new(), key, data) }
  92 }
  93
  94 /**
  95  * Modify a task-local data value. If the function returns 'None', the
  96  * data is removed (and its reference dropped).
  97  */
  98 pub fn modify<T: 'static>(key: Key<T>, f: &fn(Option<T>) -> Option<T>) {
  99     unsafe {
 100         match f(pop(::cast::unsafe_copy(&key))) {
 101             Some(next) => { set(key, next); }
 102             None => {}
 103         }
 104     }
 105 }
 106
 107 #[test]
 108 fn test_tls_multitask() {
 109     static my_key: Key<@~str> = &Key;
 110     set(my_key, @~"parent data");
 111     do task::spawn {
 112         // TLS shouldn't carry over.
 113         assert!(get(my_key, |k| k.map(|&k| *k)).is_none());
 114         set(my_key, @~"child data");
 115         assert!(*(get(my_key, |k| k.map(|&k| *k)).unwrap()) ==
 116                 ~"child data");
 117         // should be cleaned up for us
 118     }
 119     // Must work multiple times
 120     assert!(*(get(my_key, |k| k.map(|&k| *k)).unwrap()) == ~"parent data");
 121     assert!(*(get(my_key, |k| k.map(|&k| *k)).unwrap()) == ~"parent data");
 122     assert!(*(get(my_key, |k| k.map(|&k| *k)).unwrap()) == ~"parent data");
 123 }
 124
 125 #[test]
 126 fn test_tls_overwrite() {
 127     static my_key: Key<@~str> = &Key;
 128     set(my_key, @~"first data");
 129     set(my_key, @~"next data"); // Shouldn't leak.
 130     assert!(*(get(my_key, |k| k.map(|&k| *k)).unwrap()) == ~"next data");
 131 }
 132
 133 #[test]
 134 fn test_tls_pop() {
 135     static my_key: Key<@~str> = &Key;
 136     set(my_key, @~"weasel");
 137     assert!(*(pop(my_key).unwrap()) == ~"weasel");
 138     // Pop must remove the data from the map.
 139     assert!(pop(my_key).is_none());
 140 }
 141
 142 #[test]
 143 fn test_tls_modify() {
 144     static my_key: Key<@~str> = &Key;
 145     modify(my_key, |data| {
 146         match data {
 147             Some(@ref val) => fail!("unwelcome value: %s", *val),
 148             None           => Some(@~"first data")
 149         }
 150     });
 151     modify(my_key, |data| {
 152         match data {
 153             Some(@~"first data") => Some(@~"next data"),
 154             Some(@ref val)       => fail!("wrong value: %s", *val),
 155             None                 => fail!("missing value")
 156         }
 157     });
 158     assert!(*(pop(my_key).unwrap()) == ~"next data");
 159 }
 160
 161 #[test]
 162 fn test_tls_crust_automorestack_memorial_bug() {
 163     // This might result in a stack-canary clobber if the runtime fails to
 164     // set sp_limit to 0 when calling the cleanup extern - it might
 165     // automatically jump over to the rust stack, which causes next_c_sp
 166     // to get recorded as something within a rust stack segment. Then a
 167     // subsequent upcall (esp. for logging, think vsnprintf) would run on
 168     // a stack smaller than 1 MB.
 169     static my_key: Key<@~str> = &Key;
 170     do task::spawn {
 171         set(my_key, @~"hax");
 172     }
 173 }
 174
 175 #[test]
 176 fn test_tls_multiple_types() {
 177     static str_key: Key<@~str> = &Key;
 178     static box_key: Key<@@()> = &Key;
 179     static int_key: Key<@int> = &Key;
 180     do task::spawn {
 181         set(str_key, @~"string data");
 182         set(box_key, @@());
 183         set(int_key, @42);
 184     }
 185 }
 186
 187 #[test]
 188 fn test_tls_overwrite_multiple_types() {
 189     static str_key: Key<@~str> = &Key;
 190     static box_key: Key<@@()> = &Key;
 191     static int_key: Key<@int> = &Key;
 192     do task::spawn {
 193         set(str_key, @~"string data");
 194         set(int_key, @42);
 195         // This could cause a segfault if overwriting-destruction is done
 196         // with the crazy polymorphic transmute rather than the provided
 197         // finaliser.
 198         set(int_key, @31337);
 199     }
 200 }
 201
 202 #[test]
 203 #[should_fail]
 204 #[ignore(cfg(windows))]
 205 fn test_tls_cleanup_on_failure() {
 206     static str_key: Key<@~str> = &Key;
 207     static box_key: Key<@@()> = &Key;
 208     static int_key: Key<@int> = &Key;
 209     set(str_key, @~"parent data");
 210     set(box_key, @@());
 211     do task::spawn {
 212         // spawn_linked
 213         set(str_key, @~"string data");
 214         set(box_key, @@());
 215         set(int_key, @42);
 216         fail!();
 217     }
 218     // Not quite nondeterministic.
 219     set(int_key, @31337);
 220     fail!();
 221 }
 222
 223 #[test]
 224 fn test_static_pointer() {
 225     static key: Key<@&'static int> = &Key;
 226     static VALUE: int = 0;
 227     let v: @&'static int = @&VALUE;
 228     set(key, v);
 229 }
 230
 231 #[test]
 232 fn test_owned() {
 233     static key: Key<~int> = &Key;
 234     set(key, ~1);
 235
 236     do get(key) |v| {
 237         do get(key) |v| {
 238             do get(key) |v| {
 239                 assert_eq!(**v.unwrap(), 1);
 240             }
 241             assert_eq!(**v.unwrap(), 1);
 242         }
 243         assert_eq!(**v.unwrap(), 1);
 244     }
 245     set(key, ~2);
 246     do get(key) |v| {
 247         assert_eq!(**v.unwrap(), 2);
 248     }
 249 }
 250
 251 #[test]
 252 fn test_get_mut() {
 253     static key: Key<int> = &Key;
 254     set(key, 1);
 255
 256     do get_mut(key) |v| {
 257         *v.unwrap() = 2;
 258     }
 259
 260     do get(key) |v| {
 261         assert_eq!(*v.unwrap(), 2);
 262     }
 263 }
 264
 265 #[test]
 266 fn test_same_key_type() {
 267     static key1: Key<int> = &Key;
 268     static key2: Key<int> = &Key;
 269     static key3: Key<int> = &Key;
 270     static key4: Key<int> = &Key;
 271     static key5: Key<int> = &Key;
 272     set(key1, 1);
 273     set(key2, 2);
 274     set(key3, 3);
 275     set(key4, 4);
 276     set(key5, 5);
 277
 278     get(key1, |x| assert_eq!(*x.unwrap(), 1));
 279     get(key2, |x| assert_eq!(*x.unwrap(), 2));
 280     get(key3, |x| assert_eq!(*x.unwrap(), 3));
 281     get(key4, |x| assert_eq!(*x.unwrap(), 4));
 282     get(key5, |x| assert_eq!(*x.unwrap(), 5));
 283 }
 284
 285 #[test]
 286 #[should_fail]
 287 fn test_nested_get_set1() {
 288     static key: Key<int> = &Key;
 289     set(key, 4);
 290     do get(key) |_| {
 291         set(key, 4);
 292     }
 293 }
 294
 295 #[test]
 296 #[should_fail]
 297 fn test_nested_get_mut2() {
 298     static key: Key<int> = &Key;
 299     set(key, 4);
 300     do get(key) |_| {
 301         get_mut(key, |_| {})
 302     }
 303 }
 304
 305 #[test]
 306 #[should_fail]
 307 fn test_nested_get_mut3() {
 308     static key: Key<int> = &Key;
 309     set(key, 4);
 310     do get_mut(key) |_| {
 311         get(key, |_| {})
 312     }
 313 }
 314
 315 #[test]
 316 #[should_fail]
 317 fn test_nested_get_mut4() {
 318     static key: Key<int> = &Key;
 319     set(key, 4);
 320     do get_mut(key) |_| {
 321         get_mut(key, |_| {})
 322     }
 323 }