library/std/src/thread/local.rs

   1 //! Thread local storage
   2
   3 #![unstable(feature = "thread_local_internals", issue = "none")]
   4
   5 #[cfg(all(test, not(target_os = "emscripten")))]
   6 mod tests;
   7
   8 #[cfg(test)]
   9 mod dynamic_tests;
  10
  11 use crate::cell::{Cell, RefCell};
  12 use crate::error::Error;
  13 use crate::fmt;
  14
  15 /// A thread local storage key which owns its contents.
  16 ///
  17 /// This key uses the fastest possible implementation available to it for the
  18 /// target platform. It is instantiated with the [`thread_local!`] macro and the
  19 /// primary method is the [`with`] method.
  20 ///
  21 /// The [`with`] method yields a reference to the contained value which cannot be
  22 /// sent across threads or escape the given closure.
  23 ///
  24 /// [`thread_local!`]: crate::thread_local
  25 ///
  26 /// # Initialization and Destruction
  27 ///
  28 /// Initialization is dynamically performed on the first call to [`with`]
  29 /// within a thread, and values that implement [`Drop`] get destructed when a
  30 /// thread exits. Some caveats apply, which are explained below.
  31 ///
  32 /// A `LocalKey`'s initializer cannot recursively depend on itself, and using
  33 /// a `LocalKey` in this way will cause the initializer to infinitely recurse
  34 /// on the first call to `with`.
  35 ///
  36 /// # Examples
  37 ///
  38 /// ```
  39 /// use std::cell::RefCell;
  40 /// use std::thread;
  41 ///
  42 /// thread_local!(static FOO: RefCell<u32> = RefCell::new(1));
  43 ///
  44 /// FOO.with(|f| {
  45 ///     assert_eq!(*f.borrow(), 1);
  46 ///     *f.borrow_mut() = 2;
  47 /// });
  48 ///
  49 /// // each thread starts out with the initial value of 1
  50 /// let t = thread::spawn(move|| {
  51 ///     FOO.with(|f| {
  52 ///         assert_eq!(*f.borrow(), 1);
  53 ///         *f.borrow_mut() = 3;
  54 ///     });
  55 /// });
  56 ///
  57 /// // wait for the thread to complete and bail out on panic
  58 /// t.join().unwrap();
  59 ///
  60 /// // we retain our original value of 2 despite the child thread
  61 /// FOO.with(|f| {
  62 ///     assert_eq!(*f.borrow(), 2);
  63 /// });
  64 /// ```
  65 ///
  66 /// # Platform-specific behavior
  67 ///
  68 /// Note that a "best effort" is made to ensure that destructors for types
  69 /// stored in thread local storage are run, but not all platforms can guarantee
  70 /// that destructors will be run for all types in thread local storage. For
  71 /// example, there are a number of known caveats where destructors are not run:
  72 ///
  73 /// 1. On Unix systems when pthread-based TLS is being used, destructors will
  74 ///    not be run for TLS values on the main thread when it exits. Note that the
  75 ///    application will exit immediately after the main thread exits as well.
  76 /// 2. On all platforms it's possible for TLS to re-initialize other TLS slots
  77 ///    during destruction. Some platforms ensure that this cannot happen
  78 ///    infinitely by preventing re-initialization of any slot that has been
  79 ///    destroyed, but not all platforms have this guard. Those platforms that do
  80 ///    not guard typically have a synthetic limit after which point no more
  81 ///    destructors are run.
  82 /// 3. When the process exits on Windows systems, TLS destructors may only be
  83 ///    run on the thread that causes the process to exit. This is because the
  84 ///    other threads may be forcibly terminated.
  85 ///
  86 /// ## Synchronization in thread-local destructors
  87 ///
  88 /// On Windows, synchronization operations (such as [`JoinHandle::join`]) in
  89 /// thread local destructors are prone to deadlocks and so should be avoided.
  90 /// This is because the [loader lock] is held while a destructor is run. The
  91 /// lock is acquired whenever a thread starts or exits or when a DLL is loaded
  92 /// or unloaded. Therefore these events are blocked for as long as a thread
  93 /// local destructor is running.
  94 ///
  95 /// [loader lock]: https://docs.microsoft.com/en-us/windows/win32/dlls/dynamic-link-library-best-practices
  96 /// [`JoinHandle::join`]: crate::thread::JoinHandle::join
  97 /// [`with`]: LocalKey::with
  98 #[cfg_attr(not(test), rustc_diagnostic_item = "LocalKey")]
  99 #[stable(feature = "rust1", since = "1.0.0")]
 100 pub struct LocalKey<T: 'static> {
 101     // This outer `LocalKey<T>` type is what's going to be stored in statics,
 102     // but actual data inside will sometimes be tagged with #[thread_local].
 103     // It's not valid for a true static to reference a #[thread_local] static,
 104     // so we get around that by exposing an accessor through a layer of function
 105     // indirection (this thunk).
 106     //
 107     // Note that the thunk is itself unsafe because the returned lifetime of the
 108     // slot where data lives, `'static`, is not actually valid. The lifetime
 109     // here is actually slightly shorter than the currently running thread!
 110     //
 111     // Although this is an extra layer of indirection, it should in theory be
 112     // trivially devirtualizable by LLVM because the value of `inner` never
 113     // changes and the constant should be readonly within a crate. This mainly
 114     // only runs into problems when TLS statics are exported across crates.
 115     inner: unsafe fn(Option<&mut Option<T>>) -> Option<&'static T>,
 116 }
 117
 118 #[stable(feature = "std_debug", since = "1.16.0")]
 119 impl<T: 'static> fmt::Debug for LocalKey<T> {
 120     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 121         f.debug_struct("LocalKey").finish_non_exhaustive()
 122     }
 123 }
 124
 125 /// Declare a new thread local storage key of type [`std::thread::LocalKey`].
 126 ///
 127 /// # Syntax
 128 ///
 129 /// The macro wraps any number of static declarations and makes them thread local.
 130 /// Publicity and attributes for each static are allowed. Example:
 131 ///
 132 /// ```
 133 /// use std::cell::RefCell;
 134 /// thread_local! {
 135 ///     pub static FOO: RefCell<u32> = RefCell::new(1);
 136 ///
 137 ///     #[allow(unused)]
 138 ///     static BAR: RefCell<f32> = RefCell::new(1.0);
 139 /// }
 140 /// # fn main() {}
 141 /// ```
 142 ///
 143 /// See [`LocalKey` documentation][`std::thread::LocalKey`] for more
 144 /// information.
 145 ///
 146 /// [`std::thread::LocalKey`]: crate::thread::LocalKey
 147 #[macro_export]
 148 #[stable(feature = "rust1", since = "1.0.0")]
 149 #[cfg_attr(not(test), rustc_diagnostic_item = "thread_local_macro")]
 150 #[allow_internal_unstable(thread_local_internals)]
 151 macro_rules! thread_local {
 152     // empty (base case for the recursion)
 153     () => {};
 154
 155     ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty = const { $init:expr }; $($rest:tt)*) => (
 156         $crate::__thread_local_inner!($(#[$attr])* $vis $name, $t, const $init);
 157         $crate::thread_local!($($rest)*);
 158     );
 159
 160     ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty = const { $init:expr }) => (
 161         $crate::__thread_local_inner!($(#[$attr])* $vis $name, $t, const $init);
 162     );
 163
 164     // process multiple declarations
 165     ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty = $init:expr; $($rest:tt)*) => (
 166         $crate::__thread_local_inner!($(#[$attr])* $vis $name, $t, $init);
 167         $crate::thread_local!($($rest)*);
 168     );
 169
 170     // handle a single declaration
 171     ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty = $init:expr) => (
 172         $crate::__thread_local_inner!($(#[$attr])* $vis $name, $t, $init);
 173     );
 174 }
 175
 176 #[doc(hidden)]
 177 #[unstable(feature = "thread_local_internals", reason = "should not be necessary", issue = "none")]
 178 #[macro_export]
 179 #[allow_internal_unstable(thread_local_internals, cfg_target_thread_local, thread_local)]
 180 #[allow_internal_unsafe]
 181 macro_rules! __thread_local_inner {
 182     // used to generate the `LocalKey` value for const-initialized thread locals
 183     (@key $t:ty, const $init:expr) => {{
 184         #[cfg_attr(not(windows), inline)] // see comments below
 185         #[deny(unsafe_op_in_unsafe_fn)]
 186         unsafe fn __getit(
 187             _init: $crate::option::Option<&mut $crate::option::Option<$t>>,
 188         ) -> $crate::option::Option<&'static $t> {
 189             const INIT_EXPR: $t = $init;
 190
 191             // wasm without atomics maps directly to `static mut`, and dtors
 192             // aren't implemented because thread dtors aren't really a thing
 193             // on wasm right now
 194             //
 195             // FIXME(#84224) this should come after the `target_thread_local`
 196             // block.
 197             #[cfg(all(target_family = "wasm", not(target_feature = "atomics")))]
 198             {
 199                 static mut VAL: $t = INIT_EXPR;
 200                 unsafe { $crate::option::Option::Some(&VAL) }
 201             }
 202
 203             // If the platform has support for `#[thread_local]`, use it.
 204             #[cfg(all(
 205                 target_thread_local,
 206                 not(all(target_family = "wasm", not(target_feature = "atomics"))),
 207             ))]
 208             {
 209                 #[thread_local]
 210                 static mut VAL: $t = INIT_EXPR;
 211
 212                 // If a dtor isn't needed we can do something "very raw" and
 213                 // just get going.
 214                 if !$crate::mem::needs_drop::<$t>() {
 215                     unsafe {
 216                         return $crate::option::Option::Some(&VAL)
 217                     }
 218                 }
 219
 220                 // 0 == dtor not registered
 221                 // 1 == dtor registered, dtor not run
 222                 // 2 == dtor registered and is running or has run
 223                 #[thread_local]
 224                 static mut STATE: $crate::primitive::u8 = 0;
 225
 226                 unsafe extern "C" fn destroy(ptr: *mut $crate::primitive::u8) {
 227                     let ptr = ptr as *mut $t;
 228
 229                     unsafe {
 230                         $crate::debug_assert_eq!(STATE, 1);
 231                         STATE = 2;
 232                         $crate::ptr::drop_in_place(ptr);
 233                     }
 234                 }
 235
 236                 unsafe {
 237                     match STATE {
 238                         // 0 == we haven't registered a destructor, so do
 239                         //   so now.
 240                         0 => {
 241                             $crate::thread::__FastLocalKeyInner::<$t>::register_dtor(
 242                                 $crate::ptr::addr_of_mut!(VAL) as *mut $crate::primitive::u8,
 243                                 destroy,
 244                             );
 245                             STATE = 1;
 246                             $crate::option::Option::Some(&VAL)
 247                         }
 248                         // 1 == the destructor is registered and the value
 249                         //   is valid, so return the pointer.
 250                         1 => $crate::option::Option::Some(&VAL),
 251                         // otherwise the destructor has already run, so we
 252                         // can't give access.
 253                         _ => $crate::option::Option::None,
 254                     }
 255                 }
 256             }
 257
 258             // On platforms without `#[thread_local]` we fall back to the
 259             // same implementation as below for os thread locals.
 260             #[cfg(all(
 261                 not(target_thread_local),
 262                 not(all(target_family = "wasm", not(target_feature = "atomics"))),
 263             ))]
 264             {
 265                 #[inline]
 266                 const fn __init() -> $t { INIT_EXPR }
 267                 static __KEY: $crate::thread::__OsLocalKeyInner<$t> =
 268                     $crate::thread::__OsLocalKeyInner::new();
 269                 #[allow(unused_unsafe)]
 270                 unsafe {
 271                     __KEY.get(move || {
 272                         if let $crate::option::Option::Some(init) = _init {
 273                             if let $crate::option::Option::Some(value) = init.take() {
 274                                 return value;
 275                             } else if $crate::cfg!(debug_assertions) {
 276                                 $crate::unreachable!("missing initial value");
 277                             }
 278                         }
 279                         __init()
 280                     })
 281                 }
 282             }
 283         }
 284
 285         unsafe {
 286             $crate::thread::LocalKey::new(__getit)
 287         }
 288     }};
 289
 290     // used to generate the `LocalKey` value for `thread_local!`
 291     (@key $t:ty, $init:expr) => {
 292         {
 293             #[inline]
 294             fn __init() -> $t { $init }
 295
 296             // When reading this function you might ask "why is this inlined
 297             // everywhere other than Windows?", and that's a very reasonable
 298             // question to ask. The short story is that it segfaults rustc if
 299             // this function is inlined. The longer story is that Windows looks
 300             // to not support `extern` references to thread locals across DLL
 301             // boundaries. This appears to at least not be supported in the ABI
 302             // that LLVM implements.
 303             //
 304             // Because of this we never inline on Windows, but we do inline on
 305             // other platforms (where external references to thread locals
 306             // across DLLs are supported). A better fix for this would be to
 307             // inline this function on Windows, but only for "statically linked"
 308             // components. For example if two separately compiled rlibs end up
 309             // getting linked into a DLL then it's fine to inline this function
 310             // across that boundary. It's only not fine to inline this function
 311             // across a DLL boundary. Unfortunately rustc doesn't currently
 312             // have this sort of logic available in an attribute, and it's not
 313             // clear that rustc is even equipped to answer this (it's more of a
 314             // Cargo question kinda). This means that, unfortunately, Windows
 315             // gets the pessimistic path for now where it's never inlined.
 316             //
 317             // The issue of "should enable on Windows sometimes" is #84933
 318             #[cfg_attr(not(windows), inline)]
 319             unsafe fn __getit(
 320                 init: $crate::option::Option<&mut $crate::option::Option<$t>>,
 321             ) -> $crate::option::Option<&'static $t> {
 322                 #[cfg(all(target_family = "wasm", not(target_feature = "atomics")))]
 323                 static __KEY: $crate::thread::__StaticLocalKeyInner<$t> =
 324                     $crate::thread::__StaticLocalKeyInner::new();
 325
 326                 #[thread_local]
 327                 #[cfg(all(
 328                     target_thread_local,
 329                     not(all(target_family = "wasm", not(target_feature = "atomics"))),
 330                 ))]
 331                 static __KEY: $crate::thread::__FastLocalKeyInner<$t> =
 332                     $crate::thread::__FastLocalKeyInner::new();
 333
 334                 #[cfg(all(
 335                     not(target_thread_local),
 336                     not(all(target_family = "wasm", not(target_feature = "atomics"))),
 337                 ))]
 338                 static __KEY: $crate::thread::__OsLocalKeyInner<$t> =
 339                     $crate::thread::__OsLocalKeyInner::new();
 340
 341                 // FIXME: remove the #[allow(...)] marker when macros don't
 342                 // raise warning for missing/extraneous unsafe blocks anymore.
 343                 // See https://github.com/rust-lang/rust/issues/74838.
 344                 #[allow(unused_unsafe)]
 345                 unsafe {
 346                     __KEY.get(move || {
 347                         if let $crate::option::Option::Some(init) = init {
 348                             if let $crate::option::Option::Some(value) = init.take() {
 349                                 return value;
 350                             } else if $crate::cfg!(debug_assertions) {
 351                                 $crate::unreachable!("missing default value");
 352                             }
 353                         }
 354                         __init()
 355                     })
 356                 }
 357             }
 358
 359             unsafe {
 360                 $crate::thread::LocalKey::new(__getit)
 361             }
 362         }
 363     };
 364     ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty, $($init:tt)*) => {
 365         $(#[$attr])* $vis const $name: $crate::thread::LocalKey<$t> =
 366             $crate::__thread_local_inner!(@key $t, $($init)*);
 367     }
 368 }
 369
 370 /// An error returned by [`LocalKey::try_with`](struct.LocalKey.html#method.try_with).
 371 #[stable(feature = "thread_local_try_with", since = "1.26.0")]
 372 #[non_exhaustive]
 373 #[derive(Clone, Copy, Eq, PartialEq)]
 374 pub struct AccessError;
 375
 376 #[stable(feature = "thread_local_try_with", since = "1.26.0")]
 377 impl fmt::Debug for AccessError {
 378     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 379         f.debug_struct("AccessError").finish()
 380     }
 381 }
 382
 383 #[stable(feature = "thread_local_try_with", since = "1.26.0")]
 384 impl fmt::Display for AccessError {
 385     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 386         fmt::Display::fmt("already destroyed", f)
 387     }
 388 }
 389
 390 #[stable(feature = "thread_local_try_with", since = "1.26.0")]
 391 impl Error for AccessError {}
 392
 393 impl<T: 'static> LocalKey<T> {
 394     #[doc(hidden)]
 395     #[unstable(
 396         feature = "thread_local_internals",
 397         reason = "recently added to create a key",
 398         issue = "none"
 399     )]
 400     #[rustc_const_unstable(feature = "thread_local_internals", issue = "none")]
 401     pub const unsafe fn new(
 402         inner: unsafe fn(Option<&mut Option<T>>) -> Option<&'static T>,
 403     ) -> LocalKey<T> {
 404         LocalKey { inner }
 405     }
 406
 407     /// Acquires a reference to the value in this TLS key.
 408     ///
 409     /// This will lazily initialize the value if this thread has not referenced
 410     /// this key yet.
 411     ///
 412     /// # Panics
 413     ///
 414     /// This function will `panic!()` if the key currently has its
 415     /// destructor running, and it **may** panic if the destructor has
 416     /// previously been run for this thread.
 417     #[stable(feature = "rust1", since = "1.0.0")]
 418     pub fn with<F, R>(&'static self, f: F) -> R
 419     where
 420         F: FnOnce(&T) -> R,
 421     {
 422         self.try_with(f).expect(
 423             "cannot access a Thread Local Storage value \
 424              during or after destruction",
 425         )
 426     }
 427
 428     /// Acquires a reference to the value in this TLS key.
 429     ///
 430     /// This will lazily initialize the value if this thread has not referenced
 431     /// this key yet. If the key has been destroyed (which may happen if this is called
 432     /// in a destructor), this function will return an [`AccessError`].
 433     ///
 434     /// # Panics
 435     ///
 436     /// This function will still `panic!()` if the key is uninitialized and the
 437     /// key's initializer panics.
 438     #[stable(feature = "thread_local_try_with", since = "1.26.0")]
 439     #[inline]
 440     pub fn try_with<F, R>(&'static self, f: F) -> Result<R, AccessError>
 441     where
 442         F: FnOnce(&T) -> R,
 443     {
 444         unsafe {
 445             let thread_local = (self.inner)(None).ok_or(AccessError)?;
 446             Ok(f(thread_local))
 447         }
 448     }
 449
 450     /// Acquires a reference to the value in this TLS key, initializing it with
 451     /// `init` if it wasn't already initialized on this thread.
 452     ///
 453     /// If `init` was used to initialize the thread local variable, `None` is
 454     /// passed as the first argument to `f`. If it was already initialized,
 455     /// `Some(init)` is passed to `f`.
 456     ///
 457     /// # Panics
 458     ///
 459     /// This function will panic if the key currently has its destructor
 460     /// running, and it **may** panic if the destructor has previously been run
 461     /// for this thread.
 462     fn initialize_with<F, R>(&'static self, init: T, f: F) -> R
 463     where
 464         F: FnOnce(Option<T>, &T) -> R,
 465     {
 466         unsafe {
 467             let mut init = Some(init);
 468             let reference = (self.inner)(Some(&mut init)).expect(
 469                 "cannot access a Thread Local Storage value \
 470                  during or after destruction",
 471             );
 472             f(init, reference)
 473         }
 474     }
 475 }
 476
 477 impl<T: 'static> LocalKey<Cell<T>> {
 478     /// Sets or initializes the contained value.
 479     ///
 480     /// Unlike the other methods, this will *not* run the lazy initializer of
 481     /// the thread local. Instead, it will be directly initialized with the
 482     /// given value if it wasn't initialized yet.
 483     ///
 484     /// # Panics
 485     ///
 486     /// Panics if the key currently has its destructor running,
 487     /// and it **may** panic if the destructor has previously been run for this thread.
 488     ///
 489     /// # Examples
 490     ///
 491     /// ```
 492     /// #![feature(local_key_cell_methods)]
 493     /// use std::cell::Cell;
 494     ///
 495     /// thread_local! {
 496     ///     static X: Cell<i32> = panic!("!");
 497     /// }
 498     ///
 499     /// // Calling X.get() here would result in a panic.
 500     ///
 501     /// X.set(123); // But X.set() is fine, as it skips the initializer above.
 502     ///
 503     /// assert_eq!(X.get(), 123);
 504     /// ```
 505     #[unstable(feature = "local_key_cell_methods", issue = "92122")]
 506     pub fn set(&'static self, value: T) {
 507         self.initialize_with(Cell::new(value), |value, cell| {
 508             if let Some(value) = value {
 509                 // The cell was already initialized, so `value` wasn't used to
 510                 // initialize it. So we overwrite the current value with the
 511                 // new one instead.
 512                 cell.set(value.into_inner());
 513             }
 514         });
 515     }
 516
 517     /// Returns a copy of the contained value.
 518     ///
 519     /// This will lazily initialize the value if this thread has not referenced
 520     /// this key yet.
 521     ///
 522     /// # Panics
 523     ///
 524     /// Panics if the key currently has its destructor running,
 525     /// and it **may** panic if the destructor has previously been run for this thread.
 526     ///
 527     /// # Examples
 528     ///
 529     /// ```
 530     /// #![feature(local_key_cell_methods)]
 531     /// use std::cell::Cell;
 532     ///
 533     /// thread_local! {
 534     ///     static X: Cell<i32> = Cell::new(1);
 535     /// }
 536     ///
 537     /// assert_eq!(X.get(), 1);
 538     /// ```
 539     #[unstable(feature = "local_key_cell_methods", issue = "92122")]
 540     pub fn get(&'static self) -> T
 541     where
 542         T: Copy,
 543     {
 544         self.with(|cell| cell.get())
 545     }
 546
 547     /// Takes the contained value, leaving `Default::default()` in its place.
 548     ///
 549     /// This will lazily initialize the value if this thread has not referenced
 550     /// this key yet.
 551     ///
 552     /// # Panics
 553     ///
 554     /// Panics if the key currently has its destructor running,
 555     /// and it **may** panic if the destructor has previously been run for this thread.
 556     ///
 557     /// # Examples
 558     ///
 559     /// ```
 560     /// #![feature(local_key_cell_methods)]
 561     /// use std::cell::Cell;
 562     ///
 563     /// thread_local! {
 564     ///     static X: Cell<Option<i32>> = Cell::new(Some(1));
 565     /// }
 566     ///
 567     /// assert_eq!(X.take(), Some(1));
 568     /// assert_eq!(X.take(), None);
 569     /// ```
 570     #[unstable(feature = "local_key_cell_methods", issue = "92122")]
 571     pub fn take(&'static self) -> T
 572     where
 573         T: Default,
 574     {
 575         self.with(|cell| cell.take())
 576     }
 577
 578     /// Replaces the contained value, returning the old value.
 579     ///
 580     /// This will lazily initialize the value if this thread has not referenced
 581     /// this key yet.
 582     ///
 583     /// # Panics
 584     ///
 585     /// Panics if the key currently has its destructor running,
 586     /// and it **may** panic if the destructor has previously been run for this thread.
 587     ///
 588     /// # Examples
 589     ///
 590     /// ```
 591     /// #![feature(local_key_cell_methods)]
 592     /// use std::cell::Cell;
 593     ///
 594     /// thread_local! {
 595     ///     static X: Cell<i32> = Cell::new(1);
 596     /// }
 597     ///
 598     /// assert_eq!(X.replace(2), 1);
 599     /// assert_eq!(X.replace(3), 2);
 600     /// ```
 601     #[unstable(feature = "local_key_cell_methods", issue = "92122")]
 602     pub fn replace(&'static self, value: T) -> T {
 603         self.with(|cell| cell.replace(value))
 604     }
 605 }
 606
 607 impl<T: 'static> LocalKey<RefCell<T>> {
 608     /// Acquires a reference to the contained value.
 609     ///
 610     /// This will lazily initialize the value if this thread has not referenced
 611     /// this key yet.
 612     ///
 613     /// # Panics
 614     ///
 615     /// Panics if the value is currently mutably borrowed.
 616     ///
 617     /// Panics if the key currently has its destructor running,
 618     /// and it **may** panic if the destructor has previously been run for this thread.
 619     ///
 620     /// # Example
 621     ///
 622     /// ```
 623     /// #![feature(local_key_cell_methods)]
 624     /// use std::cell::RefCell;
 625     ///
 626     /// thread_local! {
 627     ///     static X: RefCell<Vec<i32>> = RefCell::new(Vec::new());
 628     /// }
 629     ///
 630     /// X.with_borrow(|v| assert!(v.is_empty()));
 631     /// ```
 632     #[unstable(feature = "local_key_cell_methods", issue = "92122")]
 633     pub fn with_borrow<F, R>(&'static self, f: F) -> R
 634     where
 635         F: FnOnce(&T) -> R,
 636     {
 637         self.with(|cell| f(&cell.borrow()))
 638     }
 639
 640     /// Acquires a mutable reference to the contained value.
 641     ///
 642     /// This will lazily initialize the value if this thread has not referenced
 643     /// this key yet.
 644     ///
 645     /// # Panics
 646     ///
 647     /// Panics if the value is currently borrowed.
 648     ///
 649     /// Panics if the key currently has its destructor running,
 650     /// and it **may** panic if the destructor has previously been run for this thread.
 651     ///
 652     /// # Example
 653     ///
 654     /// ```
 655     /// #![feature(local_key_cell_methods)]
 656     /// use std::cell::RefCell;
 657     ///
 658     /// thread_local! {
 659     ///     static X: RefCell<Vec<i32>> = RefCell::new(Vec::new());
 660     /// }
 661     ///
 662     /// X.with_borrow_mut(|v| v.push(1));
 663     ///
 664     /// X.with_borrow(|v| assert_eq!(*v, vec![1]));
 665     /// ```
 666     #[unstable(feature = "local_key_cell_methods", issue = "92122")]
 667     pub fn with_borrow_mut<F, R>(&'static self, f: F) -> R
 668     where
 669         F: FnOnce(&mut T) -> R,
 670     {
 671         self.with(|cell| f(&mut cell.borrow_mut()))
 672     }
 673
 674     /// Sets or initializes the contained value.
 675     ///
 676     /// Unlike the other methods, this will *not* run the lazy initializer of
 677     /// the thread local. Instead, it will be directly initialized with the
 678     /// given value if it wasn't initialized yet.
 679     ///
 680     /// # Panics
 681     ///
 682     /// Panics if the value is currently borrowed.
 683     ///
 684     /// Panics if the key currently has its destructor running,
 685     /// and it **may** panic if the destructor has previously been run for this thread.
 686     ///
 687     /// # Examples
 688     ///
 689     /// ```
 690     /// #![feature(local_key_cell_methods)]
 691     /// use std::cell::RefCell;
 692     ///
 693     /// thread_local! {
 694     ///     static X: RefCell<Vec<i32>> = panic!("!");
 695     /// }
 696     ///
 697     /// // Calling X.with() here would result in a panic.
 698     ///
 699     /// X.set(vec![1, 2, 3]); // But X.set() is fine, as it skips the initializer above.
 700     ///
 701     /// X.with_borrow(|v| assert_eq!(*v, vec![1, 2, 3]));
 702     /// ```
 703     #[unstable(feature = "local_key_cell_methods", issue = "92122")]
 704     pub fn set(&'static self, value: T) {
 705         self.initialize_with(RefCell::new(value), |value, cell| {
 706             if let Some(value) = value {
 707                 // The cell was already initialized, so `value` wasn't used to
 708                 // initialize it. So we overwrite the current value with the
 709                 // new one instead.
 710                 *cell.borrow_mut() = value.into_inner();
 711             }
 712         });
 713     }
 714
 715     /// Takes the contained value, leaving `Default::default()` in its place.
 716     ///
 717     /// This will lazily initialize the value if this thread has not referenced
 718     /// this key yet.
 719     ///
 720     /// # Panics
 721     ///
 722     /// Panics if the value is currently borrowed.
 723     ///
 724     /// Panics if the key currently has its destructor running,
 725     /// and it **may** panic if the destructor has previously been run for this thread.
 726     ///
 727     /// # Examples
 728     ///
 729     /// ```
 730     /// #![feature(local_key_cell_methods)]
 731     /// use std::cell::RefCell;
 732     ///
 733     /// thread_local! {
 734     ///     static X: RefCell<Vec<i32>> = RefCell::new(Vec::new());
 735     /// }
 736     ///
 737     /// X.with_borrow_mut(|v| v.push(1));
 738     ///
 739     /// let a = X.take();
 740     ///
 741     /// assert_eq!(a, vec![1]);
 742     ///
 743     /// X.with_borrow(|v| assert!(v.is_empty()));
 744     /// ```
 745     #[unstable(feature = "local_key_cell_methods", issue = "92122")]
 746     pub fn take(&'static self) -> T
 747     where
 748         T: Default,
 749     {
 750         self.with(|cell| cell.take())
 751     }
 752
 753     /// Replaces the contained value, returning the old value.
 754     ///
 755     /// # Panics
 756     ///
 757     /// Panics if the value is currently borrowed.
 758     ///
 759     /// Panics if the key currently has its destructor running,
 760     /// and it **may** panic if the destructor has previously been run for this thread.
 761     ///
 762     /// # Examples
 763     ///
 764     /// ```
 765     /// #![feature(local_key_cell_methods)]
 766     /// use std::cell::RefCell;
 767     ///
 768     /// thread_local! {
 769     ///     static X: RefCell<Vec<i32>> = RefCell::new(Vec::new());
 770     /// }
 771     ///
 772     /// let prev = X.replace(vec![1, 2, 3]);
 773     /// assert!(prev.is_empty());
 774     ///
 775     /// X.with_borrow(|v| assert_eq!(*v, vec![1, 2, 3]));
 776     /// ```
 777     #[unstable(feature = "local_key_cell_methods", issue = "92122")]
 778     pub fn replace(&'static self, value: T) -> T {
 779         self.with(|cell| cell.replace(value))
 780     }
 781 }
 782
 783 mod lazy {
 784     use crate::cell::UnsafeCell;
 785     use crate::hint;
 786     use crate::mem;
 787
 788     pub struct LazyKeyInner<T> {
 789         inner: UnsafeCell<Option<T>>,
 790     }
 791
 792     impl<T> LazyKeyInner<T> {
 793         pub const fn new() -> LazyKeyInner<T> {
 794             LazyKeyInner { inner: UnsafeCell::new(None) }
 795         }
 796
 797         pub unsafe fn get(&self) -> Option<&'static T> {
 798             // SAFETY: The caller must ensure no reference is ever handed out to
 799             // the inner cell nor mutable reference to the Option<T> inside said
 800             // cell. This make it safe to hand a reference, though the lifetime
 801             // of 'static is itself unsafe, making the get method unsafe.
 802             unsafe { (*self.inner.get()).as_ref() }
 803         }
 804
 805         /// The caller must ensure that no reference is active: this method
 806         /// needs unique access.
 807         pub unsafe fn initialize<F: FnOnce() -> T>(&self, init: F) -> &'static T {
 808             // Execute the initialization up front, *then* move it into our slot,
 809             // just in case initialization fails.
 810             let value = init();
 811             let ptr = self.inner.get();
 812
 813             // SAFETY:
 814             //
 815             // note that this can in theory just be `*ptr = Some(value)`, but due to
 816             // the compiler will currently codegen that pattern with something like:
 817             //
 818             //      ptr::drop_in_place(ptr)
 819             //      ptr::write(ptr, Some(value))
 820             //
 821             // Due to this pattern it's possible for the destructor of the value in
 822             // `ptr` (e.g., if this is being recursively initialized) to re-access
 823             // TLS, in which case there will be a `&` and `&mut` pointer to the same
 824             // value (an aliasing violation). To avoid setting the "I'm running a
 825             // destructor" flag we just use `mem::replace` which should sequence the
 826             // operations a little differently and make this safe to call.
 827             //
 828             // The precondition also ensures that we are the only one accessing
 829             // `self` at the moment so replacing is fine.
 830             unsafe {
 831                 let _ = mem::replace(&mut *ptr, Some(value));
 832             }
 833
 834             // SAFETY: With the call to `mem::replace` it is guaranteed there is
 835             // a `Some` behind `ptr`, not a `None` so `unreachable_unchecked`
 836             // will never be reached.
 837             unsafe {
 838                 // After storing `Some` we want to get a reference to the contents of
 839                 // what we just stored. While we could use `unwrap` here and it should
 840                 // always work it empirically doesn't seem to always get optimized away,
 841                 // which means that using something like `try_with` can pull in
 842                 // panicking code and cause a large size bloat.
 843                 match *ptr {
 844                     Some(ref x) => x,
 845                     None => hint::unreachable_unchecked(),
 846                 }
 847             }
 848         }
 849
 850         /// The other methods hand out references while taking &self.
 851         /// As such, callers of this method must ensure no `&` and `&mut` are
 852         /// available and used at the same time.
 853         #[allow(unused)]
 854         pub unsafe fn take(&mut self) -> Option<T> {
 855             // SAFETY: See doc comment for this method.
 856             unsafe { (*self.inner.get()).take() }
 857         }
 858     }
 859 }
 860
 861 /// On some targets like wasm there's no threads, so no need to generate
 862 /// thread locals and we can instead just use plain statics!
 863 #[doc(hidden)]
 864 #[cfg(all(target_family = "wasm", not(target_feature = "atomics")))]
 865 pub mod statik {
 866     use super::lazy::LazyKeyInner;
 867     use crate::fmt;
 868
 869     pub struct Key<T> {
 870         inner: LazyKeyInner<T>,
 871     }
 872
 873     unsafe impl<T> Sync for Key<T> {}
 874
 875     impl<T> fmt::Debug for Key<T> {
 876         fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 877             f.debug_struct("Key").finish_non_exhaustive()
 878         }
 879     }
 880
 881     impl<T> Key<T> {
 882         pub const fn new() -> Key<T> {
 883             Key { inner: LazyKeyInner::new() }
 884         }
 885
 886         pub unsafe fn get(&self, init: impl FnOnce() -> T) -> Option<&'static T> {
 887             // SAFETY: The caller must ensure no reference is ever handed out to
 888             // the inner cell nor mutable reference to the Option<T> inside said
 889             // cell. This make it safe to hand a reference, though the lifetime
 890             // of 'static is itself unsafe, making the get method unsafe.
 891             let value = unsafe {
 892                 match self.inner.get() {
 893                     Some(ref value) => value,
 894                     None => self.inner.initialize(init),
 895                 }
 896             };
 897
 898             Some(value)
 899         }
 900     }
 901 }
 902
 903 #[doc(hidden)]
 904 #[cfg(all(target_thread_local, not(all(target_family = "wasm", not(target_feature = "atomics"))),))]
 905 pub mod fast {
 906     use super::lazy::LazyKeyInner;
 907     use crate::cell::Cell;
 908     use crate::fmt;
 909     use crate::mem;
 910     use crate::sys::thread_local_dtor::register_dtor;
 911
 912     #[derive(Copy, Clone)]
 913     enum DtorState {
 914         Unregistered,
 915         Registered,
 916         RunningOrHasRun,
 917     }
 918
 919     // This data structure has been carefully constructed so that the fast path
 920     // only contains one branch on x86. That optimization is necessary to avoid
 921     // duplicated tls lookups on OSX.
 922     //
 923     // LLVM issue: https://bugs.llvm.org/show_bug.cgi?id=41722
 924     pub struct Key<T> {
 925         // If `LazyKeyInner::get` returns `None`, that indicates either:
 926         //   * The value has never been initialized
 927         //   * The value is being recursively initialized
 928         //   * The value has already been destroyed or is being destroyed
 929         // To determine which kind of `None`, check `dtor_state`.
 930         //
 931         // This is very optimizer friendly for the fast path - initialized but
 932         // not yet dropped.
 933         inner: LazyKeyInner<T>,
 934
 935         // Metadata to keep track of the state of the destructor. Remember that
 936         // this variable is thread-local, not global.
 937         dtor_state: Cell<DtorState>,
 938     }
 939
 940     impl<T> fmt::Debug for Key<T> {
 941         fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 942             f.debug_struct("Key").finish_non_exhaustive()
 943         }
 944     }
 945
 946     impl<T> Key<T> {
 947         pub const fn new() -> Key<T> {
 948             Key { inner: LazyKeyInner::new(), dtor_state: Cell::new(DtorState::Unregistered) }
 949         }
 950
 951         // note that this is just a publicly-callable function only for the
 952         // const-initialized form of thread locals, basically a way to call the
 953         // free `register_dtor` function defined elsewhere in libstd.
 954         pub unsafe fn register_dtor(a: *mut u8, dtor: unsafe extern "C" fn(*mut u8)) {
 955             unsafe {
 956                 register_dtor(a, dtor);
 957             }
 958         }
 959
 960         pub unsafe fn get<F: FnOnce() -> T>(&self, init: F) -> Option<&'static T> {
 961             // SAFETY: See the definitions of `LazyKeyInner::get` and
 962             // `try_initialize` for more information.
 963             //
 964             // The caller must ensure no mutable references are ever active to
 965             // the inner cell or the inner T when this is called.
 966             // The `try_initialize` is dependant on the passed `init` function
 967             // for this.
 968             unsafe {
 969                 match self.inner.get() {
 970                     Some(val) => Some(val),
 971                     None => self.try_initialize(init),
 972                 }
 973             }
 974         }
 975
 976         // `try_initialize` is only called once per fast thread local variable,
 977         // except in corner cases where thread_local dtors reference other
 978         // thread_local's, or it is being recursively initialized.
 979         //
 980         // Macos: Inlining this function can cause two `tlv_get_addr` calls to
 981         // be performed for every call to `Key::get`.
 982         // LLVM issue: https://bugs.llvm.org/show_bug.cgi?id=41722
 983         #[inline(never)]
 984         unsafe fn try_initialize<F: FnOnce() -> T>(&self, init: F) -> Option<&'static T> {
 985             // SAFETY: See comment above (this function doc).
 986             if !mem::needs_drop::<T>() || unsafe { self.try_register_dtor() } {
 987                 // SAFETY: See comment above (this function doc).
 988                 Some(unsafe { self.inner.initialize(init) })
 989             } else {
 990                 None
 991             }
 992         }
 993
 994         // `try_register_dtor` is only called once per fast thread local
 995         // variable, except in corner cases where thread_local dtors reference
 996         // other thread_local's, or it is being recursively initialized.
 997         unsafe fn try_register_dtor(&self) -> bool {
 998             match self.dtor_state.get() {
 999                 DtorState::Unregistered => {
1000                     // SAFETY: dtor registration happens before initialization.
1001                     // Passing `self` as a pointer while using `destroy_value<T>`
1002                     // is safe because the function will build a pointer to a
1003                     // Key<T>, which is the type of self and so find the correct
1004                     // size.
1005                     unsafe { register_dtor(self as *const _ as *mut u8, destroy_value::<T>) };
1006                     self.dtor_state.set(DtorState::Registered);
1007                     true
1008                 }
1009                 DtorState::Registered => {
1010                     // recursively initialized
1011                     true
1012                 }
1013                 DtorState::RunningOrHasRun => false,
1014             }
1015         }
1016     }
1017
1018     unsafe extern "C" fn destroy_value<T>(ptr: *mut u8) {
1019         let ptr = ptr as *mut Key<T>;
1020
1021         // SAFETY:
1022         //
1023         // The pointer `ptr` has been built just above and comes from
1024         // `try_register_dtor` where it is originally a Key<T> coming from `self`,
1025         // making it non-NUL and of the correct type.
1026         //
1027         // Right before we run the user destructor be sure to set the
1028         // `Option<T>` to `None`, and `dtor_state` to `RunningOrHasRun`. This
1029         // causes future calls to `get` to run `try_initialize_drop` again,
1030         // which will now fail, and return `None`.
1031         unsafe {
1032             let value = (*ptr).inner.take();
1033             (*ptr).dtor_state.set(DtorState::RunningOrHasRun);
1034             drop(value);
1035         }
1036     }
1037 }
1038
1039 #[doc(hidden)]
1040 #[cfg(all(
1041     not(target_thread_local),
1042     not(all(target_family = "wasm", not(target_feature = "atomics"))),
1043 ))]
1044 pub mod os {
1045     use super::lazy::LazyKeyInner;
1046     use crate::cell::Cell;
1047     use crate::fmt;
1048     use crate::marker;
1049     use crate::ptr;
1050     use crate::sys_common::thread_local_key::StaticKey as OsStaticKey;
1051
1052     /// Use a regular global static to store this key; the state provided will then be
1053     /// thread-local.
1054     pub struct Key<T> {
1055         // OS-TLS key that we'll use to key off.
1056         os: OsStaticKey,
1057         marker: marker::PhantomData<Cell<T>>,
1058     }
1059
1060     impl<T> fmt::Debug for Key<T> {
1061         fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1062             f.debug_struct("Key").finish_non_exhaustive()
1063         }
1064     }
1065
1066     unsafe impl<T> Sync for Key<T> {}
1067
1068     struct Value<T: 'static> {
1069         inner: LazyKeyInner<T>,
1070         key: &'static Key<T>,
1071     }
1072
1073     impl<T: 'static> Key<T> {
1074         #[rustc_const_unstable(feature = "thread_local_internals", issue = "none")]
1075         pub const fn new() -> Key<T> {
1076             Key { os: OsStaticKey::new(Some(destroy_value::<T>)), marker: marker::PhantomData }
1077         }
1078
1079         /// It is a requirement for the caller to ensure that no mutable
1080         /// reference is active when this method is called.
1081         pub unsafe fn get(&'static self, init: impl FnOnce() -> T) -> Option<&'static T> {
1082             // SAFETY: See the documentation for this method.
1083             let ptr = unsafe { self.os.get() as *mut Value<T> };
1084             if ptr.addr() > 1 {
1085                 // SAFETY: the check ensured the pointer is safe (its destructor
1086                 // is not running) + it is coming from a trusted source (self).
1087                 if let Some(ref value) = unsafe { (*ptr).inner.get() } {
1088                     return Some(value);
1089                 }
1090             }
1091             // SAFETY: At this point we are sure we have no value and so
1092             // initializing (or trying to) is safe.
1093             unsafe { self.try_initialize(init) }
1094         }
1095
1096         // `try_initialize` is only called once per os thread local variable,
1097         // except in corner cases where thread_local dtors reference other
1098         // thread_local's, or it is being recursively initialized.
1099         unsafe fn try_initialize(&'static self, init: impl FnOnce() -> T) -> Option<&'static T> {
1100             // SAFETY: No mutable references are ever handed out meaning getting
1101             // the value is ok.
1102             let ptr = unsafe { self.os.get() as *mut Value<T> };
1103             if ptr.addr() == 1 {
1104                 // destructor is running
1105                 return None;
1106             }
1107
1108             let ptr = if ptr.is_null() {
1109                 // If the lookup returned null, we haven't initialized our own
1110                 // local copy, so do that now.
1111                 let ptr: Box<Value<T>> = box Value { inner: LazyKeyInner::new(), key: self };
1112                 let ptr = Box::into_raw(ptr);
1113                 // SAFETY: At this point we are sure there is no value inside
1114                 // ptr so setting it will not affect anyone else.
1115                 unsafe {
1116                     self.os.set(ptr as *mut u8);
1117                 }
1118                 ptr
1119             } else {
1120                 // recursive initialization
1121                 ptr
1122             };
1123
1124             // SAFETY: ptr has been ensured as non-NUL just above an so can be
1125             // dereferenced safely.
1126             unsafe { Some((*ptr).inner.initialize(init)) }
1127         }
1128     }
1129
1130     unsafe extern "C" fn destroy_value<T: 'static>(ptr: *mut u8) {
1131         // SAFETY:
1132         //
1133         // The OS TLS ensures that this key contains a null value when this
1134         // destructor starts to run. We set it back to a sentinel value of 1 to
1135         // ensure that any future calls to `get` for this thread will return
1136         // `None`.
1137         //
1138         // Note that to prevent an infinite loop we reset it back to null right
1139         // before we return from the destructor ourselves.
1140         unsafe {
1141             let ptr = Box::from_raw(ptr as *mut Value<T>);
1142             let key = ptr.key;
1143             key.os.set(ptr::invalid_mut(1));
1144             drop(ptr);
1145             key.os.set(ptr::null_mut());
1146         }
1147     }
1148 }