tenant.openagents/omega

1use crate::{PlatformDispatcher, RunnableMeta};
2use async_task::Runnable;
3use chrono::{DateTime, Utc};
4use futures::channel::oneshot;
5use scheduler::Instant;
6use scheduler::{
7    Clock, LocalExecutor, Priority, Scheduler, SessionId, Task, TestScheduler, Timer,
8    spawn_dedicated_thread,
9};
10#[cfg(not(target_family = "wasm"))]
11use std::task::{Context, Poll};
12use std::{
13    any::Any,
14    future::Future,
15    pin::Pin,
16    sync::{
17        Arc,
18        atomic::{AtomicU16, Ordering},
19    },
20    time::Duration,
21};
22
23/// A production implementation of [`Scheduler`] that wraps a [`PlatformDispatcher`].
24///
25/// This allows GPUI to use the scheduler crate's executor types with the platform's
26/// native dispatch mechanisms (e.g., Grand Central Dispatch on macOS).
27pub struct PlatformScheduler {
28    dispatcher: Arc<dyn PlatformDispatcher>,
29    clock: Arc<PlatformClock>,
30    next_session_id: AtomicU16,
31}
32
33impl PlatformScheduler {
34    pub fn new(dispatcher: Arc<dyn PlatformDispatcher>) -> Self {
35        Self {
36            dispatcher: dispatcher.clone(),
37            clock: Arc::new(PlatformClock { dispatcher }),
38            next_session_id: AtomicU16::new(0),
39        }
40    }
41
42    pub fn foreground_executor(self: &Arc<Self>) -> LocalExecutor {
43        let session_id = self.next_session_id();
44        let scheduler = Arc::downgrade(self);
45        LocalExecutor::new(session_id, self.clone(), move |runnable| {
46            if let Some(scheduler) = scheduler.upgrade() {
47                scheduler.schedule_local(session_id, runnable);
48            }
49        })
50    }
51
52    fn next_session_id(&self) -> SessionId {
53        SessionId::new(self.next_session_id.fetch_add(1, Ordering::SeqCst))
54    }
55}
56
57impl Scheduler for PlatformScheduler {
58    fn block(
59        &self,
60        _session_id: Option<SessionId>,
61        #[cfg_attr(target_family = "wasm", allow(unused_mut))] mut future: Pin<
62            &mut dyn Future<Output = ()>,
63        >,
64        #[cfg_attr(target_family = "wasm", allow(unused_variables))] timeout: Option<Duration>,
65    ) -> bool {
66        #[cfg(target_family = "wasm")]
67        {
68            let _ = (&future, &timeout);
69            panic!("Cannot block on wasm")
70        }
71        #[cfg(not(target_family = "wasm"))]
72        {
73            use waker_fn::waker_fn;
74            let deadline = timeout.map(|t| Instant::now() + t);
75            let parker = parking::Parker::new();
76            let unparker = parker.unparker();
77            let waker = waker_fn(move || {
78                unparker.unpark();
79            });
80            let mut cx = Context::from_waker(&waker);
81            if let Poll::Ready(()) = future.as_mut().poll(&mut cx) {
82                return true;
83            }
84
85            let park_deadline = |deadline: Instant| {
86                // Timer expirations are only delivered every ~15.6 milliseconds by default on Windows.
87                // We increase the resolution during this wait so that short timeouts stay reasonably short.
88                let _timer_guard = self.dispatcher.increase_timer_resolution();
89                parker.park_deadline(deadline)
90            };
91
92            loop {
93                match deadline {
94                    Some(deadline) if !park_deadline(deadline) && deadline <= Instant::now() => {
95                        return false;
96                    }
97                    Some(_) => (),
98                    None => parker.park(),
99                }
100                if let Poll::Ready(()) = future.as_mut().poll(&mut cx) {
101                    break true;
102                }
103            }
104        }
105    }
106
107    fn schedule_local(&self, _session_id: SessionId, runnable: Runnable<RunnableMeta>) {
108        self.dispatcher
109            .dispatch_on_main_thread(runnable, Priority::default());
110    }
111
112    fn schedule_background_with_priority(
113        &self,
114        runnable: Runnable<RunnableMeta>,
115        priority: Priority,
116    ) {
117        self.dispatcher.dispatch(runnable, priority);
118    }
119
120    fn spawn_realtime(&self, f: Box<dyn FnOnce() + Send>) {
121        self.dispatcher.spawn_realtime(f);
122    }
123
124    #[track_caller]
125    fn timer(&self, duration: Duration) -> Timer {
126        let (tx, rx) = oneshot::channel();
127        let dispatcher = self.dispatcher.clone();
128
129        // Create a runnable that will send the completion signal
130        let location = std::panic::Location::caller();
131        let (runnable, _task) = async_task::Builder::new()
132            .metadata(RunnableMeta {
133                location,
134                spawned: scheduler::SpawnTime(Instant::now()),
135            })
136            .spawn(
137                move |_| async move {
138                    let _ = tx.send(());
139                },
140                move |runnable| {
141                    dispatcher.dispatch_after(duration, runnable);
142                },
143            );
144        runnable.schedule();
145
146        Timer::new(rx)
147    }
148
149    fn clock(&self) -> Arc<dyn Clock> {
150        self.clock.clone()
151    }
152
153    fn spawn_dedicated(
154        self: Arc<Self>,
155        f: Box<
156            dyn FnOnce(
157                    LocalExecutor,
158                )
159                    -> Pin<Box<dyn Future<Output = Box<dyn Any + Send + Sync>> + 'static>>
160                + Send
161                + 'static,
162        >,
163    ) -> Task<Box<dyn Any + Send + Sync>> {
164        let session_id = self.next_session_id();
165        spawn_dedicated_thread(session_id, self, move |executor| f(executor))
166    }
167
168    fn as_test(&self) -> Option<&TestScheduler> {
169        None
170    }
171}
172
173/// A production clock that uses the platform dispatcher's time.
174struct PlatformClock {
175    dispatcher: Arc<dyn PlatformDispatcher>,
176}
177
178impl Clock for PlatformClock {
179    fn utc_now(&self) -> DateTime<Utc> {
180        Utc::now()
181    }
182
183    fn now(&self) -> Instant {
184        self.dispatcher.now()
185    }
186}
187
188#[cfg(all(test, not(target_family = "wasm")))]
189mod tests {
190    use super::*;
191    use crate::RunnableVariant;
192    use scheduler::BackgroundExecutor;
193    use std::time::Instant as StdInstant;
194
195    // `spawn_dedicated` shouldn't touch the platform dispatcher at all;
196    // panicking on every method ensures the test catches it if it does.
197    struct SmokeDispatcher;
198
199    impl PlatformDispatcher for SmokeDispatcher {
200        fn is_main_thread(&self) -> bool {
201            false
202        }
203        fn dispatch(&self, _runnable: RunnableVariant, _priority: Priority) {
204            panic!("SmokeDispatcher should not be asked to dispatch in this test");
205        }
206        fn dispatch_on_main_thread(&self, _runnable: RunnableVariant, _priority: Priority) {
207            panic!("SmokeDispatcher does not implement a main thread");
208        }
209        fn dispatch_after(&self, _duration: Duration, _runnable: RunnableVariant) {
210            panic!("SmokeDispatcher does not implement timers");
211        }
212        fn spawn_realtime(&self, _f: Box<dyn FnOnce() + Send>) {
213            panic!("SmokeDispatcher does not implement realtime");
214        }
215    }
216
217    #[test]
218    fn spawn_dedicated_runs_on_a_real_separate_thread() {
219        let background =
220            BackgroundExecutor::new(Arc::new(PlatformScheduler::new(Arc::new(SmokeDispatcher))));
221        let started = StdInstant::now();
222        let task = background.spawn_dedicated(|_executor| async move {
223            // A genuine blocking syscall on the dedicated thread. If
224            // `spawn_dedicated` were running the future on any shared
225            // executor, this would stall that executor.
226            let thread_id_before = std::thread::current().id();
227            std::thread::sleep(Duration::from_millis(50));
228            let thread_id_after = std::thread::current().id();
229            assert_eq!(thread_id_before, thread_id_after);
230            (thread_id_before, "slept")
231        });
232        let (dedicated_thread_id, message) = futures::executor::block_on(task);
233        let elapsed = started.elapsed();
234        assert_eq!(message, "slept");
235        assert_ne!(
236            dedicated_thread_id,
237            std::thread::current().id(),
238            "dedicated future ran on the test thread"
239        );
240        assert!(
241            elapsed >= Duration::from_millis(40),
242            "expected the dedicated thread to genuinely sleep, elapsed = {:?}",
243            elapsed
244        );
245    }
246
247    #[test]
248    fn spawn_dedicated_returns_not_send_future_output() {
249        // The whole point of `spawn_dedicated` is that the future can be
250        // `!Send`. Constructing one with `Rc<RefCell<_>>` ensures the
251        // signature actually permits it.
252        use std::cell::RefCell;
253        use std::rc::Rc;
254
255        let background =
256            BackgroundExecutor::new(Arc::new(PlatformScheduler::new(Arc::new(SmokeDispatcher))));
257        let task = background.spawn_dedicated(|_executor| async move {
258            let state = Rc::new(RefCell::new(0_i32));
259            for _ in 0..3 {
260                *state.borrow_mut() += 1;
261            }
262            *state.borrow()
263        });
264        let output = futures::executor::block_on(task);
265        assert_eq!(output, 3);
266    }
267
268    #[test]
269    fn spawn_dedicated_dropping_task_cancels_future() {
270        use parking_lot::Mutex;
271        use std::sync::mpsc;
272
273        let background =
274            BackgroundExecutor::new(Arc::new(PlatformScheduler::new(Arc::new(SmokeDispatcher))));
275
276        let (started_tx, started_rx) = mpsc::channel::<()>();
277        let (after_park_tx, after_park_rx) = mpsc::channel::<()>();
278        let observed_post_await_write = Arc::new(Mutex::new(false));
279
280        let task = {
281            let observed_post_await_write = observed_post_await_write.clone();
282            background.spawn_dedicated(move |_executor| async move {
283                // Announce that the future is live on the dedicated thread.
284                started_tx
285                    .send(())
286                    .expect("started signal must be received");
287                // Park forever. Dropping the `Task` must cancel us here so
288                // the code below this `await` never runs.
289                futures::future::pending::<()>().await;
290                *observed_post_await_write.lock() = true;
291                after_park_tx
292                    .send(())
293                    .expect("after-park signal must be received");
294            })
295        };
296
297        // Wait until the dedicated future is actually parked at the await.
298        started_rx
299            .recv_timeout(Duration::from_secs(2))
300            .expect("dedicated future failed to start");
301
302        // Drop the root Task: this must cancel the future.
303        drop(task);
304
305        // If cancellation works, the future never advances past `pending`,
306        // so this recv must time out.
307        assert!(
308            after_park_rx
309                .recv_timeout(Duration::from_millis(100))
310                .is_err(),
311            "dedicated future advanced past the await after its Task was dropped"
312        );
313        assert!(
314            !*observed_post_await_write.lock(),
315            "dedicated future ran code past the cancellation point"
316        );
317    }
318
319    #[test]
320    fn spawn_dedicated_thread_tears_down_after_work_completes() {
321        use std::sync::mpsc;
322
323        // Fires from `Drop` so we observe teardown of the dedicated future's
324        // captured state on whichever thread runs its destructor.
325        struct DropSignal {
326            tx: Option<mpsc::Sender<std::thread::ThreadId>>,
327        }
328        impl Drop for DropSignal {
329            fn drop(&mut self) {
330                if let Some(tx) = self.tx.take() {
331                    let _ = tx.send(std::thread::current().id());
332                }
333            }
334        }
335
336        let background =
337            BackgroundExecutor::new(Arc::new(PlatformScheduler::new(Arc::new(SmokeDispatcher))));
338        let (started_tx, started_rx) = mpsc::channel::<std::thread::ThreadId>();
339        let (drop_tx, drop_rx) = mpsc::channel::<std::thread::ThreadId>();
340
341        let task = background.spawn_dedicated(move |_executor| async move {
342            // Captured by the future's state. When the future completes and
343            // its state is dropped on the dedicated thread, this guard's
344            // `Drop` fires and reports the thread id it ran on.
345            let _guard = DropSignal { tx: Some(drop_tx) };
346            started_tx
347                .send(std::thread::current().id())
348                .expect("started signal must be received");
349            // Future returns immediately. The dedicated thread should then
350            // drop the future (firing _guard), exit the recv loop, and exit.
351        });
352
353        let dedicated_thread_id = started_rx
354            .recv_timeout(Duration::from_secs(2))
355            .expect("dedicated future failed to start");
356        assert_ne!(
357            dedicated_thread_id,
358            std::thread::current().id(),
359            "dedicated future ran on the test thread"
360        );
361
362        // Drive the root task to completion so its body finishes.
363        futures::executor::block_on(task);
364
365        // The guard's drop runs from the dedicated thread as it tears down
366        // the future's captured state. If the executor/recv-loop were
367        // keeping the future alive past task completion, this would hang.
368        let drop_thread_id = drop_rx
369            .recv_timeout(Duration::from_secs(2))
370            .expect("dedicated future's captured state was not dropped after task completion");
371        assert_eq!(
372            drop_thread_id, dedicated_thread_id,
373            "dedicated future's captured state must be dropped on the dedicated thread, not elsewhere"
374        );
375    }
376
377    #[test]
378    fn spawn_dedicated_detached_child_outlives_root() {
379        use std::sync::mpsc;
380
381        let background =
382            BackgroundExecutor::new(Arc::new(PlatformScheduler::new(Arc::new(SmokeDispatcher))));
383
384        // `gate_rx` lets the detached child park until the test explicitly
385        // releases it — after we've already observed the root completing.
386        let (gate_tx, gate_rx) = mpsc::channel::<()>();
387        let (child_done_tx, child_done_rx) = mpsc::channel::<std::thread::ThreadId>();
388
389        let task = background.spawn_dedicated(move |executor| async move {
390            executor
391                .spawn(async move {
392                    // Blocking on `recv` is normally wrong inside an
393                    // executor, but the dedicated thread is exclusive to
394                    // this session, so blocking the only future on it is
395                    // fine — this is the property `spawn_dedicated` is
396                    // designed to provide.
397                    gate_rx
398                        .recv()
399                        .expect("gate sender dropped before child resumed");
400                    child_done_tx
401                        .send(std::thread::current().id())
402                        .expect("child_done receiver dropped");
403                })
404                .detach();
405            // Root finishes here. The detached child must keep the
406            // dedicated thread alive until it completes.
407        });
408
409        futures::executor::block_on(task);
410
411        // Negative assertion: the child has not finished, because the gate
412        // hasn't been released yet.
413        assert!(
414            child_done_rx
415                .recv_timeout(Duration::from_millis(50))
416                .is_err(),
417            "detached child finished before being released"
418        );
419
420        // Release the gate. The detached child should now complete on the
421        // dedicated thread.
422        gate_tx.send(()).expect("gate receiver dropped");
423
424        let child_thread_id = child_done_rx
425            .recv_timeout(Duration::from_secs(2))
426            .expect("detached child failed to complete after gate was released");
427        assert_ne!(
428            child_thread_id,
429            std::thread::current().id(),
430            "detached child ran on the test thread instead of the dedicated thread"
431        );
432    }
433}
434