tenant.openagents/omega

1use std::{
2    collections::BinaryHeap,
3    sync::Arc,
4    thread,
5    time::{Duration, Instant},
6};
7
8use parking_lot::{Condvar, Mutex};
9
10use crate::{
11    PlatformDispatcher, Priority, RunnableVariant, profiler,
12    queue::{PriorityQueueReceiver, PriorityQueueSender},
13};
14
15const MIN_THREADS: usize = 2;
16
17/// A multithreaded [`PlatformDispatcher`] for benchmarks.
18///
19/// Background tasks run in parallel on a pool of worker threads and timers fire
20/// in real time on a dedicated timer thread, mirroring the production
21/// dispatchers (see `LinuxDispatcher`). Main-thread tasks are queued until the
22/// benchmark thread drains them via [`Self::run_until_idle`], since there is no
23/// platform run loop pumping them.
24///
25/// Unlike [`TestDispatcher`](crate::TestDispatcher), which runs everything on a
26/// single thread with a virtual clock, work dispatched through this dispatcher
27/// executes with production concurrency, so wall-clock measurements reflect
28/// real parallelism.
29pub struct BenchDispatcher {
30    background_sender: PriorityQueueSender<RunnableVariant>,
31    main_sender: PriorityQueueSender<RunnableVariant>,
32    main_receiver: Mutex<PriorityQueueReceiver<RunnableVariant>>,
33    timers: Arc<TimerQueue>,
34    idle: Arc<IdleTracker>,
35    main_thread_id: thread::ThreadId,
36}
37
38/// Tracks how many background and timer runnables are queued or running so
39/// [`BenchDispatcher::run_until_idle`] knows when to stop waiting.
40#[derive(Default)]
41struct IdleTracker {
42    inflight: Mutex<usize>,
43    condvar: Condvar,
44}
45
46impl IdleTracker {
47    fn increment(&self) {
48        *self.inflight.lock() += 1;
49    }
50
51    fn decrement(&self) {
52        let mut inflight = self.inflight.lock();
53        *inflight -= 1;
54        if *inflight == 0 {
55            self.condvar.notify_all();
56        }
57    }
58
59    /// Returns a guard that decrements the in-flight count when dropped, so
60    /// the count stays correct even if the runnable being executed panics.
61    fn decrement_on_drop(&self) -> impl Drop + '_ {
62        gpui_util::defer(|| self.decrement())
63    }
64
65    /// Notifies waiters while holding the in-flight lock. `run_until_idle`
66    /// re-checks its wake conditions under this lock before waiting, so the
67    /// notification can't slip between its check and its wait and be lost.
68    fn notify_under_lock(&self) {
69        let _inflight = self.inflight.lock();
70        self.condvar.notify_all();
71    }
72}
73
74struct TimerQueue {
75    state: Mutex<TimerQueueState>,
76    condvar: Condvar,
77}
78
79struct TimerQueueState {
80    heap: BinaryHeap<TimerEntry>,
81    next_seq: u64,
82}
83
84struct TimerEntry {
85    due: Instant,
86    seq: u64,
87    runnable: RunnableVariant,
88}
89
90impl PartialEq for TimerEntry {
91    fn eq(&self, other: &Self) -> bool {
92        self.due == other.due && self.seq == other.seq
93    }
94}
95
96impl Eq for TimerEntry {}
97
98impl PartialOrd for TimerEntry {
99    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
100        Some(self.cmp(other))
101    }
102}
103
104impl Ord for TimerEntry {
105    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
106        // Reversed so that the entry with the earliest due time (breaking ties
107        // by insertion order) is at the top of the max-heap.
108        other
109            .due
110            .cmp(&self.due)
111            .then_with(|| other.seq.cmp(&self.seq))
112    }
113}
114
115impl Default for BenchDispatcher {
116    fn default() -> Self {
117        Self::new()
118    }
119}
120
121impl BenchDispatcher {
122    /// Creates a dispatcher whose main thread is the calling thread.
123    ///
124    /// Worker and timer threads live for the lifetime of the process; the
125    /// dispatcher is expected to be created once and reused across benchmarks.
126    pub fn new() -> Self {
127        let (background_sender, background_receiver) = PriorityQueueReceiver::new();
128        let (main_sender, main_receiver) = PriorityQueueReceiver::new();
129        let idle = Arc::new(IdleTracker::default());
130
131        let thread_count =
132            thread::available_parallelism().map_or(MIN_THREADS, |i| i.get().max(MIN_THREADS));
133        for i in 0..thread_count {
134            let mut receiver: PriorityQueueReceiver<RunnableVariant> = background_receiver.clone();
135            let idle = idle.clone();
136            thread::Builder::new()
137                .name(format!("BenchWorker-{i}"))
138                .spawn(move || {
139                    while let Ok(runnable) = receiver.pop() {
140                        let _decrement = idle.decrement_on_drop();
141                        let location = runnable.metadata().location;
142                        let spawned = runnable.metadata().spawned;
143                        profiler::update_running_task(spawned, location);
144                        runnable.run();
145                        profiler::save_task_timing();
146                    }
147                })
148                .expect("failed to spawn benchmark worker thread");
149        }
150        drop(background_receiver);
151
152        let timers = Arc::new(TimerQueue {
153            state: Mutex::new(TimerQueueState {
154                heap: BinaryHeap::new(),
155                next_seq: 0,
156            }),
157            condvar: Condvar::new(),
158        });
159        {
160            let timers = timers.clone();
161            let idle = idle.clone();
162            thread::Builder::new()
163                .name("BenchTimer".to_owned())
164                .spawn(move || {
165                    let mut state = timers.state.lock();
166                    loop {
167                        let Some(entry) = state.heap.peek() else {
168                            timers.condvar.wait(&mut state);
169                            continue;
170                        };
171                        let due = entry.due;
172                        if due > Instant::now() {
173                            timers.condvar.wait_until(&mut state, due);
174                            continue;
175                        }
176                        let Some(entry) = state.heap.pop() else {
177                            continue;
178                        };
179                        // Count the firing timer as in-flight before releasing
180                        // the lock so it can spawn follow-up work that
181                        // `run_until_idle` will wait for. Lock order is always
182                        // timer state, then in-flight count; `run_until_idle`
183                        // never takes them in the opposite order.
184                        idle.increment();
185                        drop(state);
186
187                        {
188                            let _decrement = idle.decrement_on_drop();
189                            let location = entry.runnable.metadata().location;
190                            let spawned = entry.runnable.metadata().spawned;
191                            profiler::update_running_task(spawned, location);
192                            entry.runnable.run();
193                            profiler::save_task_timing();
194                        }
195
196                        state = timers.state.lock();
197                    }
198                })
199                .expect("failed to spawn benchmark timer thread");
200        }
201
202        Self {
203            background_sender,
204            main_sender,
205            main_receiver: Mutex::new(main_receiver),
206            timers,
207            idle,
208            main_thread_id: thread::current().id(),
209        }
210    }
211
212    /// Runs queued main thread tasks and waits until no background or timer
213    /// work is queued, running, or already due.
214    ///
215    /// Timers that haven't reached their due time yet are *not* waited for:
216    /// the dispatcher runs in real time and cannot skip ahead like the
217    /// `TestDispatcher`'s virtual clock, so waiting on a future timer would
218    /// block for its full real duration. Tasks sleeping on such timers are
219    /// considered idle. Must be called on the thread that created this
220    /// dispatcher.
221    pub fn run_until_idle(&self) {
222        assert!(
223            self.is_main_thread(),
224            "run_until_idle must be called on the benchmark main thread"
225        );
226        loop {
227            if self.drain_main_queue() {
228                continue;
229            }
230
231            // Checked before taking the in-flight lock; the timer thread
232            // locks them in the opposite order, so nesting would deadlock.
233            if self.has_due_timer() {
234                // Poll briefly: a firing timer leaves the heap just before it
235                // registers as in-flight.
236                let mut inflight = self.idle.inflight.lock();
237                self.idle
238                    .condvar
239                    .wait_for(&mut inflight, Duration::from_millis(1));
240                continue;
241            }
242
243            let mut inflight = self.idle.inflight.lock();
244            // Re-checked under the lock that `dispatch_on_main_thread`
245            // notifies under, so the notification can't be lost.
246            if self.main_queue_has_work() {
247                continue;
248            }
249            if *inflight == 0 {
250                // Main-thread sends happen before in-flight decrements, and
251                // decrements happen under this lock, so the check above
252                // observed all completed work.
253                return;
254            }
255            // Woken when main-thread work arrives or the in-flight count
256            // reaches zero; both notify under this lock.
257            self.idle.condvar.wait(&mut inflight);
258        }
259    }
260
261    /// Runs all main-thread tasks that are queued right now, without waiting for
262    /// background work or timers to finish.
263    pub fn run_ready_main_tasks(&self) -> bool {
264        assert!(
265            self.is_main_thread(),
266            "run_ready_main_tasks must be called on the benchmark main thread"
267        );
268        self.drain_main_queue()
269    }
270
271    /// Cancels all pending timers so timers armed by one benchmark can't fire
272    /// during a later benchmark sharing this process-lifetime dispatcher.
273    ///
274    /// Dropping a timer runnable drops its completion sender, waking the task
275    /// awaiting the timer. Call [`Self::run_until_idle`] after this method to
276    /// drain any work that cancellation unblocks.
277    pub fn cancel_pending_timers(&self) -> usize {
278        let timers = {
279            let mut state = self.timers.state.lock();
280            let timers: Vec<_> = state.heap.drain().collect();
281            self.timers.condvar.notify_all();
282            timers
283        };
284        let canceled = timers.len();
285        drop(timers);
286        canceled
287    }
288
289    /// Describes the dispatcher's idle-tracking state, for diagnosing
290    /// benchmarks that fail to reach quiescence.
291    pub fn debug_state(&self) -> String {
292        let inflight = *self.idle.inflight.lock();
293        let timers = self.timers.state.lock().heap.len();
294        let main_queue_has_work = self.main_queue_has_work();
295        format!(
296            "BenchDispatcher {{ inflight: {inflight}, pending_timers: {timers}, \
297             main_queue_has_work: {main_queue_has_work} }}"
298        )
299    }
300
301    fn has_due_timer(&self) -> bool {
302        let state = self.timers.state.lock();
303        state
304            .heap
305            .peek()
306            .is_some_and(|entry| entry.due <= Instant::now())
307    }
308
309    fn main_queue_has_work(&self) -> bool {
310        !self.main_receiver.lock().is_empty()
311    }
312
313    fn drain_main_queue(&self) -> bool {
314        let mut ran_any = false;
315        loop {
316            // Lock only around the pop so runnables can re-entrantly dispatch
317            // more main-thread work through the sender while they run.
318            let runnable = self.main_receiver.lock().try_pop();
319            match runnable {
320                Ok(Some(runnable)) => {
321                    let location = runnable.metadata().location;
322                    let spawned = runnable.metadata().spawned;
323                    profiler::update_running_task(spawned, location);
324                    runnable.run();
325                    profiler::save_task_timing();
326                    ran_any = true;
327                }
328                Ok(None) | Err(_) => return ran_any,
329            }
330        }
331    }
332}
333
334impl PlatformDispatcher for BenchDispatcher {
335    fn is_main_thread(&self) -> bool {
336        thread::current().id() == self.main_thread_id
337    }
338
339    fn dispatch(&self, runnable: RunnableVariant, priority: Priority) {
340        self.idle.increment();
341        self.background_sender
342            .send(priority, runnable)
343            .unwrap_or_else(|_| panic!("benchmark worker threads are no longer running"));
344    }
345
346    fn dispatch_on_main_thread(&self, runnable: RunnableVariant, priority: Priority) {
347        if let Err(error) = self.main_sender.send(priority, runnable) {
348            // The main receiver lives as long as this dispatcher, so a failed
349            // send means we're mid-teardown. The runnable may wrap a !Send
350            // future, so forget it rather than dropping it on this thread
351            // (mirrors LinuxDispatcher).
352            std::mem::forget(error);
353            return;
354        }
355        // Wake `run_until_idle` if it's waiting for main-thread work.
356        self.idle.notify_under_lock();
357    }
358
359    fn dispatch_after(&self, duration: Duration, runnable: RunnableVariant) {
360        let mut state = self.timers.state.lock();
361        let seq = state.next_seq;
362        state.next_seq += 1;
363        state.heap.push(TimerEntry {
364            due: Instant::now() + duration,
365            seq,
366            runnable,
367        });
368        self.timers.condvar.notify_one();
369    }
370
371    fn spawn_realtime(&self, f: Box<dyn FnOnce() + Send>) {
372        // Benchmarks don't need realtime scheduling priority; a plain thread
373        // keeps this portable.
374        thread::Builder::new()
375            .name("BenchRealtime".to_owned())
376            .spawn(f)
377            .expect("failed to spawn benchmark realtime thread");
378    }
379
380    fn as_bench(&self) -> Option<&BenchDispatcher> {
381        Some(self)
382    }
383}
384
385#[cfg(test)]
386mod tests {
387    use std::sync::atomic::{AtomicBool, Ordering};
388
389    use super::*;
390    use crate::{BackgroundExecutor, ForegroundExecutor};
391
392    #[test]
393    fn run_ready_main_tasks_does_not_wait_for_background_handoffs() {
394        let dispatcher = Arc::new(BenchDispatcher::new());
395        let background = BackgroundExecutor::new(dispatcher.clone());
396        let foreground = ForegroundExecutor::new(dispatcher.clone());
397
398        let (sender, receiver) = futures::channel::oneshot::channel();
399        background
400            .spawn(async move {
401                thread::sleep(Duration::from_millis(10));
402                sender.send(()).ok();
403            })
404            .detach();
405
406        let completed = Arc::new(AtomicBool::new(false));
407        foreground
408            .spawn({
409                let completed = completed.clone();
410                async move {
411                    receiver.await.ok();
412                    completed.store(true, Ordering::SeqCst);
413                }
414            })
415            .detach();
416
417        assert!(dispatcher.run_ready_main_tasks());
418        assert!(!completed.load(Ordering::SeqCst));
419
420        dispatcher.run_until_idle();
421        assert!(completed.load(Ordering::SeqCst));
422    }
423
424    #[test]
425    fn run_until_idle_completes_background_to_main_handoffs() {
426        let dispatcher = Arc::new(BenchDispatcher::new());
427        let background = BackgroundExecutor::new(dispatcher.clone());
428        let foreground = ForegroundExecutor::new(dispatcher.clone());
429
430        let (sender, receiver) = futures::channel::oneshot::channel();
431        background
432            .spawn(async move {
433                thread::sleep(Duration::from_millis(10));
434                sender.send(()).ok();
435            })
436            .detach();
437
438        let completed = Arc::new(AtomicBool::new(false));
439        foreground
440            .spawn({
441                let completed = completed.clone();
442                async move {
443                    receiver.await.ok();
444                    completed.store(true, Ordering::SeqCst);
445                }
446            })
447            .detach();
448
449        dispatcher.run_until_idle();
450        assert!(completed.load(Ordering::SeqCst));
451    }
452
453    #[test]
454    fn timers_fire_in_real_time() {
455        let dispatcher = Arc::new(BenchDispatcher::new());
456        let background = BackgroundExecutor::new(dispatcher);
457
458        let fired = Arc::new(AtomicBool::new(false));
459        let timer = background.timer(Duration::from_millis(10));
460        background
461            .spawn({
462                let fired = fired.clone();
463                async move {
464                    timer.await;
465                    fired.store(true, Ordering::SeqCst);
466                }
467            })
468            .detach();
469
470        let deadline = Instant::now() + Duration::from_secs(10);
471        while !fired.load(Ordering::SeqCst) && Instant::now() < deadline {
472            thread::sleep(Duration::from_millis(1));
473        }
474        assert!(fired.load(Ordering::SeqCst));
475    }
476
477    #[test]
478    fn cancel_pending_timers_wakes_waiters_without_waiting_for_deadline() {
479        let dispatcher = Arc::new(BenchDispatcher::new());
480        let background = BackgroundExecutor::new(dispatcher.clone());
481
482        let fired = Arc::new(AtomicBool::new(false));
483        let timer = background.timer(Duration::from_secs(10));
484        background
485            .spawn({
486                let fired = fired.clone();
487                async move {
488                    timer.await;
489                    fired.store(true, Ordering::SeqCst);
490                }
491            })
492            .detach();
493
494        dispatcher.run_until_idle();
495        assert_eq!(dispatcher.cancel_pending_timers(), 1);
496        dispatcher.run_until_idle();
497
498        assert!(fired.load(Ordering::SeqCst));
499        assert_eq!(dispatcher.cancel_pending_timers(), 0);
500    }
501}
502