1use crate::{App, PlatformDispatcher};
2use async_task::Runnable;
3use futures::channel::mpsc;
4use smol::prelude::*;
5use std::mem::ManuallyDrop;
6use std::panic::Location;
7use std::thread::{self, ThreadId};
8use std::{
9 fmt::Debug,
10 marker::PhantomData,
11 mem,
12 num::NonZeroUsize,
13 pin::Pin,
14 rc::Rc,
15 sync::{
16 Arc,
17 atomic::{AtomicUsize, Ordering::SeqCst},
18 },
19 task::{Context, Poll},
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 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/// Task is a primitive that allows work to happen in the background.
51///
52/// It implements [`Future`] so you can `.await` on it.
53///
54/// If you drop a task it will be cancelled immediately. Calling [`Task::detach`] allows
55/// the task to continue running, but with no way to return a value.
56#[must_use]
57#[derive(Debug)]
58pub struct Task<T>(TaskState<T>);
59
60#[derive(Debug)]
61enum TaskState<T> {
62 /// A task that is ready to return a value
63 Ready(Option<T>),
64
65 /// A task that is currently running.
66 Spawned(async_task::Task<T>),
67}
68
69impl<T> Task<T> {
70 /// Creates a new task that will resolve with the value
71 pub fn ready(val: T) -> Self {
72 Task(TaskState::Ready(Some(val)))
73 }
74
75 /// Detaching a task runs it to completion in the background
76 pub fn detach(self) {
77 match self {
78 Task(TaskState::Ready(_)) => {}
79 Task(TaskState::Spawned(task)) => task.detach(),
80 }
81 }
82}
83
84impl<E, T> Task<Result<T, E>>
85where
86 T: 'static,
87 E: 'static + Debug,
88{
89 /// Run the task to completion in the background and log any
90 /// errors that occur.
91 #[track_caller]
92 pub fn detach_and_log_err(self, cx: &App) {
93 let location = core::panic::Location::caller();
94 cx.foreground_executor()
95 .spawn(self.log_tracked_err(*location))
96 .detach();
97 }
98
99 /// Convert a Task<Result<T, E>> to a Task<()> that logs all errors.
100 pub fn log_err_in_task(self, cx: &App) -> Task<Option<T>> {
101 let location = core::panic::Location::caller();
102 cx.foreground_executor()
103 .spawn(async move { self.log_tracked_err(*location).await })
104 }
105}
106
107impl<T> Future for Task<T> {
108 type Output = T;
109
110 fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
111 match unsafe { self.get_unchecked_mut() } {
112 Task(TaskState::Ready(val)) => Poll::Ready(val.take().unwrap()),
113 Task(TaskState::Spawned(task)) => task.poll(cx),
114 }
115 }
116}
117
118/// A task label is an opaque identifier that you can use to
119/// refer to a task in tests.
120#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
121pub struct TaskLabel(NonZeroUsize);
122
123impl Default for TaskLabel {
124 fn default() -> Self {
125 Self::new()
126 }
127}
128
129impl TaskLabel {
130 /// Construct a new task label.
131 pub fn new() -> Self {
132 static NEXT_TASK_LABEL: AtomicUsize = AtomicUsize::new(1);
133 Self(NEXT_TASK_LABEL.fetch_add(1, SeqCst).try_into().unwrap())
134 }
135}
136
137type AnyLocalFuture<R> = Pin<Box<dyn 'static + Future<Output = R>>>;
138
139type AnyFuture<R> = Pin<Box<dyn 'static + Send + Future<Output = R>>>;
140
141/// BackgroundExecutor lets you run things on background threads.
142/// In production this is a thread pool with no ordering guarantees.
143/// In tests this is simulated by running tasks one by one in a deterministic
144/// (but arbitrary) order controlled by the `SEED` environment variable.
145impl BackgroundExecutor {
146 #[doc(hidden)]
147 pub fn new(dispatcher: Arc<dyn PlatformDispatcher>) -> Self {
148 Self { dispatcher }
149 }
150
151 /// Enqueues the given future to be run to completion on a background thread.
152 pub fn spawn<R>(&self, future: impl Future<Output = R> + Send + 'static) -> Task<R>
153 where
154 R: Send + 'static,
155 {
156 self.spawn_internal::<R>(Box::pin(future), None)
157 }
158
159 /// Enqueues the given future to be run to completion on a background thread.
160 /// The given label can be used to control the priority of the task in tests.
161 pub fn spawn_labeled<R>(
162 &self,
163 label: TaskLabel,
164 future: impl Future<Output = R> + Send + 'static,
165 ) -> Task<R>
166 where
167 R: Send + 'static,
168 {
169 self.spawn_internal::<R>(Box::pin(future), Some(label))
170 }
171
172 fn spawn_internal<R: Send + 'static>(
173 &self,
174 future: AnyFuture<R>,
175 label: Option<TaskLabel>,
176 ) -> Task<R> {
177 let dispatcher = self.dispatcher.clone();
178 let (runnable, task) =
179 async_task::spawn(future, move |runnable| dispatcher.dispatch(runnable, label));
180 runnable.schedule();
181 Task(TaskState::Spawned(task))
182 }
183
184 /// Used by the test harness to run an async test in a synchronous fashion.
185 #[cfg(any(test, feature = "test-support"))]
186 #[track_caller]
187 pub fn block_test<R>(&self, future: impl Future<Output = R>) -> R {
188 if let Ok(value) = self.block_internal(false, future, None) {
189 value
190 } else {
191 unreachable!()
192 }
193 }
194
195 /// Block the current thread until the given future resolves.
196 /// Consider using `block_with_timeout` instead.
197 pub fn block<R>(&self, future: impl Future<Output = R>) -> R {
198 if let Ok(value) = self.block_internal(true, future, None) {
199 value
200 } else {
201 unreachable!()
202 }
203 }
204
205 #[cfg(not(any(test, feature = "test-support")))]
206 pub(crate) fn block_internal<Fut: Future>(
207 &self,
208 _background_only: bool,
209 future: Fut,
210 timeout: Option<Duration>,
211 ) -> Result<Fut::Output, impl Future<Output = Fut::Output> + use<Fut>> {
212 use std::time::Instant;
213
214 let mut future = Box::pin(future);
215 if timeout == Some(Duration::ZERO) {
216 return Err(future);
217 }
218 let deadline = timeout.map(|timeout| Instant::now() + timeout);
219
220 let unparker = self.dispatcher.unparker();
221 let waker = waker_fn(move || {
222 unparker.unpark();
223 });
224 let mut cx = std::task::Context::from_waker(&waker);
225
226 loop {
227 match future.as_mut().poll(&mut cx) {
228 Poll::Ready(result) => return Ok(result),
229 Poll::Pending => {
230 let timeout =
231 deadline.map(|deadline| deadline.saturating_duration_since(Instant::now()));
232 if !self.dispatcher.park(timeout)
233 && deadline.is_some_and(|deadline| deadline < Instant::now())
234 {
235 return Err(future);
236 }
237 }
238 }
239 }
240 }
241
242 #[cfg(any(test, feature = "test-support"))]
243 #[track_caller]
244 pub(crate) fn block_internal<Fut: Future>(
245 &self,
246 background_only: bool,
247 future: Fut,
248 timeout: Option<Duration>,
249 ) -> Result<Fut::Output, impl Future<Output = Fut::Output> + use<Fut>> {
250 use std::sync::atomic::AtomicBool;
251
252 let mut future = Box::pin(future);
253 if timeout == Some(Duration::ZERO) {
254 return Err(future);
255 }
256 let Some(dispatcher) = self.dispatcher.as_test() else {
257 return Err(future);
258 };
259
260 let mut max_ticks = if timeout.is_some() {
261 dispatcher.gen_block_on_ticks()
262 } else {
263 usize::MAX
264 };
265 let unparker = self.dispatcher.unparker();
266 let awoken = Arc::new(AtomicBool::new(false));
267 let waker = waker_fn({
268 let awoken = awoken.clone();
269 move || {
270 awoken.store(true, SeqCst);
271 unparker.unpark();
272 }
273 });
274 let mut cx = std::task::Context::from_waker(&waker);
275
276 loop {
277 match future.as_mut().poll(&mut cx) {
278 Poll::Ready(result) => return Ok(result),
279 Poll::Pending => {
280 if max_ticks == 0 {
281 return Err(future);
282 }
283 max_ticks -= 1;
284
285 if !dispatcher.tick(background_only) {
286 if awoken.swap(false, SeqCst) {
287 continue;
288 }
289
290 if !dispatcher.parking_allowed() {
291 if dispatcher.advance_clock_to_next_delayed() {
292 continue;
293 }
294 let mut backtrace_message = String::new();
295 let mut waiting_message = String::new();
296 if let Some(backtrace) = dispatcher.waiting_backtrace() {
297 backtrace_message =
298 format!("\nbacktrace of waiting future:\n{:?}", backtrace);
299 }
300 if let Some(waiting_hint) = dispatcher.waiting_hint() {
301 waiting_message = format!("\n waiting on: {}\n", waiting_hint);
302 }
303 panic!(
304 "parked with nothing left to run{waiting_message}{backtrace_message}",
305 )
306 }
307 self.dispatcher.park(None);
308 }
309 }
310 }
311 }
312 }
313
314 /// Block the current thread until the given future resolves
315 /// or `duration` has elapsed.
316 pub fn block_with_timeout<Fut: Future>(
317 &self,
318 duration: Duration,
319 future: Fut,
320 ) -> Result<Fut::Output, impl Future<Output = Fut::Output> + use<Fut>> {
321 self.block_internal(true, future, Some(duration))
322 }
323
324 /// Scoped lets you start a number of tasks and waits
325 /// for all of them to complete before returning.
326 pub async fn scoped<'scope, F>(&self, scheduler: F)
327 where
328 F: FnOnce(&mut Scope<'scope>),
329 {
330 let mut scope = Scope::new(self.clone());
331 (scheduler)(&mut scope);
332 let spawned = mem::take(&mut scope.futures)
333 .into_iter()
334 .map(|f| self.spawn(f))
335 .collect::<Vec<_>>();
336 for task in spawned {
337 task.await;
338 }
339 }
340
341 /// Get the current time.
342 ///
343 /// Calling this instead of `std::time::Instant::now` allows the use
344 /// of fake timers in tests.
345 pub fn now(&self) -> Instant {
346 self.dispatcher.now()
347 }
348
349 /// Returns a task that will complete after the given duration.
350 /// Depending on other concurrent tasks the elapsed duration may be longer
351 /// than requested.
352 pub fn timer(&self, duration: Duration) -> Task<()> {
353 if duration.is_zero() {
354 return Task::ready(());
355 }
356 let (runnable, task) = async_task::spawn(async move {}, {
357 let dispatcher = self.dispatcher.clone();
358 move |runnable| dispatcher.dispatch_after(duration, runnable)
359 });
360 runnable.schedule();
361 Task(TaskState::Spawned(task))
362 }
363
364 /// in tests, start_waiting lets you indicate which task is waiting (for debugging only)
365 #[cfg(any(test, feature = "test-support"))]
366 pub fn start_waiting(&self) {
367 self.dispatcher.as_test().unwrap().start_waiting();
368 }
369
370 /// in tests, removes the debugging data added by start_waiting
371 #[cfg(any(test, feature = "test-support"))]
372 pub fn finish_waiting(&self) {
373 self.dispatcher.as_test().unwrap().finish_waiting();
374 }
375
376 /// in tests, run an arbitrary number of tasks (determined by the SEED environment variable)
377 #[cfg(any(test, feature = "test-support"))]
378 pub fn simulate_random_delay(&self) -> impl Future<Output = ()> + use<> {
379 self.dispatcher.as_test().unwrap().simulate_random_delay()
380 }
381
382 /// in tests, indicate that a given task from `spawn_labeled` should run after everything else
383 #[cfg(any(test, feature = "test-support"))]
384 pub fn deprioritize(&self, task_label: TaskLabel) {
385 self.dispatcher.as_test().unwrap().deprioritize(task_label)
386 }
387
388 /// in tests, move time forward. This does not run any tasks, but does make `timer`s ready.
389 #[cfg(any(test, feature = "test-support"))]
390 pub fn advance_clock(&self, duration: Duration) {
391 self.dispatcher.as_test().unwrap().advance_clock(duration)
392 }
393
394 /// in tests, run one task.
395 #[cfg(any(test, feature = "test-support"))]
396 pub fn tick(&self) -> bool {
397 self.dispatcher.as_test().unwrap().tick(false)
398 }
399
400 /// in tests, run all tasks that are ready to run. If after doing so
401 /// the test still has outstanding tasks, this will panic. (See also `allow_parking`)
402 #[cfg(any(test, feature = "test-support"))]
403 pub fn run_until_parked(&self) {
404 self.dispatcher.as_test().unwrap().run_until_parked()
405 }
406
407 /// in tests, prevents `run_until_parked` from panicking if there are outstanding tasks.
408 /// This is useful when you are integrating other (non-GPUI) futures, like disk access, that
409 /// do take real async time to run.
410 #[cfg(any(test, feature = "test-support"))]
411 pub fn allow_parking(&self) {
412 self.dispatcher.as_test().unwrap().allow_parking();
413 }
414
415 /// undoes the effect of [`allow_parking`].
416 #[cfg(any(test, feature = "test-support"))]
417 pub fn forbid_parking(&self) {
418 self.dispatcher.as_test().unwrap().forbid_parking();
419 }
420
421 /// adds detail to the "parked with nothing let to run" message.
422 #[cfg(any(test, feature = "test-support"))]
423 pub fn set_waiting_hint(&self, msg: Option<String>) {
424 self.dispatcher.as_test().unwrap().set_waiting_hint(msg);
425 }
426
427 /// in tests, returns the rng used by the dispatcher and seeded by the `SEED` environment variable
428 #[cfg(any(test, feature = "test-support"))]
429 pub fn rng(&self) -> StdRng {
430 self.dispatcher.as_test().unwrap().rng()
431 }
432
433 /// How many CPUs are available to the dispatcher.
434 pub fn num_cpus(&self) -> usize {
435 #[cfg(any(test, feature = "test-support"))]
436 return 4;
437
438 #[cfg(not(any(test, feature = "test-support")))]
439 return num_cpus::get();
440 }
441
442 /// Whether we're on the main thread.
443 pub fn is_main_thread(&self) -> bool {
444 self.dispatcher.is_main_thread()
445 }
446
447 #[cfg(any(test, feature = "test-support"))]
448 /// in tests, control the number of ticks that `block_with_timeout` will run before timing out.
449 pub fn set_block_on_ticks(&self, range: std::ops::RangeInclusive<usize>) {
450 self.dispatcher.as_test().unwrap().set_block_on_ticks(range);
451 }
452}
453
454/// ForegroundExecutor runs things on the main thread.
455impl ForegroundExecutor {
456 /// Creates a new ForegroundExecutor from the given PlatformDispatcher.
457 pub fn new(dispatcher: Arc<dyn PlatformDispatcher>) -> Self {
458 Self {
459 dispatcher,
460 not_send: PhantomData,
461 }
462 }
463
464 /// Enqueues the given Task to run on the main thread at some point in the future.
465 #[track_caller]
466 pub fn spawn<R>(&self, future: impl Future<Output = R> + 'static) -> Task<R>
467 where
468 R: 'static,
469 {
470 let dispatcher = self.dispatcher.clone();
471
472 #[track_caller]
473 fn inner<R: 'static>(
474 dispatcher: Arc<dyn PlatformDispatcher>,
475 future: AnyLocalFuture<R>,
476 ) -> Task<R> {
477 let (runnable, task) = spawn_local_with_source_location(future, move |runnable| {
478 dispatcher.dispatch_on_main_thread(runnable)
479 });
480 runnable.schedule();
481 Task(TaskState::Spawned(task))
482 }
483 inner::<R>(dispatcher, Box::pin(future))
484 }
485}
486
487/// Variant of `async_task::spawn_local` that includes the source location of the spawn in panics.
488///
489/// Copy-modified from:
490/// https://github.com/smol-rs/async-task/blob/ca9dbe1db9c422fd765847fa91306e30a6bb58a9/src/runnable.rs#L405
491#[track_caller]
492fn spawn_local_with_source_location<Fut, S>(
493 future: Fut,
494 schedule: S,
495) -> (Runnable<()>, async_task::Task<Fut::Output, ()>)
496where
497 Fut: Future + 'static,
498 Fut::Output: 'static,
499 S: async_task::Schedule<()> + Send + Sync + 'static,
500{
501 #[inline]
502 fn thread_id() -> ThreadId {
503 std::thread_local! {
504 static ID: ThreadId = thread::current().id();
505 }
506 ID.try_with(|id| *id)
507 .unwrap_or_else(|_| thread::current().id())
508 }
509
510 struct Checked<F> {
511 id: ThreadId,
512 inner: ManuallyDrop<F>,
513 location: &'static Location<'static>,
514 }
515
516 impl<F> Drop for Checked<F> {
517 fn drop(&mut self) {
518 assert!(
519 self.id == thread_id(),
520 "local task dropped by a thread that didn't spawn it. Task spawned at {}",
521 self.location
522 );
523 unsafe {
524 ManuallyDrop::drop(&mut self.inner);
525 }
526 }
527 }
528
529 impl<F: Future> Future for Checked<F> {
530 type Output = F::Output;
531
532 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
533 assert!(
534 self.id == thread_id(),
535 "local task polled by a thread that didn't spawn it. Task spawned at {}",
536 self.location
537 );
538 unsafe { self.map_unchecked_mut(|c| &mut *c.inner).poll(cx) }
539 }
540 }
541
542 // Wrap the future into one that checks which thread it's on.
543 let future = Checked {
544 id: thread_id(),
545 inner: ManuallyDrop::new(future),
546 location: Location::caller(),
547 };
548
549 unsafe { async_task::spawn_unchecked(future, schedule) }
550}
551
552/// Scope manages a set of tasks that are enqueued and waited on together. See [`BackgroundExecutor::scoped`].
553pub struct Scope<'a> {
554 executor: BackgroundExecutor,
555 futures: Vec<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
556 tx: Option<mpsc::Sender<()>>,
557 rx: mpsc::Receiver<()>,
558 lifetime: PhantomData<&'a ()>,
559}
560
561impl<'a> Scope<'a> {
562 fn new(executor: BackgroundExecutor) -> Self {
563 let (tx, rx) = mpsc::channel(1);
564 Self {
565 executor,
566 tx: Some(tx),
567 rx,
568 futures: Default::default(),
569 lifetime: PhantomData,
570 }
571 }
572
573 /// How many CPUs are available to the dispatcher.
574 pub fn num_cpus(&self) -> usize {
575 self.executor.num_cpus()
576 }
577
578 /// Spawn a future into this scope.
579 pub fn spawn<F>(&mut self, f: F)
580 where
581 F: Future<Output = ()> + Send + 'a,
582 {
583 let tx = self.tx.clone().unwrap();
584
585 // SAFETY: The 'a lifetime is guaranteed to outlive any of these futures because
586 // dropping this `Scope` blocks until all of the futures have resolved.
587 let f = unsafe {
588 mem::transmute::<
589 Pin<Box<dyn Future<Output = ()> + Send + 'a>>,
590 Pin<Box<dyn Future<Output = ()> + Send + 'static>>,
591 >(Box::pin(async move {
592 f.await;
593 drop(tx);
594 }))
595 };
596 self.futures.push(f);
597 }
598}
599
600impl Drop for Scope<'_> {
601 fn drop(&mut self) {
602 self.tx.take().unwrap();
603
604 // Wait until the channel is closed, which means that all of the spawned
605 // futures have resolved.
606 self.executor.block(self.rx.next());
607 }
608}