1use crate::{Instant, Priority, RunnableMeta, Scheduler, SessionId, Timer};
2use async_task::Runnable;
3use std::{
4 any::Any,
5 future::Future,
6 marker::PhantomData,
7 mem::ManuallyDrop,
8 panic::Location,
9 pin::Pin,
10 rc::Rc,
11 sync::Arc,
12 task::{Context, Poll, Waker},
13 thread::{self, ThreadId},
14 time::Duration,
15};
16
17#[derive(Clone)]
22pub struct LocalExecutor {
23 session_id: SessionId,
24 scheduler: Arc<dyn Scheduler>,
25 dispatch: Arc<dyn Fn(Runnable<RunnableMeta>) + Send + Sync>,
29 not_send: PhantomData<Rc<()>>,
30}
31
32impl LocalExecutor {
33 pub fn new(
42 session_id: SessionId,
43 scheduler: Arc<dyn Scheduler>,
44 dispatch: impl Fn(Runnable<RunnableMeta>) + Send + Sync + 'static,
45 ) -> Self {
46 Self {
47 session_id,
48 scheduler,
49 dispatch: Arc::new(dispatch),
50 not_send: PhantomData,
51 }
52 }
53
54 pub fn session_id(&self) -> SessionId {
55 self.session_id
56 }
57
58 pub fn scheduler(&self) -> &Arc<dyn Scheduler> {
59 &self.scheduler
60 }
61
62 #[track_caller]
63 pub fn spawn<F>(&self, future: F) -> Task<F::Output>
64 where
65 F: Future + 'static,
66 F::Output: 'static,
67 {
68 let schedule = self.schedule();
69 let location = Location::caller();
70 let (runnable, task) = spawn_local_with_source_location(
71 future,
72 schedule,
73 RunnableMeta {
74 location,
75 spawned: crate::SpawnTime(Instant::now()),
76 },
77 );
78 runnable.schedule();
79 Task(TaskState::Spawned(task))
80 }
81
82 #[track_caller]
87 pub fn spawn_with_dispatch<F>(
88 &self,
89 future: F,
90 dispatch: impl Fn(Runnable<RunnableMeta>) + Send + Sync + 'static,
91 ) -> Task<F::Output>
92 where
93 F: Future + 'static,
94 F::Output: 'static,
95 {
96 let location = Location::caller();
97 let (runnable, task) = spawn_local_with_source_location(
98 future,
99 dispatch,
100 RunnableMeta {
101 location,
102 spawned: crate::SpawnTime(Instant::now()),
103 },
104 );
105 runnable.schedule();
106 Task(TaskState::Spawned(task))
107 }
108
109 #[cfg(not(target_family = "wasm"))]
110 pub fn block_on<Fut: Future>(&self, future: Fut) -> Fut::Output {
111 use std::cell::Cell;
112
113 let output = Cell::new(None);
114 let future = async {
115 output.set(Some(future.await));
116 };
117 let mut future = std::pin::pin!(future);
118
119 self.scheduler
120 .block(Some(self.session_id), future.as_mut(), None);
121
122 output.take().expect("block_on future did not complete")
123 }
124
125 #[cfg(not(target_family = "wasm"))]
128 pub fn block_with_timeout<Fut: Future>(
129 &self,
130 timeout: Duration,
131 future: Fut,
132 ) -> Result<Fut::Output, impl Future<Output = Fut::Output> + use<Fut>> {
133 use std::cell::Cell;
134
135 let output = Cell::new(None);
136 let mut future = Box::pin(future);
137
138 {
139 let future_ref = &mut future;
140 let wrapper = async {
141 output.set(Some(future_ref.await));
142 };
143 let mut wrapper = std::pin::pin!(wrapper);
144
145 self.scheduler
146 .block(Some(self.session_id), wrapper.as_mut(), Some(timeout));
147 }
148
149 match output.take() {
150 Some(value) => Ok(value),
151 None => Err(future),
152 }
153 }
154
155 #[track_caller]
156 pub fn timer(&self, duration: Duration) -> Timer {
157 self.scheduler.timer(duration)
158 }
159
160 pub fn now(&self) -> Instant {
161 self.scheduler.clock().now()
162 }
163
164 #[track_caller]
173 pub fn spawn_dedicated<F, Fut>(&self, f: F) -> Task<Fut::Output>
174 where
175 F: FnOnce(LocalExecutor) -> Fut + Send + 'static,
176 Fut: Future + 'static,
177 Fut::Output: Send + Sync + 'static,
178 {
179 self.scheduler
180 .clone()
181 .spawn_dedicated(box_dedicated(f))
182 .downcast::<Fut::Output>()
183 }
184
185 fn schedule(&self) -> impl Fn(Runnable<RunnableMeta>) + Send + Sync + 'static {
186 let dispatch = self.dispatch.clone();
187 move |runnable| dispatch(runnable)
188 }
189}
190
191fn box_dedicated<F, Fut>(
196 f: F,
197) -> Box<
198 dyn FnOnce(LocalExecutor) -> Pin<Box<dyn Future<Output = Box<dyn Any + Send + Sync>> + 'static>>
199 + Send
200 + 'static,
201>
202where
203 F: FnOnce(LocalExecutor) -> Fut + Send + 'static,
204 Fut: Future + 'static,
205 Fut::Output: Send + Sync + 'static,
206{
207 Box::new(move |executor| {
208 Box::pin(async move { Box::new(f(executor).await) as Box<dyn Any + Send + Sync> })
209 })
210}
211
212#[derive(Clone)]
213pub struct BackgroundExecutor {
214 scheduler: Arc<dyn Scheduler>,
215}
216
217impl BackgroundExecutor {
218 pub fn new(scheduler: Arc<dyn Scheduler>) -> Self {
219 Self { scheduler }
220 }
221
222 #[track_caller]
223 pub fn spawn<F>(&self, future: F) -> Task<F::Output>
224 where
225 F: Future + Send + 'static,
226 F::Output: Send + 'static,
227 {
228 self.spawn_with_priority(Priority::default(), future)
229 }
230
231 #[track_caller]
232 pub fn spawn_with_priority<F>(&self, priority: Priority, future: F) -> Task<F::Output>
233 where
234 F: Future + Send + 'static,
235 F::Output: Send + 'static,
236 {
237 let schedule = self.schedule_with_priority(priority);
238 let location = Location::caller();
239 let (runnable, task) = async_task::Builder::new()
240 .metadata(RunnableMeta {
241 location,
242 spawned: crate::SpawnTime(Instant::now()),
243 })
244 .spawn(move |_| future, schedule);
245 runnable.schedule();
246 Task(TaskState::Spawned(task))
247 }
248
249 #[track_caller]
251 pub fn spawn_realtime<F>(&self, future: F) -> Task<F::Output>
252 where
253 F: Future + Send + 'static,
254 F::Output: Send + 'static,
255 {
256 let location = Location::caller();
257 let (tx, rx) = flume::bounded::<async_task::Runnable<RunnableMeta>>(1);
258
259 self.scheduler.spawn_realtime(Box::new(move || {
260 while let Ok(runnable) = rx.recv() {
261 runnable.run();
262 }
263 }));
264
265 let (runnable, task) = async_task::Builder::new()
266 .metadata(RunnableMeta {
267 location,
268 spawned: crate::SpawnTime(Instant::now()),
269 })
270 .spawn(
271 move |_| future,
272 move |runnable| {
273 let _ = tx.send(runnable);
274 },
275 );
276 runnable.schedule();
277 Task(TaskState::Spawned(task))
278 }
279
280 #[track_caller]
281 pub fn timer(&self, duration: Duration) -> Timer {
282 self.scheduler.timer(duration)
283 }
284
285 pub fn now(&self) -> Instant {
286 self.scheduler.clock().now()
287 }
288
289 pub fn scheduler(&self) -> &Arc<dyn Scheduler> {
290 &self.scheduler
291 }
292
293 #[track_caller]
302 pub fn spawn_dedicated<F, Fut>(&self, f: F) -> Task<Fut::Output>
303 where
304 F: FnOnce(LocalExecutor) -> Fut + Send + 'static,
305 Fut: Future + 'static,
306 Fut::Output: Send + Sync + 'static,
307 {
308 self.scheduler
309 .clone()
310 .spawn_dedicated(box_dedicated(f))
311 .downcast::<Fut::Output>()
312 }
313
314 fn schedule_with_priority(
315 &self,
316 priority: Priority,
317 ) -> impl Fn(Runnable<RunnableMeta>) + Send + Sync + 'static {
318 let scheduler = Arc::downgrade(&self.scheduler);
319 move |runnable| {
320 if let Some(scheduler) = scheduler.upgrade() {
321 scheduler.schedule_background_with_priority(runnable, priority);
322 }
323 }
324 }
325}
326
327pub struct DedicatedExecutor {
338 sender: flume::Sender<Runnable<RunnableMeta>>,
339 _session: Task<()>,
340}
341
342impl DedicatedExecutor {
343 #[track_caller]
347 pub fn new(executor: &BackgroundExecutor) -> Self {
348 let (sender, receiver) = flume::unbounded::<Runnable<RunnableMeta>>();
349 let session = executor.spawn_dedicated(move |_executor| async move {
350 while let Ok(runnable) = receiver.recv_async().await {
351 runnable.run();
352 }
353 });
354 Self {
355 sender,
356 _session: session,
357 }
358 }
359
360 #[track_caller]
366 pub fn spawn<F>(&self, future: F) -> Task<F::Output>
367 where
368 F: Future + Send + 'static,
369 F::Output: Send + 'static,
370 {
371 let sender = self.sender.clone();
372 let (runnable, task) = async_task::Builder::new()
373 .metadata(RunnableMeta::new_with_callers_location())
374 .spawn(
375 move |_| future,
376 move |runnable| {
377 let _ = sender.send(runnable);
378 },
379 );
380 runnable.schedule();
381 Task(TaskState::Spawned(task))
382 }
383}
384
385#[must_use]
392pub struct Task<T>(TaskState<T>);
393
394enum TaskState<T> {
395 Ready(Option<T>),
397
398 Spawned(async_task::Task<T, RunnableMeta>),
400
401 Downcast {
405 inner: Box<Task<Box<dyn Any + Send + Sync>>>,
406 marker: PhantomData<fn() -> T>,
407 },
408
409 Rendezvous(RendezvousReceiver<T>),
413}
414
415enum RendezvousState<T> {
417 Pending(Option<Waker>),
419 Delivered(Task<T>),
421 Cancelled,
423 Detached,
425 Taken,
427}
428
429pub(crate) struct RendezvousReceiver<T> {
430 shared: Arc<parking_lot::Mutex<RendezvousState<T>>>,
431}
432
433impl<T> RendezvousReceiver<T> {
434 fn poll_take(&self, cx: &mut Context) -> Poll<Task<T>> {
435 let mut state = self.shared.lock();
436 match std::mem::replace(&mut *state, RendezvousState::Taken) {
437 RendezvousState::Delivered(task) => Poll::Ready(task),
438 RendezvousState::Pending(_) => {
439 *state = RendezvousState::Pending(Some(cx.waker().clone()));
440 Poll::Pending
441 }
442 RendezvousState::Cancelled | RendezvousState::Detached | RendezvousState::Taken => {
443 unreachable!("a rendezvous task was polled after its receiver was consumed")
444 }
445 }
446 }
447
448 fn is_ready(&self) -> bool {
449 match &*self.shared.lock() {
450 RendezvousState::Delivered(task) => task.is_ready(),
451 _ => false,
452 }
453 }
454
455 fn detach(self) {
456 let mut state = self.shared.lock();
457 match std::mem::replace(&mut *state, RendezvousState::Detached) {
458 RendezvousState::Delivered(task) => {
459 *state = RendezvousState::Taken;
460 drop(state);
461 task.detach();
462 }
463 RendezvousState::Pending(_) => {}
465 RendezvousState::Cancelled | RendezvousState::Detached | RendezvousState::Taken => {
466 unreachable!("a rendezvous task was detached after its receiver was consumed")
467 }
468 }
469 }
470}
471
472impl<T> Drop for RendezvousReceiver<T> {
473 fn drop(&mut self) {
474 let mut state = self.shared.lock();
475 match &*state {
476 RendezvousState::Pending(_) => *state = RendezvousState::Cancelled,
477 RendezvousState::Delivered(_) => {
478 let delivered = std::mem::replace(&mut *state, RendezvousState::Cancelled);
479 drop(state);
480 drop(delivered);
481 }
482 _ => {}
486 }
487 }
488}
489
490pub(crate) struct RendezvousSender<T> {
491 shared: Arc<parking_lot::Mutex<RendezvousState<T>>>,
492}
493
494impl<T> RendezvousSender<T> {
495 pub(crate) fn deliver(self, task: Task<T>) {
498 let mut state = self.shared.lock();
499 match std::mem::replace(&mut *state, RendezvousState::Delivered(task)) {
500 RendezvousState::Pending(waker) => {
501 drop(state);
502 if let Some(waker) = waker {
503 waker.wake();
504 }
505 }
506 RendezvousState::Cancelled => {
507 let delivered = std::mem::replace(&mut *state, RendezvousState::Cancelled);
508 drop(state);
509 drop(delivered);
510 }
511 RendezvousState::Detached => {
512 let RendezvousState::Delivered(task) =
513 std::mem::replace(&mut *state, RendezvousState::Detached)
514 else {
515 unreachable!("the delivered task was just stored");
516 };
517 drop(state);
518 task.detach();
519 }
520 RendezvousState::Delivered(_) | RendezvousState::Taken => {
521 unreachable!("a rendezvous task was delivered twice")
522 }
523 }
524 }
525}
526
527impl<T> Task<T> {
528 pub fn ready(val: T) -> Self {
530 Task(TaskState::Ready(Some(val)))
531 }
532
533 pub(crate) fn rendezvous() -> (Self, RendezvousSender<T>) {
538 let shared = Arc::new(parking_lot::Mutex::new(RendezvousState::Pending(None)));
539 (
540 Task(TaskState::Rendezvous(RendezvousReceiver {
541 shared: shared.clone(),
542 })),
543 RendezvousSender { shared },
544 )
545 }
546
547 pub fn from_async_task(task: async_task::Task<T, RunnableMeta>) -> Self {
549 Task(TaskState::Spawned(task))
550 }
551
552 pub fn is_ready(&self) -> bool {
553 match &self.0 {
554 TaskState::Ready(_) => true,
555 TaskState::Spawned(task) => task.is_finished(),
556 TaskState::Downcast { inner, .. } => inner.is_ready(),
557 TaskState::Rendezvous(receiver) => receiver.is_ready(),
558 }
559 }
560
561 pub fn detach(self) {
563 match self {
564 Task(TaskState::Ready(_)) => {}
565 Task(TaskState::Spawned(task)) => task.detach(),
566 Task(TaskState::Downcast { inner, .. }) => inner.detach(),
567 Task(TaskState::Rendezvous(receiver)) => receiver.detach(),
568 }
569 }
570
571 pub fn fallible(self) -> FallibleTask<T> {
573 FallibleTask(match self.0 {
574 TaskState::Ready(val) => FallibleTaskState::Ready(val),
575 TaskState::Spawned(task) => FallibleTaskState::Spawned(task.fallible()),
576 TaskState::Downcast { inner, .. } => FallibleTaskState::Downcast {
577 inner: Box::new(inner.fallible()),
578 marker: PhantomData,
579 },
580 TaskState::Rendezvous(receiver) => FallibleTaskState::Rendezvous(receiver),
581 })
582 }
583}
584
585impl Task<Box<dyn Any + Send + Sync>> {
586 pub fn downcast<T: Send + Sync + 'static>(self) -> Task<T> {
595 Task(TaskState::Downcast {
596 inner: Box::new(self),
597 marker: PhantomData,
598 })
599 }
600}
601
602impl<T> std::fmt::Debug for Task<T> {
603 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
604 match &self.0 {
605 TaskState::Ready(_) => f.debug_tuple("Task::Ready").finish(),
606 TaskState::Spawned(task) => f.debug_tuple("Task::Spawned").field(task).finish(),
607 TaskState::Downcast { inner, .. } => {
608 f.debug_tuple("Task::Downcast").field(inner).finish()
609 }
610 TaskState::Rendezvous(_) => f.debug_tuple("Task::Rendezvous").finish(),
611 }
612 }
613}
614
615#[must_use]
617pub struct FallibleTask<T>(FallibleTaskState<T>);
618
619enum FallibleTaskState<T> {
620 Ready(Option<T>),
622
623 Spawned(async_task::FallibleTask<T, RunnableMeta>),
625
626 Downcast {
628 inner: Box<FallibleTask<Box<dyn Any + Send + Sync>>>,
629 marker: PhantomData<fn() -> T>,
630 },
631
632 Rendezvous(RendezvousReceiver<T>),
634}
635
636impl<T> FallibleTask<T> {
637 pub fn ready(val: T) -> Self {
639 FallibleTask(FallibleTaskState::Ready(Some(val)))
640 }
641
642 pub fn detach(self) {
644 match self.0 {
645 FallibleTaskState::Ready(_) => {}
646 FallibleTaskState::Spawned(task) => task.detach(),
647 FallibleTaskState::Downcast { inner, .. } => inner.detach(),
648 FallibleTaskState::Rendezvous(receiver) => receiver.detach(),
649 }
650 }
651}
652
653impl<T: 'static> Future for FallibleTask<T> {
654 type Output = Option<T>;
655
656 fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
657 let this = unsafe { self.get_unchecked_mut() };
658 loop {
659 match &mut this.0 {
660 FallibleTaskState::Ready(val) => return Poll::Ready(val.take()),
661 FallibleTaskState::Spawned(task) => return Pin::new(task).poll(cx),
662 FallibleTaskState::Downcast { inner, .. } => {
663 return match Pin::new(inner.as_mut()).poll(cx) {
664 Poll::Ready(Some(boxed_any)) => Poll::Ready(Some(
665 *boxed_any
666 .downcast::<T>()
667 .expect("FallibleTask::poll: downcast type mismatch"),
668 )),
669 Poll::Ready(None) => Poll::Ready(None),
670 Poll::Pending => Poll::Pending,
671 };
672 }
673 FallibleTaskState::Rendezvous(receiver) => match receiver.poll_take(cx) {
674 Poll::Ready(task) => {
675 this.0 = task.fallible().0;
676 continue;
677 }
678 Poll::Pending => return Poll::Pending,
679 },
680 }
681 }
682 }
683}
684
685impl<T> std::fmt::Debug for FallibleTask<T> {
686 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
687 match &self.0 {
688 FallibleTaskState::Ready(_) => f.debug_tuple("FallibleTask::Ready").finish(),
689 FallibleTaskState::Spawned(task) => {
690 f.debug_tuple("FallibleTask::Spawned").field(task).finish()
691 }
692 FallibleTaskState::Downcast { inner, .. } => f
693 .debug_tuple("FallibleTask::Downcast")
694 .field(inner)
695 .finish(),
696 FallibleTaskState::Rendezvous(_) => f.debug_tuple("FallibleTask::Rendezvous").finish(),
697 }
698 }
699}
700
701impl<T: 'static> Future for Task<T> {
702 type Output = T;
703
704 fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
705 let this = unsafe { self.get_unchecked_mut() };
706 loop {
707 match &mut this.0 {
708 TaskState::Ready(val) => return Poll::Ready(val.take().unwrap()),
709 TaskState::Spawned(task) => return Pin::new(task).poll(cx),
710 TaskState::Downcast { inner, .. } => {
711 return match Pin::new(inner.as_mut()).poll(cx) {
712 Poll::Ready(boxed_any) => Poll::Ready(
713 *boxed_any
714 .downcast::<T>()
715 .expect("Task::poll: downcast type mismatch"),
716 ),
717 Poll::Pending => Poll::Pending,
718 };
719 }
720 TaskState::Rendezvous(receiver) => match receiver.poll_take(cx) {
721 Poll::Ready(task) => {
722 this.0 = task.0;
723 continue;
724 }
725 Poll::Pending => return Poll::Pending,
726 },
727 }
728 }
729 }
730}
731
732#[track_caller]
734fn spawn_local_with_source_location<Fut, S>(
735 future: Fut,
736 schedule: S,
737 metadata: RunnableMeta,
738) -> (
739 async_task::Runnable<RunnableMeta>,
740 async_task::Task<Fut::Output, RunnableMeta>,
741)
742where
743 Fut: Future + 'static,
744 Fut::Output: 'static,
745 S: async_task::Schedule<RunnableMeta> + Send + Sync + 'static,
746{
747 #[inline]
748 fn thread_id() -> ThreadId {
749 std::thread_local! {
750 static ID: ThreadId = thread::current().id();
751 }
752 ID.try_with(|id| *id)
753 .unwrap_or_else(|_| thread::current().id())
754 }
755
756 struct Checked<F> {
757 id: ThreadId,
758 inner: ManuallyDrop<F>,
759 location: &'static Location<'static>,
760 }
761
762 impl<F> Drop for Checked<F> {
763 fn drop(&mut self) {
764 assert_eq!(
765 self.id,
766 thread_id(),
767 "local task dropped by a thread that didn't spawn it. Task spawned at {}",
768 self.location
769 );
770 unsafe { ManuallyDrop::drop(&mut self.inner) };
774 }
775 }
776
777 impl<F: Future> Future for Checked<F> {
778 type Output = F::Output;
779
780 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
781 let this = unsafe { self.get_unchecked_mut() };
785 assert!(
786 this.id == thread_id(),
787 "local task polled by a thread that didn't spawn it. Task spawned at {}",
788 this.location
789 );
790 unsafe { Pin::new_unchecked(&mut *this.inner).poll(cx) }
795 }
796 }
797
798 let location = metadata.location;
799
800 let future = move |_| Checked {
801 id: thread_id(),
802 inner: ManuallyDrop::new(future),
803 location,
804 };
805
806 let builder = async_task::Builder::new().metadata(metadata);
807 unsafe { builder.spawn_unchecked(future, schedule) }
811}