freya_core/lifecycle/
task.rs1use 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
43pub 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
62pub fn spawn_in_window(future: impl Future<Output = ()> + 'static) -> TaskHandle {
73 spawn_in_scope(future, ScopeId::ROOT)
74}
75
76pub fn spawn(future: impl Future<Output = ()> + 'static) -> TaskHandle {
88 spawn_in_scope(future, None)
89}
90
91pub 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#[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 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 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 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 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#[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 pub fn cancel(&self) {
220 self.0.0.cancel();
221 }
222
223 pub fn try_cancel(&self) {
225 self.0.0.try_cancel();
226 }
227
228 pub fn is_finished(&self) -> bool {
230 self.0.0.is_finished()
231 }
232
233 pub fn downgrade(&self) -> TaskHandle {
235 self.0.0.clone()
236 }
237}
238
239pub 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
253pub struct Task {
255 pub scope_id: ScopeId,
256 pub future: LocalBoxFuture<'static, ()>,
257 pub waker: Waker,
259}
260
261#[derive(Clone, Debug, Copy, PartialEq, Eq, Hash)]
263pub struct TaskId(u64);
264
265#[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 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 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 pub fn clear(&self) {
324 let _removed_tasks = take(&mut *self.inner.tasks.borrow_mut());
325 }
326
327 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}