      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    206
    207
    208
    209
    210
    211
    212
    213
    214
    215
    216
    217
    218
    219
    220
    221
    222
    223
    224
    225
    226
    227
    228
    229
    230
    231
    232
    233
    234
    235
    236
    237
    238
    239
    240
    241
    242
    243
    244
    245
    246
    247
    248
    249
    250
    251
    252
    253
    254
    255
    256
    257
    258
    259
    260
    261
    262
    263
    264
    265
    266
    267
    268
    269
    270
    271
    272
    273
    274
    275
    276
    277
    278
    279
    280
    281
    282
    283
    284
    285
    286
    287
    288
    289
    290
    291
    292
    293
    294
    295
    296
    297
    298
    299
    300
    301

    //! OS-based thread local storage
    //!
    //! This module provides an implementation of OS-based thread local storage,
    //! using the native OS-provided facilities (think `TlsAlloc` or
    //! `pthread_setspecific`). The interface of this differs from the other types
    //! of thread-local-storage provided in this crate in that OS-based TLS can only
    //! get/set pointers,
    //!
    //! This module also provides two flavors of TLS. One is intended for static
    //! initialization, and does not contain a `Drop` implementation to deallocate
    //! the OS-TLS key. The other is a type which does implement `Drop` and hence
    //! has a safe interface.
    //!
    //! # Usage
    //!
    //! This module should likely not be used directly unless other primitives are
    //! being built on. types such as `thread_local::spawn::Key` are likely much
    //! more useful in practice than this OS-based version which likely requires
    //! unsafe code to interoperate with.
    //!
    //! # Examples
    //!
    //! Using a dynamically allocated TLS key. Note that this key can be shared
    //! among many threads via an `Arc`.
    //!
    //! ```ignore (cannot-doctest-private-modules)
    //! let key = Key::new(None);
    //! assert!(key.get().is_null());
    //! key.set(1 as *mut u8);
    //! assert!(!key.get().is_null());
    //!
    //! drop(key); // deallocate this TLS slot.
    //! ```
    //!
    //! Sometimes a statically allocated key is either required or easier to work
    //! with, however.
    //!
    //! ```ignore (cannot-doctest-private-modules)
    //! static KEY: StaticKey = INIT;
    //!
    //! unsafe {
    //!     assert!(KEY.get().is_null());
    //!     KEY.set(1 as *mut u8);
    //! }
    //! ```

    #![allow(non_camel_case_types)]
    #![unstable(feature = "thread_local_internals", issue = "0")]
    #![allow(dead_code)] // sys isn't exported yet

    use crate::ptr;
    use crate::sync::atomic::{self, AtomicUsize, Ordering};
    use crate::sys::thread_local as imp;
    use crate::sys_common::mutex::Mutex;

    /// A type for TLS keys that are statically allocated.
    ///
    /// This type is entirely `unsafe` to use as it does not protect against
    /// use-after-deallocation or use-during-deallocation.
    ///
    /// The actual OS-TLS key is lazily allocated when this is used for the first
    /// time. The key is also deallocated when the Rust runtime exits or `destroy`
    /// is called, whichever comes first.
    ///
    /// # Examples
    ///
    /// ```ignore (cannot-doctest-private-modules)
    /// use tls::os::{StaticKey, INIT};
    ///
    /// static KEY: StaticKey = INIT;
    ///
    /// unsafe {
    ///     assert!(KEY.get().is_null());
    ///     KEY.set(1 as *mut u8);
    /// }
    /// ```
    pub struct StaticKey {
        /// Inner static TLS key (internals).
        key: AtomicUsize,
        /// Destructor for the TLS value.
        ///
        /// See `Key::new` for information about when the destructor runs and how
        /// it runs.
        dtor: Option<unsafe extern fn(*mut u8)>,
    }

    /// A type for a safely managed OS-based TLS slot.
    ///
    /// This type allocates an OS TLS key when it is initialized and will deallocate
    /// the key when it falls out of scope. When compared with `StaticKey`, this
    /// type is entirely safe to use.
    ///
    /// Implementations will likely, however, contain unsafe code as this type only
    /// operates on `*mut u8`, a raw pointer.
    ///
    /// # Examples
    ///
    /// ```ignore (cannot-doctest-private-modules)
    /// use tls::os::Key;
    ///
    /// let key = Key::new(None);
    /// assert!(key.get().is_null());
    /// key.set(1 as *mut u8);
    /// assert!(!key.get().is_null());
    ///
    /// drop(key); // deallocate this TLS slot.
    /// ```
    pub struct Key {
        key: imp::Key,
    }

    /// Constant initialization value for static TLS keys.
    ///
    /// This value specifies no destructor by default.
    pub const INIT: StaticKey = StaticKey::new(None);

    impl StaticKey {
        pub const fn new(dtor: Option<unsafe extern fn(*mut u8)>) -> StaticKey {
            StaticKey {
                key: atomic::AtomicUsize::new(0),
                dtor,
            }
        }

        /// Gets the value associated with this TLS key
        ///
        /// This will lazily allocate a TLS key from the OS if one has not already
        /// been allocated.
        #[inline]
        pub unsafe fn get(&self) -> *mut u8 { imp::get(self.key()) }

        /// Sets this TLS key to a new value.
        ///
        /// This will lazily allocate a TLS key from the OS if one has not already
        /// been allocated.
        #[inline]
        pub unsafe fn set(&self, val: *mut u8) { imp::set(self.key(), val) }

        #[inline]
        unsafe fn key(&self) -> imp::Key {
            match self.key.load(Ordering::Relaxed) {
                0 => self.lazy_init() as imp::Key,
                n => n as imp::Key
            }
        }

        unsafe fn lazy_init(&self) -> usize {
            // Currently the Windows implementation of TLS is pretty hairy, and
            // it greatly simplifies creation if we just synchronize everything.
            //
            // Additionally a 0-index of a tls key hasn't been seen on windows, so
            // we just simplify the whole branch.
            if imp::requires_synchronized_create() {
                // We never call `INIT_LOCK.init()`, so it is UB to attempt to
                // acquire this mutex reentrantly!
                static INIT_LOCK: Mutex = Mutex::new();
                let _guard = INIT_LOCK.lock();
                let mut key = self.key.load(Ordering::SeqCst);
                if key == 0 {
                    key = imp::create(self.dtor) as usize;
                    self.key.store(key, Ordering::SeqCst);
                }
                rtassert!(key != 0);
                return key
            }

            // POSIX allows the key created here to be 0, but the compare_and_swap
            // below relies on using 0 as a sentinel value to check who won the
            // race to set the shared TLS key. As far as I know, there is no
            // guaranteed value that cannot be returned as a posix_key_create key,
            // so there is no value we can initialize the inner key with to
            // prove that it has not yet been set. As such, we'll continue using a
            // value of 0, but with some gyrations to make sure we have a non-0
            // value returned from the creation routine.
            // FIXME: this is clearly a hack, and should be cleaned up.
            let key1 = imp::create(self.dtor);
            let key = if key1 != 0 {
                key1
            } else {
                let key2 = imp::create(self.dtor);
                imp::destroy(key1);
                key2
            };
            rtassert!(key != 0);
            match self.key.compare_and_swap(0, key as usize, Ordering::SeqCst) {
                // The CAS succeeded, so we've created the actual key
                0 => key as usize,
                // If someone beat us to the punch, use their key instead
                n => { imp::destroy(key); n }
            }
        }
    }

    impl Key {
        /// Creates a new managed OS TLS key.
        ///
        /// This key will be deallocated when the key falls out of scope.
        ///
        /// The argument provided is an optionally-specified destructor for the
        /// value of this TLS key. When a thread exits and the value for this key
        /// is non-null the destructor will be invoked. The TLS value will be reset
        /// to null before the destructor is invoked.
        ///
        /// Note that the destructor will not be run when the `Key` goes out of
        /// scope.
        #[inline]
        pub fn new(dtor: Option<unsafe extern fn(*mut u8)>) -> Key {
            Key { key: unsafe { imp::create(dtor) } }
        }

        /// See StaticKey::get
        #[inline]
        pub fn get(&self) -> *mut u8 {
            unsafe { imp::get(self.key) }
        }

        /// See StaticKey::set
        #[inline]
        pub fn set(&self, val: *mut u8) {
            unsafe { imp::set(self.key, val) }
        }
    }

    impl Drop for Key {
        fn drop(&mut self) {
            // Right now Windows doesn't support TLS key destruction, but this also
            // isn't used anywhere other than tests, so just leak the TLS key.
            // unsafe { imp::destroy(self.key) }
        }
    }

    pub unsafe fn register_dtor_fallback(t: *mut u8,
                                         dtor: unsafe extern fn(*mut u8)) {
        // The fallback implementation uses a vanilla OS-based TLS key to track
        // the list of destructors that need to be run for this thread. The key
        // then has its own destructor which runs all the other destructors.
        //
        // The destructor for DTORS is a little special in that it has a `while`
        // loop to continuously drain the list of registered destructors. It
        // *should* be the case that this loop always terminates because we
        // provide the guarantee that a TLS key cannot be set after it is
        // flagged for destruction.

        static DTORS: StaticKey = StaticKey::new(Some(run_dtors));
        type List = Vec<(*mut u8, unsafe extern fn(*mut u8))>;
        if DTORS.get().is_null() {
            let v: Box<List> = box Vec::new();
            DTORS.set(Box::into_raw(v) as *mut u8);
        }
        let list: &mut List = &mut *(DTORS.get() as *mut List);
        list.push((t, dtor));

        unsafe extern fn run_dtors(mut ptr: *mut u8) {
            while !ptr.is_null() {
                let list: Box<List> = Box::from_raw(ptr as *mut List);
                for (ptr, dtor) in list.into_iter() {
                    dtor(ptr);
                }
                ptr = DTORS.get();
                DTORS.set(ptr::null_mut());
            }
        }
    }

    #[cfg(test)]
    mod tests {
        use super::{Key, StaticKey};

        fn assert_sync<T: Sync>() {}
        fn assert_send<T: Send>() {}

        #[test]
        fn smoke() {
            assert_sync::<Key>();
            assert_send::<Key>();

            let k1 = Key::new(None);
            let k2 = Key::new(None);
            assert!(k1.get().is_null());
            assert!(k2.get().is_null());
            k1.set(1 as *mut _);
            k2.set(2 as *mut _);
            assert_eq!(k1.get() as usize, 1);
            assert_eq!(k2.get() as usize, 2);
        }

        #[test]
        fn statik() {
            static K1: StaticKey = StaticKey::new(None);
            static K2: StaticKey = StaticKey::new(None);

            unsafe {
                assert!(K1.get().is_null());
                assert!(K2.get().is_null());
                K1.set(1 as *mut _);
                K2.set(2 as *mut _);
                assert_eq!(K1.get() as usize, 1);
                assert_eq!(K2.get() as usize, 2);
            }
        }
    }