Skip to main content

ferroui_base/threading/
dispatcher_thread_storage.rs

1use std::cell::{Cell, RefCell};
2use std::rc::Rc;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::{Arc, Mutex, MutexGuard, PoisonError, Weak};
5use std::thread::{self, ThreadId};
6
7use super::dispatcher::current_thread_id;
8
9use super::dispatcher::DispatcherLocal;
10use super::{Dispatcher, DispatcherOperationStatus, DispatcherPriority, FerroSynchronizationContext, IDispatcherImpl};
11
12struct GlobalState {
13    ui_thread: Option<Arc<Dispatcher>>,
14    dispatchers: Vec<(ThreadId, Weak<Dispatcher>)>,
15}
16
17static GLOBAL: Mutex<GlobalState> = Mutex::new(GlobalState { ui_thread: None, dispatchers: Vec::new() });
18
19/// Incremented whenever the UI thread dispatcher changes; validates the
20/// per-thread cache used by [`Dispatcher::ui_thread`].
21static UI_THREAD_GENERATION: AtomicU64 = AtomicU64::new(1);
22
23fn global() -> MutexGuard<'static, GlobalState> {
24    GLOBAL.lock().unwrap_or_else(PoisonError::into_inner)
25}
26
27/// Keeps the dispatcher of a thread alive for as long as the thread runs.
28struct CurrentThreadDispatcher(Arc<Dispatcher>);
29
30impl Drop for CurrentThreadDispatcher {
31    fn drop(&mut self) {
32        self.0.abandon();
33    }
34}
35
36crate::ferro_thread_local! {
37    static CURRENT_THREAD_DISPATCHER: RefCell<Option<CurrentThreadDispatcher>> = const { RefCell::new(None) };
38    static CURRENT_THREAD_LOCAL: RefCell<Option<Rc<DispatcherLocal>>> = const { RefCell::new(None) };
39    static UI_THREAD_CACHE: RefCell<(u64, Option<Arc<Dispatcher>>)> = const { RefCell::new((0, None)) };
40    static UNIT_TEST_ISOLATED: Cell<bool> = const { Cell::new(false) };
41}
42
43fn is_unit_test_isolated() -> bool {
44    UNIT_TEST_ISOLATED.try_with(Cell::get).unwrap_or(false)
45}
46
47impl Dispatcher {
48    pub(crate) fn register_current_thread_dispatcher(dispatcher: &Arc<Dispatcher>, local: Rc<DispatcherLocal>) {
49        let thread = current_thread_id();
50        {
51            let mut global = global();
52            if global.dispatchers.iter().any(|(id, weak)| *id == thread && weak.strong_count() > 0) {
53                panic!("The current thread already has a dispatcher");
54            }
55
56            // The first created dispatcher becomes "UI thread one"
57            if global.ui_thread.is_none() && !is_unit_test_isolated() {
58                global.ui_thread = Some(dispatcher.clone());
59                UI_THREAD_GENERATION.fetch_add(1, Ordering::AcqRel);
60            }
61
62            global.dispatchers.retain(|(id, weak)| *id != thread && weak.strong_count() > 0);
63            global.dispatchers.push((thread, Arc::downgrade(dispatcher)));
64        }
65        CURRENT_THREAD_LOCAL.with(|slot| *slot.borrow_mut() = Some(local));
66        let previous = CURRENT_THREAD_DISPATCHER
67            .with(|slot| slot.borrow_mut().replace(CurrentThreadDispatcher(dispatcher.clone())));
68        drop(previous);
69    }
70
71    pub(crate) fn current_thread_local() -> Option<Rc<DispatcherLocal>> {
72        CURRENT_THREAD_LOCAL.try_with(|slot| slot.borrow().clone()).ok().flatten()
73    }
74
75    fn try_current_dispatcher() -> Option<Arc<Dispatcher>> {
76        CURRENT_THREAD_DISPATCHER
77            .try_with(|slot| slot.borrow().as_ref().map(|current| current.0.clone()))
78            .ok()
79            .flatten()
80    }
81
82    /// The dispatcher of the calling thread. It is created when the thread
83    /// does not have one yet and then lives until the thread exits.
84    pub fn current_dispatcher() -> Arc<Dispatcher> {
85        if let Some(dispatcher) = Self::try_current_dispatcher() {
86            return dispatcher;
87        }
88
89        Dispatcher::new(None)
90    }
91
92    /// The dispatcher of `thread`, if that thread has one.
93    pub fn from_thread(thread: ThreadId) -> Option<Arc<Dispatcher>> {
94        global().dispatchers.iter().find(|(id, _)| *id == thread).and_then(|(_, weak)| weak.upgrade())
95    }
96
97    /// Gets the dispatcher for the UI thread: the first dispatcher that was
98    /// created in the process.
99    ///
100    /// Control and library authors are encouraged to use
101    /// [`current_dispatcher`](Self::current_dispatcher) and the dispatcher of
102    /// the object at hand instead.
103    pub fn ui_thread() -> Arc<Dispatcher> {
104        if is_unit_test_isolated() {
105            return Self::current_dispatcher();
106        }
107
108        let generation = UI_THREAD_GENERATION.load(Ordering::Acquire);
109        let cached = UI_THREAD_CACHE
110            .try_with(|cache| {
111                let cache = cache.borrow();
112                if cache.0 == generation {
113                    cache.1.clone()
114                } else {
115                    None
116                }
117            })
118            .ok()
119            .flatten();
120        if let Some(dispatcher) = cached {
121            return dispatcher;
122        }
123
124        Self::get_ui_thread_dispatcher_slow()
125    }
126
127    #[cold]
128    fn get_ui_thread_dispatcher_slow() -> Arc<Dispatcher> {
129        let existing = global().ui_thread.clone();
130        let current = match existing {
131            Some(_) => None,
132            // Creating the dispatcher of this thread makes it the UI thread
133            // one unless another thread got there first.
134            None => Some(Self::current_dispatcher()),
135        };
136
137        let global = global();
138        match &global.ui_thread {
139            Some(dispatcher) => {
140                let generation = UI_THREAD_GENERATION.load(Ordering::Acquire);
141                let _ = UI_THREAD_CACHE.try_with(|cache| *cache.borrow_mut() = (generation, Some(dispatcher.clone())));
142                dispatcher.clone()
143            }
144            None => {
145                drop(global);
146                current.unwrap_or_else(Self::current_dispatcher)
147            }
148        }
149    }
150
151    #[allow(dead_code)]
152    pub(crate) fn try_get_ui_thread() -> Option<Arc<Dispatcher>> {
153        if is_unit_test_isolated() {
154            return Self::try_current_dispatcher();
155        }
156        global().ui_thread.clone()
157    }
158
159    /// Installs the platform implementation of the UI thread dispatcher,
160    /// adapting a legacy threading interface.
161    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
162    pub fn initialize_ui_thread_dispatcher_legacy(impl_: Rc<dyn crate::platform::IPlatformThreadingInterface>) {
163        Self::initialize_ui_thread_dispatcher(Rc::new(super::i_dispatcher_impl::LegacyDispatcherImpl::new(impl_)));
164    }
165
166    /// Installs the platform implementation of the UI thread dispatcher.
167    ///
168    /// # Panics
169    /// Panics when not called on the UI thread or when the UI thread
170    /// dispatcher has already been initialized.
171    pub fn initialize_ui_thread_dispatcher(impl_: Rc<dyn IDispatcherImpl>) {
172        let ui_thread = Self::ui_thread();
173        ui_thread.verify_access();
174        if ui_thread.lock().initialized {
175            panic!("UI thread dispatcher is already initialized");
176        }
177        ui_thread.replace_implementation(Some(impl_));
178    }
179
180    /// Forgets the dispatcher of the calling thread.
181    fn reset_global_state() {
182        let thread = current_thread_id();
183        let current = CURRENT_THREAD_DISPATCHER.try_with(|slot| slot.borrow_mut().take()).ok().flatten();
184        if let Some(current) = &current {
185            // While the thread-local part is still reachable.
186            current.0.abandon();
187        }
188        let local = CURRENT_THREAD_LOCAL.try_with(|slot| slot.borrow_mut().take()).ok().flatten();
189        FerroSynchronizationContext::set_current(None);
190        let _ = UI_THREAD_CACHE.try_with(|cache| *cache.borrow_mut() = (0, None));
191        let ui_thread = {
192            let mut global = global();
193            global.dispatchers.retain(|(id, weak)| *id != thread && weak.strong_count() > 0);
194            let is_ui_thread = match (&global.ui_thread, &current) {
195                (Some(ui_thread), Some(current)) => Arc::ptr_eq(ui_thread, &current.0),
196                (Some(ui_thread), None) => ui_thread.thread() == thread,
197                _ => false,
198            };
199            if is_ui_thread {
200                UI_THREAD_GENERATION.fetch_add(1, Ordering::AcqRel);
201                global.ui_thread.take()
202            } else {
203                None
204            }
205        };
206        // Dropped outside of the lock: releasing the dispatcher aborts
207        // whatever is still queued.
208        drop(ui_thread);
209        drop(current);
210        drop(local);
211    }
212
213    /// Detaches the dispatcher of the calling thread without running its
214    /// pending jobs. Counterpart of [`reset_for_unit_tests`](Self::reset_for_unit_tests)
215    /// for use before a test.
216    #[doc(hidden)]
217    pub fn reset_before_unit_tests() {
218        Self::reset_global_state();
219    }
220
221    /// Runs the remaining active jobs of the calling thread's dispatcher,
222    /// shuts it down and forgets it, so that the next use of the thread gets
223    /// a fresh dispatcher.
224    ///
225    /// The reference implementation resets the process-wide UI dispatcher;
226    /// tests here run in parallel on separate threads, so the reset is
227    /// per-thread.
228    #[doc(hidden)]
229    pub fn reset_for_unit_tests() {
230        let Some(dispatcher) = Self::try_current_dispatcher() else {
231            Self::reset_global_state();
232            return;
233        };
234
235        let mut executed = 0u32;
236        loop {
237            let job = {
238                let local = dispatcher.try_local();
239                let mut st = dispatcher.lock();
240                Self::detach_from_platform(&mut st, local.as_deref());
241                st.queue.peek()
242            };
243            executed += 1;
244            if executed > 1_000_000 {
245                panic!("You've caused dispatcher loop");
246            }
247
248            match job {
249                Some(job) if job.priority() > DispatcherPriority::INACTIVE => dispatcher.execute_job(&job),
250                _ => {
251                    dispatcher.shutdown_impl();
252                    Self::reset_global_state();
253                    return;
254                }
255            }
256        }
257    }
258
259    /// Makes the calling thread self-contained for the duration of a unit
260    /// test: while the returned scope is alive,
261    /// [`ui_thread`](Self::ui_thread) resolves to the dispatcher of the
262    /// calling thread and dispatchers created on it never become the
263    /// process-wide UI thread dispatcher. Dropping the scope resets the
264    /// thread's dispatcher.
265    #[doc(hidden)]
266    pub fn unit_test_scope() -> UnitTestDispatcherScope {
267        Self::reset_for_unit_tests();
268        let previous = UNIT_TEST_ISOLATED.with(|isolated| isolated.replace(true));
269        UnitTestDispatcherScope { previous }
270    }
271
272    /// Called when the owning thread exits (or the dispatcher is reset):
273    /// nothing queued can run anymore, so everything is aborted, which also
274    /// releases threads blocked on those operations.
275    pub(crate) fn abandon(&self) {
276        let local = self.try_local();
277        let operations = {
278            let mut st = self.lock();
279            st.has_shutdown_finished = true;
280            Self::detach_from_platform(&mut st, local.as_deref());
281            st.queue.clear()
282        };
283        for operation in operations {
284            if operation.status() == DispatcherOperationStatus::Pending {
285                operation.set_status(DispatcherOperationStatus::Aborted);
286                operation.call_abort_callbacks();
287            }
288        }
289        self.cancel_tasks();
290    }
291}
292
293/// See [`Dispatcher::unit_test_scope`].
294#[doc(hidden)]
295pub struct UnitTestDispatcherScope {
296    previous: bool,
297}
298
299impl Drop for UnitTestDispatcherScope {
300    fn drop(&mut self) {
301        if thread::panicking() {
302            // Do not run more user code while unwinding.
303            Dispatcher::reset_global_state();
304        } else {
305            Dispatcher::reset_for_unit_tests();
306        }
307        let _ = UNIT_TEST_ISOLATED.try_with(|isolated| isolated.set(self.previous));
308    }
309}