ferroui_base/threading/
dispatcher_task.rs

1//! Futures running on a dispatcher.
2//!
3//! This is the executor half of what the reference implementation gets from
4//! its synchronization context: a future is stored on the dispatcher thread
5//! and polled from dispatcher jobs; its waker queues the next poll at the
6//! priority the future was started with. [`DispatcherTask`] is the handle to
7//! the outcome, the counterpart of a task object.
8
9use std::any::{Any, TypeId};
10use std::cell::Cell;
11use std::fmt;
12use std::future::Future;
13use std::marker::PhantomData;
14use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe};
15use std::pin::Pin;
16use std::sync::atomic::{AtomicBool, Ordering};
17use std::sync::{Arc, Condvar, Mutex, MutexGuard, PoisonError};
18use std::task::{Context, Poll, Wake, Waker};
19
20use super::dispatcher_operation::{next_operation_id, DispatcherException};
21use super::{Dispatcher, DispatcherPriority, OperationCanceledError};
22
23#[derive(Clone, Copy, Debug, PartialEq, Eq)]
24enum TaskStatus {
25    Running,
26    Completed,
27    Faulted,
28    Canceled,
29}
30
31struct TaskState {
32    status: TaskStatus,
33    result: Option<Box<dyn Any + Send>>,
34    panic: Option<DispatcherException>,
35    wakers: Vec<Waker>,
36}
37
38/// The shareable part of a task: its outcome and whoever waits for it.
39pub(crate) struct TaskShared {
40    id: u64,
41    dispatcher: Arc<Dispatcher>,
42    /// The result is kept in the dispatcher thread's local storage.
43    local_result: bool,
44    local_storage_used: AtomicBool,
45    state: Mutex<TaskState>,
46    finished: Condvar,
47}
48
49impl TaskShared {
50    pub(crate) fn new(dispatcher: &Arc<Dispatcher>, local_result: bool) -> Arc<Self> {
51        Arc::new(Self {
52            id: next_operation_id(),
53            dispatcher: dispatcher.clone(),
54            local_result,
55            local_storage_used: AtomicBool::new(false),
56            state: Mutex::new(TaskState { status: TaskStatus::Running, result: None, panic: None, wakers: Vec::new() }),
57            finished: Condvar::new(),
58        })
59    }
60
61    fn lock(&self) -> MutexGuard<'_, TaskState> {
62        self.state.lock().unwrap_or_else(PoisonError::into_inner)
63    }
64
65    fn finish(&self, update: impl FnOnce(&mut TaskState)) -> bool {
66        let wakers = {
67            let mut state = self.lock();
68            if state.status != TaskStatus::Running {
69                return false;
70            }
71            update(&mut state);
72            self.finished.notify_all();
73            std::mem::take(&mut state.wakers)
74        };
75        for waker in wakers {
76            waker.wake();
77        }
78        true
79    }
80
81    fn complete_send(&self, value: Box<dyn Any + Send>) {
82        self.finish(|state| {
83            state.status = TaskStatus::Completed;
84            state.result = Some(value);
85        });
86    }
87
88    /// Dispatcher thread only. `None` stands for `()`.
89    fn complete_local(&self, value: Option<Box<dyn Any>>) {
90        if let (Some(value), Some(local)) = (value, self.dispatcher.try_local()) {
91            local.results.borrow_mut().insert(self.id, value);
92            self.local_storage_used.store(true, Ordering::SeqCst);
93        }
94        self.finish(|state| state.status = TaskStatus::Completed);
95    }
96
97    fn fail(&self, payload: DispatcherException) {
98        self.finish(|state| {
99            state.status = TaskStatus::Faulted;
100            state.panic = Some(payload);
101        });
102    }
103
104    fn cancel(&self) {
105        self.finish(|state| state.status = TaskStatus::Canceled);
106    }
107
108    fn is_finished(&self) -> bool {
109        self.lock().status != TaskStatus::Running
110    }
111
112    fn take_outcome<T: 'static>(&self) -> Result<T, OperationCanceledError> {
113        let (status, panic, result) = {
114            let mut state = self.lock();
115            (state.status, state.panic.take(), state.result.take())
116        };
117        match status {
118            TaskStatus::Canceled => return Err(OperationCanceledError),
119            TaskStatus::Faulted => match panic {
120                Some(payload) => resume_unwind(payload),
121                None => panic!("The failure of the dispatcher task has already been observed."),
122            },
123            TaskStatus::Completed | TaskStatus::Running => {}
124        }
125
126        let value: Option<Box<dyn Any>> = if !self.local_result {
127            result.map(|value| value as Box<dyn Any>)
128        } else if TypeId::of::<T>() == TypeId::of::<()>() {
129            Some(Box::new(()))
130        } else {
131            if !self.dispatcher.check_access() {
132                panic!("The result of a local dispatcher task can only be read on the dispatcher thread.");
133            }
134            self.local_storage_used.store(false, Ordering::SeqCst);
135            self.dispatcher.try_local().and_then(|local| local.results.borrow_mut().remove(&self.id))
136        };
137
138        match value.and_then(|value| value.downcast::<T>().ok()) {
139            Some(value) => Ok(*value),
140            None => panic!("The result of the dispatcher task has already been taken."),
141        }
142    }
143}
144
145impl Drop for TaskShared {
146    fn drop(&mut self) {
147        if self.local_storage_used.swap(false, Ordering::SeqCst) {
148            self.dispatcher.discard_local_operation(self.id);
149        }
150    }
151}
152
153/// Completes the task as canceled when the code that was going to produce
154/// its outcome is dropped without having run to the end.
155pub(crate) struct TaskCompletion(Arc<TaskShared>);
156
157impl TaskCompletion {
158    pub(crate) fn fail(&self, payload: DispatcherException) {
159        self.0.fail(payload);
160    }
161}
162
163impl Drop for TaskCompletion {
164    fn drop(&mut self) {
165        self.0.cancel();
166    }
167}
168
169/// The outcome of a future running on a [`Dispatcher`]; what the reference
170/// implementation returns as a task object.
171///
172/// The handle can be cloned and sent to other threads. It is a [`Future`]
173/// that resolves to `Ok(result)` when the future has completed and to
174/// `Err(OperationCanceledError)` when it was dropped before completing (the
175/// dispatcher shut down, or the operation that was going to start it was
176/// aborted). A panic raised by the future is re-raised to whoever observes
177/// the outcome.
178pub struct DispatcherTask<T = ()> {
179    shared: Arc<TaskShared>,
180    _result: PhantomData<fn() -> T>,
181}
182
183impl<T> Clone for DispatcherTask<T> {
184    fn clone(&self) -> Self {
185        Self { shared: self.shared.clone(), _result: PhantomData }
186    }
187}
188
189impl<T> fmt::Debug for DispatcherTask<T> {
190    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
191        write!(f, "DispatcherTask #{} {:?}", self.shared.id, self.shared.lock().status)
192    }
193}
194
195impl<T: 'static> DispatcherTask<T> {
196    /// The dispatcher the future runs on.
197    pub fn dispatcher(&self) -> &Arc<Dispatcher> {
198        &self.shared.dispatcher
199    }
200
201    /// Whether the task has finished: completed, failed or canceled.
202    pub fn is_completed(&self) -> bool {
203        self.shared.is_finished()
204    }
205
206    /// Whether the task ran to the end without panicking.
207    pub fn is_completed_successfully(&self) -> bool {
208        self.shared.lock().status == TaskStatus::Completed
209    }
210
211    /// Whether the future panicked.
212    pub fn is_faulted(&self) -> bool {
213        self.shared.lock().status == TaskStatus::Faulted
214    }
215
216    /// Whether the future was dropped before it completed.
217    pub fn is_canceled(&self) -> bool {
218        self.shared.lock().status == TaskStatus::Canceled
219    }
220
221    fn block_until_finished(&self) {
222        if self.shared.is_finished() {
223            return;
224        }
225        if self.shared.dispatcher.check_access() {
226            // The dispatcher thread is the one that has to make progress.
227            panic!("Synchronous wait is only supported on non-UI threads");
228        }
229        let mut state = self.shared.lock();
230        while state.status == TaskStatus::Running {
231            state = self.shared.finished.wait(state).unwrap_or_else(PoisonError::into_inner);
232        }
233    }
234
235    /// Blocks until the task has finished. Only supported on threads other
236    /// than the dispatcher thread, unless the task has already finished.
237    ///
238    /// # Panics
239    /// Re-raises the panic of the future. Panics when the task is still
240    /// running and the calling thread is the dispatcher thread.
241    pub fn wait(&self) -> Result<(), OperationCanceledError> {
242        self.block_until_finished();
243        let (status, panic) = {
244            let mut state = self.shared.lock();
245            (state.status, state.panic.take())
246        };
247        if let Some(payload) = panic {
248            resume_unwind(payload);
249        }
250        match status {
251            TaskStatus::Canceled => Err(OperationCanceledError),
252            _ => Ok(()),
253        }
254    }
255
256    /// Blocks like [`wait`](Self::wait) and takes the result.
257    ///
258    /// # Panics
259    /// See [`wait`](Self::wait). Also panics when the result has already been
260    /// taken, and when a result that is not `Send` is requested from another
261    /// thread.
262    pub fn result(&self) -> Result<T, OperationCanceledError> {
263        self.block_until_finished();
264        self.shared.take_outcome()
265    }
266}
267
268impl<T: 'static> Future for DispatcherTask<T> {
269    type Output = Result<T, OperationCanceledError>;
270
271    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
272        {
273            let mut state = self.shared.lock();
274            if state.status == TaskStatus::Running {
275                if !state.wakers.iter().any(|waker| waker.will_wake(cx.waker())) {
276                    state.wakers.push(cx.waker().clone());
277                }
278                return Poll::Pending;
279            }
280        }
281        Poll::Ready(self.shared.take_outcome())
282    }
283}
284
285/// Queues polls of one task at the task's priority.
286pub(crate) struct TaskWaker {
287    dispatcher: Arc<Dispatcher>,
288    priority: DispatcherPriority,
289    task_id: u64,
290    scheduled: AtomicBool,
291}
292
293impl TaskWaker {
294    fn schedule(self: &Arc<Self>) {
295        if !self.scheduled.swap(true, Ordering::SeqCst) {
296            let this = self.clone();
297            self.dispatcher.post(
298                move || {
299                    this.dispatcher.poll_task(this.task_id);
300                },
301                self.priority,
302            );
303        }
304    }
305}
306
307impl Wake for TaskWaker {
308    fn wake(self: Arc<Self>) {
309        self.schedule();
310    }
311
312    fn wake_by_ref(self: &Arc<Self>) {
313        self.schedule();
314    }
315}
316
317/// A future stored on the dispatcher thread.
318pub(crate) struct LocalTask {
319    future: Pin<Box<dyn Future<Output = ()>>>,
320    waker: Waker,
321    wake_state: Arc<TaskWaker>,
322    completion: TaskCompletion,
323}
324
325crate::ferro_thread_local! {
326    /// The (dispatcher id, task id) of the task being polled on this thread.
327    static CURRENT_TASK: Cell<Option<(u64, u64)>> = const { Cell::new(None) };
328}
329
330/// The task of `dispatcher` that is being polled right now, if any.
331pub(crate) fn current_task_of(dispatcher: &Dispatcher) -> Option<u64> {
332    match CURRENT_TASK.try_with(Cell::get).ok().flatten() {
333        Some((dispatcher_id, task_id)) if dispatcher_id == dispatcher.id => Some(task_id),
334        _ => None,
335    }
336}
337
338/// Creates the handle of a task whose result can cross threads, and the
339/// future that produces it.
340pub(crate) fn wrap_send_future<Fut>(
341    dispatcher: &Arc<Dispatcher>,
342) -> (DispatcherTask<Fut::Output>, TaskCompletion, impl FnOnce(Fut) -> Pin<Box<dyn Future<Output = ()>>>)
343where
344    Fut: Future + 'static,
345    Fut::Output: Send + 'static,
346{
347    let shared = TaskShared::new(dispatcher, false);
348    let task = DispatcherTask { shared: shared.clone(), _result: PhantomData };
349    let completion = TaskCompletion(shared.clone());
350    let wrap = move |future: Fut| -> Pin<Box<dyn Future<Output = ()>>> {
351        Box::pin(async move {
352            let value = future.await;
353            shared.complete_send(Box::new(value));
354        })
355    };
356    (task, completion, wrap)
357}
358
359/// Creates the handle of a task whose result stays on the dispatcher thread,
360/// and the future that produces it.
361pub(crate) fn wrap_local_future<Fut>(
362    dispatcher: &Arc<Dispatcher>,
363) -> (DispatcherTask<Fut::Output>, TaskCompletion, impl FnOnce(Fut) -> Pin<Box<dyn Future<Output = ()>>>)
364where
365    Fut: Future + 'static,
366    Fut::Output: 'static,
367{
368    let shared = TaskShared::new(dispatcher, true);
369    let task = DispatcherTask { shared: shared.clone(), _result: PhantomData };
370    let completion = TaskCompletion(shared.clone());
371    let wrap = move |future: Fut| -> Pin<Box<dyn Future<Output = ()>>> {
372        Box::pin(async move {
373            let value = future.await;
374            if TypeId::of::<Fut::Output>() == TypeId::of::<()>() {
375                shared.complete_local(None);
376            } else {
377                shared.complete_local(Some(Box::new(value)));
378            }
379        })
380    };
381    (task, completion, wrap)
382}
383
384impl Dispatcher {
385    /// Stores `future` on this thread and polls it from jobs queued at
386    /// `priority`. With `poll_now` the first poll happens before returning,
387    /// otherwise it is queued.
388    ///
389    /// Dispatcher thread only.
390    pub(crate) fn start_task(
391        &self,
392        priority: DispatcherPriority,
393        future: Pin<Box<dyn Future<Output = ()>>>,
394        completion: TaskCompletion,
395        poll_now: bool,
396    ) {
397        let Some(local) = self.try_local() else {
398            // Dropping the completion cancels the task.
399            return;
400        };
401        if self.lock().has_shutdown_finished {
402            return;
403        }
404        let task_id = completion.0.id;
405        let wake_state =
406            Arc::new(TaskWaker { dispatcher: self.to_arc(), priority, task_id, scheduled: AtomicBool::new(false) });
407        let task = LocalTask { future, waker: Waker::from(wake_state.clone()), wake_state: wake_state.clone(), completion };
408        local.tasks.borrow_mut().insert(task_id, task);
409        if poll_now {
410            self.poll_task(task_id);
411        } else {
412            wake_state.schedule();
413        }
414    }
415
416    /// Polls the task once. Returns `false` when the task is not waiting to
417    /// be polled: it has finished, or it is the one being polled right now.
418    ///
419    /// Dispatcher thread only.
420    pub(crate) fn poll_task(&self, task_id: u64) -> bool {
421        let Some(local) = self.try_local() else {
422            return false;
423        };
424        let task = local.tasks.borrow_mut().remove(&task_id);
425        let Some(mut task) = task else {
426            return false;
427        };
428        task.wake_state.scheduled.store(false, Ordering::SeqCst);
429
430        let previous = CURRENT_TASK.with(|current| current.replace(Some((self.id, task_id))));
431        let result = catch_unwind(AssertUnwindSafe(|| {
432            let mut cx = Context::from_waker(&task.waker);
433            task.future.as_mut().poll(&mut cx)
434        }));
435        CURRENT_TASK.with(|current| current.set(previous));
436
437        match result {
438            Ok(Poll::Pending) => {
439                if self.lock().has_shutdown_finished {
440                    // Nothing can poll it again.
441                    drop(task);
442                } else {
443                    local.tasks.borrow_mut().insert(task_id, task);
444                }
445            }
446            Ok(Poll::Ready(())) => drop(task),
447            Err(payload) => {
448                task.completion.fail(payload);
449                drop(task);
450            }
451        }
452        true
453    }
454
455    /// Drops every future running on this dispatcher, which completes their
456    /// tasks as canceled.
457    pub(crate) fn cancel_tasks(&self) {
458        if let Some(local) = self.try_local() {
459            let tasks = std::mem::take(&mut *local.tasks.borrow_mut());
460            // Dropped outside of the borrow: futures run arbitrary code when dropped.
461            drop(tasks);
462        }
463    }
464}
ferroui_base/threading/dispatcher_task.rs

ferroui_base/threading/
dispatcher_task.rs
