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, and then asserts that the stream ends (like
1437    /// [`SimReceiver::assert_no_more`], forking the search in exhaustive mode). If the
1438    /// simulation becomes quiescent before `n` messages arrive, the test fails.
1439    ///
1440    /// Unlike [`collect_n`](SimReceiver::collect_n) on ordered streams, there is no variant
1441    /// of this API that skips the end-of-stream check. On an unordered stream, the set of
1442    /// messages that arrives *first* is not well-defined, so observing a strict prefix of
1443    /// the output would be sensitive to arrival orders that the simulator does not explore
1444    /// (delivery into the port is FIFO, with no ordering hook); sorting normalizes the
1445    /// permutation of the received messages, but not the choice of *subset*. The quiescence
1446    /// check makes the observation sound: it proves the `n` messages are *all* the messages
1447    /// the program can produce from the input so far, a set which does not depend on
1448    /// arrival order.
1449    pub async fn collect_n_sorted_only<C: Default + Extend<T> + AsMut<[T]>>(self, n: usize) -> C
1450    where
1451        T: Debug + Ord,
1452    {
1453        let out = FutureTrackingCaller {
1454            future: async move {
1455                let mut out = C::default();
1456                for i in 0..n {
1457                    // Like `next`, waiting for each message is safe mid-test; the taint on a
1458                    // forced `None` is unobservable because the test fails below.
1459                    if let Some(v) = self.try_next_impl().await {
1460                        out.extend([v]);
1461                    } else {
1462                        return Err(format!(
1463                            "Stream ended (simulation quiescent) after {} messages, but {} were expected",
1464                            i, n
1465                        ));
1466                    }
1467                }
1468                out.as_mut().sort();
1469                Ok(out)
1470            },
1471        }
1472        .await;
1473        self.assert_no_more().await;
1474        out
1475    }
1476
1477    /// Receives the next message, and then asserts that the stream ends (like
1478    /// [`SimReceiver::assert_no_more`], forking the search in exhaustive mode). If the
1479    /// simulation becomes quiescent without producing a message, the test fails.
1480    ///
1481    /// This is a shortcut for [`Self::collect_n_sorted_only`] with `n = 1`. Unlike
1482    /// [`next`](SimReceiver::next) on ordered streams, there is no variant that skips the
1483    /// end-of-stream check, because on an unordered stream *which* message arrives first is
1484    /// not well-defined; the check proves the message is the *only* one the program can
1485    /// produce from the input so far.
1486    pub async fn next_only(self) -> T
1487    where
1488        T: Debug + Ord,
1489    {
1490        let mut out: Vec<T> = self.collect_n_sorted_only(1).await;
1491        out.remove(0)
1492    }
1493
1494    /// Collects all remaining messages from the external bincode stream into a collection,
1495    /// sorting them. This will wait until no more messages can possibly arrive.
1496    ///
1497    /// If this has to force pending nondeterministic work to run, it should be the last
1498    /// observation of the test; see [`collect`](SimReceiver::collect).
1499    pub async fn collect_sorted<C: Default + Extend<T> + AsMut<[T]>>(self) -> C
1500    where
1501        T: Ord,
1502    {
1503        let mut collected = C::default();
1504        while let Some(v) = self.try_next_impl().await {
1505            collected.extend([v]);
1506        }
1507        collected.as_mut().sort();
1508        collected
1509    }
1510
1511    /// Asserts that the stream yields exactly the expected sequence of messages, in some order.
1512    /// This does not check that the stream ends, use [`Self::assert_yields_only_unordered`] for that.
1513    ///
1514    /// Like [`SimReceiver::next`], this is safe to use in the middle of a test.
1515    pub fn assert_yields_unordered<T2: Debug, I: IntoIterator<Item = T2>>(
1516        &self,
1517        expected: I,
1518    ) -> impl use<'_, T, T2, I> + Future<Output = ()>
1519    where
1520        T: Debug + PartialEq<T2>,
1521    {
1522        FutureTrackingCaller {
1523            future: async {
1524                let mut expected: Vec<T2> = expected.into_iter().collect();
1525
1526                while !expected.is_empty() {
1527                    // Like `next`, waiting for each expected message is safe mid-test; the
1528                    // taint on a forced `None` is unobservable because the test fails below.
1529                    if let Some(next) = self.try_next_impl().await {
1530                        let idx = expected.iter().enumerate().find(|(_, e)| &next == *e);
1531                        if let Some((i, _)) = idx {
1532                            expected.swap_remove(i);
1533                        } else {
1534                            return Err(format!("Stream yielded unexpected message: {:?}", next));
1535                        }
1536                    } else {
1537                        return Err(format!(
1538                            "Stream ended early, still expected: {:?}",
1539                            expected
1540                        ));
1541                    }
1542                }
1543
1544                Ok(())
1545            },
1546        }
1547    }
1548
1549    /// Asserts that the stream yields only the expected sequence of messages, in some order,
1550    /// and then ends (like [`Self::assert_no_more`], forking the search in exhaustive mode).
1551    pub fn assert_yields_only_unordered<T2: Debug, I: IntoIterator<Item = T2>>(
1552        &self,
1553        expected: I,
1554    ) -> impl use<'_, T, T2, I> + Future<Output = ()>
1555    where
1556        T: Debug + PartialEq<T2>,
1557    {
1558        ChainedFuture {
1559            first: self.assert_yields_unordered(expected),
1560            second: self.assert_no_more(),
1561            first_done: false,
1562        }
1563    }
1564}
1565
1566impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> SimSender<T, O, R> {
1567    fn with_sink<Out>(&self, thunk: impl FnOnce(&dyn Fn(T)) -> Out) -> Out {
1568        let (sender, quiescence) = CURRENT_SIM_CONNECTIONS.with(|connections| {
1569            let connections = connections.borrow();
1570            (
1571                connections
1572                    .input_senders
1573                    .get(connections.external_registered.get(&self.0).unwrap())
1574                    .unwrap()
1575                    .clone(),
1576                connections.quiescence.clone(),
1577            )
1578        });
1579
1580        thunk(&move |t| {
1581            sender
1582                .try_send(bincode::serialize(&t).unwrap().into())
1583                .unwrap();
1584            quiescence.resume();
1585        })
1586    }
1587}
1588
1589impl<T: Serialize + DeserializeOwned, O: Ordering> SimSender<T, O, ExactlyOnce> {
1590    /// Sends several messages to the external bincode sink. The messages will be asynchronously
1591    /// processed as part of the simulation, in non-deterministic order.
1592    pub fn send_many_unordered<I: IntoIterator<Item = T>>(&self, iter: I) {
1593        self.with_sink(|send| {
1594            for t in iter {
1595                send(t);
1596            }
1597        })
1598    }
1599}
1600
1601impl<T: Serialize + DeserializeOwned> SimSender<T, TotalOrder, ExactlyOnce> {
1602    /// Sends a message to the external bincode sink. The message will be asynchronously processed
1603    /// as part of the simulation.
1604    pub fn send(&self, t: T) {
1605        self.with_sink(|send| send(t));
1606    }
1607
1608    /// Sends several messages to the external bincode sink. The messages will be asynchronously
1609    /// processed as part of the simulation.
1610    pub fn send_many<I: IntoIterator<Item = T>>(&self, iter: I) {
1611        self.with_sink(|send| {
1612            for t in iter {
1613                send(t);
1614            }
1615        })
1616    }
1617}
1618
1619impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> Clone
1620    for SimClusterReceiver<T, O, R>
1621{
1622    fn clone(&self) -> Self {
1623        *self
1624    }
1625}
1626
1627impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> Copy
1628    for SimClusterReceiver<T, O, R>
1629{
1630}
1631
1632impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> SimClusterReceiver<T, O, R> {
1633    fn member_connections(
1634        &self,
1635        member_id: u32,
1636    ) -> (Rc<Mutex<UnsyncReceiver<Bytes>>>, Rc<QuiescenceState>) {
1637        CURRENT_SIM_CONNECTIONS.with(|connections| {
1638            let connections = connections.borrow();
1639            let port = connections.external_registered.get(&self.0).unwrap();
1640            let receivers = connections.cluster_output_receivers.get(port).unwrap();
1641            (
1642                receivers[&member_id].clone(),
1643                connections.quiescence.clone(),
1644            )
1645        })
1646    }
1647
1648    /// See [`try_next_bytes`].
1649    async fn try_next_impl(&self, member_id: u32) -> Option<T> {
1650        let (receiver, quiescence) = self.member_connections(member_id);
1651        try_next_bytes(&receiver, &quiescence)
1652            .await
1653            .map(|bytes| bincode::deserialize(&bytes).unwrap())
1654    }
1655
1656    /// Asserts that the stream from a specific cluster member has ended and no more messages
1657    /// can possibly arrive.
1658    ///
1659    /// If the check cannot be answered without running pending nondeterministic work (such
1660    /// as ticks with buffered inputs):
1661    /// - Under [`CompiledSim::exhaustive`], the search forks: one instance performs the
1662    ///   check and ends there, while sibling instances skip the check and continue.
1663    /// - In other modes, the pending work runs; afterwards, sending more input and then
1664    ///   attempting to receive output will panic.
1665    pub fn assert_no_more(self, member_id: u32) -> impl Future<Output = ()>
1666    where
1667        T: Debug,
1668    {
1669        QuiescenceCheckFuture::new(FutureTrackingCaller {
1670            future: async move {
1671                if let Some(next) = self.try_next_impl(member_id).await {
1672                    return Err(format!(
1673                        "Stream yielded unexpected message: {:?}, expected termination",
1674                        next
1675                    ));
1676                }
1677                Ok(())
1678            },
1679        })
1680    }
1681}
1682
1683impl<T: Serialize + DeserializeOwned> SimClusterReceiver<T, TotalOrder, ExactlyOnce> {
1684    /// Receives the next value from a specific cluster member, waiting (and letting the
1685    /// scheduler run any pending simulation work) until one is available. If the simulation
1686    /// becomes quiescent without producing a value, the test fails.
1687    ///
1688    /// This is safe to use in the middle of a test; to observe the *absence* of a value,
1689    /// use [`Self::try_next`].
1690    pub fn next(&self, member_id: u32) -> impl use<'_, T> + Future<Output = T> {
1691        // See `SimReceiver::next` for why waiting for a value never "overruns" the
1692        // simulation.
1693        FutureTrackingCaller {
1694            future: async move {
1695                self.try_next_impl(member_id).await.ok_or_else(|| {
1696                    "Stream ended (simulation quiescent), but another message was expected"
1697                        .to_owned()
1698                })
1699            },
1700        }
1701    }
1702
1703    /// Receives the next value from a specific cluster member, or returns `None` if no more
1704    /// values can possibly arrive.
1705    ///
1706    /// If answering requires forcing pending nondeterministic work to run, then afterwards,
1707    /// sending more input and then attempting to receive output will panic. Prefer
1708    /// [`Self::next`] when possible.
1709    pub async fn try_next(&self, member_id: u32) -> Option<T> {
1710        self.try_next_impl(member_id).await
1711    }
1712
1713    /// Collects all remaining values from a specific cluster member into a collection,
1714    /// waiting until no more values can possibly arrive.
1715    ///
1716    /// If this has to force pending nondeterministic work to run, it should be the last
1717    /// observation of the test; see [`SimReceiver::collect`].
1718    pub async fn collect<C: Default + Extend<T>>(self, member_id: u32) -> C {
1719        let mut out = C::default();
1720        while let Some(v) = self.try_next_impl(member_id).await {
1721            out.extend([v]);
1722        }
1723        out
1724    }
1725}
1726
1727impl<T: Serialize + DeserializeOwned> SimClusterReceiver<T, NoOrder, ExactlyOnce> {
1728    /// Receives the next `n` values from a specific cluster member, sorted, and then
1729    /// asserts that the stream ends (like [`Self::assert_no_more`], forking the search in
1730    /// exhaustive mode). If the simulation becomes quiescent before `n` values arrive, the
1731    /// test fails.
1732    ///
1733    /// There is no variant of this API that skips the end-of-stream check; see
1734    /// [`SimReceiver::collect_n_sorted_only`] for why observing a strict prefix of an
1735    /// unordered stream would be unsound.
1736    pub async fn collect_n_sorted_only<C: Default + Extend<T> + AsMut<[T]>>(
1737        self,
1738        member_id: u32,
1739        n: usize,
1740    ) -> C
1741    where
1742        T: Debug + Ord,
1743    {
1744        let out = FutureTrackingCaller {
1745            future: async move {
1746                let mut out = C::default();
1747                for i in 0..n {
1748                    // Like `SimReceiver::next`, waiting for each message is safe mid-test;
1749                    // the taint on a forced `None` is unobservable because the test fails
1750                    // below.
1751                    if let Some(v) = self.try_next_impl(member_id).await {
1752                        out.extend([v]);
1753                    } else {
1754                        return Err(format!(
1755                            "Stream ended (simulation quiescent) after {} messages, but {} were expected",
1756                            i, n
1757                        ));
1758                    }
1759                }
1760                out.as_mut().sort();
1761                Ok(out)
1762            },
1763        }
1764        .await;
1765        self.assert_no_more(member_id).await;
1766        out
1767    }
1768
1769    /// Receives the next value from a specific cluster member, and then asserts that the
1770    /// stream ends (like [`Self::assert_no_more`], forking the search in exhaustive mode).
1771    /// If the simulation becomes quiescent without producing a value, the test fails.
1772    ///
1773    /// This is a shortcut for [`Self::collect_n_sorted_only`] with `n = 1`; see
1774    /// [`SimReceiver::next_only`] for why there is no variant that skips the end-of-stream
1775    /// check.
1776    pub async fn next_only(self, member_id: u32) -> T
1777    where
1778        T: Debug + Ord,
1779    {
1780        let mut out: Vec<T> = self.collect_n_sorted_only(member_id, 1).await;
1781        out.remove(0)
1782    }
1783
1784    /// Collects all remaining values from a specific cluster member, sorted, waiting until no
1785    /// more values can possibly arrive.
1786    ///
1787    /// If this has to force pending nondeterministic work to run, it should be the last
1788    /// observation of the test; see [`SimReceiver::collect`].
1789    pub async fn collect_sorted<C: Default + Extend<T> + AsMut<[T]>>(self, member_id: u32) -> C
1790    where
1791        T: Ord,
1792    {
1793        let mut collected = C::default();
1794        while let Some(v) = self.try_next_impl(member_id).await {
1795            collected.extend([v]);
1796        }
1797        collected.as_mut().sort();
1798        collected
1799    }
1800}
1801
1802impl<T: Serialize + DeserializeOwned, O: Ordering, R: Retries> SimClusterSender<T, O, R> {
1803    fn with_sink<Out>(&self, thunk: impl FnOnce(&dyn Fn(u32, T)) -> Out) -> Out {
1804        let (senders, quiescence) = CURRENT_SIM_CONNECTIONS.with(|connections| {
1805            let connections = connections.borrow();
1806            (
1807                connections
1808                    .cluster_input_senders
1809                    .get(connections.external_registered.get(&self.0).unwrap())
1810                    .unwrap()
1811                    .clone(),
1812                connections.quiescence.clone(),
1813            )
1814        });
1815
1816        thunk(&move |member_id: u32, t: T| {
1817            let payload = bincode::serialize(&t).unwrap();
1818            senders[&member_id].try_send(Bytes::from(payload)).unwrap();
1819            quiescence.resume();
1820        })
1821    }
1822}
1823
1824impl<T: Serialize + DeserializeOwned, O: Ordering> SimClusterSender<T, O, ExactlyOnce> {
1825    /// Sends multiple values to specific cluster members. The messages will be asynchronously
1826    /// processed as part of the simulation, in non-deterministic order.
1827    pub fn send_many_unordered<I: IntoIterator<Item = (u32, T)>>(&self, iter: I) {
1828        self.with_sink(|send| {
1829            for (member_id, t) in iter {
1830                send(member_id, t);
1831            }
1832        })
1833    }
1834}
1835
1836impl<T: Serialize + DeserializeOwned> SimClusterSender<T, TotalOrder, ExactlyOnce> {
1837    /// Sends a value to a specific cluster member.
1838    pub fn send(&self, member_id: u32, t: T) {
1839        self.with_sink(|send| send(member_id, t));
1840    }
1841
1842    /// Sends multiple values to specific cluster members.
1843    pub fn send_many<I: IntoIterator<Item = (u32, T)>>(&self, iter: I) {
1844        self.with_sink(|send| {
1845            for (member_id, t) in iter {
1846                send(member_id, t);
1847            }
1848        })
1849    }
1850}
1851
1852enum LogKind<W: std::io::Write> {
1853    Null,
1854    Stderr,
1855    Custom(W),
1856}
1857
1858// via https://www.reddit.com/r/rust/comments/t69sld/is_there_a_way_to_allow_either_stdfmtwrite_or/
1859impl<W: std::io::Write> std::fmt::Write for LogKind<W> {
1860    fn write_str(&mut self, s: &str) -> Result<(), std::fmt::Error> {
1861        match self {
1862            LogKind::Null => Ok(()),
1863            LogKind::Stderr => {
1864                eprint!("{}", s);
1865                Ok(())
1866            }
1867            LogKind::Custom(w) => w.write_all(s.as_bytes()).map_err(|_| std::fmt::Error),
1868        }
1869    }
1870}
1871
1872/// A tick-scoped DFIR together with the hooks that feed it data.
1873struct SimTick {
1874    /// The location of the process/cluster the tick lives on, used to match this tick
1875    /// against the async DFIR that produces its input data.
1876    parent_location: LocationId,
1877    /// The cluster member ID, if the tick lives on a cluster.
1878    cluster_id: Option<u32>,
1879    /// The tick DFIR, executed once per tick.
1880    dfir: DfirErased,
1881    /// Hooks (e.g. from `batch`) resolved *before* the tick runs, deciding what data to
1882    /// release into it.
1883    hooks: Vec<Box<dyn SimHook>>,
1884    /// Hooks (e.g. from `assume_ordering` inside the tick) resolved *while* the tick DFIR
1885    /// is running, via a `tokio::select!` loop, for operators that block on ordering
1886    /// decisions mid-tick.
1887    inline_hooks: Vec<Box<dyn SimInlineHook>>,
1888}
1889
1890impl SimTick {
1891    /// Whether the scheduler can execute this tick right now.
1892    fn can_run(&self) -> bool {
1893        // All hooks must be ready (have received input or have a last value)...
1894        self.hooks.iter().all(|hook| hook.is_ready())
1895            // ...and at least one hook must be able to release data into the tick.
1896            && self.hooks.iter().any(|hook| hook_can_release(&**hook))
1897    }
1898}
1899
1900/// A top-level location whose hooks (e.g. from `assume_ordering` on a non-tick stream)
1901/// need scheduling decisions, but which has no tick DFIR to execute. The scheduler just
1902/// resolves the hooks.
1903struct SimObservation {
1904    /// The top-level location, used to match this observation against the async DFIR that
1905    /// produces its input data.
1906    location: LocationId,
1907    /// The cluster member ID, if the location is a cluster.
1908    cluster_id: Option<u32>,
1909    /// Hooks resolved when the scheduler selects this observation.
1910    hooks: Vec<Box<dyn SimHook>>,
1911}
1912
1913impl SimObservation {
1914    /// Whether the scheduler can resolve any of this observation's hooks right now.
1915    fn can_run(&self) -> bool {
1916        self.hooks.iter().any(|hook| hook_can_release(&**hook))
1917    }
1918}
1919
1920/// Whether the hook has already decided to release data, or has pending input that would
1921/// allow it to decide to do so.
1922fn hook_can_release(hook: &dyn SimHook) -> bool {
1923    hook.current_decision().unwrap_or(false) || hook.can_make_nontrivial_decision()
1924}
1925
1926/// A running simulation, which manages the async DFIRs, tick DFIRs, and hook-based
1927/// scheduling decisions for non-deterministic operators like `batch` and `assume_ordering`.
1928///
1929/// This struct holds all simulator state across scheduler steps. Each [`Self::step`] performs
1930/// one of three kinds of work:
1931/// - **Async DFIRs**: long-running top-level dataflows (one per process/cluster member) that
1932///   produce data consumed by ticks and observations.
1933/// - **Ticks**: tick-scoped DFIRs that execute a single tick. Before running, their associated
1934///   hooks (e.g. from `batch`) are resolved to decide what data to release into the tick.
1935/// - **Observations**: top-level locations that have hooks (e.g. from `assume_ordering` on a
1936///   non-tick stream) needing decisions, but no tick DFIR to execute. The scheduler just
1937///   resolves their hooks.
1938struct LaunchedSim<W: std::io::Write> {
1939    /// Top-level async DFIRs, one per process/cluster member. These run continuously and
1940    /// produce data that feeds into ticks and observations.
1941    async_dfirs: Vec<(LocationId, Option<u32>, DfirErased)>,
1942    /// Ticks whose parent async DFIR has made progress, so they may be ready to run.
1943    /// The scheduler further filters these by checking whether their hooks have pending decisions.
1944    possibly_ready_ticks: Vec<SimTick>,
1945    /// Ticks whose parent async DFIR has not yet made progress since they were last checked.
1946    not_ready_ticks: Vec<SimTick>,
1947    /// Observations whose async DFIR has made progress, so their hooks may have decisions
1948    /// to resolve.
1949    possibly_ready_observations: Vec<SimObservation>,
1950    /// Observations whose async DFIR has not yet made progress since they were last checked.
1951    not_ready_observations: Vec<SimObservation>,
1952    log: LogKind<W>,
1953    /// Represents quiescence state of the simulation.
1954    quiescence: Rc<QuiescenceState>,
1955}
1956
1957impl<W: std::io::Write> LaunchedSim<W> {
1958    /// Runs a single step of the simulation scheduler.
1959    ///
1960    /// A step first advances all async DFIRs; if none of them made progress, it instead runs
1961    /// one ready tick or resolves one ready observation. If nothing at all can make progress,
1962    /// the simulation is quiescent: this signals waiting receivers and returns; the driver is
1963    /// responsible for parking until new external input arrives (see
1964    /// [`QuiescenceState::resumed`]).
1965    ///
1966    /// This future is always awaited to completion by the driver, so a step is atomic: user
1967    /// code never runs (and never observes intermediate state) while a step is in flight.
1968    async fn step(&mut self) {
1969        let mut any_made_progress = false;
1970        for (loc, c_id, dfir) in &mut self.async_dfirs {
1971            if dfir.run_tick().await {
1972                any_made_progress = true;
1973
1974                // This async DFIR may have produced new data, so the ticks and observations
1975                // it feeds may now be ready.
1976                self.possibly_ready_ticks
1977                    .extend(self.not_ready_ticks.extract_if(.., |tick| {
1978                        tick.parent_location == *loc && tick.cluster_id == *c_id
1979                    }));
1980                self.possibly_ready_observations.extend(
1981                    self.not_ready_observations
1982                        .extract_if(.., |obs| obs.location == *loc && obs.cluster_id == *c_id),
1983                );
1984            }
1985        }
1986
1987        if any_made_progress {
1988            return;
1989        }
1990
1991        use bolero::generator::*;
1992
1993        // Send anything that can't make a scheduling decision back to the not-ready lists.
1994        self.not_ready_ticks.extend(
1995            self.possibly_ready_ticks
1996                .extract_if(.., |tick| !tick.can_run()),
1997        );
1998        self.not_ready_observations.extend(
1999            self.possibly_ready_observations
2000                .extract_if(.., |obs| !obs.can_run()),
2001        );
2002
2003        if self.possibly_ready_ticks.is_empty() && self.possibly_ready_observations.is_empty() {
2004            // If any tick is blocked because a hook is not ready, that's a
2005            // simulator bug — it means a singleton never received a value.
2006            for tick in &self.not_ready_ticks {
2007                abort_assert!(
2008                    tick.hooks.iter().all(|hook| hook.is_ready()),
2009                    "tick has a hook that never became ready"
2010                );
2011            }
2012
2013            // Signal quiescence, waking receivers waiting for data (their streams end). The
2014            // driver is responsible for parking until new input arrives.
2015            self.quiescence.enter_quiescence();
2016        } else if self.quiescence.pause_nondet.get() > 0 {
2017            // The test is querying whether the simulation can quiesce without
2018            // nondeterministic work (see `SettlePauseGuard::poll_settle`). Report that
2019            // ticks/observations are pending and pause; the driver parks until the test
2020            // decides how to proceed.
2021            self.quiescence.nondet_pending.set(true);
2022            self.quiescence.wake_settled();
2023        } else {
2024            let next_tick_or_obs = (0..(self.possibly_ready_ticks.len()
2025                + self.possibly_ready_observations.len()))
2026                .any();
2027
2028            if next_tick_or_obs < self.possibly_ready_ticks.len() {
2029                let mut tick = self.possibly_ready_ticks.remove(next_tick_or_obs);
2030
2031                match &mut self.log {
2032                    LogKind::Null => {}
2033                    LogKind::Stderr => {
2034                        if let Some(cid) = &tick.cluster_id {
2035                            eprintln!(
2036                                "\n{}",
2037                                format!("Running Tick (Cluster Member {})", cid)
2038                                    .color(colored::Color::Magenta)
2039                                    .bold()
2040                            )
2041                        } else {
2042                            eprintln!("\n{}", "Running Tick".color(colored::Color::Magenta).bold())
2043                        }
2044                    }
2045                    LogKind::Custom(writer) => {
2046                        writeln!(
2047                            writer,
2048                            "\n{}",
2049                            "Running Tick".color(colored::Color::Magenta).bold()
2050                        )
2051                        .unwrap();
2052                    }
2053                }
2054
2055                let mut asterisk_indenter = |_line_no, write: &mut dyn std::fmt::Write| {
2056                    write.write_str(&"*".color(colored::Color::Magenta).bold())?;
2057                    write.write_str(" ")
2058                };
2059
2060                let mut tick_decision_writer = (!matches!(self.log, LogKind::Null)).then(|| {
2061                    indenter::indented(&mut self.log).with_format(indenter::Format::Custom {
2062                        inserter: &mut asterisk_indenter,
2063                    })
2064                });
2065
2066                run_hooks(tick_decision_writer.as_mut(), &mut tick.hooks);
2067
2068                let run_tick_future = tick.dfir.run_tick();
2069                if !tick.inline_hooks.is_empty() {
2070                    let mut run_tick_future_pinned = pin!(run_tick_future);
2071
2072                    loop {
2073                        tokio::select! {
2074                            biased;
2075                            r = &mut run_tick_future_pinned => {
2076                                abort_assert!(r, "tick DFIR run_tick() returned false");
2077                                break;
2078                            }
2079                            _ = async {} => {
2080                                bolero_generator::any::scope::borrow_with(|driver| {
2081                                    for hook in tick.inline_hooks.iter_mut() {
2082                                        if hook.pending_decision() {
2083                                            if !hook.has_decision() {
2084                                                hook.autonomous_decision(driver);
2085                                            }
2086
2087                                            hook.release_decision(
2088                                                tick_decision_writer
2089                                                    .as_mut()
2090                                                    .map(|w| w as &mut dyn std::fmt::Write),
2091                                            );
2092                                        }
2093                                    }
2094                                });
2095                            }
2096                        }
2097                    }
2098                } else {
2099                    abort_assert!(run_tick_future.await, "tick DFIR run_tick() returned false");
2100                }
2101
2102                self.possibly_ready_ticks.push(tick);
2103            } else {
2104                let next_obs = next_tick_or_obs - self.possibly_ready_ticks.len();
2105                let log_writer = (!matches!(self.log, LogKind::Null)).then_some(&mut self.log);
2106                run_hooks(
2107                    log_writer,
2108                    &mut self.possibly_ready_observations[next_obs].hooks,
2109                );
2110            }
2111        }
2112    }
2113}
2114
2115fn run_hooks<W: std::fmt::Write>(
2116    mut tick_decision_writer: Option<&mut W>,
2117    hooks: &mut [Box<dyn SimHook>],
2118) {
2119    let mut remaining_decision_count = hooks.len();
2120    let mut made_nontrivial_decision = false;
2121
2122    bolero::generator::bolero_generator::any::scope::borrow_with(|driver| {
2123        // first, scan manual decisions
2124        hooks.iter_mut().for_each(|hook| {
2125            if let Some(is_nontrivial) = hook.current_decision() {
2126                made_nontrivial_decision |= is_nontrivial;
2127                remaining_decision_count -= 1;
2128            } else if !hook.can_make_nontrivial_decision() {
2129                // if no nontrivial decision is possible, make a trivial one
2130                // (we need to do this in the first pass to force nontrivial decisions
2131                // on the remaining hooks)
2132                hook.autonomous_decision(driver, false);
2133                remaining_decision_count -= 1;
2134            }
2135        });
2136
2137        hooks.iter_mut().for_each(|hook| {
2138            if hook.current_decision().is_none() {
2139                made_nontrivial_decision |= hook.autonomous_decision(
2140                    driver,
2141                    !made_nontrivial_decision && remaining_decision_count == 1,
2142                );
2143                remaining_decision_count -= 1;
2144            }
2145
2146            hook.release_decision(
2147                tick_decision_writer
2148                    .as_deref_mut()
2149                    .map(|w| w as &mut dyn std::fmt::Write),
2150            );
2151        });
2152    });
2153}