| 1 | // Copyright 2014-2015 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 | //! Scoped thread-local storage | 
|---|
| 12 | //! | 
|---|
| 13 | //! This module provides the ability to generate *scoped* thread-local | 
|---|
| 14 | //! variables. In this sense, scoped indicates that thread local storage | 
|---|
| 15 | //! actually stores a reference to a value, and this reference is only placed | 
|---|
| 16 | //! in storage for a scoped amount of time. | 
|---|
| 17 | //! | 
|---|
| 18 | //! There are no restrictions on what types can be placed into a scoped | 
|---|
| 19 | //! variable, but all scoped variables are initialized to the equivalent of | 
|---|
| 20 | //! null. Scoped thread local storage is useful when a value is present for a known | 
|---|
| 21 | //! period of time and it is not required to relinquish ownership of the | 
|---|
| 22 | //! contents. | 
|---|
| 23 | //! | 
|---|
| 24 | //! # Examples | 
|---|
| 25 | //! | 
|---|
| 26 | //! ``` | 
|---|
| 27 | //! #[macro_use] | 
|---|
| 28 | //! extern crate scoped_tls; | 
|---|
| 29 | //! | 
|---|
| 30 | //! scoped_thread_local!(static FOO: u32); | 
|---|
| 31 | //! | 
|---|
| 32 | //! # fn main() { | 
|---|
| 33 | //! // Initially each scoped slot is empty. | 
|---|
| 34 | //! assert!(!FOO.is_set()); | 
|---|
| 35 | //! | 
|---|
| 36 | //! // When inserting a value, the value is only in place for the duration | 
|---|
| 37 | //! // of the closure specified. | 
|---|
| 38 | //! FOO.set(&1, || { | 
|---|
| 39 | //!     FOO.with(|slot| { | 
|---|
| 40 | //!         assert_eq!(*slot, 1); | 
|---|
| 41 | //!     }); | 
|---|
| 42 | //! }); | 
|---|
| 43 | //! # } | 
|---|
| 44 | //! ``` | 
|---|
| 45 |  | 
|---|
| 46 | #![ deny(missing_docs, warnings)] | 
|---|
| 47 |  | 
|---|
| 48 | use std::cell::Cell; | 
|---|
| 49 | use std::marker; | 
|---|
| 50 | use std::thread::LocalKey; | 
|---|
| 51 |  | 
|---|
| 52 | /// The macro. See the module level documentation for the description and examples. | 
|---|
| 53 | #[ macro_export] | 
|---|
| 54 | macro_rules! scoped_thread_local { | 
|---|
| 55 | ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty) => ( | 
|---|
| 56 | $(#[$attrs])* | 
|---|
| 57 | $vis static $name: $crate::ScopedKey<$ty> = $crate::ScopedKey { | 
|---|
| 58 | inner: { | 
|---|
| 59 | ::std::thread_local!(static FOO: ::std::cell::Cell<*const ()> = const { | 
|---|
| 60 | ::std::cell::Cell::new(::std::ptr::null()) | 
|---|
| 61 | }); | 
|---|
| 62 | &FOO | 
|---|
| 63 | }, | 
|---|
| 64 | _marker: ::std::marker::PhantomData, | 
|---|
| 65 | }; | 
|---|
| 66 | ) | 
|---|
| 67 | } | 
|---|
| 68 |  | 
|---|
| 69 | /// Type representing a thread local storage key corresponding to a reference | 
|---|
| 70 | /// to the type parameter `T`. | 
|---|
| 71 | /// | 
|---|
| 72 | /// Keys are statically allocated and can contain a reference to an instance of | 
|---|
| 73 | /// type `T` scoped to a particular lifetime. Keys provides two methods, `set` | 
|---|
| 74 | /// and `with`, both of which currently use closures to control the scope of | 
|---|
| 75 | /// their contents. | 
|---|
| 76 | pub struct ScopedKey<T> { | 
|---|
| 77 | #[ doc(hidden)] | 
|---|
| 78 | pub inner: &'static LocalKey<Cell<*const ()>>, | 
|---|
| 79 | #[ doc(hidden)] | 
|---|
| 80 | pub _marker: marker::PhantomData<T>, | 
|---|
| 81 | } | 
|---|
| 82 |  | 
|---|
| 83 | unsafe impl<T> Sync for ScopedKey<T> {} | 
|---|
| 84 |  | 
|---|
| 85 | impl<T> ScopedKey<T> { | 
|---|
| 86 | /// Inserts a value into this scoped thread local storage slot for a | 
|---|
| 87 | /// duration of a closure. | 
|---|
| 88 | /// | 
|---|
| 89 | /// While `f` is running, the value `t` will be returned by `get` unless | 
|---|
| 90 | /// this function is called recursively inside of `f`. | 
|---|
| 91 | /// | 
|---|
| 92 | /// Upon return, this function will restore the previous value, if any | 
|---|
| 93 | /// was available. | 
|---|
| 94 | /// | 
|---|
| 95 | /// # Examples | 
|---|
| 96 | /// | 
|---|
| 97 | /// ``` | 
|---|
| 98 | /// #[macro_use] | 
|---|
| 99 | /// extern crate scoped_tls; | 
|---|
| 100 | /// | 
|---|
| 101 | /// scoped_thread_local!(static FOO: u32); | 
|---|
| 102 | /// | 
|---|
| 103 | /// # fn main() { | 
|---|
| 104 | /// FOO.set(&100, || { | 
|---|
| 105 | ///     let val = FOO.with(|v| *v); | 
|---|
| 106 | ///     assert_eq!(val, 100); | 
|---|
| 107 | /// | 
|---|
| 108 | ///     // set can be called recursively | 
|---|
| 109 | ///     FOO.set(&101, || { | 
|---|
| 110 | ///         // ... | 
|---|
| 111 | ///     }); | 
|---|
| 112 | /// | 
|---|
| 113 | ///     // Recursive calls restore the previous value. | 
|---|
| 114 | ///     let val = FOO.with(|v| *v); | 
|---|
| 115 | ///     assert_eq!(val, 100); | 
|---|
| 116 | /// }); | 
|---|
| 117 | /// # } | 
|---|
| 118 | /// ``` | 
|---|
| 119 | pub fn set<F, R>(&'static self, t: &T, f: F) -> R | 
|---|
| 120 | where F: FnOnce() -> R | 
|---|
| 121 | { | 
|---|
| 122 | struct Reset { | 
|---|
| 123 | key: &'static LocalKey<Cell<*const ()>>, | 
|---|
| 124 | val: *const (), | 
|---|
| 125 | } | 
|---|
| 126 | impl Drop for Reset { | 
|---|
| 127 | fn drop(&mut self) { | 
|---|
| 128 | self.key.with(|c| c.set(self.val)); | 
|---|
| 129 | } | 
|---|
| 130 | } | 
|---|
| 131 | let prev = self.inner.with(|c| { | 
|---|
| 132 | let prev = c.get(); | 
|---|
| 133 | c.set(t as *const T as *const ()); | 
|---|
| 134 | prev | 
|---|
| 135 | }); | 
|---|
| 136 | let _reset = Reset { key: self.inner, val: prev }; | 
|---|
| 137 | f() | 
|---|
| 138 | } | 
|---|
| 139 |  | 
|---|
| 140 | /// Gets a value out of this scoped variable. | 
|---|
| 141 | /// | 
|---|
| 142 | /// This function takes a closure which receives the value of this | 
|---|
| 143 | /// variable. | 
|---|
| 144 | /// | 
|---|
| 145 | /// # Panics | 
|---|
| 146 | /// | 
|---|
| 147 | /// This function will panic if `set` has not previously been called. | 
|---|
| 148 | /// | 
|---|
| 149 | /// # Examples | 
|---|
| 150 | /// | 
|---|
| 151 | /// ```no_run | 
|---|
| 152 | /// #[macro_use] | 
|---|
| 153 | /// extern crate scoped_tls; | 
|---|
| 154 | /// | 
|---|
| 155 | /// scoped_thread_local!(static FOO: u32); | 
|---|
| 156 | /// | 
|---|
| 157 | /// # fn main() { | 
|---|
| 158 | /// FOO.with(|slot| { | 
|---|
| 159 | ///     // work with `slot` | 
|---|
| 160 | /// # drop(slot); | 
|---|
| 161 | /// }); | 
|---|
| 162 | /// # } | 
|---|
| 163 | /// ``` | 
|---|
| 164 | pub fn with<F, R>(&'static self, f: F) -> R | 
|---|
| 165 | where F: FnOnce(&T) -> R | 
|---|
| 166 | { | 
|---|
| 167 | let val = self.inner.with(|c| c.get()); | 
|---|
| 168 | assert!(!val.is_null(), "cannot access a scoped thread local \ | 
|---|
| 169 |                                  variable without calling `set` first"); | 
|---|
| 170 | unsafe { | 
|---|
| 171 | f(&*(val as *const T)) | 
|---|
| 172 | } | 
|---|
| 173 | } | 
|---|
| 174 |  | 
|---|
| 175 | /// Test whether this TLS key has been `set` for the current thread. | 
|---|
| 176 | pub fn is_set(&'static self) -> bool { | 
|---|
| 177 | self.inner.with(|c| !c.get().is_null()) | 
|---|
| 178 | } | 
|---|
| 179 | } | 
|---|
| 180 |  | 
|---|
| 181 | #[ cfg(test)] | 
|---|
| 182 | mod tests { | 
|---|
| 183 | use std::cell::Cell; | 
|---|
| 184 | use std::sync::mpsc::{channel, Sender}; | 
|---|
| 185 | use std::thread; | 
|---|
| 186 |  | 
|---|
| 187 | scoped_thread_local!(static FOO: u32); | 
|---|
| 188 |  | 
|---|
| 189 | #[ test] | 
|---|
| 190 | fn smoke() { | 
|---|
| 191 | scoped_thread_local!(static BAR: u32); | 
|---|
| 192 |  | 
|---|
| 193 | assert!(!BAR.is_set()); | 
|---|
| 194 | BAR.set(&1, || { | 
|---|
| 195 | assert!(BAR.is_set()); | 
|---|
| 196 | BAR.with(|slot| { | 
|---|
| 197 | assert_eq!(*slot, 1); | 
|---|
| 198 | }); | 
|---|
| 199 | }); | 
|---|
| 200 | assert!(!BAR.is_set()); | 
|---|
| 201 | } | 
|---|
| 202 |  | 
|---|
| 203 | #[ test] | 
|---|
| 204 | fn cell_allowed() { | 
|---|
| 205 | scoped_thread_local!(static BAR: Cell<u32>); | 
|---|
| 206 |  | 
|---|
| 207 | BAR.set(&Cell::new(1), || { | 
|---|
| 208 | BAR.with(|slot| { | 
|---|
| 209 | assert_eq!(slot.get(), 1); | 
|---|
| 210 | }); | 
|---|
| 211 | }); | 
|---|
| 212 | } | 
|---|
| 213 |  | 
|---|
| 214 | #[ test] | 
|---|
| 215 | fn scope_item_allowed() { | 
|---|
| 216 | assert!(!FOO.is_set()); | 
|---|
| 217 | FOO.set(&1, || { | 
|---|
| 218 | assert!(FOO.is_set()); | 
|---|
| 219 | FOO.with(|slot| { | 
|---|
| 220 | assert_eq!(*slot, 1); | 
|---|
| 221 | }); | 
|---|
| 222 | }); | 
|---|
| 223 | assert!(!FOO.is_set()); | 
|---|
| 224 | } | 
|---|
| 225 |  | 
|---|
| 226 | #[ test] | 
|---|
| 227 | fn panic_resets() { | 
|---|
| 228 | struct Check(Sender<u32>); | 
|---|
| 229 | impl Drop for Check { | 
|---|
| 230 | fn drop(&mut self) { | 
|---|
| 231 | FOO.with(|r| { | 
|---|
| 232 | self.0.send(*r).unwrap(); | 
|---|
| 233 | }) | 
|---|
| 234 | } | 
|---|
| 235 | } | 
|---|
| 236 |  | 
|---|
| 237 | let (tx, rx) = channel(); | 
|---|
| 238 | let t = thread::spawn(|| { | 
|---|
| 239 | FOO.set(&1, || { | 
|---|
| 240 | let _r = Check(tx); | 
|---|
| 241 |  | 
|---|
| 242 | FOO.set(&2, || { | 
|---|
| 243 | panic!() | 
|---|
| 244 | }); | 
|---|
| 245 | }); | 
|---|
| 246 | }); | 
|---|
| 247 |  | 
|---|
| 248 | assert_eq!(rx.recv().unwrap(), 1); | 
|---|
| 249 | assert!(t.join().is_err()); | 
|---|
| 250 | } | 
|---|
| 251 |  | 
|---|
| 252 | #[ test] | 
|---|
| 253 | fn attrs_allowed() { | 
|---|
| 254 | scoped_thread_local!( | 
|---|
| 255 | /// Docs | 
|---|
| 256 | static BAZ: u32 | 
|---|
| 257 | ); | 
|---|
| 258 |  | 
|---|
| 259 | scoped_thread_local!( | 
|---|
| 260 | #[allow(non_upper_case_globals)] | 
|---|
| 261 | static quux: u32 | 
|---|
| 262 | ); | 
|---|
| 263 |  | 
|---|
| 264 | let _ = BAZ; | 
|---|
| 265 | let _ = quux; | 
|---|
| 266 | } | 
|---|
| 267 | } | 
|---|
| 268 |  | 
|---|