Skip to main content

scheduler/
executor.rs

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/// A `!Send` executor pinned to a single session. Tasks spawned on it run in
18/// order on whichever thread drains the dispatch destination supplied at
19/// construction time — typically the main thread for the default session, or
20/// a dedicated OS thread for sessions created by `spawn_dedicated_thread`.
21#[derive(Clone)]
22pub struct LocalExecutor {
23    session_id: SessionId,
24    scheduler: Arc<dyn Scheduler>,
25    // Spawned tasks' schedule callbacks each hold an `Arc` clone of this
26    // closure, so the destination it captures stays alive as long as work
27    // could still land on it.
28    dispatch: Arc<dyn Fn(Runnable<RunnableMeta>) + Send + Sync>,
29    not_send: PhantomData<Rc<()>>,
30}
31
32impl LocalExecutor {
33    /// Constructs a local executor that runs spawned tasks by sending their
34    /// runnables through `dispatch`. The `scheduler` is retained for access to
35    /// clocks, timers, and other scheduler-level services.
36    ///
37    /// For the common case of routing runnables through
38    /// `Scheduler::schedule_local`, callers pass a closure that does exactly
39    /// that. `spawn_dedicated_thread` instead passes a closure that sends to
40    /// the dedicated thread's channel.
41    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    /// Like [`Self::spawn`], but routes the task's runnables through the
83    /// given dispatch destination instead of this executor's own. The
84    /// destination must deliver runnables to this executor's thread: the
85    /// future is polled and dropped on the spawning thread.
86    #[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    /// Block until the future completes or timeout occurs.
126    /// Returns Ok(output) if completed, Err(future) if timed out.
127    #[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    /// Spawn a closure on a fresh session pinned to its own [`LocalExecutor`].
165    /// The closure runs on a new OS thread under `PlatformScheduler`, or on
166    /// the test scheduler's loop under `TestScheduler`.
167    ///
168    /// The returned `Task` represents the dedicated work: dropping it cancels
169    /// the dedicated closure, `.await`ing it yields the closure's return
170    /// value, `.detach()`ing it lets the dedicated work run independently of
171    /// the caller.
172    #[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
191/// Boxes the user-supplied dedicated closure into the type-erased shape
192/// expected by [`Scheduler::spawn_dedicated`]. The user's `Fut::Output` is
193/// boxed as `Box<dyn Any + Send + Sync>` on the dedicated side and downcast
194/// back to `Fut::Output` by [`Task::downcast`] in the wrapper.
195fn 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    /// Spawns a future on a dedicated realtime thread for audio processing.
250    #[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    /// Spawn a closure on a fresh session pinned to its own [`LocalExecutor`].
294    /// The closure runs on a new OS thread under `PlatformScheduler`, or on
295    /// the test scheduler's loop under `TestScheduler`.
296    ///
297    /// The returned `Task` represents the dedicated work: dropping it cancels
298    /// the dedicated closure, `.await`ing it yields the closure's return
299    /// value, `.detach()`ing it lets the dedicated work run independently of
300    /// the caller.
301    #[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
327/// A long-lived handle to one dedicated session: every future spawned on it
328/// runs on that session's single thread, never on the background pool.
329///
330/// Use this when tasks rely on thread-local state being coherent across
331/// spawns, or when the per-thread setup cost of the work is high enough that
332/// it should be paid once. Creating the session costs a thread (on the web, a
333/// worker — expensive); each spawn afterwards is just a channel send.
334///
335/// Dropping the handle cancels the session: queued and future spawns resolve
336/// to cancelled tasks.
337pub struct DedicatedExecutor {
338    sender: flume::Sender<Runnable<RunnableMeta>>,
339    _session: Task<()>,
340}
341
342impl DedicatedExecutor {
343    /// Starts a dedicated session on `executor` and returns a handle that
344    /// spawns futures onto it. Under `TestScheduler` the session runs on the
345    /// deterministic test loop; no real thread is created.
346    #[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    /// Spawns a future onto the dedicated session's thread.
361    ///
362    /// The returned task has the usual semantics: dropping it cancels the
363    /// future, `.await`ing it yields the output, `.detach()`ing it lets it
364    /// run to completion on its own.
365    #[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/// Task is a primitive that allows work to happen in the background.
386///
387/// It implements [`Future`] so you can `.await` on it.
388///
389/// If you drop a task it will be cancelled immediately. Calling [`Task::detach`] allows
390/// the task to continue running, but with no way to return a value.
391#[must_use]
392pub struct Task<T>(TaskState<T>);
393
394enum TaskState<T> {
395    /// A task that is ready to return a value
396    Ready(Option<T>),
397
398    /// A task that is currently running.
399    Spawned(async_task::Task<T, RunnableMeta>),
400
401    /// A typed view of a [`Task<Box<dyn Any + Send + Sync>>`] obtained via
402    /// [`Task::downcast`]. The inner task drives the actual work; the
403    /// downcast layer just unwraps the `Box<dyn Any + Send + Sync>` on poll.
404    Downcast {
405        inner: Box<Task<Box<dyn Any + Send + Sync>>>,
406        marker: PhantomData<fn() -> T>,
407    },
408
409    /// A task whose real handle is delivered later by another thread (see
410    /// [`Task::rendezvous`]). Once delivered, polling replaces this state
411    /// with the delivered task's state.
412    Rendezvous(RendezvousReceiver<T>),
413}
414
415/// State shared between the two halves of a [`Task::rendezvous`] pair.
416enum RendezvousState<T> {
417    /// No task delivered yet; holds the consumer's waker if it polled.
418    Pending(Option<Waker>),
419    /// The producer delivered before the consumer consumed the task.
420    Delivered(Task<T>),
421    /// The consumer was dropped before delivery; delivery cancels the task.
422    Cancelled,
423    /// The consumer was detached before delivery; delivery detaches the task.
424    Detached,
425    /// The consumer took the delivered task.
426    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            // The producer applies the detached disposition on delivery.
464            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            // `Taken`, `Detached`, and `Cancelled` record dispositions that
483            // this drop must not overwrite: the receiver is also dropped as a
484            // normal side effect of taking or detaching.
485            _ => {}
486        }
487    }
488}
489
490pub(crate) struct RendezvousSender<T> {
491    shared: Arc<parking_lot::Mutex<RendezvousState<T>>>,
492}
493
494impl<T> RendezvousSender<T> {
495    /// Hands the real task to the rendezvous, applying the consumer's
496    /// disposition if it was dropped or detached before delivery.
497    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    /// Creates a new task that will resolve with the value
529    pub fn ready(val: T) -> Self {
530        Task(TaskState::Ready(Some(val)))
531    }
532
533    /// Creates a task whose real handle arrives later through the returned
534    /// sender, typically from another thread. Until delivery the task is
535    /// pending; afterwards it behaves exactly like the delivered task.
536    /// Dropping or detaching the task before delivery is honored on delivery.
537    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    /// Creates a Task from an async_task::Task
548    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    /// Detaching a task runs it to completion in the background
562    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    /// Converts this task into a fallible task that returns `Option<T>`.
572    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    /// Reinterprets the boxed output as a concrete `T` via downcast on
587    /// completion. Used by [`LocalExecutor::spawn_dedicated`] and
588    /// [`BackgroundExecutor::spawn_dedicated`] to recover the user closure's
589    /// `Fut::Output` from the dyn-safe [`Scheduler::spawn_dedicated`].
590    ///
591    /// Panics on poll if the inner output is not in fact a `T` -- a logic
592    /// error in whatever produced the inner task, since the downcast type is
593    /// chosen by the caller of `downcast`.
594    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/// A task that returns `Option<T>` instead of panicking when cancelled.
616#[must_use]
617pub struct FallibleTask<T>(FallibleTaskState<T>);
618
619enum FallibleTaskState<T> {
620    /// A task that is ready to return a value
621    Ready(Option<T>),
622
623    /// A task that is currently running (wraps async_task::FallibleTask).
624    Spawned(async_task::FallibleTask<T, RunnableMeta>),
625
626    /// Mirror of [`TaskState::Downcast`] for fallible tasks.
627    Downcast {
628        inner: Box<FallibleTask<Box<dyn Any + Send + Sync>>>,
629        marker: PhantomData<fn() -> T>,
630    },
631
632    /// Mirror of [`TaskState::Rendezvous`] for fallible tasks.
633    Rendezvous(RendezvousReceiver<T>),
634}
635
636impl<T> FallibleTask<T> {
637    /// Creates a new fallible task that will resolve with the value.
638    pub fn ready(val: T) -> Self {
639        FallibleTask(FallibleTaskState::Ready(Some(val)))
640    }
641
642    /// Detaching a task runs it to completion in the background.
643    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/// Variant of `async_task::spawn_local` that includes the source location of the spawn in panics.
733#[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            // SAFETY: `inner` is wrapped in `ManuallyDrop`, so this is the only
771            // place it is dropped. The thread check above ensures local futures
772            // are dropped on the thread that created them.
773            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            // SAFETY: We don't move any fields out of `self`; this mutable
782            // reference is only used to check metadata and to project the pin to
783            // `inner` below.
784            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            // SAFETY: `inner` is structurally pinned by `Checked`; after
791            // `Checked` is pinned, `inner` is never moved. The thread check
792            // above ensures the local future is only polled by its spawning
793            // thread.
794            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    // SAFETY: `Checked` enforces the invariants required by `spawn_unchecked`:
808    // the non-`Send` future is only polled and dropped on the thread that
809    // spawned it.
810    unsafe { builder.spawn_unchecked(future, schedule) }
811}