|
| 1 | +//! Remote-write load generator for throughput benchmarks of the standalone |
| 2 | +//! `precompute_engine` binary. Run it as a separate process on separate cores |
| 3 | +//! so its own CPU never counts against the engine. |
| 4 | +//! |
| 5 | +//! Series `i` carries labels `l0..l4` = the base-10 digits of `i` (`l4` takes |
| 6 | +//! the remainder), so grouping by a label prefix gives a known fan-in. Each |
| 7 | +//! pass sends one sample per series and advances event time by `step_ms`, so |
| 8 | +//! window closes per sample are fixed regardless of send rate. |
| 9 | +//! |
| 10 | +//! Prints one JSON line: acked samples/sec and engine CPU cores over the |
| 11 | +//! measurement window (after warmup). |
| 12 | +
|
| 13 | +use clap::Parser; |
| 14 | +use prost::Message; |
| 15 | +use query_engine_rust::drivers::ingest::prometheus_remote_write::{ |
| 16 | + Label, Sample, TimeSeries, WriteRequest, |
| 17 | +}; |
| 18 | +use std::ops::Range; |
| 19 | +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; |
| 20 | +use std::sync::Arc; |
| 21 | +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; |
| 22 | +use tokio::sync::mpsc; |
| 23 | + |
| 24 | +const METRIC_NAME: &str = "bench_metric"; |
| 25 | +const NUM_LABELS: usize = 5; |
| 26 | +const LABEL_RADIX: usize = 10; |
| 27 | +/// Sample values are uniform in `[0, VALUE_RANGE)`. |
| 28 | +const VALUE_RANGE: f64 = 10_000.0; |
| 29 | +/// Bodies an encoder thread may build ahead of its sender. |
| 30 | +const ENCODE_AHEAD: usize = 8; |
| 31 | +/// Max passes a sender may lead the slowest one. Groups span senders and the |
| 32 | +/// engine keeps one watermark per group, so `MAX_PASS_SKEW * step_ms` must |
| 33 | +/// stay below the engine's `--allowed_lateness_ms`, or the lagging sender's |
| 34 | +/// samples are dropped as late. |
| 35 | +const MAX_PASS_SKEW: u64 = 2; |
| 36 | +const SKEW_POLL_INTERVAL: Duration = Duration::from_micros(200); |
| 37 | + |
| 38 | +#[derive(Parser, Debug)] |
| 39 | +#[command(name = "bench_ingest_sender")] |
| 40 | +struct Args { |
| 41 | + /// PID of the precompute_engine process to measure. |
| 42 | + #[arg(long)] |
| 43 | + engine_pid: u32, |
| 44 | + |
| 45 | + #[arg(long)] |
| 46 | + port: u16, |
| 47 | + |
| 48 | + #[arg(long)] |
| 49 | + num_series: usize, |
| 50 | + |
| 51 | + /// Target samples/sec across all senders; 0 sends open-loop (saturation). |
| 52 | + #[arg(long)] |
| 53 | + target_rate: f64, |
| 54 | + |
| 55 | + /// Concurrent senders (capped at --num-series). Each owns a disjoint slice |
| 56 | + /// of the series and keeps one request in flight (preserves per-series |
| 57 | + /// order). The engine handles a request mostly serially, so this must be |
| 58 | + /// at least the engine's core count to saturate it, even if that makes |
| 59 | + /// requests smaller at low cardinality. |
| 60 | + #[arg(long)] |
| 61 | + concurrency: usize, |
| 62 | + |
| 63 | + /// Max samples per remote-write request (Prometheus' default |
| 64 | + /// `max_samples_per_send` is 2000). Smaller when a sender's slice is. |
| 65 | + #[arg(long)] |
| 66 | + samples_per_request: usize, |
| 67 | + |
| 68 | + /// Event-time advance per pass (the simulated scrape interval). |
| 69 | + #[arg(long)] |
| 70 | + step_ms: i64, |
| 71 | + |
| 72 | + #[arg(long)] |
| 73 | + warmup_secs: u64, |
| 74 | + |
| 75 | + #[arg(long)] |
| 76 | + measure_secs: u64, |
| 77 | +} |
| 78 | + |
| 79 | +type Body = (u64, u64, Vec<u8>); // (pass, samples, snappy-compressed body) |
| 80 | + |
| 81 | +fn make_series(i: usize) -> TimeSeries { |
| 82 | + let mut labels = vec![Label { |
| 83 | + name: "__name__".into(), |
| 84 | + value: METRIC_NAME.into(), |
| 85 | + }]; |
| 86 | + let mut rest = i; |
| 87 | + for l in 0..NUM_LABELS { |
| 88 | + let digit = if l + 1 == NUM_LABELS { |
| 89 | + rest |
| 90 | + } else { |
| 91 | + rest % LABEL_RADIX |
| 92 | + }; |
| 93 | + rest /= LABEL_RADIX; |
| 94 | + labels.push(Label { |
| 95 | + name: format!("l{l}"), |
| 96 | + value: format!("v{digit}"), |
| 97 | + }); |
| 98 | + } |
| 99 | + TimeSeries { |
| 100 | + labels, |
| 101 | + samples: vec![Sample::default()], |
| 102 | + } |
| 103 | +} |
| 104 | + |
| 105 | +/// splitmix64; seeded per encoder so senders don't emit identical values. |
| 106 | +fn next_value(state: &mut u64) -> f64 { |
| 107 | + *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15); |
| 108 | + let mut z = *state; |
| 109 | + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); |
| 110 | + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); |
| 111 | + z ^= z >> 31; |
| 112 | + (z >> 11) as f64 / (1u64 << 53) as f64 * VALUE_RANGE |
| 113 | +} |
| 114 | + |
| 115 | +/// Builds request bodies pass after pass until the sender hangs up. Labels are |
| 116 | +/// built once; each pass only rewrites the sample. |
| 117 | +fn encode_loop( |
| 118 | + series: Range<usize>, |
| 119 | + samples_per_request: usize, |
| 120 | + start_ms: i64, |
| 121 | + step_ms: i64, |
| 122 | + tx: mpsc::Sender<Body>, |
| 123 | +) { |
| 124 | + let ids: Vec<usize> = series.clone().collect(); |
| 125 | + let mut requests: Vec<WriteRequest> = ids |
| 126 | + .chunks(samples_per_request) |
| 127 | + .map(|chunk| WriteRequest { |
| 128 | + timeseries: chunk.iter().map(|&i| make_series(i)).collect(), |
| 129 | + }) |
| 130 | + .collect(); |
| 131 | + let mut rng = series.start as u64; |
| 132 | + let mut encoder = snap::raw::Encoder::new(); |
| 133 | + for pass in 0u64.. { |
| 134 | + let timestamp = start_ms + pass as i64 * step_ms; |
| 135 | + for request in &mut requests { |
| 136 | + for ts in &mut request.timeseries { |
| 137 | + ts.samples[0] = Sample { |
| 138 | + value: next_value(&mut rng), |
| 139 | + timestamp, |
| 140 | + }; |
| 141 | + } |
| 142 | + let body = encoder |
| 143 | + .compress_vec(&request.encode_to_vec()) |
| 144 | + .expect("snappy compression should succeed"); |
| 145 | + if tx |
| 146 | + .blocking_send((pass, request.timeseries.len() as u64, body)) |
| 147 | + .is_err() |
| 148 | + { |
| 149 | + return; |
| 150 | + } |
| 151 | + } |
| 152 | + } |
| 153 | +} |
| 154 | + |
| 155 | +struct Shared { |
| 156 | + client: reqwest::Client, |
| 157 | + url: String, |
| 158 | + /// Pass each sender is currently on, for the skew bound. |
| 159 | + passes: Vec<AtomicU64>, |
| 160 | + acked_samples: AtomicU64, |
| 161 | + stop: AtomicBool, |
| 162 | +} |
| 163 | + |
| 164 | +async fn send_loop( |
| 165 | + shared: Arc<Shared>, |
| 166 | + idx: usize, |
| 167 | + rate: f64, |
| 168 | + mut rx: mpsc::Receiver<Body>, |
| 169 | +) -> Result<(), String> { |
| 170 | + let start = Instant::now(); |
| 171 | + let mut sent_samples = 0u64; |
| 172 | + while let Some((pass, samples, body)) = rx.recv().await { |
| 173 | + shared.passes[idx].store(pass, Ordering::Relaxed); |
| 174 | + loop { |
| 175 | + if shared.stop.load(Ordering::Relaxed) { |
| 176 | + return Ok(()); |
| 177 | + } |
| 178 | + let slowest = shared |
| 179 | + .passes |
| 180 | + .iter() |
| 181 | + .map(|p| p.load(Ordering::Relaxed)) |
| 182 | + .min() |
| 183 | + .expect("concurrency >= 1"); |
| 184 | + if pass <= slowest + MAX_PASS_SKEW { |
| 185 | + break; |
| 186 | + } |
| 187 | + tokio::time::sleep(SKEW_POLL_INTERVAL).await; |
| 188 | + } |
| 189 | + if rate > 0.0 { |
| 190 | + let due = start + Duration::from_secs_f64(sent_samples as f64 / rate); |
| 191 | + tokio::time::sleep_until(due.into()).await; |
| 192 | + } |
| 193 | + let response = shared |
| 194 | + .client |
| 195 | + .post(&shared.url) |
| 196 | + .header("Content-Type", "application/x-protobuf") |
| 197 | + .header("Content-Encoding", "snappy") |
| 198 | + .body(body) |
| 199 | + .send() |
| 200 | + .await |
| 201 | + .map_err(|e| format!("sender {idx}: request failed: {e}"))?; |
| 202 | + if response.status() != reqwest::StatusCode::NO_CONTENT { |
| 203 | + return Err(format!("sender {idx}: HTTP {}", response.status())); |
| 204 | + } |
| 205 | + sent_samples += samples; |
| 206 | + shared.acked_samples.fetch_add(samples, Ordering::Relaxed); |
| 207 | + } |
| 208 | + Ok(()) |
| 209 | +} |
| 210 | + |
| 211 | +fn clock_ticks_per_sec() -> u64 { |
| 212 | + let out = std::process::Command::new("getconf") |
| 213 | + .arg("CLK_TCK") |
| 214 | + .output() |
| 215 | + .expect("getconf CLK_TCK should run"); |
| 216 | + String::from_utf8(out.stdout) |
| 217 | + .expect("getconf output is ASCII") |
| 218 | + .trim() |
| 219 | + .parse() |
| 220 | + .expect("CLK_TCK is an integer") |
| 221 | +} |
| 222 | + |
| 223 | +/// utime + stime of the whole process, in clock ticks. |
| 224 | +fn process_cpu_ticks(pid: u32) -> u64 { |
| 225 | + let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")) |
| 226 | + .unwrap_or_else(|e| panic!("pid {pid} not readable: {e}")); |
| 227 | + // Fields after the parenthesised comm; utime/stime are fields 14/15. |
| 228 | + let fields: Vec<&str> = stat[stat.rfind(')').expect("stat has comm") + 2..] |
| 229 | + .split_whitespace() |
| 230 | + .collect(); |
| 231 | + let field = |i: usize| -> u64 { fields[i].parse().expect("stat field is an integer") }; |
| 232 | + field(11) + field(12) |
| 233 | +} |
| 234 | + |
| 235 | +fn engine_peak_rss_mb(pid: u32) -> f64 { |
| 236 | + let status = std::fs::read_to_string(format!("/proc/{pid}/status")) |
| 237 | + .unwrap_or_else(|e| panic!("engine pid {pid} not readable: {e}")); |
| 238 | + let kb: f64 = status |
| 239 | + .lines() |
| 240 | + .find_map(|l| l.strip_prefix("VmHWM:")) |
| 241 | + .expect("status has VmHWM") |
| 242 | + .trim() |
| 243 | + .trim_end_matches("kB") |
| 244 | + .trim() |
| 245 | + .parse() |
| 246 | + .expect("VmHWM is a number"); |
| 247 | + kb / 1024.0 |
| 248 | +} |
| 249 | + |
| 250 | +#[tokio::main] |
| 251 | +async fn main() { |
| 252 | + let args = Args::parse(); |
| 253 | + assert!(args.num_series > 0 && args.concurrency > 0 && args.samples_per_request > 0); |
| 254 | + let concurrency = args.concurrency.min(args.num_series); |
| 255 | + let ticks_per_sec = clock_ticks_per_sec(); |
| 256 | + let start_ms = SystemTime::now() |
| 257 | + .duration_since(UNIX_EPOCH) |
| 258 | + .expect("clock after epoch") |
| 259 | + .as_millis() as i64; |
| 260 | + |
| 261 | + let shared = Arc::new(Shared { |
| 262 | + client: reqwest::Client::new(), |
| 263 | + url: format!("http://127.0.0.1:{}/api/v1/write", args.port), |
| 264 | + passes: (0..concurrency).map(|_| AtomicU64::new(0)).collect(), |
| 265 | + acked_samples: AtomicU64::new(0), |
| 266 | + stop: AtomicBool::new(false), |
| 267 | + }); |
| 268 | + |
| 269 | + let mut senders = Vec::with_capacity(concurrency); |
| 270 | + for idx in 0..concurrency { |
| 271 | + let series = idx * args.num_series / concurrency..(idx + 1) * args.num_series / concurrency; |
| 272 | + let rate = args.target_rate * series.len() as f64 / args.num_series as f64; |
| 273 | + let (tx, rx) = mpsc::channel(ENCODE_AHEAD); |
| 274 | + let step_ms = args.step_ms; |
| 275 | + let samples_per_request = args.samples_per_request; |
| 276 | + std::thread::spawn(move || encode_loop(series, samples_per_request, start_ms, step_ms, tx)); |
| 277 | + senders.push(tokio::spawn(send_loop(shared.clone(), idx, rate, rx))); |
| 278 | + } |
| 279 | + |
| 280 | + tokio::time::sleep(Duration::from_secs(args.warmup_secs)).await; |
| 281 | + let acked_0 = shared.acked_samples.load(Ordering::Relaxed); |
| 282 | + let ticks_0 = process_cpu_ticks(args.engine_pid); |
| 283 | + let sender_ticks_0 = process_cpu_ticks(std::process::id()); |
| 284 | + let t0 = Instant::now(); |
| 285 | + assert!(acked_0 > 0, "engine acked no samples during warmup"); |
| 286 | + |
| 287 | + tokio::time::sleep(Duration::from_secs(args.measure_secs)).await; |
| 288 | + let acked_1 = shared.acked_samples.load(Ordering::Relaxed); |
| 289 | + let ticks_1 = process_cpu_ticks(args.engine_pid); |
| 290 | + let sender_ticks_1 = process_cpu_ticks(std::process::id()); |
| 291 | + let elapsed = t0.elapsed().as_secs_f64(); |
| 292 | + let cores = |ticks: u64| ticks as f64 / ticks_per_sec as f64 / elapsed; |
| 293 | + |
| 294 | + shared.stop.store(true, Ordering::Relaxed); |
| 295 | + for sender in senders { |
| 296 | + if let Err(e) = sender.await.expect("sender task panicked") { |
| 297 | + panic!("{e}"); |
| 298 | + } |
| 299 | + } |
| 300 | + |
| 301 | + println!( |
| 302 | + "{}", |
| 303 | + serde_json::json!({ |
| 304 | + "num_series": args.num_series, |
| 305 | + "concurrency": concurrency, |
| 306 | + "target_rate": args.target_rate, |
| 307 | + "achieved_rate": (acked_1 - acked_0) as f64 / elapsed, |
| 308 | + "engine_cpu_cores": cores(ticks_1 - ticks_0), |
| 309 | + // Near the sender's core count means the sender, not the engine, |
| 310 | + // capped an open-loop run. |
| 311 | + "sender_cpu_cores": cores(sender_ticks_1 - sender_ticks_0), |
| 312 | + "engine_peak_rss_mb": engine_peak_rss_mb(args.engine_pid), |
| 313 | + "measure_secs": elapsed, |
| 314 | + }) |
| 315 | + ); |
| 316 | + // Encoder threads may be parked on a full channel; don't wait for them. |
| 317 | + std::process::exit(0); |
| 318 | +} |
0 commit comments