Skip to main content

freya_core/lifecycle/
task.rs

1use std::{
2    cell::{
3        Cell,
4        RefCell,
5    },
6    hash::{
7        Hash,
8        Hasher,
9    },
10    mem::take,
11    rc::{
12        Rc,
13        Weak,
14    },
15    sync::{
16        Arc,
17        atomic::Ordering,
18    },
19    task::{
20        Context,
21        Waker,
22    },
23};
24
25use futures_channel::mpsc::UnboundedSender;
26use futures_util::{
27    future::LocalBoxFuture,
28    task::{
29        ArcWake,
30        waker,
31    },
32};
33use rustc_hash::FxHashMap;
34
35use crate::{
36    current_context::CurrentContext,
37    lifecycle::global_context::GlobalContexts,
38    prelude::current_scope_id,
39    runner::Message,
40    scope_id::ScopeId,
41};
42
43/// Spawn an event-loop task with global context that survives window closure.
44/// It has no component context and must not retain window-owned state.
45///
46/// ```rust,no_run
47/// # use freya::prelude::*;
48/// # async fn my_task() {}
49/// # fn app() -> impl IntoElement {
50/// use_hook(|| spawn_global(my_task()));
51/// rect()
52/// # }
53/// ```
54pub fn spawn_global(future: impl Future<Output = ()> + 'static) -> TaskHandle {
55    if let Some(tasks) = GlobalTasks::try_get() {
56        tasks.spawn(future)
57    } else {
58        spawn_in_window(future)
59    }
60}
61
62/// Spawn a task that survives its component but is cancelled when its window closes.
63///
64/// ```rust,no_run
65/// # use freya::prelude::*;
66/// # async fn my_task() {}
67/// # fn app() -> impl IntoElement {
68/// use_hook(|| spawn_in_window(my_task()));
69/// rect()
70/// # }
71/// ```
72pub fn spawn_in_window(future: impl Future<Output = ()> + 'static) -> TaskHandle {
73    spawn_in_scope(future, ScopeId::ROOT)
74}
75
76/// Spawn a task cancelled when its current component unmounts.
77/// The returned [`TaskHandle`] can cancel it earlier.
78///
79/// ```rust,no_run
80/// # use freya::prelude::*;
81/// # async fn load_data() {}
82/// # fn app() -> impl IntoElement {
83/// use_hook(|| spawn(load_data()));
84/// rect()
85/// # }
86/// ```
87pub fn spawn(future: impl Future<Output = ()> + 'static) -> TaskHandle {
88    spawn_in_scope(future, None)
89}
90
91/// Spawn a task attached to the given scope, or to the current one when passing `None`.
92pub fn spawn_in_scope(
93    future: impl Future<Output = ()> + 'static,
94    scope_id: impl Into<Option<ScopeId>>,
95) -> TaskHandle {
96    let scope_id = scope_id.into().unwrap_or_else(current_scope_id);
97    CurrentContext::with(|context| {
98        let task_id = TaskId(context.task_id_counter.fetch_add(1, Ordering::Relaxed));
99        context.tasks.borrow_mut().insert(
100            task_id,
101            Rc::new(RefCell::new(Task {
102                scope_id,
103                future: Box::pin(future),
104                waker: waker(Arc::new(TaskWaker {
105                    task_id,
106                    sender: context.sender.clone(),
107                })),
108            })),
109        );
110        context
111            .sender
112            .unbounded_send(Message::PollTask(task_id))
113            .unwrap();
114        TaskHandle {
115            task_id,
116            owner: TaskOwner::Window,
117            global_tasks: Weak::new(),
118        }
119    })
120}
121
122/// A non-owning task handle whose drop does not cancel the task.
123/// Use [`TaskHandle::owned`] for cancellation on drop.
124#[derive(Clone, Debug)]
125pub struct TaskHandle {
126    task_id: TaskId,
127    owner: TaskOwner,
128    global_tasks: Weak<InnerGlobalTasks>,
129}
130
131impl PartialEq for TaskHandle {
132    fn eq(&self, other: &Self) -> bool {
133        self.task_id == other.task_id
134            && self.owner == other.owner
135            && Weak::ptr_eq(&self.global_tasks, &other.global_tasks)
136    }
137}
138
139impl Eq for TaskHandle {}
140
141impl Hash for TaskHandle {
142    fn hash<H: Hasher>(&self, state: &mut H) {
143        self.task_id.hash(state);
144        self.owner.hash(state);
145        self.global_tasks.as_ptr().hash(state);
146    }
147}
148
149#[derive(Clone, Debug, Copy, PartialEq, Eq, Hash)]
150enum TaskOwner {
151    Window,
152    Global,
153}
154
155impl TaskHandle {
156    /// Cancel the task. Component and window tasks require Freya's current context.
157    pub fn cancel(&self) {
158        if self.owner == TaskOwner::Window {
159            CurrentContext::with(|context| context.tasks.borrow_mut().remove(&self.task_id));
160        } else {
161            self.try_cancel();
162        }
163    }
164
165    /// Try to cancel the task if Freya's current context is available.
166    ///
167    /// Unlike [`TaskHandle::cancel`], this does nothing when called outside
168    /// Freya's context. Prefer it in destructors and other cleanup paths where a
169    /// context might no longer exist.
170    pub fn try_cancel(&self) {
171        if self.owner == TaskOwner::Window {
172            CurrentContext::try_with(|context| context.tasks.borrow_mut().remove(&self.task_id));
173        } else if let Some(tasks) = self.global_tasks.upgrade() {
174            tasks.cancel(self.task_id);
175        }
176    }
177
178    /// Check whether the task is no longer scheduled.
179    pub fn is_finished(&self) -> bool {
180        if self.owner == TaskOwner::Window {
181            CurrentContext::with(|context| !context.tasks.borrow().contains_key(&self.task_id))
182        } else {
183            self.global_tasks
184                .upgrade()
185                .is_none_or(|tasks| !tasks.tasks.borrow().contains_key(&self.task_id))
186        }
187    }
188
189    /// Upgrade to an [`OwnedTaskHandle`] that cancels the task when its last
190    /// clone is dropped.
191    /// Retain the returned handle for as long as the task should run. Useful for
192    /// a task owned by another long-lived value rather than by a component scope.
193    pub fn owned(self) -> OwnedTaskHandle {
194        OwnedTaskHandle(Rc::new(InnerOwnedTaskHandle(self)))
195    }
196}
197
198struct InnerOwnedTaskHandle(TaskHandle);
199
200impl Drop for InnerOwnedTaskHandle {
201    fn drop(&mut self) {
202        self.0.try_cancel();
203    }
204}
205
206/// An owning handle that cancels its task when the last clone is dropped.
207/// Created with [`TaskHandle::owned`].
208#[derive(Clone)]
209pub struct OwnedTaskHandle(Rc<InnerOwnedTaskHandle>);
210
211impl PartialEq for OwnedTaskHandle {
212    fn eq(&self, other: &Self) -> bool {
213        Rc::ptr_eq(&self.0, &other.0)
214    }
215}
216
217impl OwnedTaskHandle {
218    /// Cancel the owned task. Component and window tasks require Freya's current context.
219    pub fn cancel(&self) {
220        self.0.0.cancel();
221    }
222
223    /// Try to cancel the task.
224    pub fn try_cancel(&self) {
225        self.0.0.try_cancel();
226    }
227
228    /// Check whether the task is no longer scheduled.
229    pub fn is_finished(&self) -> bool {
230        self.0.0.is_finished()
231    }
232
233    /// Get a non-owning handle.
234    pub fn downgrade(&self) -> TaskHandle {
235        self.0.0.clone()
236    }
237}
238
239/// Wakes a Freya task by asking the runner to poll it again.
240pub struct TaskWaker {
241    task_id: TaskId,
242    sender: UnboundedSender<Message>,
243}
244
245impl ArcWake for TaskWaker {
246    fn wake_by_ref(arc_self: &Arc<Self>) {
247        _ = arc_self
248            .sender
249            .unbounded_send(Message::PollTask(arc_self.task_id));
250    }
251}
252
253/// A component or window task scheduled by the runner.
254pub struct Task {
255    pub scope_id: ScopeId,
256    pub future: LocalBoxFuture<'static, ()>,
257    /// Used to notify the runner that this task needs progress.
258    pub waker: Waker,
259}
260
261/// The opaque identifier of a task scheduled by Freya's async runtime.
262#[derive(Clone, Debug, Copy, PartialEq, Eq, Hash)]
263pub struct TaskId(u64);
264
265/// Tasks owned by the event loop rather than a window runner.
266#[derive(Clone)]
267pub struct GlobalTasks {
268    inner: Rc<InnerGlobalTasks>,
269}
270
271struct InnerGlobalTasks {
272    tasks: RefCell<FxHashMap<TaskId, Rc<RefCell<LocalBoxFuture<'static, ()>>>>>,
273    next_id: Cell<u64>,
274    waker: Waker,
275}
276
277impl InnerGlobalTasks {
278    fn cancel(&self, task_id: TaskId) {
279        let _removed_task = self.tasks.borrow_mut().remove(&task_id);
280    }
281}
282
283impl GlobalTasks {
284    #[track_caller]
285    pub fn get() -> Self {
286        Self::try_get().expect("Global tasks are unavailable outside the event loop.")
287    }
288
289    pub fn try_get() -> Option<Self> {
290        GlobalContexts::try_get().and_then(|contexts| contexts.try_get_context::<Self>())
291    }
292
293    /// Create a task pool woken by the event loop.
294    pub fn new(waker: Waker) -> Self {
295        Self {
296            inner: Rc::new(InnerGlobalTasks {
297                tasks: RefCell::default(),
298                next_id: Cell::default(),
299                waker,
300            }),
301        }
302    }
303
304    /// Schedule a future on the event loop.
305    pub fn spawn(&self, future: impl Future<Output = ()> + 'static) -> TaskHandle {
306        let task_id = TaskId(self.inner.next_id.get());
307        self.inner.next_id.set(self.inner.next_id.get() + 1);
308
309        self.inner
310            .tasks
311            .borrow_mut()
312            .insert(task_id, Rc::new(RefCell::new(Box::pin(future))));
313        self.inner.waker.wake_by_ref();
314
315        TaskHandle {
316            task_id,
317            owner: TaskOwner::Global,
318            global_tasks: Rc::downgrade(&self.inner),
319        }
320    }
321
322    /// Cancel all tasks before the event loop is dropped.
323    pub fn clear(&self) {
324        let _removed_tasks = take(&mut *self.inner.tasks.borrow_mut());
325    }
326
327    /// Poll all event-loop tasks that were scheduled before this call.
328    pub fn poll(&self) {
329        let task_ids: Vec<_> = self.inner.tasks.borrow().keys().copied().collect();
330        let mut context = Context::from_waker(&self.inner.waker);
331
332        for task_id in task_ids {
333            let Some(future) = self.inner.tasks.borrow().get(&task_id).cloned() else {
334                continue;
335            };
336
337            let poll_result = future.borrow_mut().as_mut().poll(&mut context);
338            if poll_result.is_ready() {
339                self.inner.cancel(task_id);
340            }
341        }
342    }
343}