tenant.openagents/omega

1use crate::{App, PlatformDispatcher, PlatformScheduler};
2use futures::channel::mpsc;
3use futures::prelude::*;
4use gpui_util::{TryFutureExt, TryFutureExtBacktrace};
5use scheduler::Instant;
6use scheduler::Scheduler;
7use std::{future::Future, marker::PhantomData, mem, pin::Pin, rc::Rc, sync::Arc, time::Duration};
8
9pub use scheduler::{FallibleTask, LocalExecutor as SchedulerLocalExecutor, Priority, Task};
10
11/// A pointer to the executor that is currently running,
12/// for spawning background tasks.
13#[derive(Clone)]
14pub struct BackgroundExecutor {
15    inner: scheduler::BackgroundExecutor,
16    dispatcher: Arc<dyn PlatformDispatcher>,
17}
18
19/// A pointer to the executor that is currently running,
20/// for spawning tasks on the main thread.
21#[derive(Clone)]
22pub struct ForegroundExecutor {
23    inner: scheduler::LocalExecutor,
24    dispatcher: Arc<dyn PlatformDispatcher>,
25    not_send: PhantomData<Rc<()>>,
26}
27
28/// Extension trait for `Task<Result<T, E>>` that adds `detach_and_log_err` with an `&App` context.
29///
30/// This trait is automatically implemented for all `Task<Result<T, E>>` types.
31pub trait TaskExt<T, E> {
32    /// Run the task to completion in the background and log any errors that occur.
33    fn detach_and_log_err(self, cx: &App);
34    /// Like [`Self::detach_and_log_err`], but uses `{:?}` formatting on failure so `anyhow::Error`
35    /// values emit their full backtrace. Prefer `detach_and_log_err` unless a backtrace is wanted.
36    fn detach_and_log_err_with_backtrace(self, cx: &App);
37}
38
39impl<T, E> TaskExt<T, E> for Task<Result<T, E>>
40where
41    T: 'static,
42    E: 'static + std::fmt::Display + std::fmt::Debug,
43{
44    #[track_caller]
45    fn detach_and_log_err(self, cx: &App) {
46        let location = core::panic::Location::caller();
47        cx.foreground_executor()
48            .spawn(self.log_tracked_err(*location))
49            .detach();
50    }
51
52    #[track_caller]
53    fn detach_and_log_err_with_backtrace(self, cx: &App) {
54        let location = *core::panic::Location::caller();
55        cx.foreground_executor()
56            .spawn(self.log_tracked_err_with_backtrace(location))
57            .detach();
58    }
59}
60
61impl BackgroundExecutor {
62    /// Creates a new BackgroundExecutor from the given PlatformDispatcher.
63    pub fn new(dispatcher: Arc<dyn PlatformDispatcher>) -> Self {
64        #[cfg(any(test, feature = "test-support"))]
65        let scheduler: Arc<dyn Scheduler> = if let Some(test_dispatcher) = dispatcher.as_test() {
66            test_dispatcher.scheduler().clone()
67        } else {
68            Arc::new(PlatformScheduler::new(dispatcher.clone()))
69        };
70
71        #[cfg(not(any(test, feature = "test-support")))]
72        let scheduler: Arc<dyn Scheduler> = Arc::new(PlatformScheduler::new(dispatcher.clone()));
73
74        Self {
75            inner: scheduler::BackgroundExecutor::new(scheduler),
76            dispatcher,
77        }
78    }
79
80    /// Returns the underlying scheduler::BackgroundExecutor.
81    ///
82    /// This is used by Ex to pass the executor to thread/worktree code.
83    pub fn scheduler_executor(&self) -> scheduler::BackgroundExecutor {
84        self.inner.clone()
85    }
86
87    /// Enqueues the given future to be run to completion on a background thread.
88    #[track_caller]
89    pub fn spawn<R>(&self, future: impl Future<Output = R> + Send + 'static) -> Task<R>
90    where
91        R: Send + 'static,
92    {
93        self.spawn_with_priority(Priority::default(), future.boxed())
94    }
95
96    /// Enqueues the given future to be run to completion on a background thread with the given priority.
97    ///
98    /// When `Priority::RealtimeAudio` is used, the task runs on a dedicated thread with
99    /// realtime scheduling priority, suitable for audio processing.
100    #[track_caller]
101    pub fn spawn_with_priority<R>(
102        &self,
103        priority: Priority,
104        future: impl Future<Output = R> + Send + 'static,
105    ) -> Task<R>
106    where
107        R: Send + 'static,
108    {
109        if priority == Priority::RealtimeAudio {
110            self.inner.spawn_realtime(future)
111        } else {
112            self.inner.spawn_with_priority(priority, future)
113        }
114    }
115
116    /// Scoped lets you start a number of tasks and waits
117    /// for all of them to complete before returning.
118    pub async fn scoped<'scope, F>(&self, scheduler: F)
119    where
120        F: FnOnce(&mut Scope<'scope>),
121    {
122        let mut scope = Scope::new(self.clone(), Priority::default());
123        (scheduler)(&mut scope);
124        let spawned = mem::take(&mut scope.futures)
125            .into_iter()
126            .map(|f| self.spawn_with_priority(scope.priority, f))
127            .collect::<Vec<_>>();
128        for task in spawned {
129            task.await;
130        }
131    }
132
133    /// Scoped lets you start a number of tasks and waits
134    /// for all of them to complete before returning.
135    pub async fn scoped_priority<'scope, F>(&self, priority: Priority, scheduler: F)
136    where
137        F: FnOnce(&mut Scope<'scope>),
138    {
139        let mut scope = Scope::new(self.clone(), priority);
140        (scheduler)(&mut scope);
141        let spawned = mem::take(&mut scope.futures)
142            .into_iter()
143            .map(|f| self.spawn_with_priority(scope.priority, f))
144            .collect::<Vec<_>>();
145        for task in spawned {
146            task.await;
147        }
148    }
149
150    /// Get the current time.
151    ///
152    /// Calling this instead of `std::time::Instant::now` allows the use
153    /// of fake timers in tests.
154    pub fn now(&self) -> Instant {
155        self.inner.scheduler().clock().now()
156    }
157
158    /// Returns a task that will complete after the given duration.
159    /// Depending on other concurrent tasks the elapsed duration may be longer
160    /// than requested.
161    #[track_caller]
162    pub fn timer(&self, duration: Duration) -> Task<()> {
163        if duration.is_zero() {
164            return Task::ready(());
165        }
166        self.spawn(self.inner.scheduler().timer(duration))
167    }
168
169    /// In tests, run an arbitrary number of tasks (determined by the SEED environment variable)
170    #[cfg(any(test, feature = "test-support"))]
171    pub fn simulate_random_delay(&self) -> impl Future<Output = ()> + use<> {
172        self.dispatcher.as_test().unwrap().simulate_random_delay()
173    }
174
175    /// In tests, move time forward. This does not run any tasks, but does make `timer`s ready.
176    #[cfg(any(test, feature = "test-support"))]
177    pub fn advance_clock(&self, duration: Duration) {
178        self.dispatcher.as_test().unwrap().advance_clock(duration)
179    }
180
181    /// In tests, run one task.
182    #[cfg(any(test, feature = "test-support"))]
183    pub fn tick(&self) -> bool {
184        self.dispatcher.as_test().unwrap().scheduler().tick()
185    }
186
187    /// In tests, run tasks until the scheduler would park.
188    ///
189    /// Under the scheduler-backed test dispatcher, `tick()` will not advance the clock, so a pending
190    /// timer can keep `has_pending_tasks()` true even after all currently-runnable tasks have been
191    /// drained. To preserve the historical semantics that tests relied on (drain all work that can
192    /// make progress), we advance the clock to the next timer when no runnable tasks remain.
193    #[cfg(any(test, feature = "test-support"))]
194    pub fn run_until_parked(&self) {
195        let scheduler = self.dispatcher.as_test().unwrap().scheduler();
196        scheduler.run();
197    }
198
199    /// In tests, prevents `run_until_parked` from panicking if there are outstanding tasks.
200    #[cfg(any(test, feature = "test-support"))]
201    pub fn allow_parking(&self) {
202        self.dispatcher
203            .as_test()
204            .unwrap()
205            .scheduler()
206            .allow_parking();
207
208        if std::env::var("GPUI_RUN_UNTIL_PARKED_LOG").ok().as_deref() == Some("1") {
209            log::warn!("[gpui::executor] allow_parking: enabled");
210        }
211    }
212
213    /// Sets the range of ticks to run before timing out in block_on.
214    #[cfg(any(test, feature = "test-support"))]
215    pub fn set_block_on_ticks(&self, range: std::ops::RangeInclusive<usize>) {
216        self.dispatcher
217            .as_test()
218            .unwrap()
219            .scheduler()
220            .set_timeout_ticks(range);
221    }
222
223    /// Undoes the effect of [`Self::allow_parking`].
224    #[cfg(any(test, feature = "test-support"))]
225    pub fn forbid_parking(&self) {
226        self.dispatcher
227            .as_test()
228            .unwrap()
229            .scheduler()
230            .forbid_parking();
231    }
232
233    /// In tests, returns the rng used by the dispatcher.
234    #[cfg(any(test, feature = "test-support"))]
235    pub fn rng(&self) -> scheduler::SharedRng {
236        self.dispatcher.as_test().unwrap().scheduler().rng()
237    }
238
239    /// How many CPUs are available to the dispatcher.
240    pub fn num_cpus(&self) -> usize {
241        #[cfg(any(test, feature = "test-support"))]
242        if let Some(test) = self.dispatcher.as_test() {
243            return test.num_cpus_override().unwrap_or(4);
244        }
245        num_cpus::get()
246    }
247
248    /// Override the number of CPUs reported by this executor in tests.
249    /// Panics if not called on a test executor.
250    #[cfg(any(test, feature = "test-support"))]
251    pub fn set_num_cpus(&self, count: usize) {
252        self.dispatcher
253            .as_test()
254            .expect("set_num_cpus can only be called on a test executor")
255            .set_num_cpus(count);
256    }
257
258    /// Whether we're on the main thread.
259    pub fn is_main_thread(&self) -> bool {
260        self.dispatcher.is_main_thread()
261    }
262
263    #[doc(hidden)]
264    pub fn dispatcher(&self) -> &Arc<dyn PlatformDispatcher> {
265        &self.dispatcher
266    }
267}
268
269impl ForegroundExecutor {
270    /// Creates a new ForegroundExecutor from the given PlatformDispatcher.
271    pub fn new(dispatcher: Arc<dyn PlatformDispatcher>) -> Self {
272        #[cfg(any(test, feature = "test-support"))]
273        let (scheduler, session_id): (Arc<dyn Scheduler>, _) =
274            if let Some(test_dispatcher) = dispatcher.as_test() {
275                (
276                    test_dispatcher.scheduler().clone(),
277                    test_dispatcher.session_id(),
278                )
279            } else {
280                let platform_scheduler = Arc::new(PlatformScheduler::new(dispatcher.clone()));
281                let inner = platform_scheduler.foreground_executor();
282                return Self {
283                    inner,
284                    dispatcher,
285                    not_send: PhantomData,
286                };
287            };
288
289        #[cfg(not(any(test, feature = "test-support")))]
290        let inner = {
291            let platform_scheduler = Arc::new(PlatformScheduler::new(dispatcher.clone()));
292            platform_scheduler.foreground_executor()
293        };
294
295        #[cfg(any(test, feature = "test-support"))]
296        let inner = {
297            let scheduler_for_dispatch = Arc::downgrade(&scheduler);
298            scheduler::LocalExecutor::new(session_id, scheduler, move |runnable| {
299                if let Some(scheduler) = scheduler_for_dispatch.upgrade() {
300                    scheduler.schedule_local(session_id, runnable);
301                }
302            })
303        };
304
305        Self {
306            inner,
307            dispatcher,
308            not_send: PhantomData,
309        }
310    }
311
312    /// Enqueues the given Task to run on the main thread.
313    #[track_caller]
314    pub fn spawn<R>(&self, future: impl Future<Output = R> + 'static) -> Task<R>
315    where
316        R: 'static,
317    {
318        self.inner.spawn(future.boxed_local())
319    }
320
321    /// Enqueues the given Task to run on the main thread with the given priority.
322    #[track_caller]
323    pub fn spawn_with_priority<R>(
324        &self,
325        _priority: Priority,
326        future: impl Future<Output = R> + 'static,
327    ) -> Task<R>
328    where
329        R: 'static,
330    {
331        // Priority is ignored for foreground tasks - they run in order on the main thread
332        self.inner.spawn(future)
333    }
334
335    /// Used by the test harness to run an async test in a synchronous fashion.
336    #[cfg(any(test, feature = "test-support"))]
337    #[track_caller]
338    pub fn block_test<R>(&self, future: impl Future<Output = R>) -> R {
339        use std::cell::Cell;
340
341        let scheduler = self.inner.scheduler();
342
343        let output = Cell::new(None);
344        let future = async {
345            output.set(Some(future.await));
346        };
347        let mut future = std::pin::pin!(future);
348
349        // In async GPUI tests, we must allow foreground tasks scheduled by the test itself
350        // (which are associated with the test session) to make progress while we block.
351        // Otherwise, awaiting futures that depend on same-session foreground work can deadlock.
352        scheduler.block(None, future.as_mut(), None);
353
354        output.take().expect("block_test future did not complete")
355    }
356
357    /// Block the current thread until the given future resolves.
358    /// Consider using `block_with_timeout` instead.
359    pub fn block_on<R>(&self, future: impl Future<Output = R>) -> R {
360        self.inner.block_on(future)
361    }
362
363    /// Block the current thread until the given future resolves or the timeout elapses.
364    pub fn block_with_timeout<R, Fut: Future<Output = R>>(
365        &self,
366        duration: Duration,
367        future: Fut,
368    ) -> Result<R, impl Future<Output = R> + use<R, Fut>> {
369        self.inner.block_with_timeout(duration, future)
370    }
371
372    #[doc(hidden)]
373    pub fn dispatcher(&self) -> &Arc<dyn PlatformDispatcher> {
374        &self.dispatcher
375    }
376
377    #[doc(hidden)]
378    pub fn scheduler_executor(&self) -> SchedulerLocalExecutor {
379        self.inner.clone()
380    }
381}
382
383/// Scope manages a set of tasks that are enqueued and waited on together. See [`BackgroundExecutor::scoped`].
384pub struct Scope<'a> {
385    executor: BackgroundExecutor,
386    priority: Priority,
387    futures: Vec<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
388    tx: Option<mpsc::Sender<()>>,
389    rx: mpsc::Receiver<()>,
390    lifetime: PhantomData<&'a ()>,
391}
392
393impl<'a> Scope<'a> {
394    fn new(executor: BackgroundExecutor, priority: Priority) -> Self {
395        let (tx, rx) = mpsc::channel(1);
396        Self {
397            executor,
398            priority,
399            tx: Some(tx),
400            rx,
401            futures: Default::default(),
402            lifetime: PhantomData,
403        }
404    }
405
406    /// How many CPUs are available to the dispatcher.
407    pub fn num_cpus(&self) -> usize {
408        self.executor.num_cpus()
409    }
410
411    /// Spawn a future into this scope.
412    #[track_caller]
413    pub fn spawn<F>(&mut self, f: F)
414    where
415        F: Future<Output = ()> + Send + 'a,
416    {
417        let tx = self.tx.clone().unwrap();
418
419        // SAFETY: The 'a lifetime is guaranteed to outlive any of these futures because
420        // dropping this `Scope` blocks until all of the futures have resolved.
421        let f = unsafe {
422            mem::transmute::<
423                Pin<Box<dyn Future<Output = ()> + Send + 'a>>,
424                Pin<Box<dyn Future<Output = ()> + Send + 'static>>,
425            >(Box::pin(async move {
426                f.await;
427                drop(tx);
428            }))
429        };
430        self.futures.push(f);
431    }
432}
433
434impl Drop for Scope<'_> {
435    fn drop(&mut self) {
436        self.tx.take().unwrap();
437
438        // Wait until the channel is closed, which means that all of the spawned
439        // futures have resolved.
440        let future = async {
441            self.rx.next().await;
442        };
443        let mut future = std::pin::pin!(future);
444        self.executor
445            .inner
446            .scheduler()
447            .block(None, future.as_mut(), None);
448    }
449}
450
451#[cfg(test)]
452mod test {
453    use super::*;
454    use crate::{App, TestDispatcher, TestPlatform};
455    use std::cell::RefCell;
456
457    /// Helper to create test infrastructure.
458    /// Returns (dispatcher, background_executor, app).
459    fn create_test_app() -> (TestDispatcher, BackgroundExecutor, Rc<crate::AppCell>) {
460        let dispatcher = TestDispatcher::new(0);
461        let arc_dispatcher = Arc::new(dispatcher.clone());
462        let background_executor = BackgroundExecutor::new(arc_dispatcher.clone());
463        let foreground_executor = ForegroundExecutor::new(arc_dispatcher);
464
465        let platform = TestPlatform::new(background_executor.clone(), foreground_executor);
466        let asset_source = Arc::new(());
467        let http_client = http_client::FakeHttpClient::with_404_response();
468
469        let app = App::new_app(platform, asset_source, http_client);
470        (dispatcher, background_executor, app)
471    }
472
473    #[test]
474    fn sanity_test_tasks_run() {
475        let (dispatcher, _background_executor, app) = create_test_app();
476        let foreground_executor = app.borrow().foreground_executor.clone();
477
478        let task_ran = Rc::new(RefCell::new(false));
479
480        foreground_executor
481            .spawn({
482                let task_ran = Rc::clone(&task_ran);
483                async move {
484                    *task_ran.borrow_mut() = true;
485                }
486            })
487            .detach();
488
489        // Run dispatcher while app is still alive
490        dispatcher.run_until_parked();
491
492        // Task should have run
493        assert!(
494            *task_ran.borrow(),
495            "Task should run normally when app is alive"
496        );
497    }
498}
499