1use crate::{Priority, RunnableMeta, Scheduler, SessionId, Timer};
2use std::{
3 future::Future,
4 marker::PhantomData,
5 mem::ManuallyDrop,
6 panic::Location,
7 pin::Pin,
8 rc::Rc,
9 sync::{
10 Arc,
11 atomic::{AtomicBool, Ordering},
12 },
13 task::{Context, Poll},
14 thread::{self, ThreadId},
15 time::{Duration, Instant},
16};
17
18#[derive(Clone)]
19pub struct ForegroundExecutor {
20 session_id: SessionId,
21 scheduler: Arc<dyn Scheduler>,
22 closed: Arc<AtomicBool>,
23 not_send: PhantomData<Rc<()>>,
24}
25
26impl ForegroundExecutor {
27 pub fn new(session_id: SessionId, scheduler: Arc<dyn Scheduler>) -> Self {
28 Self {
29 session_id,
30 scheduler,
31 closed: Arc::new(AtomicBool::new(false)),
32 not_send: PhantomData,
33 }
34 }
35
36 pub fn session_id(&self) -> SessionId {
37 self.session_id
38 }
39
40 pub fn scheduler(&self) -> &Arc<dyn Scheduler> {
41 &self.scheduler
42 }
43
44 /// Returns the closed flag for this executor.
45 pub fn closed(&self) -> &Arc<AtomicBool> {
46 &self.closed
47 }
48
49 /// Close this executor. Tasks will not run after this is called.
50 pub fn close(&self) {
51 self.closed.store(true, Ordering::SeqCst);
52 }
53
54 #[track_caller]
55 pub fn spawn<F>(&self, future: F) -> Task<F::Output>
56 where
57 F: Future + 'static,
58 F::Output: 'static,
59 {
60 let session_id = self.session_id;
61 let scheduler = Arc::clone(&self.scheduler);
62 let location = Location::caller();
63 let closed = self.closed.clone();
64 let (runnable, task) = spawn_local_with_source_location(
65 future,
66 move |runnable| {
67 scheduler.schedule_foreground(session_id, runnable);
68 },
69 RunnableMeta { location, closed },
70 );
71 runnable.schedule();
72 Task(TaskState::Spawned(task))
73 }
74
75 pub fn block_on<Fut: Future>(&self, future: Fut) -> Fut::Output {
76 use std::cell::Cell;
77
78 let output = Cell::new(None);
79 let future = async {
80 output.set(Some(future.await));
81 };
82 let mut future = std::pin::pin!(future);
83
84 self.scheduler
85 .block(Some(self.session_id), future.as_mut(), None);
86
87 output.take().expect("block_on future did not complete")
88 }
89
90 /// Block until the future completes or timeout occurs.
91 /// Returns Ok(output) if completed, Err(future) if timed out.
92 pub fn block_with_timeout<Fut: Future>(
93 &self,
94 timeout: Duration,
95 future: Fut,
96 ) -> Result<Fut::Output, impl Future<Output = Fut::Output> + use<Fut>> {
97 use std::cell::Cell;
98
99 let output = Cell::new(None);
100 let mut future = Box::pin(future);
101
102 {
103 let future_ref = &mut future;
104 let wrapper = async {
105 output.set(Some(future_ref.await));
106 };
107 let mut wrapper = std::pin::pin!(wrapper);
108
109 self.scheduler
110 .block(Some(self.session_id), wrapper.as_mut(), Some(timeout));
111 }
112
113 match output.take() {
114 Some(value) => Ok(value),
115 None => Err(future),
116 }
117 }
118
119 pub fn timer(&self, duration: Duration) -> Timer {
120 self.scheduler.timer(duration)
121 }
122
123 pub fn now(&self) -> Instant {
124 self.scheduler.clock().now()
125 }
126}
127
128#[derive(Clone)]
129pub struct BackgroundExecutor {
130 scheduler: Arc<dyn Scheduler>,
131 closed: Arc<AtomicBool>,
132}
133
134impl BackgroundExecutor {
135 pub fn new(scheduler: Arc<dyn Scheduler>) -> Self {
136 Self {
137 scheduler,
138 closed: Arc::new(AtomicBool::new(false)),
139 }
140 }
141
142 /// Returns the closed flag for this executor.
143 pub fn closed(&self) -> &Arc<AtomicBool> {
144 &self.closed
145 }
146
147 /// Close this executor. Tasks will not run after this is called.
148 pub fn close(&self) {
149 self.closed.store(true, Ordering::SeqCst);
150 }
151
152 #[track_caller]
153 pub fn spawn<F>(&self, future: F) -> Task<F::Output>
154 where
155 F: Future + Send + 'static,
156 F::Output: Send + 'static,
157 {
158 self.spawn_with_priority(Priority::default(), future)
159 }
160
161 #[track_caller]
162 pub fn spawn_with_priority<F>(&self, priority: Priority, future: F) -> Task<F::Output>
163 where
164 F: Future + Send + 'static,
165 F::Output: Send + 'static,
166 {
167 let scheduler = Arc::clone(&self.scheduler);
168 let location = Location::caller();
169 let closed = self.closed.clone();
170 let (runnable, task) = async_task::Builder::new()
171 .metadata(RunnableMeta { location, closed })
172 .spawn(
173 move |_| future,
174 move |runnable| {
175 scheduler.schedule_background_with_priority(runnable, priority);
176 },
177 );
178 runnable.schedule();
179 Task(TaskState::Spawned(task))
180 }
181
182 /// Spawns a future on a dedicated realtime thread for audio processing.
183 #[track_caller]
184 pub fn spawn_realtime<F>(&self, future: F) -> Task<F::Output>
185 where
186 F: Future + Send + 'static,
187 F::Output: Send + 'static,
188 {
189 let location = Location::caller();
190 let closed = self.closed.clone();
191 let (tx, rx) = flume::bounded::<async_task::Runnable<RunnableMeta>>(1);
192
193 self.scheduler.spawn_realtime(Box::new(move || {
194 while let Ok(runnable) = rx.recv() {
195 if runnable.metadata().is_closed() {
196 continue;
197 }
198 runnable.run();
199 }
200 }));
201
202 let (runnable, task) = async_task::Builder::new()
203 .metadata(RunnableMeta { location, closed })
204 .spawn(
205 move |_| future,
206 move |runnable| {
207 let _ = tx.send(runnable);
208 },
209 );
210 runnable.schedule();
211 Task(TaskState::Spawned(task))
212 }
213
214 pub fn timer(&self, duration: Duration) -> Timer {
215 self.scheduler.timer(duration)
216 }
217
218 pub fn now(&self) -> Instant {
219 self.scheduler.clock().now()
220 }
221
222 pub fn scheduler(&self) -> &Arc<dyn Scheduler> {
223 &self.scheduler
224 }
225}
226
227/// Task is a primitive that allows work to happen in the background.
228///
229/// It implements [`Future`] so you can `.await` on it.
230///
231/// If you drop a task it will be cancelled immediately. Calling [`Task::detach`] allows
232/// the task to continue running, but with no way to return a value.
233#[must_use]
234#[derive(Debug)]
235pub struct Task<T>(TaskState<T>);
236
237#[derive(Debug)]
238enum TaskState<T> {
239 /// A task that is ready to return a value
240 Ready(Option<T>),
241
242 /// A task that is currently running.
243 Spawned(async_task::Task<T, RunnableMeta>),
244}
245
246impl<T> Task<T> {
247 /// Creates a new task that will resolve with the value
248 pub fn ready(val: T) -> Self {
249 Task(TaskState::Ready(Some(val)))
250 }
251
252 /// Creates a Task from an async_task::Task
253 pub fn from_async_task(task: async_task::Task<T, RunnableMeta>) -> Self {
254 Task(TaskState::Spawned(task))
255 }
256
257 pub fn is_ready(&self) -> bool {
258 match &self.0 {
259 TaskState::Ready(_) => true,
260 TaskState::Spawned(task) => task.is_finished(),
261 }
262 }
263
264 /// Detaching a task runs it to completion in the background
265 pub fn detach(self) {
266 match self {
267 Task(TaskState::Ready(_)) => {}
268 Task(TaskState::Spawned(task)) => task.detach(),
269 }
270 }
271
272 /// Converts this task into a fallible task that returns `Option<T>`.
273 pub fn fallible(self) -> FallibleTask<T> {
274 FallibleTask(match self.0 {
275 TaskState::Ready(val) => FallibleTaskState::Ready(val),
276 TaskState::Spawned(task) => FallibleTaskState::Spawned(task.fallible()),
277 })
278 }
279}
280
281/// A task that returns `Option<T>` instead of panicking when cancelled.
282#[must_use]
283pub struct FallibleTask<T>(FallibleTaskState<T>);
284
285enum FallibleTaskState<T> {
286 /// A task that is ready to return a value
287 Ready(Option<T>),
288
289 /// A task that is currently running (wraps async_task::FallibleTask).
290 Spawned(async_task::FallibleTask<T, RunnableMeta>),
291}
292
293impl<T> FallibleTask<T> {
294 /// Creates a new fallible task that will resolve with the value.
295 pub fn ready(val: T) -> Self {
296 FallibleTask(FallibleTaskState::Ready(Some(val)))
297 }
298
299 /// Detaching a task runs it to completion in the background.
300 pub fn detach(self) {
301 match self.0 {
302 FallibleTaskState::Ready(_) => {}
303 FallibleTaskState::Spawned(task) => task.detach(),
304 }
305 }
306}
307
308impl<T> Future for FallibleTask<T> {
309 type Output = Option<T>;
310
311 fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
312 match unsafe { self.get_unchecked_mut() } {
313 FallibleTask(FallibleTaskState::Ready(val)) => Poll::Ready(val.take()),
314 FallibleTask(FallibleTaskState::Spawned(task)) => Pin::new(task).poll(cx),
315 }
316 }
317}
318
319impl<T> std::fmt::Debug for FallibleTask<T> {
320 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
321 match &self.0 {
322 FallibleTaskState::Ready(_) => f.debug_tuple("FallibleTask::Ready").finish(),
323 FallibleTaskState::Spawned(task) => {
324 f.debug_tuple("FallibleTask::Spawned").field(task).finish()
325 }
326 }
327 }
328}
329
330impl<T> Future for Task<T> {
331 type Output = T;
332
333 fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
334 match unsafe { self.get_unchecked_mut() } {
335 Task(TaskState::Ready(val)) => Poll::Ready(val.take().unwrap()),
336 Task(TaskState::Spawned(task)) => Pin::new(task).poll(cx),
337 }
338 }
339}
340
341/// Variant of `async_task::spawn_local` that includes the source location of the spawn in panics.
342#[track_caller]
343fn spawn_local_with_source_location<Fut, S>(
344 future: Fut,
345 schedule: S,
346 metadata: RunnableMeta,
347) -> (
348 async_task::Runnable<RunnableMeta>,
349 async_task::Task<Fut::Output, RunnableMeta>,
350)
351where
352 Fut: Future + 'static,
353 Fut::Output: 'static,
354 S: async_task::Schedule<RunnableMeta> + Send + Sync + 'static,
355{
356 #[inline]
357 fn thread_id() -> ThreadId {
358 std::thread_local! {
359 static ID: ThreadId = thread::current().id();
360 }
361 ID.try_with(|id| *id)
362 .unwrap_or_else(|_| thread::current().id())
363 }
364
365 struct Checked<F> {
366 id: ThreadId,
367 inner: ManuallyDrop<F>,
368 location: &'static Location<'static>,
369 }
370
371 impl<F> Drop for Checked<F> {
372 fn drop(&mut self) {
373 assert!(
374 self.id == thread_id(),
375 "local task dropped by a thread that didn't spawn it. Task spawned at {}",
376 self.location
377 );
378 unsafe {
379 ManuallyDrop::drop(&mut self.inner);
380 }
381 }
382 }
383
384 impl<F: Future> Future for Checked<F> {
385 type Output = F::Output;
386
387 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
388 assert!(
389 self.id == thread_id(),
390 "local task polled by a thread that didn't spawn it. Task spawned at {}",
391 self.location
392 );
393 unsafe { self.map_unchecked_mut(|c| &mut *c.inner).poll(cx) }
394 }
395 }
396
397 let location = metadata.location;
398
399 unsafe {
400 async_task::Builder::new()
401 .metadata(metadata)
402 .spawn_unchecked(
403 move |_| Checked {
404 id: thread_id(),
405 inner: ManuallyDrop::new(future),
406 location,
407 },
408 schedule,
409 )
410 }
411}