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