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