forge/os/
tls.rs

1use core::{ffi::c_void, marker::PhantomData};
2
3use sys::os::tls::*;
4
5use alloc::boxed::Box;
6
7/// A thread-local storage slot holding an independent value of type `T` per thread.
8///
9/// The value for each thread is constructed lazily on first access using the `init`
10/// function supplied at creation time, and dropped when the thread exits.
11pub struct ThreadLocal<T> {
12    slot: TlsSlot,
13    init: fn() -> T,
14    _marker: PhantomData<*mut T>,
15}
16
17unsafe impl<T: Send> Send for ThreadLocal<T> {}
18unsafe impl<T: Send> Sync for ThreadLocal<T> {}
19
20impl<T> ThreadLocal<T> {
21    /// Allocates a TLS slot. `init` is called once per thread to construct its value.
22    pub fn new(init: fn() -> T) -> Self {
23        let mut slot = TlsSlot::new();
24        let result = unsafe { nnosAllocateTlsSlot(&mut slot, Self::destructor) };
25
26        if result != 0 || !slot.valid() {
27            panic!("Failed to allocate TLS Slot");
28        }
29
30        Self {
31            slot,
32            init,
33            _marker: Default::default(),
34        }
35    }
36
37    /// Calls `f` with a shared reference to this thread's value, initializing it on first access.
38    pub fn with<R>(&self, f: impl FnOnce(&T) -> R) -> R {
39        f(self.get_or_init())
40    }
41
42    /// Calls `f` with a mutable reference to this thread's value, initializing it on first access.
43    pub fn with_mut<R>(&self, f: impl FnOnce(&mut T) -> R) -> R {
44        f(self.get_or_init())
45    }
46
47    fn get_or_init(&self) -> &mut T {
48        let ptr = unsafe { nnosGetTlsValue(self.slot) };
49
50        if !ptr.is_null() {
51            unsafe { &mut *(ptr as *mut T) }
52        } else {
53            let val = Box::leak(Box::new((self.init)()));
54            unsafe { nnosSetTlsValue(self.slot, val as *mut T as *mut c_void) };
55            val
56        }
57    }
58
59    unsafe extern "C" fn destructor(value: *mut c_void) {
60        if !value.is_null() {
61            drop(unsafe { Box::from_raw(value as *mut T) })
62        }
63    }
64}
65
66impl<T> Drop for ThreadLocal<T> {
67    fn drop(&mut self) {
68        unsafe { nnosFreeTlsSlot(self.slot) };
69    }
70}
forge/os/tls.rs

forge/os/
tls.rs
