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 pub fn timer(&self, duration: Duration) -> Timer {
183 self.scheduler.timer(duration)
184 }
185
186 pub fn now(&self) -> Instant {
187 self.scheduler.clock().now()
188 }
189
190 pub fn scheduler(&self) -> &Arc<dyn Scheduler> {
191 &self.scheduler
192 }
193}
194
195/// Task is a primitive that allows work to happen in the background.
196///
197/// It implements [`Future`] so you can `.await` on it.
198///
199/// If you drop a task it will be cancelled immediately. Calling [`Task::detach`] allows
200/// the task to continue running, but with no way to return a value.
201#[must_use]
202#[derive(Debug)]
203pub struct Task<T>(TaskState<T>);
204
205#[derive(Debug)]
206enum TaskState<T> {
207 /// A task that is ready to return a value
208 Ready(Option<T>),
209
210 /// A task that is currently running.
211 Spawned(async_task::Task<T, RunnableMeta>),
212}
213
214impl<T> Task<T> {
215 /// Creates a new task that will resolve with the value
216 pub fn ready(val: T) -> Self {
217 Task(TaskState::Ready(Some(val)))
218 }
219
220 /// Creates a Task from an async_task::Task
221 pub fn from_async_task(task: async_task::Task<T, RunnableMeta>) -> Self {
222 Task(TaskState::Spawned(task))
223 }
224
225 pub fn is_ready(&self) -> bool {
226 match &self.0 {
227 TaskState::Ready(_) => true,
228 TaskState::Spawned(task) => task.is_finished(),
229 }
230 }
231
232 /// Detaching a task runs it to completion in the background
233 pub fn detach(self) {
234 match self {
235 Task(TaskState::Ready(_)) => {}
236 Task(TaskState::Spawned(task)) => task.detach(),
237 }
238 }
239
240 /// Converts this task into a fallible task that returns `Option<T>`.
241 pub fn fallible(self) -> FallibleTask<T> {
242 FallibleTask(match self.0 {
243 TaskState::Ready(val) => FallibleTaskState::Ready(val),
244 TaskState::Spawned(task) => FallibleTaskState::Spawned(task.fallible()),
245 })
246 }
247}
248
249/// A task that returns `Option<T>` instead of panicking when cancelled.
250#[must_use]
251pub struct FallibleTask<T>(FallibleTaskState<T>);
252
253enum FallibleTaskState<T> {
254 /// A task that is ready to return a value
255 Ready(Option<T>),
256
257 /// A task that is currently running (wraps async_task::FallibleTask).
258 Spawned(async_task::FallibleTask<T, RunnableMeta>),
259}
260
261impl<T> FallibleTask<T> {
262 /// Creates a new fallible task that will resolve with the value.
263 pub fn ready(val: T) -> Self {
264 FallibleTask(FallibleTaskState::Ready(Some(val)))
265 }
266
267 /// Detaching a task runs it to completion in the background.
268 pub fn detach(self) {
269 match self.0 {
270 FallibleTaskState::Ready(_) => {}
271 FallibleTaskState::Spawned(task) => task.detach(),
272 }
273 }
274}
275
276impl<T> Future for FallibleTask<T> {
277 type Output = Option<T>;
278
279 fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
280 match unsafe { self.get_unchecked_mut() } {
281 FallibleTask(FallibleTaskState::Ready(val)) => Poll::Ready(val.take()),
282 FallibleTask(FallibleTaskState::Spawned(task)) => Pin::new(task).poll(cx),
283 }
284 }
285}
286
287impl<T> std::fmt::Debug for FallibleTask<T> {
288 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
289 match &self.0 {
290 FallibleTaskState::Ready(_) => f.debug_tuple("FallibleTask::Ready").finish(),
291 FallibleTaskState::Spawned(task) => {
292 f.debug_tuple("FallibleTask::Spawned").field(task).finish()
293 }
294 }
295 }
296}
297
298impl<T> Future for Task<T> {
299 type Output = T;
300
301 fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
302 match unsafe { self.get_unchecked_mut() } {
303 Task(TaskState::Ready(val)) => Poll::Ready(val.take().unwrap()),
304 Task(TaskState::Spawned(task)) => Pin::new(task).poll(cx),
305 }
306 }
307}
308
309/// Variant of `async_task::spawn_local` that includes the source location of the spawn in panics.
310#[track_caller]
311fn spawn_local_with_source_location<Fut, S>(
312 future: Fut,
313 schedule: S,
314 metadata: RunnableMeta,
315) -> (
316 async_task::Runnable<RunnableMeta>,
317 async_task::Task<Fut::Output, RunnableMeta>,
318)
319where
320 Fut: Future + 'static,
321 Fut::Output: 'static,
322 S: async_task::Schedule<RunnableMeta> + Send + Sync + 'static,
323{
324 #[inline]
325 fn thread_id() -> ThreadId {
326 std::thread_local! {
327 static ID: ThreadId = thread::current().id();
328 }
329 ID.try_with(|id| *id)
330 .unwrap_or_else(|_| thread::current().id())
331 }
332
333 struct Checked<F> {
334 id: ThreadId,
335 inner: ManuallyDrop<F>,
336 location: &'static Location<'static>,
337 }
338
339 impl<F> Drop for Checked<F> {
340 fn drop(&mut self) {
341 assert!(
342 self.id == thread_id(),
343 "local task dropped by a thread that didn't spawn it. Task spawned at {}",
344 self.location
345 );
346 unsafe {
347 ManuallyDrop::drop(&mut self.inner);
348 }
349 }
350 }
351
352 impl<F: Future> Future for Checked<F> {
353 type Output = F::Output;
354
355 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
356 assert!(
357 self.id == thread_id(),
358 "local task polled by a thread that didn't spawn it. Task spawned at {}",
359 self.location
360 );
361 unsafe { self.map_unchecked_mut(|c| &mut *c.inner).poll(cx) }
362 }
363 }
364
365 let location = metadata.location;
366
367 unsafe {
368 async_task::Builder::new()
369 .metadata(metadata)
370 .spawn_unchecked(
371 move |_| Checked {
372 id: thread_id(),
373 inner: ManuallyDrop::new(future),
374 location,
375 },
376 schedule,
377 )
378 }
379}