1use crate::{App, PlatformDispatcher, RunnableMeta, RunnableVariant, TaskTiming, profiler};
2use async_task::Runnable;
3use futures::channel::mpsc;
4use parking_lot::{Condvar, Mutex};
5use smol::prelude::*;
6use std::{
7 fmt::Debug,
8 marker::PhantomData,
9 mem::{self, ManuallyDrop},
10 num::NonZeroUsize,
11 panic::Location,
12 pin::Pin,
13 rc::Rc,
14 sync::{
15 Arc,
16 atomic::{AtomicUsize, Ordering},
17 },
18 task::{Context, Poll},
19 thread::{self, ThreadId},
20 time::{Duration, Instant},
21};
22use util::TryFutureExt;
23use waker_fn::waker_fn;
24
25#[cfg(any(test, feature = "test-support"))]
26use rand::rngs::StdRng;
27
28/// A pointer to the executor that is currently running,
29/// for spawning background tasks.
30#[derive(Clone)]
31pub struct BackgroundExecutor {
32 #[doc(hidden)]
33 pub dispatcher: Arc<dyn PlatformDispatcher>,
34}
35
36/// A pointer to the executor that is currently running,
37/// for spawning tasks on the main thread.
38///
39/// This is intentionally `!Send` via the `not_send` marker field. This is because
40/// `ForegroundExecutor::spawn` does not require `Send` but checks at runtime that the future is
41/// only polled from the same thread it was spawned from. These checks would fail when spawning
42/// foreground tasks from background threads.
43#[derive(Clone)]
44pub struct ForegroundExecutor {
45 #[doc(hidden)]
46 pub dispatcher: Arc<dyn PlatformDispatcher>,
47 not_send: PhantomData<Rc<()>>,
48}
49
50/// Realtime task priority
51#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
52#[repr(u8)]
53pub enum RealtimePriority {
54 /// Audio task
55 Audio,
56 /// Other realtime task
57 #[default]
58 Other,
59}
60
61/// Task priority
62#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
63#[repr(u8)]
64pub enum Priority {
65 /// Realtime priority
66 ///
67 /// Spawning a task with this priority will spin it off on a separate thread dedicated just to that task.
68 Realtime(RealtimePriority),
69 /// High priority
70 ///
71 /// Only use for tasks that are critical to the user experience / responsiveness of the editor.
72 High,
73 /// Medium priority, probably suits most of your use cases.
74 #[default]
75 Medium,
76 /// Low priority
77 ///
78 /// Prioritize this for background work that can come in large quantities
79 /// to not starve the executor of resources for high priority tasks
80 Low,
81}
82
83impl Priority {
84 #[allow(dead_code)]
85 pub(crate) const fn probability(&self) -> u32 {
86 match self {
87 // realtime priorities are not considered for probability scheduling
88 Priority::Realtime(_) => 0,
89 Priority::High => 60,
90 Priority::Medium => 30,
91 Priority::Low => 10,
92 }
93 }
94}
95
96/// Task is a primitive that allows work to happen in the background.
97///
98/// It implements [`Future`] so you can `.await` on it.
99///
100/// If you drop a task it will be cancelled immediately. Calling [`Task::detach`] allows
101/// the task to continue running, but with no way to return a value.
102#[must_use]
103#[derive(Debug)]
104pub struct Task<T>(TaskState<T>);
105
106#[derive(Debug)]
107enum TaskState<T> {
108 /// A task that is ready to return a value
109 Ready(Option<T>),
110
111 /// A task that is currently running.
112 Spawned(async_task::Task<T, RunnableMeta>),
113}
114
115impl<T> Task<T> {
116 /// Creates a new task that will resolve with the value
117 pub fn ready(val: T) -> Self {
118 Task(TaskState::Ready(Some(val)))
119 }
120
121 /// Detaching a task runs it to completion in the background
122 pub fn detach(self) {
123 match self {
124 Task(TaskState::Ready(_)) => {}
125 Task(TaskState::Spawned(task)) => task.detach(),
126 }
127 }
128}
129
130impl<E, T> Task<Result<T, E>>
131where
132 T: 'static,
133 E: 'static + Debug,
134{
135 /// Run the task to completion in the background and log any
136 /// errors that occur.
137 #[track_caller]
138 pub fn detach_and_log_err(self, cx: &App) {
139 let location = core::panic::Location::caller();
140 cx.foreground_executor()
141 .spawn(self.log_tracked_err(*location))
142 .detach();
143 }
144}
145
146impl<T> Future for Task<T> {
147 type Output = T;
148
149 fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
150 match unsafe { self.get_unchecked_mut() } {
151 Task(TaskState::Ready(val)) => Poll::Ready(val.take().unwrap()),
152 Task(TaskState::Spawned(task)) => task.poll(cx),
153 }
154 }
155}
156
157/// A task label is an opaque identifier that you can use to
158/// refer to a task in tests.
159#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
160pub struct TaskLabel(NonZeroUsize);
161
162impl Default for TaskLabel {
163 fn default() -> Self {
164 Self::new()
165 }
166}
167
168impl TaskLabel {
169 /// Construct a new task label.
170 pub fn new() -> Self {
171 static NEXT_TASK_LABEL: AtomicUsize = AtomicUsize::new(1);
172 Self(
173 NEXT_TASK_LABEL
174 .fetch_add(1, Ordering::SeqCst)
175 .try_into()
176 .unwrap(),
177 )
178 }
179}
180
181type AnyLocalFuture<R> = Pin<Box<dyn 'static + Future<Output = R>>>;
182
183type AnyFuture<R> = Pin<Box<dyn 'static + Send + Future<Output = R>>>;
184
185/// BackgroundExecutor lets you run things on background threads.
186/// In production this is a thread pool with no ordering guarantees.
187/// In tests this is simulated by running tasks one by one in a deterministic
188/// (but arbitrary) order controlled by the `SEED` environment variable.
189impl BackgroundExecutor {
190 #[doc(hidden)]
191 pub fn new(dispatcher: Arc<dyn PlatformDispatcher>) -> Self {
192 Self { dispatcher }
193 }
194
195 /// Enqueues the given future to be run to completion on a background thread.
196 #[track_caller]
197 pub fn spawn<R>(&self, future: impl Future<Output = R> + Send + 'static) -> Task<R>
198 where
199 R: Send + 'static,
200 {
201 self.spawn_with_priority(Priority::default(), future)
202 }
203
204 /// Enqueues the given future to be run to completion on a background thread.
205 #[track_caller]
206 pub fn spawn_with_priority<R>(
207 &self,
208 priority: Priority,
209 future: impl Future<Output = R> + Send + 'static,
210 ) -> Task<R>
211 where
212 R: Send + 'static,
213 {
214 self.spawn_internal::<R>(Box::pin(future), None, priority)
215 }
216
217 /// Enqueues the given future to be run to completion on a background thread and blocking the current task on it.
218 ///
219 /// This allows to spawn background work that borrows from its scope. Note that the supplied future will run to
220 /// completion before the current task is resumed, even if the current task is slated for cancellation.
221 pub async fn await_on_background<R>(&self, future: impl Future<Output = R> + Send) -> R
222 where
223 R: Send,
224 {
225 // We need to ensure that cancellation of the parent task does not drop the environment
226 // before the our own task has completed or got cancelled.
227 struct NotifyOnDrop<'a>(&'a (Condvar, Mutex<bool>));
228
229 impl Drop for NotifyOnDrop<'_> {
230 fn drop(&mut self) {
231 *self.0.1.lock() = true;
232 self.0.0.notify_all();
233 }
234 }
235
236 struct WaitOnDrop<'a>(&'a (Condvar, Mutex<bool>));
237
238 impl Drop for WaitOnDrop<'_> {
239 fn drop(&mut self) {
240 let mut done = self.0.1.lock();
241 if !*done {
242 self.0.0.wait(&mut done);
243 }
244 }
245 }
246
247 let dispatcher = self.dispatcher.clone();
248 let location = core::panic::Location::caller();
249
250 let pair = &(Condvar::new(), Mutex::new(false));
251 let _wait_guard = WaitOnDrop(pair);
252
253 let (runnable, task) = unsafe {
254 async_task::Builder::new()
255 .metadata(RunnableMeta {
256 location,
257 app: None,
258 })
259 .spawn_unchecked(
260 move |_| async {
261 let _notify_guard = NotifyOnDrop(pair);
262 future.await
263 },
264 move |runnable| {
265 dispatcher.dispatch(
266 RunnableVariant::Meta(runnable),
267 None,
268 Priority::default(),
269 )
270 },
271 )
272 };
273 runnable.schedule();
274 task.await
275 }
276
277 /// Enqueues the given future to be run to completion on a background thread.
278 /// The given label can be used to control the priority of the task in tests.
279 #[track_caller]
280 pub fn spawn_labeled<R>(
281 &self,
282 label: TaskLabel,
283 future: impl Future<Output = R> + Send + 'static,
284 ) -> Task<R>
285 where
286 R: Send + 'static,
287 {
288 self.spawn_internal::<R>(Box::pin(future), Some(label), Priority::default())
289 }
290
291 #[track_caller]
292 fn spawn_internal<R: Send + 'static>(
293 &self,
294 future: AnyFuture<R>,
295 label: Option<TaskLabel>,
296 #[cfg_attr(
297 target_os = "windows",
298 expect(
299 unused_variables,
300 reason = "Multi priority scheduler is broken on windows"
301 )
302 )]
303 priority: Priority,
304 ) -> Task<R> {
305 let dispatcher = self.dispatcher.clone();
306 #[cfg(target_os = "windows")]
307 let priority = Priority::Medium; // multi-prio scheduler is broken on windows
308
309 let (runnable, task) = if let Priority::Realtime(realtime) = priority {
310 let location = core::panic::Location::caller();
311 let (mut tx, rx) = flume::bounded::<Runnable<RunnableMeta>>(1);
312
313 dispatcher.spawn_realtime(
314 realtime,
315 Box::new(move || {
316 while let Ok(runnable) = rx.recv() {
317 let start = Instant::now();
318 let location = runnable.metadata().location;
319 let mut timing = TaskTiming {
320 location,
321 start,
322 end: None,
323 };
324 profiler::add_task_timing(timing);
325
326 runnable.run();
327
328 let end = Instant::now();
329 timing.end = Some(end);
330 profiler::add_task_timing(timing);
331 }
332 }),
333 );
334
335 async_task::Builder::new()
336 .metadata(RunnableMeta {
337 location,
338 app: None,
339 })
340 .spawn(
341 move |_| future,
342 move |runnable| {
343 let _ = tx.send(runnable);
344 },
345 )
346 } else {
347 let location = core::panic::Location::caller();
348 async_task::Builder::new()
349 .metadata(RunnableMeta {
350 location,
351 app: None,
352 })
353 .spawn(
354 move |_| future,
355 move |runnable| {
356 dispatcher.dispatch(RunnableVariant::Meta(runnable), label, priority)
357 },
358 )
359 };
360
361 runnable.schedule();
362 Task(TaskState::Spawned(task))
363 }
364
365 /// Used by the test harness to run an async test in a synchronous fashion.
366 #[cfg(any(test, feature = "test-support"))]
367 #[track_caller]
368 pub fn block_test<R>(&self, future: impl Future<Output = R>) -> R {
369 if let Ok(value) = self.block_internal(false, future, None) {
370 value
371 } else {
372 unreachable!()
373 }
374 }
375
376 /// Block the current thread until the given future resolves.
377 /// Consider using `block_with_timeout` instead.
378 pub fn block<R>(&self, future: impl Future<Output = R>) -> R {
379 if let Ok(value) = self.block_internal(true, future, None) {
380 value
381 } else {
382 unreachable!()
383 }
384 }
385
386 #[cfg(not(any(test, feature = "test-support")))]
387 pub(crate) fn block_internal<Fut: Future>(
388 &self,
389 _background_only: bool,
390 future: Fut,
391 timeout: Option<Duration>,
392 ) -> Result<Fut::Output, impl Future<Output = Fut::Output> + use<Fut>> {
393 use std::time::Instant;
394
395 let mut future = Box::pin(future);
396 if timeout == Some(Duration::ZERO) {
397 return Err(future);
398 }
399 let deadline = timeout.map(|timeout| Instant::now() + timeout);
400
401 let parker = parking::Parker::new();
402 let unparker = parker.unparker();
403 let waker = waker_fn(move || {
404 unparker.unpark();
405 });
406 let mut cx = std::task::Context::from_waker(&waker);
407
408 loop {
409 match future.as_mut().poll(&mut cx) {
410 Poll::Ready(result) => return Ok(result),
411 Poll::Pending => {
412 let timeout =
413 deadline.map(|deadline| deadline.saturating_duration_since(Instant::now()));
414 if let Some(timeout) = timeout {
415 if !parker.park_timeout(timeout)
416 && deadline.is_some_and(|deadline| deadline < Instant::now())
417 {
418 return Err(future);
419 }
420 } else {
421 parker.park();
422 }
423 }
424 }
425 }
426 }
427
428 #[cfg(any(test, feature = "test-support"))]
429 #[track_caller]
430 pub(crate) fn block_internal<Fut: Future>(
431 &self,
432 background_only: bool,
433 future: Fut,
434 timeout: Option<Duration>,
435 ) -> Result<Fut::Output, impl Future<Output = Fut::Output> + use<Fut>> {
436 use std::sync::atomic::AtomicBool;
437
438 use parking::Parker;
439
440 let mut future = Box::pin(future);
441 if timeout == Some(Duration::ZERO) {
442 return Err(future);
443 }
444 let Some(dispatcher) = self.dispatcher.as_test() else {
445 return Err(future);
446 };
447
448 let mut max_ticks = if timeout.is_some() {
449 dispatcher.gen_block_on_ticks()
450 } else {
451 usize::MAX
452 };
453
454 let parker = Parker::new();
455 let unparker = parker.unparker();
456
457 let awoken = Arc::new(AtomicBool::new(false));
458 let waker = waker_fn({
459 let awoken = awoken.clone();
460 let unparker = unparker.clone();
461 move || {
462 awoken.store(true, Ordering::SeqCst);
463 unparker.unpark();
464 }
465 });
466 let mut cx = std::task::Context::from_waker(&waker);
467
468 let duration = Duration::from_secs(
469 option_env!("GPUI_TEST_TIMEOUT")
470 .and_then(|s| s.parse::<u64>().ok())
471 .unwrap_or(180),
472 );
473 let mut test_should_end_by = Instant::now() + duration;
474
475 loop {
476 match future.as_mut().poll(&mut cx) {
477 Poll::Ready(result) => return Ok(result),
478 Poll::Pending => {
479 if max_ticks == 0 {
480 return Err(future);
481 }
482 max_ticks -= 1;
483
484 if !dispatcher.tick(background_only) {
485 if awoken.swap(false, Ordering::SeqCst) {
486 continue;
487 }
488
489 if !dispatcher.parking_allowed() {
490 if dispatcher.advance_clock_to_next_delayed() {
491 continue;
492 }
493 let mut backtrace_message = String::new();
494 let mut waiting_message = String::new();
495 if let Some(backtrace) = dispatcher.waiting_backtrace() {
496 backtrace_message =
497 format!("\nbacktrace of waiting future:\n{:?}", backtrace);
498 }
499 if let Some(waiting_hint) = dispatcher.waiting_hint() {
500 waiting_message = format!("\n waiting on: {}\n", waiting_hint);
501 }
502 panic!(
503 "parked with nothing left to run{waiting_message}{backtrace_message}",
504 )
505 }
506 dispatcher.push_unparker(unparker.clone());
507 parker.park_timeout(Duration::from_millis(1));
508 if Instant::now() > test_should_end_by {
509 panic!("test timed out after {duration:?} with allow_parking")
510 }
511 }
512 }
513 }
514 }
515 }
516
517 /// Block the current thread until the given future resolves
518 /// or `duration` has elapsed.
519 pub fn block_with_timeout<Fut: Future>(
520 &self,
521 duration: Duration,
522 future: Fut,
523 ) -> Result<Fut::Output, impl Future<Output = Fut::Output> + use<Fut>> {
524 self.block_internal(true, future, Some(duration))
525 }
526
527 /// Scoped lets you start a number of tasks and waits
528 /// for all of them to complete before returning.
529 pub async fn scoped<'scope, F>(&self, scheduler: F)
530 where
531 F: FnOnce(&mut Scope<'scope>),
532 {
533 let mut scope = Scope::new(self.clone(), Priority::default());
534 (scheduler)(&mut scope);
535 let spawned = mem::take(&mut scope.futures)
536 .into_iter()
537 .map(|f| self.spawn_with_priority(scope.priority, f))
538 .collect::<Vec<_>>();
539 for task in spawned {
540 task.await;
541 }
542 }
543
544 /// Scoped lets you start a number of tasks and waits
545 /// for all of them to complete before returning.
546 pub async fn scoped_priority<'scope, F>(&self, priority: Priority, scheduler: F)
547 where
548 F: FnOnce(&mut Scope<'scope>),
549 {
550 let mut scope = Scope::new(self.clone(), priority);
551 (scheduler)(&mut scope);
552 let spawned = mem::take(&mut scope.futures)
553 .into_iter()
554 .map(|f| self.spawn_with_priority(scope.priority, f))
555 .collect::<Vec<_>>();
556 for task in spawned {
557 task.await;
558 }
559 }
560
561 /// Get the current time.
562 ///
563 /// Calling this instead of `std::time::Instant::now` allows the use
564 /// of fake timers in tests.
565 pub fn now(&self) -> Instant {
566 self.dispatcher.now()
567 }
568
569 /// Returns a task that will complete after the given duration.
570 /// Depending on other concurrent tasks the elapsed duration may be longer
571 /// than requested.
572 pub fn timer(&self, duration: Duration) -> Task<()> {
573 if duration.is_zero() {
574 return Task::ready(());
575 }
576 let location = core::panic::Location::caller();
577 let (runnable, task) = async_task::Builder::new()
578 .metadata(RunnableMeta {
579 location,
580 app: None,
581 })
582 .spawn(move |_| async move {}, {
583 let dispatcher = self.dispatcher.clone();
584 move |runnable| dispatcher.dispatch_after(duration, RunnableVariant::Meta(runnable))
585 });
586 runnable.schedule();
587 Task(TaskState::Spawned(task))
588 }
589
590 /// in tests, start_waiting lets you indicate which task is waiting (for debugging only)
591 #[cfg(any(test, feature = "test-support"))]
592 pub fn start_waiting(&self) {
593 self.dispatcher.as_test().unwrap().start_waiting();
594 }
595
596 /// in tests, removes the debugging data added by start_waiting
597 #[cfg(any(test, feature = "test-support"))]
598 pub fn finish_waiting(&self) {
599 self.dispatcher.as_test().unwrap().finish_waiting();
600 }
601
602 /// in tests, run an arbitrary number of tasks (determined by the SEED environment variable)
603 #[cfg(any(test, feature = "test-support"))]
604 pub fn simulate_random_delay(&self) -> impl Future<Output = ()> + use<> {
605 self.dispatcher.as_test().unwrap().simulate_random_delay()
606 }
607
608 /// in tests, indicate that a given task from `spawn_labeled` should run after everything else
609 #[cfg(any(test, feature = "test-support"))]
610 pub fn deprioritize(&self, task_label: TaskLabel) {
611 self.dispatcher.as_test().unwrap().deprioritize(task_label)
612 }
613
614 /// in tests, move time forward. This does not run any tasks, but does make `timer`s ready.
615 #[cfg(any(test, feature = "test-support"))]
616 pub fn advance_clock(&self, duration: Duration) {
617 self.dispatcher.as_test().unwrap().advance_clock(duration)
618 }
619
620 /// in tests, run one task.
621 #[cfg(any(test, feature = "test-support"))]
622 pub fn tick(&self) -> bool {
623 self.dispatcher.as_test().unwrap().tick(false)
624 }
625
626 /// in tests, run all tasks that are ready to run. If after doing so
627 /// the test still has outstanding tasks, this will panic. (See also [`Self::allow_parking`])
628 #[cfg(any(test, feature = "test-support"))]
629 pub fn run_until_parked(&self) {
630 self.dispatcher.as_test().unwrap().run_until_parked()
631 }
632
633 /// in tests, prevents `run_until_parked` from panicking if there are outstanding tasks.
634 /// This is useful when you are integrating other (non-GPUI) futures, like disk access, that
635 /// do take real async time to run.
636 #[cfg(any(test, feature = "test-support"))]
637 pub fn allow_parking(&self) {
638 self.dispatcher.as_test().unwrap().allow_parking();
639 }
640
641 /// undoes the effect of [`Self::allow_parking`].
642 #[cfg(any(test, feature = "test-support"))]
643 pub fn forbid_parking(&self) {
644 self.dispatcher.as_test().unwrap().forbid_parking();
645 }
646
647 /// adds detail to the "parked with nothing let to run" message.
648 #[cfg(any(test, feature = "test-support"))]
649 pub fn set_waiting_hint(&self, msg: Option<String>) {
650 self.dispatcher.as_test().unwrap().set_waiting_hint(msg);
651 }
652
653 /// in tests, returns the rng used by the dispatcher and seeded by the `SEED` environment variable
654 #[cfg(any(test, feature = "test-support"))]
655 pub fn rng(&self) -> StdRng {
656 self.dispatcher.as_test().unwrap().rng()
657 }
658
659 /// How many CPUs are available to the dispatcher.
660 pub fn num_cpus(&self) -> usize {
661 #[cfg(any(test, feature = "test-support"))]
662 return 4;
663
664 #[cfg(not(any(test, feature = "test-support")))]
665 return num_cpus::get();
666 }
667
668 /// Whether we're on the main thread.
669 pub fn is_main_thread(&self) -> bool {
670 self.dispatcher.is_main_thread()
671 }
672
673 #[cfg(any(test, feature = "test-support"))]
674 /// in tests, control the number of ticks that `block_with_timeout` will run before timing out.
675 pub fn set_block_on_ticks(&self, range: std::ops::RangeInclusive<usize>) {
676 self.dispatcher.as_test().unwrap().set_block_on_ticks(range);
677 }
678}
679
680/// ForegroundExecutor runs things on the main thread.
681impl ForegroundExecutor {
682 /// Creates a new ForegroundExecutor from the given PlatformDispatcher.
683 pub fn new(dispatcher: Arc<dyn PlatformDispatcher>) -> Self {
684 Self {
685 dispatcher,
686 not_send: PhantomData,
687 }
688 }
689
690 /// Enqueues the given Task to run on the main thread at some point in the future.
691 #[track_caller]
692 pub fn spawn<R>(&self, future: impl Future<Output = R> + 'static) -> Task<R>
693 where
694 R: 'static,
695 {
696 self.spawn_with_app_and_priority(None, Priority::default(), future)
697 }
698
699 /// Enqueues the given Task to run on the main thread at some point in the future.
700 #[track_caller]
701 pub fn spawn_with_priority<R>(
702 &self,
703 priority: Priority,
704 future: impl Future<Output = R> + 'static,
705 ) -> Task<R>
706 where
707 R: 'static,
708 {
709 self.spawn_with_app_and_priority(None, priority, future)
710 }
711
712 /// Enqueues the given Task to run on the main thread at some point in the future,
713 /// with a weak reference to the app for cancellation checking.
714 ///
715 /// When the app is dropped, pending tasks spawned with this method will be cancelled
716 /// before they run, rather than panicking when they try to access the dropped app.
717 #[track_caller]
718 pub fn spawn_with_app<R>(
719 &self,
720 app: std::rc::Weak<crate::AppCell>,
721 future: impl Future<Output = R> + 'static,
722 ) -> Task<R>
723 where
724 R: 'static,
725 {
726 self.spawn_with_app_and_priority(Some(app), Priority::default(), future)
727 }
728
729 /// Enqueues the given Task to run on the main thread at some point in the future,
730 /// with an optional weak reference to the app for cancellation checking and a specific priority.
731 #[track_caller]
732 pub fn spawn_with_app_and_priority<R>(
733 &self,
734 app: Option<std::rc::Weak<crate::AppCell>>,
735 priority: Priority,
736 future: impl Future<Output = R> + 'static,
737 ) -> Task<R>
738 where
739 R: 'static,
740 {
741 let dispatcher = self.dispatcher.clone();
742 let location = core::panic::Location::caller();
743
744 #[track_caller]
745 fn inner<R: 'static>(
746 dispatcher: Arc<dyn PlatformDispatcher>,
747 future: AnyLocalFuture<R>,
748 location: &'static core::panic::Location<'static>,
749 app: Option<std::rc::Weak<crate::AppCell>>,
750 priority: Priority,
751 ) -> Task<R> {
752 // SAFETY: We are on the main thread (ForegroundExecutor is !Send), and the
753 // MainThreadWeak will only be accessed on the main thread in the trampoline.
754 let app_weak = app.map(|weak| unsafe { crate::MainThreadWeak::new(weak) });
755 let (runnable, task) = spawn_local_with_source_location(
756 future,
757 move |runnable| {
758 dispatcher.dispatch_on_main_thread(RunnableVariant::Meta(runnable), priority)
759 },
760 RunnableMeta {
761 location,
762 app: app_weak,
763 },
764 );
765 runnable.schedule();
766 Task(TaskState::Spawned(task))
767 }
768 inner::<R>(dispatcher, Box::pin(future), location, app, priority)
769 }
770}
771
772/// Variant of `async_task::spawn_local` that includes the source location of the spawn in panics.
773///
774/// Copy-modified from:
775/// <https://github.com/smol-rs/async-task/blob/ca9dbe1db9c422fd765847fa91306e30a6bb58a9/src/runnable.rs#L405>
776#[track_caller]
777fn spawn_local_with_source_location<Fut, S, M>(
778 future: Fut,
779 schedule: S,
780 metadata: M,
781) -> (Runnable<M>, async_task::Task<Fut::Output, M>)
782where
783 Fut: Future + 'static,
784 Fut::Output: 'static,
785 S: async_task::Schedule<M> + Send + Sync + 'static,
786 M: 'static,
787{
788 #[inline]
789 fn thread_id() -> ThreadId {
790 std::thread_local! {
791 static ID: ThreadId = thread::current().id();
792 }
793 ID.try_with(|id| *id)
794 .unwrap_or_else(|_| thread::current().id())
795 }
796
797 struct Checked<F> {
798 id: ThreadId,
799 inner: ManuallyDrop<F>,
800 location: &'static Location<'static>,
801 }
802
803 impl<F> Drop for Checked<F> {
804 fn drop(&mut self) {
805 assert!(
806 self.id == thread_id(),
807 "local task dropped by a thread that didn't spawn it. Task spawned at {}",
808 self.location
809 );
810 unsafe { ManuallyDrop::drop(&mut self.inner) };
811 }
812 }
813
814 impl<F: Future> Future for Checked<F> {
815 type Output = F::Output;
816
817 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
818 assert!(
819 self.id == thread_id(),
820 "local task polled by a thread that didn't spawn it. Task spawned at {}",
821 self.location
822 );
823 unsafe { self.map_unchecked_mut(|c| &mut *c.inner).poll(cx) }
824 }
825 }
826
827 // Wrap the future into one that checks which thread it's on.
828 let future = Checked {
829 id: thread_id(),
830 inner: ManuallyDrop::new(future),
831 location: Location::caller(),
832 };
833
834 unsafe {
835 async_task::Builder::new()
836 .metadata(metadata)
837 .spawn_unchecked(move |_| future, schedule)
838 }
839}
840
841/// Scope manages a set of tasks that are enqueued and waited on together. See [`BackgroundExecutor::scoped`].
842pub struct Scope<'a> {
843 executor: BackgroundExecutor,
844 priority: Priority,
845 futures: Vec<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
846 tx: Option<mpsc::Sender<()>>,
847 rx: mpsc::Receiver<()>,
848 lifetime: PhantomData<&'a ()>,
849}
850
851impl<'a> Scope<'a> {
852 fn new(executor: BackgroundExecutor, priority: Priority) -> Self {
853 let (tx, rx) = mpsc::channel(1);
854 Self {
855 executor,
856 priority,
857 tx: Some(tx),
858 rx,
859 futures: Default::default(),
860 lifetime: PhantomData,
861 }
862 }
863
864 /// How many CPUs are available to the dispatcher.
865 pub fn num_cpus(&self) -> usize {
866 self.executor.num_cpus()
867 }
868
869 /// Spawn a future into this scope.
870 #[track_caller]
871 pub fn spawn<F>(&mut self, f: F)
872 where
873 F: Future<Output = ()> + Send + 'a,
874 {
875 let tx = self.tx.clone().unwrap();
876
877 // SAFETY: The 'a lifetime is guaranteed to outlive any of these futures because
878 // dropping this `Scope` blocks until all of the futures have resolved.
879 let f = unsafe {
880 mem::transmute::<
881 Pin<Box<dyn Future<Output = ()> + Send + 'a>>,
882 Pin<Box<dyn Future<Output = ()> + Send + 'static>>,
883 >(Box::pin(async move {
884 f.await;
885 drop(tx);
886 }))
887 };
888 self.futures.push(f);
889 }
890}
891
892impl Drop for Scope<'_> {
893 fn drop(&mut self) {
894 self.tx.take().unwrap();
895
896 // Wait until the channel is closed, which means that all of the spawned
897 // futures have resolved.
898 self.executor.block(self.rx.next());
899 }
900}