harness/
driver.rs

1//! Shared `DogStatsD` load-driver engine.
2//!
3//! A producer thread samples lines into a bounded channel; a consumer thread
4//! fans each line out to every socket and tallies per-socket sends. Drivers
5//! differ only in how many sockets they target and which anchors they fire, so
6//! both the single-socket and differential drivers run on this one engine.
7
8use std::os::unix::net::UnixDatagram;
9use std::path::Path;
10use std::sync::mpsc::sync_channel;
11use std::thread::{self, sleep};
12use std::time::{Duration, Instant};
13
14use antithesis_sdk::random::{random_choice, AntithesisRng};
15use rand::{rand_core::UnwrapErr, RngExt};
16
17use crate::payload::dogstatsd;
18
19/// Per-batch composition: 50% clean, 25% feral, 25% mixed.
20#[derive(Clone, Copy, Debug)]
21pub enum Batch {
22    /// Every line clean.
23    Clean,
24    /// Every line feral.
25    Feral,
26    /// A per-line clean-or-feral mix.
27    Mixed,
28}
29
30impl Batch {
31    /// Sample a batch composition: half clean, a quarter feral, a quarter mixed.
32    #[must_use]
33    pub fn sample() -> Self {
34        match random_choice(&[Batch::Clean, Batch::Clean, Batch::Feral, Batch::Mixed]) {
35            Some(Batch::Feral) => Batch::Feral,
36            Some(Batch::Mixed) => Batch::Mixed,
37            _ => Batch::Clean,
38        }
39    }
40
41    /// The vibe for one line drawn from this batch.
42    fn vibe(self) -> dogstatsd::Vibe {
43        match self {
44            Batch::Clean => dogstatsd::Vibe::Clean,
45            Batch::Feral => dogstatsd::Vibe::Feral,
46            Batch::Mixed => dogstatsd::sample_vibe(),
47        }
48    }
49}
50
51/// A generated dogstatsd line queued for the sockets.
52enum Line {
53    /// A single-value line.
54    Single { bytes: Vec<u8> },
55    /// A multi-value `:`-packed metric.
56    Multi {
57        /// The encoded line.
58        bytes: Vec<u8>,
59        /// The number of values in the packed run.
60        count: usize,
61    },
62}
63
64impl Line {
65    /// The encoded bytes to ship over a socket.
66    fn bytes(&self) -> &[u8] {
67        match self {
68            Line::Single { bytes } | Line::Multi { bytes, .. } => bytes,
69        }
70    }
71}
72
73/// What a driver run shipped, for anchoring assertions.
74#[derive(Clone, Debug)]
75pub struct Stats {
76    /// Lines pulled from the channel, whether or not any send succeeded.
77    pub received: usize,
78    /// Successful sends per socket, indexed as the sockets were passed to [`run`].
79    pub sent: Vec<usize>,
80    /// Largest packed run that reached each socket, indexed likewise. Zero when
81    /// no multi-value line reached that socket.
82    pub max_packed: Vec<usize>,
83}
84
85/// Drive one batch of sampled `DogStatsD` lines to every socket.
86///
87/// A producer samples up to ~10k lines at the batch's vibe and queues them; a
88/// consumer ships each line to every socket and tallies per-socket sends. The
89/// returned [`Stats`] anchor the caller's assertions.
90///
91/// # Panics
92///
93/// Panics if the producer or consumer thread panics.
94#[must_use]
95pub fn run(batch: Batch, sockets: Vec<UnixDatagram>) -> Stats {
96    let count = {
97        let mut rng = UnwrapErr(AntithesisRng);
98        rng.random_range(0..=10_000u64)
99    };
100
101    let (tx, rx) = sync_channel::<Line>(2024);
102
103    let producer = thread::spawn(move || {
104        let mut rng = UnwrapErr(AntithesisRng);
105        for _ in 0..count {
106            let mut bytes = Vec::new();
107            let line = match dogstatsd::send(&mut rng, &mut bytes, batch.vibe()) {
108                None => Line::Single { bytes },
109                Some(count) => Line::Multi { bytes, count },
110            };
111            if tx.send(line).is_err() {
112                break;
113            }
114        }
115    });
116
117    let consumer = thread::spawn(move || {
118        let mut received = 0usize;
119        let mut sent = vec![0usize; sockets.len()];
120        let mut max_packed = vec![0usize; sockets.len()];
121        while let Ok(line) = rx.recv() {
122            received += 1;
123            for (i, socket) in sockets.iter().enumerate() {
124                if socket.send(line.bytes()).is_ok() {
125                    sent[i] += 1;
126                    if let Line::Multi { count, .. } = &line {
127                        max_packed[i] = max_packed[i].max(*count);
128                    }
129                }
130            }
131        }
132        Stats {
133            received,
134            sent,
135            max_packed,
136        }
137    });
138
139    producer.join().expect("producer thread panicked");
140    consumer.join().expect("consumer thread panicked")
141}
142
143/// Wait for the remote process to bind `path`, intentionally naive. Returns
144/// `None` if the socket is still unavailable after 30 seconds.
145#[must_use]
146pub fn connect_with_retry(path: &Path) -> Option<UnixDatagram> {
147    let deadline = Instant::now() + Duration::from_secs(30);
148    loop {
149        if let Ok(socket) = UnixDatagram::unbound() {
150            if socket.connect(path).is_ok() {
151                return Some(socket);
152            }
153        }
154        if Instant::now() >= deadline {
155            return None;
156        }
157        sleep(Duration::from_millis(250));
158    }
159}
harness/driver.rs

harness/
driver.rs
