shuttle_engine/
thread_support.rs

1use crate::runtime::execution::ExecutionState;
2use crate::runtime::thread;
3use std::marker::PhantomData;
4
5/// Cooperatively gives up a timeslice to the Shuttle scheduler.
6pub fn yield_now() {
7    let waker = ExecutionState::with(|state| state.current().waker());
8    waker.wake_by_ref();
9    ExecutionState::request_yield();
10    thread::switch();
11}
12
13/// The body of a spawned thread. Runs `f`, drops thread locals, publishes result.
14pub fn thread_fn<F, T>(
15    f: F,
16    switch_before_exit: bool,
17    result: std::sync::Arc<std::sync::Mutex<Option<std::thread::Result<T>>>>,
18) where
19    F: FnOnce() -> T,
20{
21    let ret = f();
22
23    if switch_before_exit && ExecutionState::with(|s| s.exit_current_truncates_execution()) {
24        thread::switch();
25    }
26
27    tracing::trace!("thread finished, dropping thread locals");
28
29    while let Some(local) = ExecutionState::with(|state| state.current_mut().pop_local()) {
30        tracing::trace!("dropping thread local {:p}", local);
31        drop(local);
32    }
33
34    tracing::trace!("done dropping thread locals");
35
36    *result.lock().unwrap() = Some(Ok(ret));
37    ExecutionState::with(|state| {
38        if let Some(waiter) = state.current_mut().take_waiter() {
39            state.get_mut(waiter).unblock();
40        }
41    });
42}
43
44/// A key into Shuttle's thread-local storage.
45pub struct LocalKey<T: 'static> {
46    #[doc(hidden)]
47    pub init: fn() -> T,
48    #[doc(hidden)]
49    pub _p: PhantomData<T>,
50}
51
52// Safety: `LocalKey` implements thread-local storage; each thread sees its own value of the type T.
53unsafe impl<T> Send for LocalKey<T> {}
54unsafe impl<T> Sync for LocalKey<T> {}
55
56impl<T: 'static> std::fmt::Debug for LocalKey<T> {
57    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        f.debug_struct("LocalKey").finish_non_exhaustive()
59    }
60}
61
62impl<T: 'static> LocalKey<T> {
63    /// Acquires a reference to the value in this TLS key.
64    ///
65    /// This will lazily initialize the value if this thread has not referenced this key yet.
66    pub fn with<F, R>(&'static self, f: F) -> R
67    where
68        F: FnOnce(&T) -> R,
69    {
70        self.try_with(f).expect(
71            "cannot access a Thread Local Storage value \
72            during or after destruction",
73        )
74    }
75
76    /// Acquires a reference to the value in this TLS key.
77    ///
78    /// This will lazily initialize the value if this thread has not referenced this key yet. If the
79    /// key has been destroyed (which may happen if this is called in a destructor), this function
80    /// will return an AccessError.
81    pub fn try_with<F, R>(&'static self, f: F) -> std::result::Result<R, AccessError>
82    where
83        F: FnOnce(&T) -> R,
84    {
85        let value = self.get().unwrap_or_else(|| {
86            let value = (self.init)();
87
88            ExecutionState::with(move |state| {
89                state.current_mut().init_local(self, value);
90            });
91
92            self.get().unwrap()
93        })?;
94
95        Ok(f(value))
96    }
97
98    fn get(&'static self) -> Option<std::result::Result<&'static T, AccessError>> {
99        // Safety: see the usage below
100        unsafe fn extend_lt<'b, T>(t: &'_ T) -> &'b T {
101            std::mem::transmute(t)
102        }
103
104        ExecutionState::with(|state| {
105            if let Ok(value) = state.current().local(self)? {
106                // Safety: the `ExecutionState` outlives any thread, including the caller, and so
107                // it's safe to give the caller the lifetime it's asking for here.
108                Some(Ok(unsafe { extend_lt(value) }))
109            } else {
110                Some(Err(AccessError))
111            }
112        })
113    }
114}
115
116/// An error returned by [`LocalKey::try_with`]
117#[derive(Clone, Copy, PartialEq, Eq, Debug)]
118#[non_exhaustive]
119pub struct AccessError;
120
121impl std::fmt::Display for AccessError {
122    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
123        std::fmt::Display::fmt("already destroyed", f)
124    }
125}
126
127impl std::error::Error for AccessError {}
shuttle_engine/thread_support.rs

shuttle_engine/
thread_support.rs
