mockforge_bench/preflight.rs
1//! Pre-flight latency probe for sizing `--vus` against `--rps`.
2//!
3//! Issue #79 round 8 — Srikanth's reply on 0.3.137 flagged that the
4//! `--vus * 10 < --rps` warning's "1 VU sustains ~10 req/s at 100ms latency"
5//! rule of thumb is wrong for fast targets (~2ms). At 2ms response time, one
6//! VU can drive ~500 req/s, so `--vus 5` is enough for `--rps 1000` but the
7//! static rule incorrectly warns "bump to --vus 100".
8//!
9//! Instead, do a tiny (1-3 request) HTTP probe of the actual target to
10//! measure baseline latency, then derive a more accurate "VUs needed to
11//! sustain rate" estimate. Skip the warning entirely if the measured rate
12//! comfortably covers the requested rate.
13
14use std::time::{Duration, Instant};
15
16/// Result of a pre-flight latency probe.
17#[derive(Debug, Clone, Copy)]
18pub struct ProbeResult {
19 /// Observed average request latency.
20 pub avg_latency: Duration,
21 /// Number of successful probe requests.
22 pub samples: u32,
23}
24
25impl ProbeResult {
26 /// Required VUs to sustain `target_rps` end-to-end with `num_operations`
27 /// per iteration.
28 ///
29 /// Why `num_operations` matters: k6's `constant-arrival-rate` executor
30 /// (set by `--rps`) targets ITERATIONS per second, not requests. Each
31 /// iteration runs every operation in the generated script sequentially.
32 /// So sustaining `--rps R` with a spec of `N` operations actually needs
33 /// `R × N × latency_secs` VUs, not `R × latency_secs`.
34 ///
35 /// Issue #79 round 9 — Srikanth saw "Pre-flight probe: --vus 5 is
36 /// sufficient" followed by k6 emitting "Insufficient VUs" mid-run.
37 /// Cause: he had a ~12-operation spec, so the real iteration time
38 /// was ~12 × measured latency, not 1 ×.
39 ///
40 /// Formula: `rps × num_operations × latency_secs + 1` (safety margin).
41 /// Returns at least 1.
42 pub fn required_vus(&self, target_rps: u32, num_operations: u32) -> u32 {
43 let lat_secs = self.avg_latency.as_secs_f64().max(0.001);
44 let ops = num_operations.max(1) as f64;
45 let raw = (target_rps as f64 * ops * lat_secs).ceil() as u32;
46 raw.saturating_add(1).max(1)
47 }
48}
49
50/// Probe `target` with up to `samples` quick HEAD/GET requests and report
51/// the average successful-response latency. Each probe has a 5s timeout;
52/// failed probes are excluded from the average. Returns `None` if no
53/// probe succeeded.
54///
55/// Used pre-flight to size the `--vus` warning. We deliberately *don't*
56/// fail the bench if probes fail — the target might require auth or be
57/// strict about HEADs. Falling back to the static heuristic is fine.
58pub async fn probe_target_latency(
59 target: &str,
60 samples: u32,
61 skip_tls_verify: bool,
62) -> Option<ProbeResult> {
63 probe_target_latency_guarded(target, samples, skip_tls_verify, crate::ssrf::process_guard())
64 .await
65}
66
67/// [`probe_target_latency`] with an explicit SSRF guard; `None` is unguarded.
68pub async fn probe_target_latency_guarded(
69 target: &str,
70 samples: u32,
71 skip_tls_verify: bool,
72 guard: Option<crate::ssrf::Policy>,
73) -> Option<ProbeResult> {
74 if let Some(policy) = guard {
75 let parsed = url::Url::parse(target).ok()?;
76 crate::ssrf::check_url(&parsed, policy).ok()?;
77 }
78 let client = crate::ssrf::client_builder_for(guard)
79 .timeout(Duration::from_secs(5))
80 .danger_accept_invalid_certs(skip_tls_verify)
81 .build()
82 .ok()?;
83
84 let mut total = Duration::ZERO;
85 let mut count: u32 = 0;
86 for _ in 0..samples {
87 let start = Instant::now();
88 // HEAD first — cheaper for the target. Fall back to GET if HEAD
89 // fails (some servers / WAFs reject HEAD).
90 let head = client.head(target).send().await;
91 let elapsed = match head {
92 Ok(_) => start.elapsed(),
93 Err(_) => {
94 let start = Instant::now();
95 match client.get(target).send().await {
96 Ok(_) => start.elapsed(),
97 Err(_) => continue,
98 }
99 }
100 };
101 total += elapsed;
102 count = count.saturating_add(1);
103 }
104
105 if count == 0 {
106 return None;
107 }
108
109 Some(ProbeResult {
110 avg_latency: total / count,
111 samples: count,
112 })
113}
114
115#[cfg(test)]
116mod tests {
117 use super::*;
118
119 #[test]
120 fn required_vus_scales_with_latency() {
121 let fast = ProbeResult {
122 avg_latency: Duration::from_millis(2),
123 samples: 3,
124 };
125 // Single-op spec: 1000 rps × 1 op × 2ms = 2 VU + 1 margin = 3
126 assert_eq!(fast.required_vus(1000, 1), 3);
127
128 let slow = ProbeResult {
129 avg_latency: Duration::from_millis(100),
130 samples: 3,
131 };
132 // Single-op: 100 rps × 1 op × 100ms = 10 VU + 1 margin = 11
133 assert_eq!(slow.required_vus(100, 1), 11);
134 }
135
136 #[test]
137 fn required_vus_scales_with_operation_count() {
138 // Issue #79 round 9 — Srikanth saw "vus 5 is sufficient" then k6
139 // hit Insufficient VUs because his spec had ~12 operations per
140 // iteration. With 15ms baseline × 12 ops × 100 rps = 18 VUs.
141 let probe = ProbeResult {
142 avg_latency: Duration::from_millis(15),
143 samples: 3,
144 };
145 assert_eq!(probe.required_vus(100, 1), 3); // single op
146 assert_eq!(probe.required_vus(100, 12), 19); // 12 ops + 1 margin
147 }
148
149 #[test]
150 fn required_vus_clamps_to_one() {
151 let fast = ProbeResult {
152 avg_latency: Duration::from_micros(50),
153 samples: 1,
154 };
155 // (1 × 0.001s clamp) × 1 op × 1 rps = 1, +1 margin = 2
156 assert!(fast.required_vus(1, 1) >= 1);
157 }
158
159 #[test]
160 fn required_vus_treats_zero_operations_as_one() {
161 let probe = ProbeResult {
162 avg_latency: Duration::from_millis(10),
163 samples: 1,
164 };
165 // num_operations=0 should clamp to 1 so we never divide-by-zero
166 // upstream or report an impossible "0 VUs" recommendation.
167 assert_eq!(probe.required_vus(100, 0), probe.required_vus(100, 1));
168 }
169
170 #[tokio::test]
171 async fn probe_returns_none_for_unreachable() {
172 // Reserved-for-docs port on a non-routable address — should error
173 // fast (no DNS lookup) without hanging.
174 let result = probe_target_latency("http://127.0.0.1:1/", 1, false).await;
175 assert!(result.is_none());
176 }
177
178 #[tokio::test]
179 async fn guarded_probe_never_reaches_an_internal_target() {
180 use crate::ssrf::{test_server, Policy};
181 use std::sync::atomic::Ordering;
182
183 let (port, hits) = test_server::serve(test_server::OK).await;
184 // Literal loopback and a name that resolves to it.
185 for target in [
186 format!("http://127.0.0.1:{port}/"),
187 format!("http://localhost:{port}/"),
188 ] {
189 let r = probe_target_latency_guarded(&target, 2, false, Some(Policy::strict())).await;
190 assert!(r.is_none(), "{target}");
191 }
192 assert_eq!(hits.load(Ordering::SeqCst), 0);
193
194 // A first hop that redirects inward is refused too. 0.0.0.0 reaches
195 // local listeners on Linux and is blocked even when loopback is
196 // allowed, so the inner server proves whether the hop was followed.
197 let (inner, inner_hits) = test_server::serve(test_server::OK).await;
198 let (outer, _) =
199 test_server::serve(test_server::redirect_to(&format!("http://0.0.0.0:{inner}/"))).await;
200 let outer_url = format!("http://127.0.0.1:{outer}/");
201 let r = probe_target_latency_guarded(&outer_url, 1, false, Some(Policy::for_test())).await;
202 assert!(r.is_none());
203 assert_eq!(inner_hits.load(Ordering::SeqCst), 0);
204 // Control: unguarded, the same redirect does reach the inner server.
205 probe_target_latency_guarded(&outer_url, 1, false, None).await;
206 assert!(inner_hits.load(Ordering::SeqCst) > 0);
207
208 // Unguarded (the CLI) still probes localhost.
209 let r = probe_target_latency_guarded(&format!("http://127.0.0.1:{port}/"), 1, false, None)
210 .await;
211 assert!(r.is_some());
212 }
213}