Skip to main content

ferroui_base/threading/
dispatcher_priority_awaitable.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::sync::Arc;
5use std::task::{Context, Poll};
6
7use super::dispatcher_task::current_task_of;
8use super::{Dispatcher, DispatcherPriority};
9
10/// Resumes the awaiting code from a dispatcher job of a given priority.
11struct Resumption {
12    dispatcher: Arc<Dispatcher>,
13    priority: DispatcherPriority,
14    posted: bool,
15    resumed: Arc<AtomicBool>,
16}
17
18impl Resumption {
19    fn new(dispatcher: Arc<Dispatcher>, priority: DispatcherPriority) -> Self {
20        Self { dispatcher, priority, posted: false, resumed: Arc::new(AtomicBool::new(false)) }
21    }
22
23    /// Never ready the first time it is polled.
24    fn poll(&mut self, cx: &mut Context<'_>) -> Poll<()> {
25        if self.resumed.load(Ordering::SeqCst) {
26            return Poll::Ready(());
27        }
28        if !self.posted {
29            self.posted = true;
30            let resumed = self.resumed.clone();
31            let waker = cx.waker().clone();
32            // When the awaiting code is a task of this dispatcher, the job
33            // polls it directly, so that it continues inside the job of the
34            // requested priority. Any other executor is woken up instead.
35            let task = current_task_of(&self.dispatcher).map(|task_id| (self.dispatcher.clone(), task_id));
36            self.dispatcher.post(
37                move || {
38                    resumed.store(true, Ordering::SeqCst);
39                    let polled = match &task {
40                        Some((dispatcher, task_id)) => dispatcher.poll_task(*task_id),
41                        None => false,
42                    };
43                    if !polled {
44                        waker.wake();
45                    }
46                },
47                self.priority,
48            );
49        }
50        Poll::Pending
51    }
52}
53
54/// A future that completes after the dispatcher has processed a job at the
55/// requested priority, which resumes the awaiting code from the dispatcher
56/// thread. Created by [`Dispatcher::resume`] and [`Dispatcher::yield_now`].
57///
58/// Like its counterpart in the reference implementation it is never
59/// complete when first awaited.
60pub struct DispatcherPriorityAwaitable {
61    resumption: Resumption,
62}
63
64impl DispatcherPriorityAwaitable {
65    pub(crate) fn new(dispatcher: Arc<Dispatcher>, priority: DispatcherPriority) -> Self {
66        Self { resumption: Resumption::new(dispatcher, priority) }
67    }
68}
69
70impl Future for DispatcherPriorityAwaitable {
71    type Output = ();
72
73    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
74        self.resumption.poll(cx)
75    }
76}
77
78/// A future that awaits another future and then resumes the awaiting code
79/// from a dispatcher job of the requested priority, yielding the inner
80/// future's output. Created by [`Dispatcher::await_with_priority`].
81///
82/// The awaiting code is resumed through the dispatcher even when the inner
83/// future is already complete.
84pub struct DispatcherPriorityTaskAwaitable<F: Future> {
85    task: Pin<Box<F>>,
86    output: Option<F::Output>,
87    resumption: Resumption,
88}
89
90// The inner future is boxed and the output is never pinned.
91impl<F: Future> Unpin for DispatcherPriorityTaskAwaitable<F> {}
92
93impl<F: Future> DispatcherPriorityTaskAwaitable<F> {
94    pub(crate) fn new(dispatcher: Arc<Dispatcher>, task: F, priority: DispatcherPriority) -> Self {
95        Self { task: Box::pin(task), output: None, resumption: Resumption::new(dispatcher, priority) }
96    }
97}
98
99impl<F: Future> Future for DispatcherPriorityTaskAwaitable<F> {
100    type Output = F::Output;
101
102    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
103        let this = &mut *self;
104        if this.output.is_none() && !this.resumption.posted {
105            match this.task.as_mut().poll(cx) {
106                Poll::Ready(output) => this.output = Some(output),
107                Poll::Pending => return Poll::Pending,
108            }
109        }
110        match this.resumption.poll(cx) {
111            Poll::Ready(()) => match this.output.take() {
112                Some(output) => Poll::Ready(output),
113                None => panic!("DispatcherPriorityTaskAwaitable polled after completion"),
114            },
115            Poll::Pending => Poll::Pending,
116        }
117    }
118}