Skip to main content

hydro_lang/sim/
compiled.rs

1//! Interfaces for compiled Hydro simulators and concrete simulation instances.
2//!
3//! # Quiescence and observation soundness
4//!
5//! The scheduler distinguishes two kinds of simulation work:
6//! - **Deterministic work**: running the top-level async dataflows, which simply propagate
7//!   whatever data is already in flight. This makes no `nondet!` decisions, so running it can
8//!   never change which executions are explored.
9//! - **Nondeterministic work**: running ticks and observations, whose behavior depends on
10//!   decisions drawn from the bolero driver (batch boundaries, snapshot versions, message
11//!   orderings). Each decision forks the space of possible executions.
12//!
13//! The simulation is **quiescent** when neither kind of work can make progress without new
14//! external input. Test-side observations (the methods on [`SimReceiver`] /
15//! [`SimClusterReceiver`]) interact with the scheduler while waiting, and the key soundness
16//! question is: *when is it okay for an observation to let nondeterministic work run?*
17//!
18//! **Waiting for a message is always sound.** If the message eventually arrives, the work
19//! that ran was necessary to produce it (schedules that run *extra* work are also valid
20//! executions and are explored separately). If the simulation instead quiesces without
21//! producing the message, the assertion fails and the instance ends, so nothing can observe
22//! the overrun. This is why [`SimReceiver::next`], [`SimReceiver::collect_n`], and the
23//! `assert_yields*` prefix checks are safe to use in the middle of a test.
24//!
25//! **Observing the *absence* of a message is dangerous.** Proving that "no more messages can
26//! arrive" requires driving the simulation all the way to quiescence, running *all* pending
27//! nondeterministic work. A later assertion may have needed to observe a state where that
28//! work had not yet run — e.g., `assert_yields_only([1, 2])` followed by reading a counter
29//! must be able to see the counter *before* the ticks that count `1` and `2` have fired.
30//! Forcing quiescence at the first assertion would make some executions unobservable, and
31//! extra messages produced by the forced work could surface at a *later* assertion,
32//! misattributing the failure. Absence-observing APIs therefore proceed in phases:
33//!
34//! 1. **Settle** (see `SettlePauseGuard::poll_settle`): the scheduler runs only deterministic work, pausing
35//!    just before nondeterministic work. If the simulation reaches quiescence this way, the
36//!    end-of-stream check is *free* — no decision was forced, no execution was cut off — and
37//!    the test simply continues.
38//! 2. If nondeterministic work is pending, the check would overrun. What happens next depends
39//!    on the API and engine:
40//!    - The assertion APIs ([`SimReceiver::assert_no_more`], `assert_yields_only*`,
41//!      `collect_n_only`) under [`CompiledSim::exhaustive`] **fork** the search on a bolero
42//!      decision: one instance performs the check and then ends (via a discard panic, like
43//!      `sim::continue_if!`), while sibling instances skip the check entirely and continue. The
44//!      exhaustive driver enumerates the checking instance *first*, so a failing check is
45//!      found before any instance runs past it — with a decision trace that leads exactly to
46//!      the failing assertion. Since nothing after the check runs in the checking instance,
47//!      the overrun it performs is unobservable, and the continuing instances never quiesce,
48//!      so every downstream state remains reachable.
49//!    - Otherwise (fuzz / RNG / replay engines, or the drain-everything APIs
50//!      [`SimReceiver::try_next`], [`SimReceiver::collect`], and `collect_sorted` in every
51//!      mode), the pending work runs and the instance is **tainted**
52//!      (`QuiescenceState::tainted`). Reads of the now-quiescent state remain sound (they
53//!      observe a fully-drained simulation that can no longer advance), so tests may drain
54//!      multiple output ports at the end. But once new input is sent, the instance is
55//!      **poisoned** (`QuiescenceState::poisoned`): any further receive panics (see
56//!      `guard_not_poisoned`), because a failure observed after the forced overrun could
57//!      have been caused by it and attributed to the wrong assertion.
58//!
59//! NOTE: This module runs inside bolero's `catch_unwind` scope, which silently
60//! swallows panics. Internal invariant checks should use `abort_assert!`
61//! rather than `panic!`/`assert!`.
62//!
63//! TODO(mingwei): Panics inside the tick DFIR (generated code in the dylib) are
64//! also caught by bolero's `catch_unwind`. Consider a mechanism to detect and
65//! propagate those as well.
66
67/// Like `assert!`, but calls `std::process::abort()` instead of `panic!()`.
68/// Use for internal invariants that must not be silently caught by bolero.
69macro_rules! abort_assert {
70    ($cond:expr, $($arg:tt)*) => {
71        if !$cond {
72            eprintln!("Simulator internal error: {}", format!($($arg)*));
73            std::process::abort();
74        }
75    };
76}
77
78use core::{fmt, panic};
79use std::cell::{Cell, RefCell};
80use std::collections::{HashMap, VecDeque};
81use std::fmt::Debug;
82use std::panic::RefUnwindSafe;
83use std::path::Path;
84use std::pin::{Pin, pin};
85use std::rc::Rc;
86use std::task::{Poll, ready};
87
88use bytes::Bytes;
89use colored::Colorize;
90use dfir_rs::scheduled::context::DfirErased;
91use dfir_rs::util::unsync::mpsc::{Receiver as UnsyncReceiver, Sender as UnsyncSender};
92use futures::StreamExt;
93use libloading::Library;
94use serde::Serialize;
95use serde::de::DeserializeOwned;
96use tempfile::TempPath;
97use tokio::sync::{Mutex, Notify};
98
99use super::runtime::{Hooks, InlineHooks};
100use super::{SimClusterReceiver, SimClusterSender, SimReceiver, SimSender};
101use crate::compile::builder::ExternalPortId;
102use crate::live_collections::stream::{ExactlyOnce, NoOrder, Ordering, Retries, TotalOrder};
103use crate::location::dynamic::LocationId;
104use crate::sim::graph::{SimExternalPort, SimExternalPortRegistry};
105use crate::sim::runtime::{SimHook, SimInlineHook};
106
107struct QuiescenceState {
108    /// Set to true when the scheduler reaches quiescence; reset to false when new input is sent.
109    quiescent: Cell<bool>,
110    /// Notified when the scheduler reaches quiescence (wakes receivers waiting for data).
111    quiescence_notify: Notify,
112    /// Notified when new input is sent, signaling the scheduler to resume.
113    resume_notify: Notify,
114    /// When nonzero, the scheduler must not start nondeterministic work (ticks /
115    /// observations): once only such work remains, it sets `nondet_pending` and pauses until
116    /// resumed. Used by receivers to query whether the simulation can quiesce
117    /// deterministically. This is a count (not a bool) because multiple settling futures can
118    /// be in flight at once (e.g. `select!`/`join!` between two receiver awaits): the
119    /// scheduler must stay paused until *every* one of them has finished settling.
120    pause_nondet: Cell<usize>,
121    /// Set while the scheduler is paused because nondeterministic work is ready to run but
122    /// `pause_nondet` is set.
123    nondet_pending: Cell<bool>,
124    /// Wakers for test-side tasks waiting for the scheduler to settle (either quiesce or set
125    /// `nondet_pending`) while `pause_nondet` is set.
126    settle_wakers: RefCell<Vec<std::task::Waker>>,
127    /// Set when an observation *forced* the simulation to quiesce (running pending
128    /// nondeterministic work) outside of exhaustive mode's forking. Further observations of
129    /// the quiescent state remain sound, but once new input is sent (see `poisoned`), later
130    /// observations could misattribute failures caused by the forced overrun.
131    tainted: Cell<bool>,
132    /// Set when new input is sent after `tainted`; all further receives panic.
133    poisoned: Cell<bool>,
134}
135
136impl QuiescenceState {
137    /// Signal that new input has been sent, waking the scheduler if it was quiescent.
138    fn resume(&self) {
139        if self.tainted.get() {
140            self.poisoned.set(true);
141        }
142        self.quiescent.set(false);
143        // `notify_one` (rather than `notify_waiters`) stores a permit if the scheduler driver
144        // is not currently parked on [`Self::resumed`], so a resume that fires before the
145        // driver parks (e.g. input sent while the driver is polling the thunk) is not lost.
146        self.resume_notify.notify_one();
147    }
148
149    /// Whether the scheduler is currently quiescent (no more progress possible without input).
150    fn is_quiescent(&self) -> bool {
151        self.quiescent.get()
152    }
153
154    /// Returns a future that completes when the scheduler next reaches quiescence.
155    fn notified(&self) -> tokio::sync::futures::Notified<'_> {
156        self.quiescence_notify.notified()
157    }
158
159    /// Wakes test-side tasks waiting for the scheduler to settle.
160    fn wake_settled(&self) {
161        for waker in self.settle_wakers.borrow_mut().drain(..) {
162            waker.wake();
163        }
164    }
165
166    /// Enter quiescence, waking receivers waiting for data (their streams end). The scheduler
167    /// driver is responsible for parking until [`Self::resume`] is called with new input.
168    fn enter_quiescence(&self) {
169        self.quiescent.set(true);
170        self.quiescence_notify.notify_waiters();
171        self.wake_settled();
172    }
173
174    /// Completes when new input arrives (via [`Self::resume`]).
175    async fn resumed(&self) {
176        self.resume_notify.notified().await;
177    }
178}
179
180/// Tracks a pending "settle" pause request to the scheduler (see
181/// [`QuiescenceState::pause_nondet`]), releasing it if the requesting future is dropped
182/// mid-settle (e.g. by `select!`) so the scheduler is not left paused forever. Pause
183/// requests are counted, so concurrent settling futures each hold their own request.
184struct SettlePauseGuard {
185    quiescence: Rc<QuiescenceState>,
186    active: bool,
187}
188
189impl SettlePauseGuard {
190    fn new(quiescence: Rc<QuiescenceState>) -> Self {
191        SettlePauseGuard {
192            quiescence,
193            active: false,
194        }
195    }
196
197    fn acquire(&mut self) {
198        abort_assert!(!self.active, "settle pause acquired twice");
199        self.quiescence
200            .pause_nondet
201            .set(self.quiescence.pause_nondet.get() + 1);
202        self.active = true;
203    }
204
205    fn release(&mut self) {
206        abort_assert!(self.active, "settle pause released without being acquired");
207        self.active = false;
208        self.quiescence
209            .pause_nondet
210            .set(self.quiescence.pause_nondet.get() - 1);
211    }
212
213    /// Polls the "settle" handshake with the scheduler: deterministic (non-tick) work is
214    /// allowed to run, but the scheduler pauses instead of starting nondeterministic work
215    /// (ticks / observations). Resolves to `true` if the simulation reached quiescence
216    /// deterministically, or `false` if nondeterministic work is pending (in which case the
217    /// scheduler is resumed).
218    fn poll_settle(&mut self, cx: &mut std::task::Context<'_>) -> Poll<bool> {
219        let quiescence = self.quiescence.clone();
220        if !self.active {
221            if quiescence.is_quiescent() {
222                return Poll::Ready(true);
223            }
224            self.acquire();
225        }
226
227        if quiescence.is_quiescent() {
228            self.release();
229            Poll::Ready(true)
230        } else if quiescence.nondet_pending.get() {
231            self.release();
232            // `notify_one` (permit-based): the driver only parks *between* thunk polls, so it
233            // is not parked right now — the permit ensures this resume is not lost.
234            quiescence.resume_notify.notify_one();
235            Poll::Ready(false)
236        } else {
237            // This may push a duplicate waker if we are re-polled without an intervening
238            // `wake_settled` (e.g. a `join!` sibling waking the shared task), but duplicates
239            // are harmless (waking is idempotent) and are cleared at the next `wake_settled`,
240            // so deduplicating here isn't worth the scan on every poll.
241            quiescence
242                .settle_wakers
243                .borrow_mut()
244                .push(cx.waker().clone());
245            Poll::Pending
246        }
247    }
248}
249
250impl Drop for SettlePauseGuard {
251    fn drop(&mut self) {
252        if self.active {
253            self.release();
254            // Resume the scheduler in case this was the last pause request (otherwise it
255            // would stay parked forever with nobody left to resume it). `notify_one`
256            // (permit-based) so the resume is not lost if the driver has not parked yet. If
257            // other settlers still hold requests, this wakeup is spurious but harmless: the
258            // scheduler re-checks `pause_nondet > 0` before starting any nondeterministic
259            // work, so it immediately re-parks without running anything.
260            self.quiescence.resume_notify.notify_one();
261        }
262    }
263}
264
265/// Panics if the simulation has been poisoned: an earlier observation forced the simulation
266/// to quiesce (running pending nondeterministic work), and new input has been sent since, so
267/// further observations could misattribute failures caused by the forced overrun.
268fn guard_not_poisoned(quiescence: &QuiescenceState) {
269    if quiescence.poisoned.get() {
270        panic!(
271            "cannot receive more simulator output: an earlier observation (such as `try_next`, `collect`, or a quiescence assertion outside exhaustive mode) forced the simulation to quiesce by running pending nondeterministic work, and new input has been sent since. Failures observed now could be misattributed, so either restructure the test to make quiescence-forcing observations its last step, or insert an explicit `sim::quiesce().await` phase barrier before sending more input."
272        );
273    }
274}
275
276/// Runs the simulation to quiescence, as an explicit *phase barrier* between rounds of a
277/// multi-phase test.
278///
279/// All pending nondeterministic work (ticks / observations) is forced to run until no more
280/// progress is possible without new input. This deliberately narrows the explored executions:
281/// inputs sent after the barrier will never interleave with work from before it, modeling
282/// scenarios where new stimuli (such as timer ticks) arrive long after the system settles.
283/// Pair such tests with a separate barrier-free test if interleaved executions should also be
284/// explored.
285///
286/// Because the barrier is explicit, observations after it are *intended* to see the fully
287/// settled state, so — unlike [`SimReceiver::try_next`] / [`SimReceiver::collect`] forcing
288/// quiescence implicitly — it does not restrict what the test may do afterwards: receives
289/// after the barrier observe only buffered output (plus whatever later input produces), and
290/// failures cannot be misattributed across it.
291pub async fn quiesce() {
292    let quiescence =
293        CURRENT_SIM_CONNECTIONS.with(|connections| connections.borrow().quiescence.clone());
294    guard_not_poisoned(&quiescence);
295
296    let mut notified_fut = pin!(None);
297    std::future::poll_fn(|cx| {
298        if quiescence.is_quiescent() {
299            return Poll::Ready(());
300        }
301        // Registered before the scheduler can run (single-threaded), so the quiescence
302        // notification cannot be missed.
303        if notified_fut.is_none() {
304            notified_fut.set(Some(quiescence.notified()));
305        }
306        let () = ready!(notified_fut.as_mut().as_pin_mut().unwrap().poll(cx));
307        Poll::Ready(())
308    })
309    .await;
310
311    // The barrier subsumes any quiescence forced by earlier observations in this phase:
312    // everything before it has fully settled, and the test has explicitly opted into
313    // observing only post-quiescence states from here on.
314    quiescence.tainted.set(false);
315}
316
317/// Receives the next message from `receiver` while trying not to overrun the simulation:
318/// first the simulation *settles* (deterministic work runs, but the scheduler pauses before
319/// nondeterministic work). If a message arrives, it is returned; if the simulation settles to
320/// quiescence, returns `None` without having run any nondeterministic work. Otherwise the
321/// scheduler is resumed and pending nondeterministic work runs until a message arrives or the
322/// simulation quiesces; quiescing this way *taints* the simulation (see
323/// [`QuiescenceState::tainted`]).
324async fn try_next_bytes(
325    receiver: &Mutex<UnsyncReceiver<Bytes>>,
326    quiescence: &Rc<QuiescenceState>,
327) -> Option<Bytes> {
328    guard_not_poisoned(quiescence);
329
330    let mut receiver_stream = receiver.lock().await;
331    let mut settle_guard = SettlePauseGuard::new(quiescence.clone());
332    // `Some` once the settle phase has concluded that nondeterministic work is pending and
333    // we have started forcing it to run.
334    let mut notified_fut = pin!(None);
335
336    std::future::poll_fn(|cx| {
337        // A message may become available at any point (including from deterministic work
338        // while settling), so always check the stream first.
339        match receiver_stream.poll_next_unpin(cx) {
340            Poll::Ready(Some(bytes)) => return Poll::Ready(Some(bytes)),
341            Poll::Ready(None) => return Poll::Ready(None),
342            Poll::Pending => {}
343        }
344
345        if notified_fut.is_none() {
346            match settle_guard.poll_settle(cx) {
347                // Deterministically quiescent: no more messages, and nothing was overrun.
348                Poll::Ready(true) => return Poll::Ready(None),
349                // Nondeterministic work is pending; start forcing it to run. The `Notified`
350                // is created here and polled (registered) below in this same synchronous
351                // poll — before the scheduler can run — and the simulation is not currently
352                // quiescent, so the quiescence notification cannot be missed.
353                Poll::Ready(false) => notified_fut.set(Some(quiescence.notified())),
354                Poll::Pending => return Poll::Pending,
355            }
356        }
357
358        // Let the scheduler run nondeterministic work until a message arrives or the
359        // simulation quiesces. Note that merely entering this phase does not taint: if a
360        // message arrives (the `Some` exit at the top), waiting was sound for the same
361        // reason as `SimReceiver::next` — the work that ran was needed to produce it. Only
362        // *observing quiescence* after forcing the pending work taints, since that is the
363        // overrun a later observation could misattribute.
364        let () = ready!(notified_fut.as_mut().as_pin_mut().unwrap().poll(cx));
365        quiescence.tainted.set(true);
366        Poll::Ready(None)
367    })
368    .await
369}
370
371struct SimConnections {
372    input_senders: HashMap<SimExternalPort, UnsyncSender<Bytes>>,
373    output_receivers: HashMap<SimExternalPort, Rc<Mutex<UnsyncReceiver<Bytes>>>>,
374    cluster_input_senders: HashMap<SimExternalPort, HashMap<u32, UnsyncSender<Bytes>>>,
375    cluster_output_receivers:
376        HashMap<SimExternalPort, HashMap<u32, Rc<Mutex<UnsyncReceiver<Bytes>>>>>,
377    external_registered: HashMap<ExternalPortId, SimExternalPort>,
378    quiescence: Rc<QuiescenceState>,
379    log: bool,
380    /// Whether this instance is being executed by the exhaustive engine (see
381    /// [`CompiledSim::exhaustive`]), which affects how `assert_yields_only` explores
382    /// quiescence checks.
383    exhaustive: bool,
384}
385
386/// Implementation detail of [`crate::sim::continue_if!`](crate::continue_if); do not call directly.
387///
388/// If `condition` is false, aborts the current simulation instance by panicking with a special
389/// payload ([`bolero::generator::bolero_generator::any::Error`]) that bolero recognizes as an
390/// "invalid input" marker: the instance is discarded (not treated as a test failure, and never
391/// recorded as a reproducer) and exploration moves on to the next instance. If logging is
392/// enabled for the current instance, the failed assumption is logged first.
393#[doc(hidden)]
394#[track_caller]
395pub fn continue_if_impl(condition: bool, message: fmt::Arguments<'_>) {
396    if condition {
397        return;
398    }
399
400    let log = CURRENT_SIM_CONNECTIONS
401        .try_with(|connections| connections.borrow().log)
402        .unwrap_or(true);
403    if log {
404        eprintln!(
405            "{}",
406            render_continue_if_failure(std::panic::Location::caller(), message)
407        );
408    }
409
410    // Panics with `bolero_generator::any::Error`, which bolero's engines treat as an invalid
411    // input rather than a test failure. Both this function and bolero's `assume` are
412    // `#[track_caller]`, so the recorded location is the user's `continue_if!` call site.
413    bolero::generator::bolero_generator::any::assume(false, "simulation assumption failed");
414}
415
416/// Renders the log message for a failed assumption, echoing the source line with a caret
417/// pointing at the `continue_if!` call site, in the same style as the other simulator logs.
418fn render_continue_if_failure(
419    location: &std::panic::Location<'_>,
420    message: fmt::Arguments<'_>,
421) -> String {
422    use std::fmt::Write;
423
424    // `Location::file()` is relative to the directory the crate was compiled from (e.g. the
425    // workspace root), which may not match the current working directory (e.g. the crate
426    // root when running `cargo test`), so walk up from the current directory to find it.
427    let source_line = std::env::current_dir()
428        .ok()
429        .and_then(|cwd| {
430            cwd.ancestors()
431                .find_map(|base| std::fs::read_to_string(base.join(location.file())).ok())
432        })
433        .and_then(|content| {
434            content
435                .lines()
436                .nth((location.line() as usize).saturating_sub(1))
437                .map(|line| line.to_owned())
438        })
439        .unwrap_or_default();
440
441    let caret_indent = " ".repeat((location.column() as usize).saturating_sub(1));
442
443    let mut out = String::new();
444    let _ = writeln!(
445        out,
446        "\n{}",
447        "Condition failed (discarding simulation instance):"
448            .color(colored::Color::Yellow)
449            .bold()
450    );
451    let _ = writeln!(out, "{} {}", "-->".color(colored::Color::Blue), location);
452    let _ = writeln!(out, " {}{}", "|".color(colored::Color::Blue), source_line);
453    let _ = write!(
454        out,
455        " {}{}{}",
456        "|".color(colored::Color::Blue),
457        caret_indent,
458        format!("^ {}", message).color(colored::Color::Yellow)
459    );
460    out
461}
462
463tokio::task_local! {
464    static CURRENT_SIM_CONNECTIONS: RefCell<SimConnections>;
465}
466
467/// A handle to a compiled Hydro simulation, which can be instantiated and run.
468pub struct CompiledSim {
469    pub(super) _path: TempPath,
470    pub(super) lib: Library,
471    pub(super) externals_port_registry: SimExternalPortRegistry,
472    pub(super) unit_test_fuzz_iterations: usize,
473}
474
475#[sealed::sealed]
476/// A trait implemented by closures that can instantiate a compiled simulation.
477///
478/// This is needed to ensure [`RefUnwindSafe`] so instances can be created during fuzzing.
479pub trait Instantiator<'a>: RefUnwindSafe + Fn() -> CompiledSimInstance<'a> {}
480#[sealed::sealed]
481impl<'a, T: RefUnwindSafe + Fn() -> CompiledSimInstance<'a>> Instantiator<'a> for T {}
482
483fn null_handler(_args: fmt::Arguments<'_>) {}
484
485fn println_handler(args: fmt::Arguments<'_>) {
486    println!("{}", args);
487}
488
489fn eprintln_handler(args: fmt::Arguments<'_>) {
490    eprintln!("{}", args);
491}
492
493/// Creates a simulation instance, returning:
494/// - A list of async DFIRs to run (all process / cluster logic outside a tick)
495/// - A list of tick DFIRs to run (where the &'static str is for the tick location id)
496/// - A mapping of hooks for non-deterministic decisions at tick-input boundaries
497/// - A mapping of inline hooks for non-deterministic decisions inside ticks
498type SimLoaded<'a> = libloading::Symbol<
499    'a,
500    unsafe extern "Rust" fn(
501        should_color: bool,
502        external_out: &mut HashMap<usize, UnsyncReceiver<Bytes>>,
503        external_in: &mut HashMap<usize, UnsyncSender<Bytes>>,
504        cluster_external_out: &mut HashMap<usize, HashMap<u32, UnsyncReceiver<Bytes>>>,
505        cluster_external_in: &mut HashMap<usize, HashMap<u32, UnsyncSender<Bytes>>>,
506        println_handler: fn(fmt::Arguments<'_>),
507        eprintln_handler: fn(fmt::Arguments<'_>),
508    ) -> (
509        Vec<(&'static str, Option<u32>, DfirErased)>,
510        Vec<(&'static str, Option<u32>, DfirErased)>,
511        Hooks<&'static str>,
512        InlineHooks<&'static str>,
513    ),
514>;
515
516impl CompiledSim {
517    /// Executes the given closure with a single instance of the compiled simulation.
518    pub fn with_instance<T>(&self, thunk: impl FnOnce(CompiledSimInstance<'_>) -> T) -> T {
519        self.with_instantiator(|instantiator| thunk(instantiator()), true)
520    }
521
522    /// Executes the given closure with an [`Instantiator`], which can be called to create
523    /// independent instances of the simulation. This is useful for fuzzing, where we need to
524    /// re-execute the simulation several times with different decisions.
525    ///
526    /// The `always_log` parameter controls whether to log tick executions and stream releases. If
527    /// it is `true`, logging will always be enabled. If it is `false`, logging will only be
528    /// enabled if the `HYDRO_SIM_LOG` environment variable is set to `1`.
529    pub fn with_instantiator<T>(
530        &self,
531        thunk: impl FnOnce(&dyn Instantiator<'_>) -> T,
532        always_log: bool,
533    ) -> T {
534        let func: SimLoaded<'_> = unsafe { self.lib.get(b"__hydro_runtime").unwrap() };
535        let log = always_log || std::env::var("HYDRO_SIM_LOG").is_ok_and(|v| v == "1");
536        thunk(
537            &(|| CompiledSimInstance {
538                func: func.clone(),
539                externals_port_registry: self.externals_port_registry.clone(),
540                dylib_result: None,
541                log,
542                exhaustive: false,
543            }),
544        )
545    }
546
547    /// Uses a fuzzing strategy to explore possible executions of the simulation. The provided
548    /// closure will be repeatedly executed with instances of the Hydro program where the
549    /// batching boundaries, order of messages, and retries are varied.
550    ///
551    /// During development, you should run the test that invokes this function with the `cargo sim`
552    /// command, which will use `libfuzzer` to intelligently explore the execution space. If a
553    /// failure is found, a minimized test case will be produced in a `sim-failures` directory.
554    /// When running the test with `cargo test` (such as in CI), if a reproducer is found it will
555    /// be executed, and if no reproducer is found a small number of random executions will be
556    /// performed.
557    pub fn fuzz(&self, mut thunk: impl AsyncFnMut() + RefUnwindSafe) {
558        let caller_fn = crate::compile::ir::backtrace::Backtrace::get_backtrace(0)
559            .elements()
560            .into_iter()
561            .find(|e| {
562                !e.fn_name.starts_with("hydro_lang::sim::compiled")
563                    && !e.fn_name.starts_with("hydro_lang::sim::flow")
564                    && !e.fn_name.starts_with("fuzz<")
565                    && !e.fn_name.starts_with("<hydro_lang::sim")
566            })
567            .unwrap();
568
569        let caller_path = Path::new(&caller_fn.filename.unwrap()).to_path_buf();
570        let repro_folder = caller_path.parent().unwrap().join("sim-failures");
571
572        let caller_fuzz_repro_path = repro_folder
573            .join(caller_fn.fn_name.replace("::", "__"))
574            .with_extension("bin");
575
576        if std::env::var("BOLERO_FUZZER").is_ok() {
577            let corpus_dir = std::env::current_dir().unwrap().join(".fuzz-corpus");
578            std::fs::create_dir_all(&corpus_dir).unwrap();
579            let libfuzzer_args = format!(
580                "{} {} -artifact_prefix={}/ -handle_abrt=0",
581                corpus_dir.to_str().unwrap(),
582                corpus_dir.to_str().unwrap(),
583                corpus_dir.to_str().unwrap(),
584            );
585
586            std::fs::create_dir_all(&repro_folder).unwrap();
587
588            if !std::env::var("HYDRO_NO_FAILURE_OUTPUT").is_ok_and(|v| v == "1") {
589                unsafe {
590                    std::env::set_var(
591                        "BOLERO_FAILURE_OUTPUT",
592                        caller_fuzz_repro_path.to_str().unwrap(),
593                    );
594                }
595            }
596
597            unsafe {
598                std::env::set_var("BOLERO_LIBFUZZER_ARGS", libfuzzer_args);
599            }
600
601            self.with_instantiator(
602                |instantiator| {
603                    bolero::test(bolero::TargetLocation {
604                        package_name: "",
605                        manifest_dir: "",
606                        module_path: "",
607                        file: "",
608                        line: 0,
609                        item_path: "<unknown>::__bolero_item_path__",
610                        test_name: None,
611                    })
612                    .run_with_replay(move |is_replay| {
613                        let mut instance = instantiator();
614
615                        if instance.log {
616                            eprintln!(
617                                "{}",
618                                "\n==== New Simulation Instance ===="
619                                    .color(colored::Color::Cyan)
620                                    .bold()
621                            );
622                        }
623
624                        if is_replay {
625                            instance.log = true;
626                        }
627
628                        tokio::runtime::Builder::new_current_thread()
629                            .build()
630                            .unwrap()
631                            .block_on(async { instance.run(&mut thunk).await })
632                    })
633                },
634                false,
635            );
636        } else if let Ok(existing_bytes) = std::fs::read(&caller_fuzz_repro_path) {
637            self.fuzz_repro(existing_bytes, async |compiled| {
638                compiled.run_with_scheduler(thunk()).await
639            });
640        } else {
641            eprintln!(
642                "Running a fuzz test without `cargo sim` and no reproducer found at {}, using {} iterations with random inputs.",
643                caller_fuzz_repro_path.display(),
644                self.unit_test_fuzz_iterations,
645            );
646            self.with_instantiator(
647                |instantiator| {
648                    bolero::test(bolero::TargetLocation {
649                        package_name: "",
650                        manifest_dir: "",
651                        module_path: "",
652                        file: ".",
653                        line: 0,
654                        item_path: "<unknown>::__bolero_item_path__",
655                        test_name: None,
656                    })
657                    .with_iterations(self.unit_test_fuzz_iterations)
658                    .run_with_replay(move |is_replay| {
659                        let mut instance = instantiator();
660
661                        if instance.log {
662                            eprintln!(
663                                "{}",
664                                "\n==== New Simulation Instance ===="
665                                    .color(colored::Color::Cyan)
666                                    .bold()
667                            );
668                        }
669
670                        if is_replay {
671                            instance.log = true;
672                        }
673
674                        tokio::runtime::Builder::new_current_thread()
675                            .build()
676                            .unwrap()
677                            .block_on(async { instance.run(&mut thunk).await })
678                    })
679                },
680                false,
681            );
682        }
683    }
684
685    /// Executes the given closure with a single instance of the compiled simulation, using the
686    /// provided bytes as the source of fuzzing decisions. This can be used to manually reproduce a
687    /// failure found during fuzzing.
688    pub fn fuzz_repro<'a>(
689        &'a self,
690        bytes: Vec<u8>,
691        thunk: impl AsyncFnOnce(CompiledSimInstance<'_>) + RefUnwindSafe,
692    ) {
693        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
694            self.with_instance(|instance| {
695                bolero::bolero_engine::any::scope::with(
696                    Box::new(bolero::bolero_engine::driver::object::Object(
697                        bolero::bolero_engine::driver::bytes::Driver::new(
698                            bytes,
699                            &Default::default(),
700                        ),
701                    )),
702                    || {
703                        tokio::runtime::Builder::new_current_thread()
704                            .build()
705                            .unwrap()
706                            .block_on(async { instance.run_without_launching(thunk).await })
707                    },
708                )
709            })
710        }));
711
712        if let Err(payload) = result {
713            if payload
714                .downcast_ref::<bolero::generator::bolero_generator::any::Error>()
715                .is_some()
716            {
717                // A `continue_if!` failed (or the driver ran out of entropy) while replaying the
718                // recorded bytes. Instances that fail an assumption are never recorded as
719                // failures, so this means the reproducer is stale or does not correspond to
720                // this program.
721                panic!(
722                    "simulation assumption failed while replaying recorded fuzz decisions; the reproducer may be stale or may not correspond to this program"
723                );
724            }
725            std::panic::resume_unwind(payload);
726        }
727    }
728
729    /// Exhaustively searches all possible executions of the simulation. The provided
730    /// closure will be repeatedly executed with instances of the Hydro program where the
731    /// batching boundaries, order of messages, and retries are varied.
732    ///
733    /// Exhaustive searching is feasible when the inputs to the Hydro program are finite and there
734    /// are no dataflow loops that generate infinite messages. Exhaustive searching provides a
735    /// stronger guarantee of correctness than fuzzing, but may take a long time to complete.
736    /// Because no fuzzer is involved, you can run exhaustive tests with `cargo test`.
737    ///
738    /// Returns the number of distinct executions explored.
739    pub fn exhaustive(&self, mut thunk: impl AsyncFnMut() + RefUnwindSafe) -> usize {
740        if std::env::var("BOLERO_FUZZER").is_ok() {
741            eprintln!(
742                "Cannot run exhaustive tests with a fuzzer. Please use `cargo test` instead of `cargo sim`."
743            );
744            std::process::abort();
745        }
746
747        let mut count = 0;
748        let count_mut = &mut count;
749
750        let _span = tracing::debug_span!(target: "hydro_build", "sim_exhaustive").entered();
751
752        self.with_instantiator(
753            |instantiator| {
754                bolero::test(bolero::TargetLocation {
755                    package_name: "",
756                    manifest_dir: "",
757                    module_path: "",
758                    file: "",
759                    line: 0,
760                    item_path: "<unknown>::__bolero_item_path__",
761                    test_name: None,
762                })
763                .exhaustive()
764                .run_with_replay(move |is_replay| {
765                    *count_mut += 1;
766
767                    let mut instance = instantiator();
768                    instance.exhaustive = true;
769                    if instance.log {
770                        eprintln!(
771                            "{}",
772                            "\n==== New Simulation Instance ===="
773                                .color(colored::Color::Cyan)
774                                .bold()
775                        );
776                    }
777
778                    if is_replay {
779                        instance.log = true;
780                    }
781
782                    tokio::runtime::Builder::new_current_thread()
783                        .build()
784                        .unwrap()
785                        .block_on(async { instance.run(&mut thunk).await })
786                })
787            },
788            false,
789        );
790
791        count
792    }
793}
794
795// This must be a tuple because it is referenced from generated code in `graph.rs`.
796type DylibResult = (
797    Vec<(&'static str, Option<u32>, DfirErased)>,
798    Vec<(&'static str, Option<u32>, DfirErased)>,
799    Hooks<&'static str>,
800    InlineHooks<&'static str>,
801);
802
803/// A single instance of a compiled Hydro simulation, which provides methods to interactively
804/// execute the simulation, feed inputs, and receive outputs.
805pub struct CompiledSimInstance<'a> {
806    func: SimLoaded<'a>,
807    externals_port_registry: SimExternalPortRegistry,
808    dylib_result: Option<DylibResult>,
809    log: bool,
810    exhaustive: bool,
811}
812
813impl<'a> CompiledSimInstance<'a> {
814    async fn run(self, thunk: impl AsyncFnOnce() + RefUnwindSafe) {
815        self.run_without_launching(async |instance| {
816            instance.run_with_scheduler(thunk()).await;
817        })
818        .await;
819    }
820
821    async fn run_without_launching(
822        mut self,
823        thunk: impl AsyncFnOnce(CompiledSimInstance<'_>) + RefUnwindSafe,
824    ) {
825        let mut external_out: HashMap<usize, UnsyncReceiver<Bytes>> = HashMap::new();
826        let mut external_in: HashMap<usize, UnsyncSender<Bytes>> = HashMap::new();
827        let mut cluster_external_out: HashMap<usize, HashMap<u32, UnsyncReceiver<Bytes>>> =
828            HashMap::new();
829        let mut cluster_external_in: HashMap<usize, HashMap<u32, UnsyncSender<Bytes>>> =
830            HashMap::new();
831
832        let dylib_result = unsafe {
833            (self.func)(
834                colored::control::SHOULD_COLORIZE.should_colorize(),
835                &mut external_out,
836                &mut external_in,
837                &mut cluster_external_out,
838                &mut cluster_external_in,
839                if self.log {
840                    println_handler
841                } else {
842                    null_handler
843                },
844                if self.log {
845                    eprintln_handler
846                } else {
847                    null_handler
848                },
849            )
850        };
851
852        let registered = &self.externals_port_registry.registered;
853
854        let quiescence = Rc::new(QuiescenceState {
855            quiescent: Cell::new(false),
856            quiescence_notify: Notify::new(),
857            resume_notify: Notify::new(),
858            pause_nondet: Cell::new(0),
859            nondet_pending: Cell::new(false),
860            settle_wakers: RefCell::new(vec![]),
861            tainted: Cell::new(false),
862            poisoned: Cell::new(false),
863        });
864
865        let mut input_senders = HashMap::new();
866        let mut output_receivers = HashMap::new();
867        let mut cluster_input_senders = HashMap::new();
868        let mut cluster_output_receivers = HashMap::new();
869
870        #[expect(
871            clippy::disallowed_methods,
872            reason = "inserts into maps also unordered"
873        )]
874        for sim_port in registered.values() {
875            let usize_key = sim_port.into_inner();
876            if let Some(sender) = external_in.remove(&usize_key) {
877                input_senders.insert(*sim_port, sender);
878            }
879            if let Some(receiver) = external_out.remove(&usize_key) {
880                output_receivers.insert(*sim_port, Rc::new(Mutex::new(receiver)));
881            }
882            if let Some(senders) = cluster_external_in.remove(&usize_key) {
883                cluster_input_senders.insert(*sim_port, senders);
884            }
885            if let Some(receivers) = cluster_external_out.remove(&usize_key) {
886                cluster_output_receivers.insert(
887                    *sim_port,
888                    receivers
889                        .into_iter()
890                        .map(|(member, r)| (member, Rc::new(Mutex::new(r))))
891                        .collect(),
892                );
893            }
894        }
895
896        self.dylib_result = Some(dylib_result);
897
898        CURRENT_SIM_CONNECTIONS
899            .scope(
900                RefCell::new(SimConnections {
901                    input_senders,
902                    output_receivers,
903                    cluster_input_senders,
904                    cluster_output_receivers,
905                    external_registered: self.externals_port_registry.registered.clone(),
906                    quiescence: quiescence.clone(),
907                    log: self.log,
908                    exhaustive: self.exhaustive,
909                }),
910                async move {
911                    thunk(self).await;
912                },
913            )
914            .await;
915    }
916
917    /// Runs the simulation scheduler alongside the given future, until the future completes.
918    ///
919    /// The future always gets to run first; whenever it is blocked (e.g. waiting to receive
920    /// simulation outputs), the scheduler runs a single step to completion. Steps are atomic
921    /// with respect to the future: it is re-polled between every pair of scheduler steps, but
922    /// never while a step is in flight. The [`LaunchedSim`] state struct lives across steps,
923    /// in this function's frame.
924    async fn run_with_scheduler(self, thunk: impl Future<Output = ()>) {
925        self.run_with_scheduler_and_maybe_logger::<std::io::Empty>(None, thunk)
926            .await;
927    }
928
929    /// Runs the simulation scheduler alongside the given future, until the future completes,
930    /// reporting the simulation trace to the given logger.
931    ///
932    /// The future always gets to run first; whenever it is blocked (e.g. waiting to receive
933    /// simulation outputs), the scheduler runs a single step to completion. Steps are atomic
934    /// with respect to the future: it is re-polled between every pair of scheduler steps, but
935    /// never while a step is in flight.
936    pub async fn run_with_scheduler_and_logger<W: std::io::Write>(
937        self,
938        log_writer: W,
939        thunk: impl Future<Output = ()>,
940    ) {
941        self.run_with_scheduler_and_maybe_logger(Some(log_writer), thunk)
942            .await;
943    }
944
945    async fn run_with_scheduler_and_maybe_logger<W: std::io::Write>(
946        self,
947        log_override: Option<W>,
948        thunk: impl Future<Output = ()>,
949    ) {
950        let mut sim = self.start(log_override);
951        let mut thunk_fut = pin!(thunk);
952        loop {
953            // The thunk always gets to run first.
954            if futures::poll!(thunk_fut.as_mut()).is_ready() {
955                break;
956            }
957
958            if sim.quiescence.is_quiescent() || sim.quiescence.nondet_pending.get() {
959                // The scheduler is parked: either no step can make progress until the thunk
960                // sends new input (quiescent), or nondeterministic work is ready but a
961                // settling test-side observation has paused the scheduler (nondet_pending).
962                // Park until either the thunk is woken independently or the scheduler is
963                // resumed. (`resumed()` is permit-based, so a resume that fired while polling
964                // the thunk above is not lost.)
965                tokio::select! {
966                    biased;
967                    () = &mut thunk_fut => break,
968                    () = sim.quiescence.resumed() => {}
969                }
970                sim.quiescence.nondet_pending.set(false);
971            } else {
972                // Run a single scheduler step to completion. This is awaited directly (not
973                // raced against the thunk), so a step is atomic: the thunk is never polled
974                // while a step is in flight, and a step is never cancelled mid-execution.
975                sim.step().await;
976            }
977        }
978    }
979
980    /// Consumes this instance and constructs the [`LaunchedSim`] state struct, which is
981    /// advanced incrementally via [`LaunchedSim::step`].
982    fn start<W: std::io::Write>(mut self, log_override: Option<W>) -> LaunchedSim<W> {
983        let (async_dfirs, tick_dfirs, mut hooks, mut inline_hooks) =
984            self.dylib_result.take().unwrap();
985
986        // The generated code keys hooks and tick DFIRs by the same serialized location
987        // strings, so we can move each tick's / observation's hooks out of the maps and
988        // attach them directly. This lets the scheduler's hot paths avoid keyed lookups
989        // (which would clone `LocationId`s) entirely.
990        let not_ready_ticks = tick_dfirs
991            .into_iter()
992            .map(|(lid, cluster_id, dfir)| {
993                let location: LocationId = serde_json::from_str(lid).unwrap();
994                let LocationId::Tick(_, parent_location) = location else {
995                    unreachable!("tick DFIRs are always keyed by a tick location")
996                };
997                SimTick {
998                    parent_location: *parent_location,
999                    cluster_id,
1000                    dfir,
1001                    hooks: hooks
1002                        .remove(&(lid, cluster_id))
1003                        .expect("every tick DFIR must have at least one hook"),
1004                    inline_hooks: inline_hooks.remove(&(lid, cluster_id)).unwrap_or_default(),
1005                }
1006            })
1007            .collect();
1008
1009        let quiescence = CURRENT_SIM_CONNECTIONS.with(|connections| {
1010            let connections = connections.borrow();
1011            connections.quiescence.clone()
1012        });
1013
1014        let not_ready_observations = async_dfirs
1015            .iter()
1016            .map(|(lid, cluster_id, _)| SimObservation {
1017                location: serde_json::from_str(lid).unwrap(),
1018                cluster_id: *cluster_id,
1019                hooks: hooks.remove(&(*lid, *cluster_id)).unwrap_or_default(),
1020            })
1021            .collect();
1022
1023        debug_assert!(
1024            hooks.is_empty() && inline_hooks.is_empty(),
1025            "all hooks should belong to either a tick DFIR or a top-level location"
1026        );
1027
1028        LaunchedSim {
1029            async_dfirs: async_dfirs
1030                .into_iter()
1031                .map(|(lid, c_id, dfir)| (serde_json::from_str(lid).unwrap(), c_id, dfir))
1032                .collect(),
1033            possibly_ready_ticks: vec![],
1034            not_ready_ticks,
1035            possibly_ready_observations: vec![],
1036            not_ready_observations,
1037            log: if self.log {
1038                if let Some(w) = log_override {
1039                    LogKind::Custom(w)
1040                } else {
1041                    LogKind::Stderr
1042                }
1043            } else {
1044                LogKind::Null
1045            },
1046            quiescence,
1047        }
1048    }
1049}
1050
1051impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> Clone for SimReceiver<T, O, R> {
1052    fn clone(&self) -> Self {
1053        *self
1054    }
1055}
1056
1057impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> Copy for SimReceiver<T, O, R> {}
1058
1059/// How a [`QuiescenceCheckFuture`] resolves the "did the stream end?" check of
1060/// `assert_no_more`. Decided once the simulation has settled (run out of deterministic
1061/// work).
1062#[derive(Clone, Copy)]
1063enum QuiescenceBranch {
1064    /// Skip the check and continue the test. Only taken in exhaustive mode, where a
1065    /// sibling instance performs the check instead.
1066    Continue,
1067    /// Perform the check, then end this simulation instance (exhaustive mode), letting
1068    /// sibling instances continue past this point without forcing quiescence.
1069    CheckThenEnd,
1070    /// Perform the check and keep running. Taken when the simulation is already quiescent
1071    /// (the check is free) and in non-exhaustive modes.
1072    CheckAndKeepRunning,
1073}
1074
1075/// Decides how to run the quiescence check when the simulation has pending nondeterministic
1076/// work (ticks / observations) that the check would force to run.
1077fn decide_quiescence_branch() -> QuiescenceBranch {
1078    let (exhaustive, log) = CURRENT_SIM_CONNECTIONS.with(|connections| {
1079        let connections = connections.borrow();
1080        (connections.exhaustive, connections.log)
1081    });
1082
1083    if !exhaustive {
1084        return QuiescenceBranch::CheckAndKeepRunning;
1085    }
1086
1087    // In exhaustive mode, fork the search on a bolero decision. The exhaustive driver
1088    // enumerates `false` first, so the instance that performs the quiescence check is
1089    // explored *before* any instance that continues past this assertion. This ensures that
1090    // if the stream has extra output, the failure is attributed to this assertion (with a
1091    // decision trace leading exactly to the check) rather than leaking the extra messages
1092    // into a later assertion.
1093    let continue_without_check: bool = bolero::any();
1094    if continue_without_check {
1095        if log {
1096            eprintln!(
1097                "\n{}",
1098                "Continuing past quiescence assertion without checking (checked by an earlier instance)"
1099                    .color(colored::Color::Cyan)
1100                    .bold()
1101            );
1102        }
1103        QuiescenceBranch::Continue
1104    } else {
1105        if log {
1106            eprintln!(
1107                "\n{}",
1108                "Checking that no more messages arrive (this instance will end after the check)"
1109                    .color(colored::Color::Cyan)
1110                    .bold()
1111            );
1112        }
1113        QuiescenceBranch::CheckThenEnd
1114    }
1115}
1116
1117/// Ends the current simulation instance after a passing quiescence check, by panicking with
1118/// [`bolero::generator::bolero_generator::any::Error`], which bolero's engines treat as an
1119/// invalid input rather than a test failure. The instance has verified everything up to and
1120/// including the quiescence check; sibling instances continue past the check instead.
1121fn end_instance_after_quiescence_check() -> ! {
1122    bolero::generator::bolero_generator::any::assume(
1123        false,
1124        "simulation instance ended after quiescence check",
1125    );
1126    unreachable!()
1127}
1128
1129pin_project_lite::pin_project! {
1130    // The "and then the stream ends" half of `assert_no_more` (and thus of
1131    // `assert_yields_only*` / `collect_n_only`). First lets the simulation *settle* (see
1132    // `poll_settle`): if it settles to quiescence, the check is free and the test simply
1133    // continues. Otherwise, in exhaustive mode the search forks into a checking instance and
1134    // continuing instances (see `SimReceiver::assert_no_more` and
1135    // `decide_quiescence_branch`); in non-exhaustive modes the check runs, forcing the
1136    // pending work (which taints the simulation, via `try_next_bytes`).
1137    //
1138    // See [`FutureTrackingCaller`] for why `poll` is `#[track_caller]`.
1139    struct QuiescenceCheckFuture<F: Future<Output = ()>> {
1140        #[pin]
1141        check: F,
1142        settle: SettlePauseGuard,
1143        branch: Option<QuiescenceBranch>,
1144    }
1145}
1146
1147impl<F: Future<Output = ()>> QuiescenceCheckFuture<F> {
1148    fn new(check: F) -> Self {
1149        QuiescenceCheckFuture {
1150            check,
1151            settle: SettlePauseGuard::new(
1152                CURRENT_SIM_CONNECTIONS.with(|connections| connections.borrow().quiescence.clone()),
1153            ),
1154            branch: None,
1155        }
1156    }
1157}
1158
1159impl<F: Future<Output = ()>> Future for QuiescenceCheckFuture<F> {
1160    type Output = ();
1161
1162    #[track_caller]
1163    fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
1164        let this = self.as_mut().project();
1165
1166        if this.branch.is_none() {
1167            *this.branch = Some(if ready!(this.settle.poll_settle(cx)) {
1168                // Settled to quiescence deterministically, so the check is free.
1169                QuiescenceBranch::CheckAndKeepRunning
1170            } else {
1171                // The check would force nondeterministic work to run.
1172                decide_quiescence_branch()
1173            });
1174        }
1175
1176        match this.branch.unwrap() {
1177            QuiescenceBranch::Continue => Poll::Ready(()),
1178            QuiescenceBranch::CheckAndKeepRunning => this.check.poll(cx),
1179            QuiescenceBranch::CheckThenEnd => {
1180                ready!(this.check.poll(cx));
1181                end_instance_after_quiescence_check()
1182            }
1183        }
1184    }
1185}
1186
1187impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> SimReceiver<T, O, R> {
1188    fn connections(&self) -> (Rc<Mutex<UnsyncReceiver<Bytes>>>, Rc<QuiescenceState>) {
1189        CURRENT_SIM_CONNECTIONS.with(|connections| {
1190            let connections = connections.borrow();
1191            let port = connections.external_registered.get(&self.0).unwrap();
1192            (
1193                connections.output_receivers.get(port).unwrap().clone(),
1194                connections.quiescence.clone(),
1195            )
1196        })
1197    }
1198
1199    /// See [`try_next_bytes`].
1200    async fn try_next_impl(&self) -> Option<T> {
1201        let (receiver, quiescence) = self.connections();
1202        try_next_bytes(&receiver, &quiescence)
1203            .await
1204            .map(|bytes| bincode::deserialize(&bytes).unwrap())
1205    }
1206
1207    /// Asserts that the stream has ended and no more messages can possibly arrive.
1208    ///
1209    /// If the check cannot be answered without running pending nondeterministic work (such
1210    /// as ticks with buffered inputs):
1211    /// - Under [`CompiledSim::exhaustive`], the search forks: one instance performs the
1212    ///   check and ends there, while sibling instances skip the check and continue.
1213    /// - In other modes, the pending work runs; afterwards, sending more input and then
1214    ///   attempting to receive output will panic.
1215    pub fn assert_no_more(self) -> impl Future<Output = ()>
1216    where
1217        T: Debug,
1218    {
1219        QuiescenceCheckFuture::new(FutureTrackingCaller {
1220            future: async move {
1221                if let Some(next) = self.try_next_impl().await {
1222                    return Err(format!(
1223                        "Stream yielded unexpected message: {:?}, expected termination",
1224                        next
1225                    ));
1226                }
1227                Ok(())
1228            },
1229        })
1230    }
1231}
1232
1233impl<T: Serialize + DeserializeOwned> SimReceiver<T, TotalOrder, ExactlyOnce> {
1234    /// Receives the next message from the external bincode stream, waiting (and letting the
1235    /// scheduler run any pending simulation work) until one is available. If the simulation
1236    /// becomes quiescent without producing a message, the test fails.
1237    ///
1238    /// This is safe to use in the middle of a test; to observe the *absence* of a message,
1239    /// use [`Self::try_next`] or [`Self::assert_no_more`].
1240    pub fn next(&self) -> impl use<'_, T> + Future<Output = T> {
1241        // Waiting for a message never "overruns" the simulation, even though the scheduler
1242        // may run nondeterministic ticks while we wait: if a message arrives, some pending
1243        // work was necessary to produce it (schedules that run *extra* work are also valid
1244        // executions, explored separately), and if the simulation quiesces instead, the test
1245        // fails right here — so no later observation can be affected by the overrun (the
1246        // taint set by `try_next_impl` is unobservable). See the module docs for the full
1247        // soundness reasoning.
1248        FutureTrackingCaller {
1249            future: async move {
1250                self.try_next_impl().await.ok_or_else(|| {
1251                    "Stream ended (simulation quiescent), but another message was expected"
1252                        .to_owned()
1253                })
1254            },
1255        }
1256    }
1257
1258    /// Receives the next message from the external bincode stream, or returns `None` if no
1259    /// more messages can possibly arrive.
1260    ///
1261    /// If answering requires forcing pending nondeterministic work to run, then afterwards,
1262    /// sending more input and then attempting to receive output will panic. Prefer
1263    /// [`Self::next`] (or [`Self::assert_no_more`]) when possible.
1264    pub async fn try_next(&self) -> Option<T> {
1265        self.try_next_impl().await
1266    }
1267
1268    /// Receives the next `n` messages from the external bincode stream, waiting (and letting
1269    /// the scheduler run any pending simulation work) until they are available. If the
1270    /// simulation becomes quiescent before `n` messages arrive, the test fails.
1271    ///
1272    /// Like [`Self::next`], this is safe to use in the middle of a test. It does not check
1273    /// that the stream ends afterwards; use [`Self::collect_n_only`] for that.
1274    pub fn collect_n<C: Default + Extend<T>>(
1275        &self,
1276        n: usize,
1277    ) -> impl use<'_, T, C> + Future<Output = C> {
1278        FutureTrackingCaller {
1279            future: async move {
1280                let mut out = C::default();
1281                for i in 0..n {
1282                    // Like `next`, waiting for each message is safe mid-test; the taint on a
1283                    // forced `None` is unobservable because the test fails below.
1284                    if let Some(v) = self.try_next_impl().await {
1285                        out.extend([v]);
1286                    } else {
1287                        return Err(format!(
1288                            "Stream ended (simulation quiescent) after {} messages, but {} were expected",
1289                            i, n
1290                        ));
1291                    }
1292                }
1293                Ok(out)
1294            },
1295        }
1296    }
1297
1298    /// Receives the next `n` messages (like [`Self::collect_n`]) and then asserts that the
1299    /// stream ends (like [`Self::assert_no_more`], forking the search in exhaustive mode).
1300    pub async fn collect_n_only<C: Default + Extend<T>>(self, n: usize) -> C
1301    where
1302        T: Debug,
1303    {
1304        let out = self.collect_n(n).await;
1305        self.assert_no_more().await;
1306        out
1307    }
1308
1309    /// Collects all remaining messages from the external bincode stream into a collection,
1310    /// waiting until no more messages can possibly arrive.
1311    ///
1312    /// If this has to force pending nondeterministic work to run, it should be the last
1313    /// observation of the test: afterwards, sending more input and then attempting to
1314    /// receive output will panic. When the number of expected messages is known, prefer
1315    /// [`Self::collect_n`] / [`Self::collect_n_only`].
1316    pub async fn collect<C: Default + Extend<T>>(self) -> C {
1317        let mut out = C::default();
1318        while let Some(v) = self.try_next_impl().await {
1319            out.extend([v]);
1320        }
1321        out
1322    }
1323
1324    /// Asserts that the stream yields exactly the expected sequence of messages, in order.
1325    /// This does not check that the stream ends, use [`Self::assert_yields_only`] for that.
1326    ///
1327    /// Like [`Self::next`], this is safe to use in the middle of a test.
1328    pub fn assert_yields<T2: Debug, I: IntoIterator<Item = T2>>(
1329        &self,
1330        expected: I,
1331    ) -> impl use<'_, T, T2, I> + Future<Output = ()>
1332    where
1333        T: Debug + PartialEq<T2>,
1334    {
1335        FutureTrackingCaller {
1336            future: async {
1337                let mut expected: VecDeque<T2> = expected.into_iter().collect();
1338
1339                while !expected.is_empty() {
1340                    // Like `next`, waiting for each expected message is safe mid-test; the
1341                    // taint on a forced `None` is unobservable because the test fails below.
1342                    if let Some(next) = self.try_next_impl().await {
1343                        let next_expected = expected.pop_front().unwrap();
1344                        if next != next_expected {
1345                            return Err(format!(
1346                                "Stream yielded unexpected message: {:?}, expected: {:?}",
1347                                next, next_expected
1348                            ));
1349                        }
1350                    } else {
1351                        return Err(format!(
1352                            "Stream ended early, still expected: {:?}",
1353                            expected
1354                        ));
1355                    }
1356                }
1357
1358                Ok(())
1359            },
1360        }
1361    }
1362
1363    /// Asserts that the stream yields only the expected sequence of messages, in order,
1364    /// and then ends (like [`Self::assert_no_more`], forking the search in exhaustive mode).
1365    pub fn assert_yields_only<T2: Debug, I: IntoIterator<Item = T2>>(
1366        &self,
1367        expected: I,
1368    ) -> impl use<'_, T, T2, I> + Future<Output = ()>
1369    where
1370        T: Debug + PartialEq<T2>,
1371    {
1372        ChainedFuture {
1373            first: self.assert_yields(expected),
1374            second: self.assert_no_more(),
1375            first_done: false,
1376        }
1377    }
1378}
1379
1380pin_project_lite::pin_project! {
1381    // A future that tracks the location of the `.await` call for better panic messages.
1382    //
1383    // `#[track_caller]` is important for us to create assertion methods because it makes
1384    // the panic backtrace show up at that method (instead of inside the call tree within
1385    // that method). This is e.g. what `Option::unwrap` uses. Unfortunately, `#[track_caller]`
1386    // does not work correctly for async methods (or `dyn Future` either), so we have to
1387    // create these concrete future types that (1) have `#[track_caller]` on their `poll()`
1388    // method and (2) have the `panic!` triggered in their `poll()` method (or in a directly
1389    // nested concrete future).
1390    struct FutureTrackingCaller<F> {
1391        #[pin]
1392        future: F,
1393    }
1394}
1395
1396impl<T, F: Future<Output = Result<T, String>>> Future for FutureTrackingCaller<F> {
1397    type Output = T;
1398
1399    #[track_caller]
1400    fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
1401        match ready!(self.as_mut().project().future.poll(cx)) {
1402            Ok(v) => Poll::Ready(v),
1403            Err(e) => panic!("{}", e),
1404        }
1405    }
1406}
1407
1408pin_project_lite::pin_project! {
1409    // A future that first awaits the first future, then the second, propagating caller info.
1410    //
1411    // See [`FutureTrackingCaller`] for context.
1412    struct ChainedFuture<F1: Future<Output = ()>, F2: Future<Output = ()>> {
1413        #[pin]
1414        first: F1,
1415        #[pin]
1416        second: F2,
1417        first_done: bool,
1418    }
1419}
1420
1421impl<F1: Future<Output = ()>, F2: Future<Output = ()>> Future for ChainedFuture<F1, F2> {
1422    type Output = ();
1423
1424    #[track_caller]
1425    fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
1426        if !self.first_done {
1427            ready!(self.as_mut().project().first.poll(cx));
1428            *self.as_mut().project().first_done = true;
1429        }
1430
1431        self.as_mut().project().second.poll(cx)
1432    }
1433}
1434
1435impl<T: Serialize + DeserializeOwned> SimReceiver<T, NoOrder, ExactlyOnce> {
1436    /// Receives the next `n` messages, sorted, waiting (and letting the scheduler run any
1437    /// pending simulation work) until they are available. If the simulation becomes quiescent
1438    /// before `n` messages arrive, the test fails.
1439    ///
1440    /// Like [`SimReceiver::next`], this is safe to use in the middle of a test.
1441    pub fn collect_n_sorted<C: Default + Extend<T> + AsMut<[T]>>(
1442        &self,
1443        n: usize,
1444    ) -> impl use<'_, T, C> + Future<Output = C>
1445    where
1446        T: Ord,
1447    {
1448        FutureTrackingCaller {
1449            future: async move {
1450                let mut out = C::default();
1451                for i in 0..n {
1452                    // Like `next`, waiting for each message is safe mid-test; the taint on a
1453                    // forced `None` is unobservable because the test fails below.
1454                    if let Some(v) = self.try_next_impl().await {
1455                        out.extend([v]);
1456                    } else {
1457                        return Err(format!(
1458                            "Stream ended (simulation quiescent) after {} messages, but {} were expected",
1459                            i, n
1460                        ));
1461                    }
1462                }
1463                out.as_mut().sort();
1464                Ok(out)
1465            },
1466        }
1467    }
1468
1469    /// Collects all remaining messages from the external bincode stream into a collection,
1470    /// sorting them. This will wait until no more messages can possibly arrive.
1471    ///
1472    /// If this has to force pending nondeterministic work to run, it should be the last
1473    /// observation of the test; see [`collect`](SimReceiver::collect).
1474    pub async fn collect_sorted<C: Default + Extend<T> + AsMut<[T]>>(self) -> C
1475    where
1476        T: Ord,
1477    {
1478        let mut collected = C::default();
1479        while let Some(v) = self.try_next_impl().await {
1480            collected.extend([v]);
1481        }
1482        collected.as_mut().sort();
1483        collected
1484    }
1485
1486    /// Asserts that the stream yields exactly the expected sequence of messages, in some order.
1487    /// This does not check that the stream ends, use [`Self::assert_yields_only_unordered`] for that.
1488    ///
1489    /// Like [`SimReceiver::next`], this is safe to use in the middle of a test.
1490    pub fn assert_yields_unordered<T2: Debug, I: IntoIterator<Item = T2>>(
1491        &self,
1492        expected: I,
1493    ) -> impl use<'_, T, T2, I> + Future<Output = ()>
1494    where
1495        T: Debug + PartialEq<T2>,
1496    {
1497        FutureTrackingCaller {
1498            future: async {
1499                let mut expected: Vec<T2> = expected.into_iter().collect();
1500
1501                while !expected.is_empty() {
1502                    // Like `next`, waiting for each expected message is safe mid-test; the
1503                    // taint on a forced `None` is unobservable because the test fails below.
1504                    if let Some(next) = self.try_next_impl().await {
1505                        let idx = expected.iter().enumerate().find(|(_, e)| &next == *e);
1506                        if let Some((i, _)) = idx {
1507                            expected.swap_remove(i);
1508                        } else {
1509                            return Err(format!("Stream yielded unexpected message: {:?}", next));
1510                        }
1511                    } else {
1512                        return Err(format!(
1513                            "Stream ended early, still expected: {:?}",
1514                            expected
1515                        ));
1516                    }
1517                }
1518
1519                Ok(())
1520            },
1521        }
1522    }
1523
1524    /// Asserts that the stream yields only the expected sequence of messages, in some order,
1525    /// and then ends (like [`Self::assert_no_more`], forking the search in exhaustive mode).
1526    pub fn assert_yields_only_unordered<T2: Debug, I: IntoIterator<Item = T2>>(
1527        &self,
1528        expected: I,
1529    ) -> impl use<'_, T, T2, I> + Future<Output = ()>
1530    where
1531        T: Debug + PartialEq<T2>,
1532    {
1533        ChainedFuture {
1534            first: self.assert_yields_unordered(expected),
1535            second: self.assert_no_more(),
1536            first_done: false,
1537        }
1538    }
1539}
1540
1541impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> SimSender<T, O, R> {
1542    fn with_sink<Out>(&self, thunk: impl FnOnce(&dyn Fn(T)) -> Out) -> Out {
1543        let (sender, quiescence) = CURRENT_SIM_CONNECTIONS.with(|connections| {
1544            let connections = connections.borrow();
1545            (
1546                connections
1547                    .input_senders
1548                    .get(connections.external_registered.get(&self.0).unwrap())
1549                    .unwrap()
1550                    .clone(),
1551                connections.quiescence.clone(),
1552            )
1553        });
1554
1555        thunk(&move |t| {
1556            sender
1557                .try_send(bincode::serialize(&t).unwrap().into())
1558                .unwrap();
1559            quiescence.resume();
1560        })
1561    }
1562}
1563
1564impl<T: Serialize + DeserializeOwned, O: Ordering> SimSender<T, O, ExactlyOnce> {
1565    /// Sends several messages to the external bincode sink. The messages will be asynchronously
1566    /// processed as part of the simulation, in non-deterministic order.
1567    pub fn send_many_unordered<I: IntoIterator<Item = T>>(&self, iter: I) {
1568        self.with_sink(|send| {
1569            for t in iter {
1570                send(t);
1571            }
1572        })
1573    }
1574}
1575
1576impl<T: Serialize + DeserializeOwned> SimSender<T, TotalOrder, ExactlyOnce> {
1577    /// Sends a message to the external bincode sink. The message will be asynchronously processed
1578    /// as part of the simulation.
1579    pub fn send(&self, t: T) {
1580        self.with_sink(|send| send(t));
1581    }
1582
1583    /// Sends several messages to the external bincode sink. The messages will be asynchronously
1584    /// processed as part of the simulation.
1585    pub fn send_many<I: IntoIterator<Item = T>>(&self, iter: I) {
1586        self.with_sink(|send| {
1587            for t in iter {
1588                send(t);
1589            }
1590        })
1591    }
1592}
1593
1594impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> Clone
1595    for SimClusterReceiver<T, O, R>
1596{
1597    fn clone(&self) -> Self {
1598        *self
1599    }
1600}
1601
1602impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> Copy
1603    for SimClusterReceiver<T, O, R>
1604{
1605}
1606
1607impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> SimClusterReceiver<T, O, R> {
1608    fn member_connections(
1609        &self,
1610        member_id: u32,
1611    ) -> (Rc<Mutex<UnsyncReceiver<Bytes>>>, Rc<QuiescenceState>) {
1612        CURRENT_SIM_CONNECTIONS.with(|connections| {
1613            let connections = connections.borrow();
1614            let port = connections.external_registered.get(&self.0).unwrap();
1615            let receivers = connections.cluster_output_receivers.get(port).unwrap();
1616            (
1617                receivers[&member_id].clone(),
1618                connections.quiescence.clone(),
1619            )
1620        })
1621    }
1622
1623    /// See [`try_next_bytes`].
1624    async fn try_next_impl(&self, member_id: u32) -> Option<T> {
1625        let (receiver, quiescence) = self.member_connections(member_id);
1626        try_next_bytes(&receiver, &quiescence)
1627            .await
1628            .map(|bytes| bincode::deserialize(&bytes).unwrap())
1629    }
1630}
1631
1632impl<T: Serialize + DeserializeOwned> SimClusterReceiver<T, TotalOrder, ExactlyOnce> {
1633    /// Receives the next value from a specific cluster member, waiting (and letting the
1634    /// scheduler run any pending simulation work) until one is available. If the simulation
1635    /// becomes quiescent without producing a value, the test fails.
1636    ///
1637    /// This is safe to use in the middle of a test; to observe the *absence* of a value,
1638    /// use [`Self::try_next`].
1639    pub fn next(&self, member_id: u32) -> impl use<'_, T> + Future<Output = T> {
1640        // See `SimReceiver::next` for why waiting for a value never "overruns" the
1641        // simulation.
1642        FutureTrackingCaller {
1643            future: async move {
1644                self.try_next_impl(member_id).await.ok_or_else(|| {
1645                    "Stream ended (simulation quiescent), but another message was expected"
1646                        .to_owned()
1647                })
1648            },
1649        }
1650    }
1651
1652    /// Receives the next value from a specific cluster member, or returns `None` if no more
1653    /// values can possibly arrive.
1654    ///
1655    /// If answering requires forcing pending nondeterministic work to run, then afterwards,
1656    /// sending more input and then attempting to receive output will panic. Prefer
1657    /// [`Self::next`] when possible.
1658    pub async fn try_next(&self, member_id: u32) -> Option<T> {
1659        self.try_next_impl(member_id).await
1660    }
1661
1662    /// Collects all remaining values from a specific cluster member into a collection,
1663    /// waiting until no more values can possibly arrive.
1664    ///
1665    /// If this has to force pending nondeterministic work to run, it should be the last
1666    /// observation of the test; see [`SimReceiver::collect`].
1667    pub async fn collect<C: Default + Extend<T>>(self, member_id: u32) -> C {
1668        let mut out = C::default();
1669        while let Some(v) = self.try_next_impl(member_id).await {
1670            out.extend([v]);
1671        }
1672        out
1673    }
1674}
1675
1676impl<T: Serialize + DeserializeOwned> SimClusterReceiver<T, NoOrder, ExactlyOnce> {
1677    /// Receives the next `n` values from a specific cluster member, sorted, waiting (and
1678    /// letting the scheduler run any pending simulation work) until they are available. If
1679    /// the simulation becomes quiescent before `n` values arrive, the test fails.
1680    ///
1681    /// Like [`SimReceiver::next`], this is safe to use in the middle of a test.
1682    pub fn collect_n_sorted<C: Default + Extend<T> + AsMut<[T]>>(
1683        &self,
1684        member_id: u32,
1685        n: usize,
1686    ) -> impl use<'_, T, C> + Future<Output = C>
1687    where
1688        T: Ord,
1689    {
1690        FutureTrackingCaller {
1691            future: async move {
1692                let mut out = C::default();
1693                for i in 0..n {
1694                    // Like `SimReceiver::next`, waiting for each message is safe mid-test;
1695                    // the taint on a forced `None` is unobservable because the test fails
1696                    // below.
1697                    if let Some(v) = self.try_next_impl(member_id).await {
1698                        out.extend([v]);
1699                    } else {
1700                        return Err(format!(
1701                            "Stream ended (simulation quiescent) after {} messages, but {} were expected",
1702                            i, n
1703                        ));
1704                    }
1705                }
1706                out.as_mut().sort();
1707                Ok(out)
1708            },
1709        }
1710    }
1711
1712    /// Collects all remaining values from a specific cluster member, sorted, waiting until no
1713    /// more values can possibly arrive.
1714    ///
1715    /// If this has to force pending nondeterministic work to run, it should be the last
1716    /// observation of the test; see [`SimReceiver::collect`].
1717    pub async fn collect_sorted<C: Default + Extend<T> + AsMut<[T]>>(self, member_id: u32) -> C
1718    where
1719        T: Ord,
1720    {
1721        let mut collected = C::default();
1722        while let Some(v) = self.try_next_impl(member_id).await {
1723            collected.extend([v]);
1724        }
1725        collected.as_mut().sort();
1726        collected
1727    }
1728}
1729
1730impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> SimClusterSender<T, O, R> {
1731    fn with_sink<Out>(&self, thunk: impl FnOnce(&dyn Fn(u32, T)) -> Out) -> Out {
1732        let (senders, quiescence) = CURRENT_SIM_CONNECTIONS.with(|connections| {
1733            let connections = connections.borrow();
1734            (
1735                connections
1736                    .cluster_input_senders
1737                    .get(connections.external_registered.get(&self.0).unwrap())
1738                    .unwrap()
1739                    .clone(),
1740                connections.quiescence.clone(),
1741            )
1742        });
1743
1744        thunk(&move |member_id: u32, t: T| {
1745            let payload = bincode::serialize(&t).unwrap();
1746            senders[&member_id].try_send(Bytes::from(payload)).unwrap();
1747            quiescence.resume();
1748        })
1749    }
1750}
1751
1752impl<T: Serialize + DeserializeOwned, O: Ordering> SimClusterSender<T, O, ExactlyOnce> {
1753    /// Sends multiple values to specific cluster members. The messages will be asynchronously
1754    /// processed as part of the simulation, in non-deterministic order.
1755    pub fn send_many_unordered<I: IntoIterator<Item = (u32, T)>>(&self, iter: I) {
1756        self.with_sink(|send| {
1757            for (member_id, t) in iter {
1758                send(member_id, t);
1759            }
1760        })
1761    }
1762}
1763
1764impl<T: Serialize + DeserializeOwned> SimClusterSender<T, TotalOrder, ExactlyOnce> {
1765    /// Sends a value to a specific cluster member.
1766    pub fn send(&self, member_id: u32, t: T) {
1767        self.with_sink(|send| send(member_id, t));
1768    }
1769
1770    /// Sends multiple values to specific cluster members.
1771    pub fn send_many<I: IntoIterator<Item = (u32, T)>>(&self, iter: I) {
1772        self.with_sink(|send| {
1773            for (member_id, t) in iter {
1774                send(member_id, t);
1775            }
1776        })
1777    }
1778}
1779
1780enum LogKind<W: std::io::Write> {
1781    Null,
1782    Stderr,
1783    Custom(W),
1784}
1785
1786// via https://www.reddit.com/r/rust/comments/t69sld/is_there_a_way_to_allow_either_stdfmtwrite_or/
1787impl<W: std::io::Write> std::fmt::Write for LogKind<W> {
1788    fn write_str(&mut self, s: &str) -> Result<(), std::fmt::Error> {
1789        match self {
1790            LogKind::Null => Ok(()),
1791            LogKind::Stderr => {
1792                eprint!("{}", s);
1793                Ok(())
1794            }
1795            LogKind::Custom(w) => w.write_all(s.as_bytes()).map_err(|_| std::fmt::Error),
1796        }
1797    }
1798}
1799
1800/// A tick-scoped DFIR together with the hooks that feed it data.
1801struct SimTick {
1802    /// The location of the process/cluster the tick lives on, used to match this tick
1803    /// against the async DFIR that produces its input data.
1804    parent_location: LocationId,
1805    /// The cluster member ID, if the tick lives on a cluster.
1806    cluster_id: Option<u32>,
1807    /// The tick DFIR, executed once per tick.
1808    dfir: DfirErased,
1809    /// Hooks (e.g. from `batch`) resolved *before* the tick runs, deciding what data to
1810    /// release into it.
1811    hooks: Vec<Box<dyn SimHook>>,
1812    /// Hooks (e.g. from `assume_ordering` inside the tick) resolved *while* the tick DFIR
1813    /// is running, via a `tokio::select!` loop, for operators that block on ordering
1814    /// decisions mid-tick.
1815    inline_hooks: Vec<Box<dyn SimInlineHook>>,
1816}
1817
1818impl SimTick {
1819    /// Whether the scheduler can execute this tick right now.
1820    fn can_run(&self) -> bool {
1821        // All hooks must be ready (have received input or have a last value)...
1822        self.hooks.iter().all(|hook| hook.is_ready())
1823            // ...and at least one hook must be able to release data into the tick.
1824            && self.hooks.iter().any(|hook| hook_can_release(&**hook))
1825    }
1826}
1827
1828/// A top-level location whose hooks (e.g. from `assume_ordering` on a non-tick stream)
1829/// need scheduling decisions, but which has no tick DFIR to execute. The scheduler just
1830/// resolves the hooks.
1831struct SimObservation {
1832    /// The top-level location, used to match this observation against the async DFIR that
1833    /// produces its input data.
1834    location: LocationId,
1835    /// The cluster member ID, if the location is a cluster.
1836    cluster_id: Option<u32>,
1837    /// Hooks resolved when the scheduler selects this observation.
1838    hooks: Vec<Box<dyn SimHook>>,
1839}
1840
1841impl SimObservation {
1842    /// Whether the scheduler can resolve any of this observation's hooks right now.
1843    fn can_run(&self) -> bool {
1844        self.hooks.iter().any(|hook| hook_can_release(&**hook))
1845    }
1846}
1847
1848/// Whether the hook has already decided to release data, or has pending input that would
1849/// allow it to decide to do so.
1850fn hook_can_release(hook: &dyn SimHook) -> bool {
1851    hook.current_decision().unwrap_or(false) || hook.can_make_nontrivial_decision()
1852}
1853
1854/// A running simulation, which manages the async DFIRs, tick DFIRs, and hook-based
1855/// scheduling decisions for non-deterministic operators like `batch` and `assume_ordering`.
1856///
1857/// This struct holds all simulator state across scheduler steps. Each [`Self::step`] performs
1858/// one of three kinds of work:
1859/// - **Async DFIRs**: long-running top-level dataflows (one per process/cluster member) that
1860///   produce data consumed by ticks and observations.
1861/// - **Ticks**: tick-scoped DFIRs that execute a single tick. Before running, their associated
1862///   hooks (e.g. from `batch`) are resolved to decide what data to release into the tick.
1863/// - **Observations**: top-level locations that have hooks (e.g. from `assume_ordering` on a
1864///   non-tick stream) needing decisions, but no tick DFIR to execute. The scheduler just
1865///   resolves their hooks.
1866struct LaunchedSim<W: std::io::Write> {
1867    /// Top-level async DFIRs, one per process/cluster member. These run continuously and
1868    /// produce data that feeds into ticks and observations.
1869    async_dfirs: Vec<(LocationId, Option<u32>, DfirErased)>,
1870    /// Ticks whose parent async DFIR has made progress, so they may be ready to run.
1871    /// The scheduler further filters these by checking whether their hooks have pending decisions.
1872    possibly_ready_ticks: Vec<SimTick>,
1873    /// Ticks whose parent async DFIR has not yet made progress since they were last checked.
1874    not_ready_ticks: Vec<SimTick>,
1875    /// Observations whose async DFIR has made progress, so their hooks may have decisions
1876    /// to resolve.
1877    possibly_ready_observations: Vec<SimObservation>,
1878    /// Observations whose async DFIR has not yet made progress since they were last checked.
1879    not_ready_observations: Vec<SimObservation>,
1880    log: LogKind<W>,
1881    /// Represents quiescence state of the simulation.
1882    quiescence: Rc<QuiescenceState>,
1883}
1884
1885impl<W: std::io::Write> LaunchedSim<W> {
1886    /// Runs a single step of the simulation scheduler.
1887    ///
1888    /// A step first advances all async DFIRs; if none of them made progress, it instead runs
1889    /// one ready tick or resolves one ready observation. If nothing at all can make progress,
1890    /// the simulation is quiescent: this signals waiting receivers and returns; the driver is
1891    /// responsible for parking until new external input arrives (see
1892    /// [`QuiescenceState::resumed`]).
1893    ///
1894    /// This future is always awaited to completion by the driver, so a step is atomic: user
1895    /// code never runs (and never observes intermediate state) while a step is in flight.
1896    async fn step(&mut self) {
1897        let mut any_made_progress = false;
1898        for (loc, c_id, dfir) in &mut self.async_dfirs {
1899            if dfir.run_tick().await {
1900                any_made_progress = true;
1901
1902                // This async DFIR may have produced new data, so the ticks and observations
1903                // it feeds may now be ready.
1904                self.possibly_ready_ticks
1905                    .extend(self.not_ready_ticks.extract_if(.., |tick| {
1906                        tick.parent_location == *loc && tick.cluster_id == *c_id
1907                    }));
1908                self.possibly_ready_observations.extend(
1909                    self.not_ready_observations
1910                        .extract_if(.., |obs| obs.location == *loc && obs.cluster_id == *c_id),
1911                );
1912            }
1913        }
1914
1915        if any_made_progress {
1916            return;
1917        }
1918
1919        use bolero::generator::*;
1920
1921        // Send anything that can't make a scheduling decision back to the not-ready lists.
1922        self.not_ready_ticks.extend(
1923            self.possibly_ready_ticks
1924                .extract_if(.., |tick| !tick.can_run()),
1925        );
1926        self.not_ready_observations.extend(
1927            self.possibly_ready_observations
1928                .extract_if(.., |obs| !obs.can_run()),
1929        );
1930
1931        if self.possibly_ready_ticks.is_empty() && self.possibly_ready_observations.is_empty() {
1932            // If any tick is blocked because a hook is not ready, that's a
1933            // simulator bug — it means a singleton never received a value.
1934            for tick in &self.not_ready_ticks {
1935                abort_assert!(
1936                    tick.hooks.iter().all(|hook| hook.is_ready()),
1937                    "tick has a hook that never became ready"
1938                );
1939            }
1940
1941            // Signal quiescence, waking receivers waiting for data (their streams end). The
1942            // driver is responsible for parking until new input arrives.
1943            self.quiescence.enter_quiescence();
1944        } else if self.quiescence.pause_nondet.get() > 0 {
1945            // The test is querying whether the simulation can quiesce without
1946            // nondeterministic work (see `SettlePauseGuard::poll_settle`). Report that
1947            // ticks/observations are pending and pause; the driver parks until the test
1948            // decides how to proceed.
1949            self.quiescence.nondet_pending.set(true);
1950            self.quiescence.wake_settled();
1951        } else {
1952            let next_tick_or_obs = (0..(self.possibly_ready_ticks.len()
1953                + self.possibly_ready_observations.len()))
1954                .any();
1955
1956            if next_tick_or_obs < self.possibly_ready_ticks.len() {
1957                let mut tick = self.possibly_ready_ticks.remove(next_tick_or_obs);
1958
1959                match &mut self.log {
1960                    LogKind::Null => {}
1961                    LogKind::Stderr => {
1962                        if let Some(cid) = &tick.cluster_id {
1963                            eprintln!(
1964                                "\n{}",
1965                                format!("Running Tick (Cluster Member {})", cid)
1966                                    .color(colored::Color::Magenta)
1967                                    .bold()
1968                            )
1969                        } else {
1970                            eprintln!("\n{}", "Running Tick".color(colored::Color::Magenta).bold())
1971                        }
1972                    }
1973                    LogKind::Custom(writer) => {
1974                        writeln!(
1975                            writer,
1976                            "\n{}",
1977                            "Running Tick".color(colored::Color::Magenta).bold()
1978                        )
1979                        .unwrap();
1980                    }
1981                }
1982
1983                let mut asterisk_indenter = |_line_no, write: &mut dyn std::fmt::Write| {
1984                    write.write_str(&"*".color(colored::Color::Magenta).bold())?;
1985                    write.write_str(" ")
1986                };
1987
1988                let mut tick_decision_writer = (!matches!(self.log, LogKind::Null)).then(|| {
1989                    indenter::indented(&mut self.log).with_format(indenter::Format::Custom {
1990                        inserter: &mut asterisk_indenter,
1991                    })
1992                });
1993
1994                run_hooks(tick_decision_writer.as_mut(), &mut tick.hooks);
1995
1996                let run_tick_future = tick.dfir.run_tick();
1997                if !tick.inline_hooks.is_empty() {
1998                    let mut run_tick_future_pinned = pin!(run_tick_future);
1999
2000                    loop {
2001                        tokio::select! {
2002                            biased;
2003                            r = &mut run_tick_future_pinned => {
2004                                abort_assert!(r, "tick DFIR run_tick() returned false");
2005                                break;
2006                            }
2007                            _ = async {} => {
2008                                bolero_generator::any::scope::borrow_with(|driver| {
2009                                    for hook in tick.inline_hooks.iter_mut() {
2010                                        if hook.pending_decision() {
2011                                            if !hook.has_decision() {
2012                                                hook.autonomous_decision(driver);
2013                                            }
2014
2015                                            hook.release_decision(
2016                                                tick_decision_writer
2017                                                    .as_mut()
2018                                                    .map(|w| w as &mut dyn std::fmt::Write),
2019                                            );
2020                                        }
2021                                    }
2022                                });
2023                            }
2024                        }
2025                    }
2026                } else {
2027                    abort_assert!(run_tick_future.await, "tick DFIR run_tick() returned false");
2028                }
2029
2030                self.possibly_ready_ticks.push(tick);
2031            } else {
2032                let next_obs = next_tick_or_obs - self.possibly_ready_ticks.len();
2033                let log_writer = (!matches!(self.log, LogKind::Null)).then_some(&mut self.log);
2034                run_hooks(
2035                    log_writer,
2036                    &mut self.possibly_ready_observations[next_obs].hooks,
2037                );
2038            }
2039        }
2040    }
2041}
2042
2043fn run_hooks<W: std::fmt::Write>(
2044    mut tick_decision_writer: Option<&mut W>,
2045    hooks: &mut [Box<dyn SimHook>],
2046) {
2047    let mut remaining_decision_count = hooks.len();
2048    let mut made_nontrivial_decision = false;
2049
2050    bolero::generator::bolero_generator::any::scope::borrow_with(|driver| {
2051        // first, scan manual decisions
2052        hooks.iter_mut().for_each(|hook| {
2053            if let Some(is_nontrivial) = hook.current_decision() {
2054                made_nontrivial_decision |= is_nontrivial;
2055                remaining_decision_count -= 1;
2056            } else if !hook.can_make_nontrivial_decision() {
2057                // if no nontrivial decision is possible, make a trivial one
2058                // (we need to do this in the first pass to force nontrivial decisions
2059                // on the remaining hooks)
2060                hook.autonomous_decision(driver, false);
2061                remaining_decision_count -= 1;
2062            }
2063        });
2064
2065        hooks.iter_mut().for_each(|hook| {
2066            if hook.current_decision().is_none() {
2067                made_nontrivial_decision |= hook.autonomous_decision(
2068                    driver,
2069                    !made_nontrivial_decision && remaining_decision_count == 1,
2070                );
2071                remaining_decision_count -= 1;
2072            }
2073
2074            hook.release_decision(
2075                tick_decision_writer
2076                    .as_deref_mut()
2077                    .map(|w| w as &mut dyn std::fmt::Write),
2078            );
2079        });
2080    });
2081}