std\sys\thread_local\key/
windows.rs

1//! Implementation of `LazyKey` for Windows.
2//!
3//! Windows has no native support for running destructors so we manage our own
4//! list of destructors to keep track of how to destroy keys. We then install a
5//! callback later to get invoked whenever a thread exits, running all
6//! appropriate destructors (see the [`guard`](guard) module documentation).
7//!
8//! This will likely need to be improved over time, but this module attempts a
9//! "poor man's" destructor callback system. Once we've got a list of what to
10//! run, we iterate over all keys, check their values, and then run destructors
11//! if the values turn out to be non null (setting them to null just beforehand).
12//! We do this a few times in a loop to basically match Unix semantics. If we
13//! don't reach a fixed point after a short while then we just inevitably leak
14//! something.
15//!
16//! The list is implemented as an atomic single-linked list of `LazyKey`s and
17//! does not support unregistration. Unfortunately, this means that we cannot
18//! use racy initialization for creating the keys in `LazyKey`, as that could
19//! result in destructors being missed. Hence, we synchronize the creation of
20//! keys with destructors through [`INIT_ONCE`](c::INIT_ONCE) (`std`'s
21//! [`Once`](crate::sync::Once) cannot be used since it might use TLS itself).
22//! For keys without destructors, racy initialization suffices.
23
24// FIXME: investigate using a fixed-size array instead, as the maximum number
25//        of keys is [limited to 1088](https://learn.microsoft.com/en-us/windows/win32/ProcThread/thread-local-storage).
26
27use crate::cell::UnsafeCell;
28use crate::ptr;
29use crate::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release};
30use crate::sync::atomic::{Atomic, AtomicPtr, AtomicU32};
31use crate::sys::c;
32use crate::sys::thread_local::guard;
33
34pub type Key = u32;
35type Dtor = unsafe extern "C" fn(*mut u8);
36
37pub struct LazyKey {
38    /// The key value shifted up by one. Since TLS_OUT_OF_INDEXES == u32::MAX
39    /// is not a valid key value, this allows us to use zero as sentinel value
40    /// without risking overflow.
41    key: Atomic<Key>,
42    dtor: Option<Dtor>,
43    next: Atomic<*mut LazyKey>,
44    /// Currently, destructors cannot be unregistered, so we cannot use racy
45    /// initialization for keys. Instead, we need synchronize initialization.
46    /// Use the Windows-provided `Once` since it does not require TLS.
47    once: UnsafeCell<c::INIT_ONCE>,
48}
49
50impl LazyKey {
51    #[inline]
52    pub const fn new(dtor: Option<Dtor>) -> LazyKey {
53        LazyKey {
54            key: AtomicU32::new(0),
55            dtor,
56            next: AtomicPtr::new(ptr::null_mut()),
57            once: UnsafeCell::new(c::INIT_ONCE_STATIC_INIT),
58        }
59    }
60
61    #[inline]
62    pub fn force(&'static self) -> Key {
63        match self.key.load(Acquire) {
64            0 => unsafe { self.init() },
65            key => key - 1,
66        }
67    }
68
69    #[cold]
70    unsafe fn init(&'static self) -> Key {
71        if self.dtor.is_some() {
72            let mut pending = c::FALSE;
73            let r = unsafe {
74                c::InitOnceBeginInitialize(self.once.get(), 0, &mut pending, ptr::null_mut())
75            };
76            assert_eq!(r, c::TRUE);
77
78            if pending == c::FALSE {
79                // Some other thread initialized the key, load it.
80                self.key.load(Relaxed) - 1
81            } else {
82                let key = unsafe { c::TlsAlloc() };
83                if key == c::TLS_OUT_OF_INDEXES {
84                    // Wakeup the waiting threads before panicking to avoid deadlock.
85                    unsafe {
86                        c::InitOnceComplete(
87                            self.once.get(),
88                            c::INIT_ONCE_INIT_FAILED,
89                            ptr::null_mut(),
90                        );
91                    }
92                    panic!("out of TLS indexes");
93                }
94
95                unsafe {
96                    register_dtor(self);
97                }
98
99                // Release-storing the key needs to be the last thing we do.
100                // This is because in `fn key()`, other threads will do an acquire load of the key,
101                // and if that sees this write then it will entirely bypass the `InitOnce`. We thus
102                // need to establish synchronization through `key`. In particular that acquire load
103                // must happen-after the register_dtor above, to ensure the dtor actually runs!
104                self.key.store(key + 1, Release);
105
106                let r = unsafe { c::InitOnceComplete(self.once.get(), 0, ptr::null_mut()) };
107                debug_assert_eq!(r, c::TRUE);
108
109                key
110            }
111        } else {
112            // If there is no destructor to clean up, we can use racy initialization.
113
114            let key = unsafe { c::TlsAlloc() };
115            assert_ne!(key, c::TLS_OUT_OF_INDEXES, "out of TLS indexes");
116
117            match self.key.compare_exchange(0, key + 1, AcqRel, Acquire) {
118                Ok(_) => key,
119                Err(new) => unsafe {
120                    // Some other thread completed initialization first, so destroy
121                    // our key and use theirs.
122                    let r = c::TlsFree(key);
123                    debug_assert_eq!(r, c::TRUE);
124                    new - 1
125                },
126            }
127        }
128    }
129}
130
131unsafe impl Send for LazyKey {}
132unsafe impl Sync for LazyKey {}
133
134#[inline]
135pub unsafe fn set(key: Key, val: *mut u8) {
136    let r = unsafe { c::TlsSetValue(key, val.cast()) };
137    debug_assert_eq!(r, c::TRUE);
138}
139
140#[inline]
141pub unsafe fn get(key: Key) -> *mut u8 {
142    unsafe { c::TlsGetValue(key).cast() }
143}
144
145static DTORS: Atomic<*mut LazyKey> = AtomicPtr::new(ptr::null_mut());
146
147/// Should only be called once per key, otherwise loops or breaks may occur in
148/// the linked list.
149unsafe fn register_dtor(key: &'static LazyKey) {
150    guard::enable();
151
152    let this = <*const LazyKey>::cast_mut(key);
153    // Use acquire ordering to pass along the changes done by the previously
154    // registered keys when we store the new head with release ordering.
155    let mut head = DTORS.load(Acquire);
156    loop {
157        key.next.store(head, Relaxed);
158        match DTORS.compare_exchange_weak(head, this, Release, Acquire) {
159            Ok(_) => break,
160            Err(new) => head = new,
161        }
162    }
163}
164
165/// This will and must only be run by the destructor callback in [`guard`].
166pub unsafe fn run_dtors() {
167    for _ in 0..5 {
168        let mut any_run = false;
169
170        // Use acquire ordering to observe key initialization.
171        let mut cur = DTORS.load(Acquire);
172        while !cur.is_null() {
173            let pre_key = unsafe { (*cur).key.load(Acquire) };
174            let dtor = unsafe { (*cur).dtor.unwrap() };
175            cur = unsafe { (*cur).next.load(Relaxed) };
176
177            // In LazyKey::init, we register the dtor before setting `key`.
178            // So if one thread's `run_dtors` races with another thread executing `init` on the same
179            // `LazyKey`, we can encounter a key of 0 here. That means this key was never
180            // initialized in this thread so we can safely skip it.
181            if pre_key == 0 {
182                continue;
183            }
184            // If this is non-zero, then via the `Acquire` load above we synchronized with
185            // everything relevant for this key. (It's not clear that this is needed, since the
186            // release-acquire pair on DTORS also establishes synchronization, but better safe than
187            // sorry.)
188            let key = pre_key - 1;
189
190            let ptr = unsafe { c::TlsGetValue(key) };
191            if !ptr.is_null() {
192                unsafe {
193                    c::TlsSetValue(key, ptr::null_mut());
194                    dtor(ptr as *mut _);
195                    any_run = true;
196                }
197            }
198        }
199
200        if !any_run {
201            break;
202        }
203    }
204}