1use crate::{AppContext, PlatformDispatcher};
2use futures::{channel::mpsc, pin_mut, FutureExt};
3use smol::prelude::*;
4use std::{
5 fmt::Debug,
6 marker::PhantomData,
7 mem,
8 pin::Pin,
9 rc::Rc,
10 sync::{
11 atomic::{AtomicBool, Ordering::SeqCst},
12 Arc,
13 },
14 task::{Context, Poll},
15 time::Duration,
16};
17use util::TryFutureExt;
18use waker_fn::waker_fn;
19
20#[cfg(any(test, feature = "test-support"))]
21use rand::rngs::StdRng;
22
23#[derive(Clone)]
24pub struct BackgroundExecutor {
25 dispatcher: Arc<dyn PlatformDispatcher>,
26}
27
28#[derive(Clone)]
29pub struct ForegroundExecutor {
30 dispatcher: Arc<dyn PlatformDispatcher>,
31 not_send: PhantomData<Rc<()>>,
32}
33
34#[must_use]
35#[derive(Debug)]
36pub enum Task<T> {
37 Ready(Option<T>),
38 Spawned(async_task::Task<T>),
39}
40
41impl<T> Task<T> {
42 pub fn ready(val: T) -> Self {
43 Task::Ready(Some(val))
44 }
45
46 pub fn detach(self) {
47 match self {
48 Task::Ready(_) => {}
49 Task::Spawned(task) => task.detach(),
50 }
51 }
52}
53
54impl<E, T> Task<Result<T, E>>
55where
56 T: 'static,
57 E: 'static + Debug,
58{
59 pub fn detach_and_log_err(self, cx: &mut AppContext) {
60 cx.foreground_executor().spawn(self.log_err()).detach();
61 }
62}
63
64impl<T> Future for Task<T> {
65 type Output = T;
66
67 fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
68 match unsafe { self.get_unchecked_mut() } {
69 Task::Ready(val) => Poll::Ready(val.take().unwrap()),
70 Task::Spawned(task) => task.poll(cx),
71 }
72 }
73}
74type AnyLocalFuture<R> = Pin<Box<dyn 'static + Future<Output = R>>>;
75type AnyFuture<R> = Pin<Box<dyn 'static + Send + Future<Output = R>>>;
76impl BackgroundExecutor {
77 pub fn new(dispatcher: Arc<dyn PlatformDispatcher>) -> Self {
78 Self { dispatcher }
79 }
80
81 /// Enqueues the given closure to be run on any thread. The closure returns
82 /// a future which will be run to completion on any available thread.
83 pub fn spawn<R>(&self, future: impl Future<Output = R> + Send + 'static) -> Task<R>
84 where
85 R: Send + 'static,
86 {
87 let dispatcher = self.dispatcher.clone();
88 fn inner<R: Send + 'static>(
89 dispatcher: Arc<dyn PlatformDispatcher>,
90 future: AnyFuture<R>,
91 ) -> Task<R> {
92 let (runnable, task) =
93 async_task::spawn(future, move |runnable| dispatcher.dispatch(runnable));
94 runnable.schedule();
95 Task::Spawned(task)
96 }
97 inner::<R>(dispatcher, Box::pin(future))
98 }
99
100 #[cfg(any(test, feature = "test-support"))]
101 #[track_caller]
102 pub fn block_test<R>(&self, future: impl Future<Output = R>) -> R {
103 self.block_internal(false, future)
104 }
105
106 pub fn block<R>(&self, future: impl Future<Output = R>) -> R {
107 self.block_internal(true, future)
108 }
109
110 #[track_caller]
111 pub(crate) fn block_internal<R>(
112 &self,
113 background_only: bool,
114 future: impl Future<Output = R>,
115 ) -> R {
116 pin_mut!(future);
117 let unparker = self.dispatcher.unparker();
118 let awoken = Arc::new(AtomicBool::new(false));
119
120 let waker = waker_fn({
121 let awoken = awoken.clone();
122 move || {
123 awoken.store(true, SeqCst);
124 unparker.unpark();
125 }
126 });
127 let mut cx = std::task::Context::from_waker(&waker);
128
129 loop {
130 match future.as_mut().poll(&mut cx) {
131 Poll::Ready(result) => return result,
132 Poll::Pending => {
133 if !self.dispatcher.poll(background_only) {
134 if awoken.swap(false, SeqCst) {
135 continue;
136 }
137
138 #[cfg(any(test, feature = "test-support"))]
139 if let Some(test) = self.dispatcher.as_test() {
140 if !test.parking_allowed() {
141 let mut backtrace_message = String::new();
142 if let Some(backtrace) = test.waiting_backtrace() {
143 backtrace_message =
144 format!("\nbacktrace of waiting future:\n{:?}", backtrace);
145 }
146 panic!("parked with nothing left to run\n{:?}", backtrace_message)
147 }
148 }
149
150 self.dispatcher.park();
151 }
152 }
153 }
154 }
155 }
156
157 pub fn block_with_timeout<R>(
158 &self,
159 duration: Duration,
160 future: impl Future<Output = R>,
161 ) -> Result<R, impl Future<Output = R>> {
162 let mut future = Box::pin(future.fuse());
163 if duration.is_zero() {
164 return Err(future);
165 }
166
167 let mut timer = self.timer(duration).fuse();
168 let timeout = async {
169 futures::select_biased! {
170 value = future => Ok(value),
171 _ = timer => Err(()),
172 }
173 };
174 match self.block(timeout) {
175 Ok(value) => Ok(value),
176 Err(_) => Err(future),
177 }
178 }
179
180 pub async fn scoped<'scope, F>(&self, scheduler: F)
181 where
182 F: FnOnce(&mut Scope<'scope>),
183 {
184 let mut scope = Scope::new(self.clone());
185 (scheduler)(&mut scope);
186 let spawned = mem::take(&mut scope.futures)
187 .into_iter()
188 .map(|f| self.spawn(f))
189 .collect::<Vec<_>>();
190 for task in spawned {
191 task.await;
192 }
193 }
194
195 pub fn timer(&self, duration: Duration) -> Task<()> {
196 let (runnable, task) = async_task::spawn(async move {}, {
197 let dispatcher = self.dispatcher.clone();
198 move |runnable| dispatcher.dispatch_after(duration, runnable)
199 });
200 runnable.schedule();
201 Task::Spawned(task)
202 }
203
204 #[cfg(any(test, feature = "test-support"))]
205 pub fn start_waiting(&self) {
206 self.dispatcher.as_test().unwrap().start_waiting();
207 }
208
209 #[cfg(any(test, feature = "test-support"))]
210 pub fn finish_waiting(&self) {
211 self.dispatcher.as_test().unwrap().finish_waiting();
212 }
213
214 #[cfg(any(test, feature = "test-support"))]
215 pub fn simulate_random_delay(&self) -> impl Future<Output = ()> {
216 self.dispatcher.as_test().unwrap().simulate_random_delay()
217 }
218
219 #[cfg(any(test, feature = "test-support"))]
220 pub fn advance_clock(&self, duration: Duration) {
221 self.dispatcher.as_test().unwrap().advance_clock(duration)
222 }
223
224 #[cfg(any(test, feature = "test-support"))]
225 pub fn run_until_parked(&self) {
226 self.dispatcher.as_test().unwrap().run_until_parked()
227 }
228
229 #[cfg(any(test, feature = "test-support"))]
230 pub fn allow_parking(&self) {
231 self.dispatcher.as_test().unwrap().allow_parking();
232 }
233
234 #[cfg(any(test, feature = "test-support"))]
235 pub fn rng(&self) -> StdRng {
236 self.dispatcher.as_test().unwrap().rng()
237 }
238
239 pub fn num_cpus(&self) -> usize {
240 num_cpus::get()
241 }
242
243 pub fn is_main_thread(&self) -> bool {
244 self.dispatcher.is_main_thread()
245 }
246}
247
248impl ForegroundExecutor {
249 pub fn new(dispatcher: Arc<dyn PlatformDispatcher>) -> Self {
250 Self {
251 dispatcher,
252 not_send: PhantomData,
253 }
254 }
255
256 /// Enqueues the given closure to be run on any thread. The closure returns
257 /// a future which will be run to completion on any available thread.
258 pub fn spawn<R>(&self, future: impl Future<Output = R> + 'static) -> Task<R>
259 where
260 R: 'static,
261 {
262 let dispatcher = self.dispatcher.clone();
263 fn inner<R: 'static>(
264 dispatcher: Arc<dyn PlatformDispatcher>,
265 future: AnyLocalFuture<R>,
266 ) -> Task<R> {
267 let (runnable, task) = async_task::spawn_local(future, move |runnable| {
268 dispatcher.dispatch_on_main_thread(runnable)
269 });
270 runnable.schedule();
271 Task::Spawned(task)
272 }
273 inner::<R>(dispatcher, Box::pin(future))
274 }
275}
276
277pub struct Scope<'a> {
278 executor: BackgroundExecutor,
279 futures: Vec<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
280 tx: Option<mpsc::Sender<()>>,
281 rx: mpsc::Receiver<()>,
282 lifetime: PhantomData<&'a ()>,
283}
284
285impl<'a> Scope<'a> {
286 fn new(executor: BackgroundExecutor) -> Self {
287 let (tx, rx) = mpsc::channel(1);
288 Self {
289 executor,
290 tx: Some(tx),
291 rx,
292 futures: Default::default(),
293 lifetime: PhantomData,
294 }
295 }
296
297 pub fn spawn<F>(&mut self, f: F)
298 where
299 F: Future<Output = ()> + Send + 'a,
300 {
301 let tx = self.tx.clone().unwrap();
302
303 // Safety: The 'a lifetime is guaranteed to outlive any of these futures because
304 // dropping this `Scope` blocks until all of the futures have resolved.
305 let f = unsafe {
306 mem::transmute::<
307 Pin<Box<dyn Future<Output = ()> + Send + 'a>>,
308 Pin<Box<dyn Future<Output = ()> + Send + 'static>>,
309 >(Box::pin(async move {
310 f.await;
311 drop(tx);
312 }))
313 };
314 self.futures.push(f);
315 }
316}
317
318impl<'a> Drop for Scope<'a> {
319 fn drop(&mut self) {
320 self.tx.take().unwrap();
321
322 // Wait until the channel is closed, which means that all of the spawned
323 // futures have resolved.
324 self.executor.block(self.rx.next());
325 }
326}