tenant.openagents/omega

1use crate::{
2    BackgroundExecutor, Clock, Instant, LocalExecutor, Priority, RunnableMeta, Scheduler,
3    SessionId, Task, TestClock, Timer,
4};
5use async_task::Runnable;
6use backtrace::{Backtrace, BacktraceFrame};
7use futures::channel::oneshot;
8use parking_lot::{Mutex, MutexGuard};
9use rand::{
10    distr::{StandardUniform, uniform::SampleRange, uniform::SampleUniform},
11    prelude::*,
12};
13use std::any::Any;
14use std::{
15    any::type_name_of_val,
16    collections::{BTreeMap, HashSet, VecDeque},
17    env,
18    fmt::Write,
19    future::Future,
20    mem,
21    ops::RangeInclusive,
22    panic::{self, AssertUnwindSafe},
23    pin::Pin,
24    sync::{
25        Arc, Weak,
26        atomic::{AtomicBool, Ordering::SeqCst},
27    },
28    task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
29    thread::{self, Thread},
30    time::Duration,
31};
32
33const PENDING_TRACES_VAR_NAME: &str = "PENDING_TRACES";
34
35pub struct TestScheduler {
36    clock: Arc<TestClock>,
37    rng: Arc<Mutex<StdRng>>,
38    state: Arc<Mutex<SchedulerState>>,
39    thread: Thread,
40}
41
42impl TestScheduler {
43    /// Run a test once with default configuration (seed 0)
44    pub fn once<R>(f: impl AsyncFnOnce(Arc<TestScheduler>) -> R) -> R {
45        Self::with_seed(0, f)
46    }
47
48    /// Run a test multiple times with sequential seeds (0, 1, 2, ...)
49    pub fn many<R>(
50        default_iterations: usize,
51        mut f: impl AsyncFnMut(Arc<TestScheduler>) -> R,
52    ) -> Vec<R> {
53        let num_iterations = std::env::var("ITERATIONS")
54            .map(|iterations| iterations.parse().unwrap())
55            .unwrap_or(default_iterations);
56
57        let seed = std::env::var("SEED")
58            .map(|seed| seed.parse().unwrap())
59            .unwrap_or(0);
60
61        let interactive = !std::env::var("SCHEDULER_NONINTERACTIVE").is_ok();
62
63        (seed..seed + num_iterations as u64)
64            .map(|seed| {
65                let mut unwind_safe_f = AssertUnwindSafe(&mut f);
66                if interactive {
67                    eprintln!("Running seed: {seed}");
68                }
69                match panic::catch_unwind(move || Self::with_seed(seed, &mut *unwind_safe_f)) {
70                    Ok(result) => result,
71                    Err(error) => {
72                        eprintln!("\x1b[31mFailing Seed: {seed}\x1b[0m");
73                        panic::resume_unwind(error);
74                    }
75                }
76            })
77            .collect()
78    }
79
80    fn with_seed<R>(seed: u64, f: impl AsyncFnOnce(Arc<TestScheduler>) -> R) -> R {
81        let scheduler = Arc::new(TestScheduler::new(TestSchedulerConfig::with_seed(seed)));
82        let future = f(scheduler.clone());
83        let result = scheduler.foreground().block_on(future);
84        scheduler.run(); // Ensure spawned tasks finish up before returning in tests
85        result
86    }
87
88    pub fn new(config: TestSchedulerConfig) -> Self {
89        Self {
90            rng: Arc::new(Mutex::new(StdRng::seed_from_u64(config.seed))),
91            state: Arc::new(Mutex::new(SchedulerState {
92                runnables: VecDeque::new(),
93                timers: Vec::new(),
94                blocked_sessions: Vec::new(),
95                randomize_order: config.randomize_order,
96                allow_parking: config.allow_parking,
97                timeout_ticks: config.timeout_ticks,
98                next_session_id: SessionId(0),
99                capture_pending_traces: config.capture_pending_traces,
100                pending_traces: BTreeMap::new(),
101                next_trace_id: TraceId(0),
102                is_main_thread: true,
103                non_determinism_error: None,
104                finished: false,
105                parking_allowed_once: false,
106                unparked: false,
107            })),
108            clock: Arc::new(TestClock::new()),
109            thread: thread::current(),
110        }
111    }
112
113    pub fn end_test(&self) {
114        let mut state = self.state.lock();
115        if let Some((message, backtrace)) = &state.non_determinism_error {
116            if cfg!(miri) {
117                // miri cannot debug print backtraces with `miri-disable-isolation` enabled
118                panic!("{}", message)
119            } else {
120                panic!("{}\n{:?}", message, backtrace)
121            }
122        }
123        state.finished = true;
124    }
125
126    pub fn clock(&self) -> Arc<TestClock> {
127        self.clock.clone()
128    }
129
130    pub fn rng(&self) -> SharedRng {
131        SharedRng(self.rng.clone())
132    }
133
134    pub fn set_timeout_ticks(&self, timeout_ticks: RangeInclusive<usize>) {
135        self.state.lock().timeout_ticks = timeout_ticks;
136    }
137
138    pub fn allow_parking(&self) {
139        let mut state = self.state.lock();
140        state.allow_parking = true;
141        state.parking_allowed_once = true;
142    }
143
144    pub fn forbid_parking(&self) {
145        self.state.lock().allow_parking = false;
146    }
147
148    pub fn parking_allowed(&self) -> bool {
149        self.state.lock().allow_parking
150    }
151
152    pub fn is_main_thread(&self) -> bool {
153        self.state.lock().is_main_thread
154    }
155
156    pub fn allocate_session_id(&self) -> SessionId {
157        let mut state = self.state.lock();
158        state.next_session_id.0 += 1;
159        state.next_session_id
160    }
161
162    /// Create a local executor for this scheduler.
163    pub fn foreground(self: &Arc<Self>) -> LocalExecutor {
164        let session_id = self.allocate_session_id();
165        LocalExecutor::new(
166            session_id,
167            self.clone(),
168            schedule_local_dispatch(Arc::downgrade(self), session_id),
169        )
170    }
171
172    /// Create a background executor for this scheduler
173    pub fn background(self: &Arc<Self>) -> BackgroundExecutor {
174        BackgroundExecutor::new(self.clone())
175    }
176
177    pub fn yield_random(&self) -> Yield {
178        let rng = &mut *self.rng.lock();
179        if rng.random_bool(0.1) {
180            Yield(rng.random_range(10..20))
181        } else {
182            Yield(rng.random_range(0..2))
183        }
184    }
185
186    pub fn run(&self) {
187        while self.step() {
188            // Continue until no work remains
189        }
190    }
191
192    pub fn run_with_clock_advancement(&self) {
193        while self.step() || self.advance_clock_to_next_timer() {
194            // Continue until no work remains
195        }
196    }
197
198    /// Execute one tick of the scheduler, processing expired timers and running
199    /// at most one task. Returns true if any work was done.
200    ///
201    /// This is the public interface for GPUI's TestDispatcher to drive task execution.
202    pub fn tick(&self) -> bool {
203        self.step_filtered(false)
204    }
205
206    /// Execute one tick, but only run background tasks (no foreground/session tasks).
207    /// Returns true if any work was done.
208    pub fn tick_background_only(&self) -> bool {
209        self.step_filtered(true)
210    }
211
212    /// Check if there are any pending tasks or timers that could run.
213    pub fn has_pending_tasks(&self) -> bool {
214        let state = self.state.lock();
215        !state.runnables.is_empty() || !state.timers.is_empty()
216    }
217
218    /// Returns counts of (foreground_tasks, background_tasks) currently queued.
219    /// Foreground tasks are those with a session_id, background tasks have none.
220    pub fn pending_task_counts(&self) -> (usize, usize) {
221        let state = self.state.lock();
222        let foreground = state
223            .runnables
224            .iter()
225            .filter(|r| r.session_id.is_some())
226            .count();
227        let background = state
228            .runnables
229            .iter()
230            .filter(|r| r.session_id.is_none())
231            .count();
232        (foreground, background)
233    }
234
235    fn step(&self) -> bool {
236        self.step_filtered(false)
237    }
238
239    fn step_filtered(&self, background_only: bool) -> bool {
240        let (elapsed_count, runnables_before) = {
241            let mut state = self.state.lock();
242            let end_ix = state
243                .timers
244                .partition_point(|timer| timer.expiration <= self.clock.now());
245            let elapsed: Vec<_> = state.timers.drain(..end_ix).collect();
246            let count = elapsed.len();
247            let runnables = state.runnables.len();
248            drop(state);
249            // Dropping elapsed timers here wakes the waiting futures
250            drop(elapsed);
251            (count, runnables)
252        };
253
254        if elapsed_count > 0 {
255            let runnables_after = self.state.lock().runnables.len();
256            if std::env::var("DEBUG_SCHEDULER").is_ok() {
257                eprintln!(
258                    "[scheduler] Expired {} timers at {:?}, runnables: {} -> {}",
259                    elapsed_count,
260                    self.clock.now(),
261                    runnables_before,
262                    runnables_after
263                );
264            }
265            return true;
266        }
267
268        let runnable = {
269            let state = &mut *self.state.lock();
270
271            // Find candidate tasks:
272            // - For foreground tasks (with session_id), only the first task from each session
273            //   is a candidate (to preserve intra-session ordering)
274            // - For background tasks (no session_id), all are candidates
275            // - Tasks from blocked sessions are excluded
276            // - If background_only is true, skip foreground tasks entirely
277            let mut seen_sessions = HashSet::new();
278            let candidate_indices: Vec<usize> = state
279                .runnables
280                .iter()
281                .enumerate()
282                .filter(|(_, runnable)| {
283                    if let Some(session_id) = runnable.session_id {
284                        // Skip foreground tasks if background_only mode
285                        if background_only {
286                            return false;
287                        }
288                        // Exclude tasks from blocked sessions
289                        if state.blocked_sessions.contains(&session_id) {
290                            return false;
291                        }
292                        // Only include first task from each session (insert returns true if new)
293                        seen_sessions.insert(session_id)
294                    } else {
295                        // Background tasks are always candidates
296                        true
297                    }
298                })
299                .map(|(ix, _)| ix)
300                .collect();
301
302            if candidate_indices.is_empty() {
303                None
304            } else if state.randomize_order {
305                // Use priority-weighted random selection
306                let weights: Vec<u32> = candidate_indices
307                    .iter()
308                    .map(|&ix| state.runnables[ix].priority.weight())
309                    .collect();
310                let total_weight: u32 = weights.iter().sum();
311
312                if total_weight == 0 {
313                    // Fallback to uniform random if all weights are zero
314                    let choice = self.rng.lock().random_range(0..candidate_indices.len());
315                    state.runnables.remove(candidate_indices[choice])
316                } else {
317                    let mut target = self.rng.lock().random_range(0..total_weight);
318                    let mut selected_idx = 0;
319                    for (i, &weight) in weights.iter().enumerate() {
320                        if target < weight {
321                            selected_idx = i;
322                            break;
323                        }
324                        target -= weight;
325                    }
326                    state.runnables.remove(candidate_indices[selected_idx])
327                }
328            } else {
329                // Non-randomized: just take the first candidate task
330                state.runnables.remove(candidate_indices[0])
331            }
332        };
333
334        if let Some(runnable) = runnable {
335            let is_foreground = runnable.session_id.is_some();
336            let was_main_thread = self.state.lock().is_main_thread;
337            self.state.lock().is_main_thread = is_foreground;
338            runnable.run();
339            self.state.lock().is_main_thread = was_main_thread;
340            return true;
341        }
342
343        false
344    }
345
346    /// Drops all runnable tasks from the scheduler.
347    ///
348    /// This is used by the leak detector to ensure that all tasks have been dropped as tasks may keep entities alive otherwise.
349    /// Why do we even have tasks left when tests finish you may ask. The reason for that is simple, the scheduler itself is the executor and it retains the scheduled runnables.
350    /// A lot of tasks, including every foreground task contain an executor handle that keeps the test scheduler alive, causing a reference cycle, thus the need for this function right now.
351    pub fn drain_tasks(&self) {
352        // dropping runnables may reschedule tasks
353        // due to drop impls with executors in them
354        // so drop until we reach a fixpoint
355        loop {
356            let mut state = self.state.lock();
357            if state.runnables.is_empty() && state.timers.is_empty() {
358                break;
359            }
360            let runnables = std::mem::take(&mut state.runnables);
361            let timers = std::mem::take(&mut state.timers);
362            drop(state);
363            drop(timers);
364            drop(runnables);
365        }
366    }
367
368    pub fn advance_clock_to_next_timer(&self) -> bool {
369        if let Some(timer) = self.state.lock().timers.first() {
370            self.clock.advance(timer.expiration - self.clock.now());
371            true
372        } else {
373            false
374        }
375    }
376
377    pub fn advance_clock(&self, duration: Duration) {
378        let debug = std::env::var("DEBUG_SCHEDULER").is_ok();
379        let start = self.clock.now();
380        let next_now = start + duration;
381        if debug {
382            let timer_count = self.state.lock().timers.len();
383            eprintln!(
384                "[scheduler] advance_clock({:?}) from {:?}, {} pending timers",
385                duration, start, timer_count
386            );
387        }
388        loop {
389            self.run();
390            if let Some(timer) = self.state.lock().timers.first()
391                && timer.expiration <= next_now
392            {
393                let advance_to = timer.expiration;
394                if debug {
395                    eprintln!(
396                        "[scheduler] Advancing clock {:?} -> {:?} for timer",
397                        self.clock.now(),
398                        advance_to
399                    );
400                }
401                self.clock.advance(advance_to - self.clock.now());
402            } else {
403                break;
404            }
405        }
406        self.clock.advance(next_now - self.clock.now());
407        if debug {
408            eprintln!(
409                "[scheduler] advance_clock done, now at {:?}",
410                self.clock.now()
411            );
412        }
413    }
414
415    fn park(&self, deadline: Option<Instant>) -> bool {
416        if self.state.lock().allow_parking {
417            let start = Instant::now();
418            // Enforce a hard timeout to prevent tests from hanging indefinitely
419            let hard_deadline = start + Duration::from_secs(15);
420
421            // Use the earlier of the provided deadline or the hard timeout deadline
422            let effective_deadline = deadline
423                .map(|d| d.min(hard_deadline))
424                .unwrap_or(hard_deadline);
425
426            // Park in small intervals to allow checking both deadlines
427            const PARK_INTERVAL: Duration = Duration::from_millis(100);
428            loop {
429                let now = Instant::now();
430                if now >= effective_deadline {
431                    // Check if we hit the hard timeout
432                    if now >= hard_deadline {
433                        panic!(
434                            "Test timed out after 15 seconds while parking. \
435                            This may indicate a deadlock or missing waker.",
436                        );
437                    }
438                    // Hit the provided deadline
439                    return false;
440                }
441
442                let remaining = effective_deadline.saturating_duration_since(now);
443                let park_duration = remaining.min(PARK_INTERVAL);
444                let before_park = Instant::now();
445                thread::park_timeout(park_duration);
446                let elapsed = before_park.elapsed();
447
448                // Advance the test clock by the real elapsed time while parking
449                self.clock.advance(elapsed);
450
451                // Check if any timers have expired after advancing the clock.
452                // If so, return so the caller can process them.
453                if self
454                    .state
455                    .lock()
456                    .timers
457                    .first()
458                    .map_or(false, |t| t.expiration <= self.clock.now())
459                {
460                    return true;
461                }
462
463                // Check if we were woken up by a different thread.
464                // We use a flag because timing-based detection is unreliable:
465                // OS scheduling delays can cause elapsed >= park_duration even when
466                // we were woken early by unpark().
467                if std::mem::take(&mut self.state.lock().unparked) {
468                    return true;
469                }
470            }
471        } else if deadline.is_some() {
472            false
473        } else if cfg!(miri) {
474            // miri cannot debug print backtraces with `miri-disable-isolation` enabled
475            panic!("Parking forbidden.");
476        } else if self.state.lock().capture_pending_traces {
477            let mut pending_traces = String::new();
478            for (_, trace) in mem::take(&mut self.state.lock().pending_traces) {
479                writeln!(pending_traces, "{:?}", exclude_wakers_from_trace(trace)).unwrap();
480            }
481            panic!("Parking forbidden. Pending traces:\n{}", pending_traces);
482        } else {
483            panic!(
484                "Parking forbidden. Re-run with {PENDING_TRACES_VAR_NAME}=1 to show pending traces"
485            );
486        }
487    }
488}
489
490fn assert_correct_thread(expected: &Thread, state: &Arc<Mutex<SchedulerState>>) {
491    let current_thread = thread::current();
492    let mut state = state.lock();
493    if state.parking_allowed_once {
494        return;
495    }
496    if current_thread.id() == expected.id() {
497        return;
498    }
499
500    let message = format!(
501        "Detected activity on thread {:?} {:?}, but test scheduler is running on {:?} {:?}. Your test is not deterministic.",
502        current_thread.name(),
503        current_thread.id(),
504        expected.name(),
505        expected.id(),
506    );
507    let backtrace = Backtrace::new();
508    if state.finished {
509        panic!("{}", message);
510    } else {
511        state.non_determinism_error = Some((message, backtrace))
512    }
513}
514
515impl Scheduler for TestScheduler {
516    /// Block until the given future completes, with an optional timeout. If the
517    /// future is unable to make progress at any moment before the timeout and
518    /// no other tasks or timers remain, we panic unless parking is allowed. If
519    /// parking is allowed, we block up to the timeout or indefinitely if none
520    /// is provided. This is to allow testing a mix of deterministic and
521    /// non-deterministic async behavior, such as when interacting with I/O in
522    /// an otherwise deterministic test.
523    fn block(
524        &self,
525        session_id: Option<SessionId>,
526        mut future: Pin<&mut dyn Future<Output = ()>>,
527        timeout: Option<Duration>,
528    ) -> bool {
529        if let Some(session_id) = session_id {
530            self.state.lock().blocked_sessions.push(session_id);
531        }
532
533        let deadline = timeout.map(|timeout| Instant::now() + timeout);
534        let awoken = Arc::new(AtomicBool::new(false));
535        let waker = Box::new(TracingWaker {
536            id: None,
537            awoken: awoken.clone(),
538            thread: self.thread.clone(),
539            state: self.state.clone(),
540        });
541        let waker = unsafe { Waker::new(Box::into_raw(waker) as *const (), &WAKER_VTABLE) };
542        let max_ticks = if timeout.is_some() {
543            self.rng
544                .lock()
545                .random_range(self.state.lock().timeout_ticks.clone())
546        } else {
547            usize::MAX
548        };
549        let mut cx = Context::from_waker(&waker);
550
551        let mut completed = false;
552        for _ in 0..max_ticks {
553            match future.as_mut().poll(&mut cx) {
554                Poll::Ready(()) => {
555                    completed = true;
556                    break;
557                }
558                Poll::Pending => {}
559            }
560
561            let mut stepped = None;
562            while self.rng.lock().random() {
563                let stepped = stepped.get_or_insert(false);
564                if self.step() {
565                    *stepped = true;
566                } else {
567                    break;
568                }
569            }
570
571            let stepped = stepped.unwrap_or(true);
572            let awoken = awoken.swap(false, SeqCst);
573            if !stepped && !awoken {
574                let parking_allowed = self.state.lock().allow_parking;
575                // In deterministic mode (parking forbidden), instantly jump to the next timer.
576                // In non-deterministic mode (parking allowed), let real time pass instead.
577                let advanced_to_timer = !parking_allowed && self.advance_clock_to_next_timer();
578                if !advanced_to_timer && !self.park(deadline) {
579                    break;
580                }
581            }
582        }
583
584        if session_id.is_some() {
585            self.state.lock().blocked_sessions.pop();
586        }
587
588        completed
589    }
590
591    fn schedule_local(&self, session_id: SessionId, runnable: Runnable<RunnableMeta>) {
592        assert_correct_thread(&self.thread, &self.state);
593        let mut state = self.state.lock();
594        let ix = if state.randomize_order {
595            let start_ix = state
596                .runnables
597                .iter()
598                .rposition(|task| task.session_id == Some(session_id))
599                .map_or(0, |ix| ix + 1);
600            self.rng
601                .lock()
602                .random_range(start_ix..=state.runnables.len())
603        } else {
604            state.runnables.len()
605        };
606        state.runnables.insert(
607            ix,
608            ScheduledRunnable {
609                session_id: Some(session_id),
610                priority: Priority::default(),
611                runnable,
612            },
613        );
614        state.unparked = true;
615        drop(state);
616        self.thread.unpark();
617    }
618
619    fn schedule_background_with_priority(
620        &self,
621        runnable: Runnable<RunnableMeta>,
622        priority: Priority,
623    ) {
624        assert_correct_thread(&self.thread, &self.state);
625        let mut state = self.state.lock();
626        let ix = if state.randomize_order {
627            self.rng.lock().random_range(0..=state.runnables.len())
628        } else {
629            state.runnables.len()
630        };
631        state.runnables.insert(
632            ix,
633            ScheduledRunnable {
634                session_id: None,
635                priority,
636                runnable,
637            },
638        );
639        state.unparked = true;
640        drop(state);
641        self.thread.unpark();
642    }
643
644    fn spawn_realtime(&self, f: Box<dyn FnOnce() + Send>) {
645        std::thread::spawn(move || {
646            f();
647        });
648    }
649
650    #[track_caller]
651    fn timer(&self, duration: Duration) -> Timer {
652        let (tx, rx) = oneshot::channel();
653        let state = &mut *self.state.lock();
654        state.timers.push(ScheduledTimer {
655            expiration: self.clock.now() + duration,
656            _notify: tx,
657        });
658        state.timers.sort_by_key(|timer| timer.expiration);
659        Timer(rx)
660    }
661
662    fn clock(&self) -> Arc<dyn Clock> {
663        self.clock.clone()
664    }
665
666    /// In the test world, dedicated work is just a fresh local session driven
667    /// by the test scheduler's run loop alongside everything else. No real
668    /// thread is spawned, so determinism under `TestScheduler::many` is
669    /// preserved.
670    fn spawn_dedicated(
671        self: Arc<Self>,
672        f: Box<
673            dyn FnOnce(
674                    LocalExecutor,
675                )
676                    -> Pin<Box<dyn Future<Output = Box<dyn Any + Send + Sync>> + 'static>>
677                + Send
678                + 'static,
679        >,
680    ) -> Task<Box<dyn Any + Send + Sync>> {
681        let session_id = self.allocate_session_id();
682        let weak_scheduler = Arc::downgrade(&self);
683        let spawner = LocalExecutor::new(
684            session_id,
685            self,
686            schedule_local_dispatch(weak_scheduler.clone(), session_id),
687        );
688        // The dedicated future sits in this scheduler's own queue until its
689        // first poll. A `LocalExecutor` holds the scheduler strongly, so
690        // capturing one here would create a reference cycle
691        // (scheduler -> queued runnable -> future -> executor -> scheduler)
692        // that leaks both if the scheduler is dropped before ever being run.
693        // Construct the executor lazily at first poll, when the runnable is
694        // already out of the queue.
695        spawner.spawn(async move {
696            let scheduler = weak_scheduler
697                .upgrade()
698                .expect("dedicated tasks are only polled by their scheduler, which is still alive when polled");
699            let executor = LocalExecutor::new(
700                session_id,
701                scheduler,
702                schedule_local_dispatch(weak_scheduler, session_id),
703            );
704            f(executor).await
705        })
706    }
707
708    fn as_test(&self) -> Option<&TestScheduler> {
709        Some(self)
710    }
711}
712
713/// Dispatch closure for `LocalExecutor`s backed by a `TestScheduler`. Holds
714/// the scheduler weakly so queued runnables don't keep it alive.
715fn schedule_local_dispatch(
716    scheduler: Weak<TestScheduler>,
717    session_id: SessionId,
718) -> impl Fn(Runnable<RunnableMeta>) + Send + Sync + 'static {
719    move |runnable| {
720        if let Some(scheduler) = scheduler.upgrade() {
721            scheduler.schedule_local(session_id, runnable);
722        }
723    }
724}
725
726#[derive(Clone, Debug)]
727pub struct TestSchedulerConfig {
728    pub seed: u64,
729    pub randomize_order: bool,
730    pub allow_parking: bool,
731    pub capture_pending_traces: bool,
732    pub timeout_ticks: RangeInclusive<usize>,
733}
734
735impl TestSchedulerConfig {
736    pub fn with_seed(seed: u64) -> Self {
737        Self {
738            seed,
739            ..Default::default()
740        }
741    }
742}
743
744impl Default for TestSchedulerConfig {
745    fn default() -> Self {
746        Self {
747            seed: 0,
748            randomize_order: true,
749            allow_parking: false,
750            capture_pending_traces: env::var(PENDING_TRACES_VAR_NAME)
751                .map_or(false, |var| var == "1" || var == "true"),
752            timeout_ticks: 1..=1000,
753        }
754    }
755}
756
757struct ScheduledRunnable {
758    session_id: Option<SessionId>,
759    priority: Priority,
760    runnable: Runnable<RunnableMeta>,
761}
762
763impl ScheduledRunnable {
764    fn run(self) {
765        self.runnable.run();
766    }
767}
768
769struct ScheduledTimer {
770    expiration: Instant,
771    _notify: oneshot::Sender<()>,
772}
773
774struct SchedulerState {
775    runnables: VecDeque<ScheduledRunnable>,
776    timers: Vec<ScheduledTimer>,
777    blocked_sessions: Vec<SessionId>,
778    randomize_order: bool,
779    allow_parking: bool,
780    timeout_ticks: RangeInclusive<usize>,
781    next_session_id: SessionId,
782    capture_pending_traces: bool,
783    next_trace_id: TraceId,
784    pending_traces: BTreeMap<TraceId, Backtrace>,
785    is_main_thread: bool,
786    non_determinism_error: Option<(String, Backtrace)>,
787    parking_allowed_once: bool,
788    finished: bool,
789    unparked: bool,
790}
791
792const WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
793    TracingWaker::clone_raw,
794    TracingWaker::wake_raw,
795    TracingWaker::wake_by_ref_raw,
796    TracingWaker::drop_raw,
797);
798
799#[derive(Copy, Clone, Eq, PartialEq, PartialOrd, Ord)]
800struct TraceId(usize);
801
802struct TracingWaker {
803    id: Option<TraceId>,
804    awoken: Arc<AtomicBool>,
805    thread: Thread,
806    state: Arc<Mutex<SchedulerState>>,
807}
808
809impl Clone for TracingWaker {
810    fn clone(&self) -> Self {
811        let mut state = self.state.lock();
812        let id = if state.capture_pending_traces {
813            let id = state.next_trace_id;
814            state.next_trace_id.0 += 1;
815            state.pending_traces.insert(id, Backtrace::new_unresolved());
816            Some(id)
817        } else {
818            None
819        };
820        Self {
821            id,
822            awoken: self.awoken.clone(),
823            thread: self.thread.clone(),
824            state: self.state.clone(),
825        }
826    }
827}
828
829impl Drop for TracingWaker {
830    fn drop(&mut self) {
831        assert_correct_thread(&self.thread, &self.state);
832
833        if let Some(id) = self.id {
834            self.state.lock().pending_traces.remove(&id);
835        }
836    }
837}
838
839impl TracingWaker {
840    fn wake(self) {
841        self.wake_by_ref();
842    }
843
844    fn wake_by_ref(&self) {
845        assert_correct_thread(&self.thread, &self.state);
846
847        let mut state = self.state.lock();
848        if let Some(id) = self.id {
849            state.pending_traces.remove(&id);
850        }
851        state.unparked = true;
852        drop(state);
853        self.awoken.store(true, SeqCst);
854        self.thread.unpark();
855    }
856
857    fn clone_raw(waker: *const ()) -> RawWaker {
858        let waker = waker as *const TracingWaker;
859        let waker = unsafe { &*waker };
860        RawWaker::new(
861            Box::into_raw(Box::new(waker.clone())) as *const (),
862            &WAKER_VTABLE,
863        )
864    }
865
866    fn wake_raw(waker: *const ()) {
867        let waker = unsafe { Box::from_raw(waker as *mut TracingWaker) };
868        waker.wake();
869    }
870
871    fn wake_by_ref_raw(waker: *const ()) {
872        let waker = waker as *const TracingWaker;
873        let waker = unsafe { &*waker };
874        waker.wake_by_ref();
875    }
876
877    fn drop_raw(waker: *const ()) {
878        let waker = unsafe { Box::from_raw(waker as *mut TracingWaker) };
879        drop(waker);
880    }
881}
882
883pub struct Yield(usize);
884
885/// A wrapper around `Arc<Mutex<StdRng>>` that provides convenient methods
886/// for random number generation without requiring explicit locking.
887#[derive(Clone)]
888pub struct SharedRng(Arc<Mutex<StdRng>>);
889
890impl SharedRng {
891    /// Lock the inner RNG for direct access. Use this when you need multiple
892    /// random operations without re-locking between each one.
893    pub fn lock(&self) -> MutexGuard<'_, StdRng> {
894        self.0.lock()
895    }
896
897    /// Generate a random value in the given range.
898    pub fn random_range<T, R>(&self, range: R) -> T
899    where
900        T: SampleUniform,
901        R: SampleRange<T>,
902    {
903        self.0.lock().random_range(range)
904    }
905
906    /// Generate a random boolean with the given probability of being true.
907    pub fn random_bool(&self, p: f64) -> bool {
908        self.0.lock().random_bool(p)
909    }
910
911    /// Generate a random value of the given type.
912    pub fn random<T>(&self) -> T
913    where
914        StandardUniform: Distribution<T>,
915    {
916        self.0.lock().random()
917    }
918
919    /// Generate a random ratio - true with probability `numerator/denominator`.
920    pub fn random_ratio(&self, numerator: u32, denominator: u32) -> bool {
921        self.0.lock().random_ratio(numerator, denominator)
922    }
923}
924
925impl Future for Yield {
926    type Output = ();
927
928    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
929        if self.0 == 0 {
930            Poll::Ready(())
931        } else {
932            self.0 -= 1;
933            cx.waker().wake_by_ref();
934            Poll::Pending
935        }
936    }
937}
938
939fn exclude_wakers_from_trace(mut trace: Backtrace) -> Backtrace {
940    trace.resolve();
941    let mut frames: Vec<BacktraceFrame> = trace.into();
942    let waker_clone_frame_ix = frames.iter().position(|frame| {
943        frame.symbols().iter().any(|symbol| {
944            symbol
945                .name()
946                .is_some_and(|name| format!("{name:#?}") == type_name_of_val(&Waker::clone))
947        })
948    });
949
950    if let Some(waker_clone_frame_ix) = waker_clone_frame_ix {
951        frames.drain(..waker_clone_frame_ix + 1);
952    }
953
954    Backtrace::from(frames)
955}
956