bench_latency/
bench_latency.rs

1//! Locked produce-ack and fetch-request latency example.
2//!
3//! Sequential `Producer::send` (enqueue to Produce ack) and non-empty
4//! `Consumer::fetch` (already-on-log Fetch RPC). Prints one JSON object
5//! per kind with p50/p99 in microseconds. Not a throughput bench.
6
7use std::time::{Duration, Instant};
8
9use bytes::Bytes;
10use partitionline::{Consumer, ConsumerConfig, ProduceRecord, Producer, ProducerConfig};
11
12fn env_or(key: &str, default: &str) -> String {
13    std::env::var(key).unwrap_or_else(|_| default.into())
14}
15
16fn env_parse<T: std::str::FromStr>(key: &str, default: T) -> T {
17    std::env::var(key)
18        .ok()
19        .and_then(|s| s.parse().ok())
20        .unwrap_or(default)
21}
22
23fn percentile_us(sorted: &[u64], p: u32) -> partitionline::Result<u64> {
24    if sorted.is_empty() {
25        return Err(partitionline::Error::protocol("no latency samples"));
26    }
27    let n = sorted.len();
28    let rank = n
29        .saturating_mul(p as usize)
30        .div_ceil(100)
31        .saturating_sub(1)
32        .min(n.saturating_sub(1));
33    sorted
34        .get(rank)
35        .copied()
36        .ok_or_else(|| partitionline::Error::protocol("percentile index"))
37}
38
39fn print_latency(kind: &str, mut samples: Vec<u64>, extra: &str) -> partitionline::Result<()> {
40    if samples.is_empty() {
41        return Err(partitionline::Error::protocol(format!(
42            "{kind}: no latency samples"
43        )));
44    }
45    samples.sort_unstable();
46    let n = samples.len();
47    let min_us = samples.first().copied().unwrap_or(0);
48    let max_us = samples.last().copied().unwrap_or(0);
49    let sum: u128 = samples.iter().map(|v| u128::from(*v)).sum();
50    let mean_us = u64::try_from(sum / u128::from(n as u64))
51        .map_err(|_| partitionline::Error::protocol("mean overflow"))?;
52    let p50_us = percentile_us(&samples, 50)?;
53    let p99_us = percentile_us(&samples, 99)?;
54    println!(
55        "{{\"kind\":\"{kind}\",\"samples\":{n},\"p50_us\":{p50_us},\"p99_us\":{p99_us},\"min_us\":{min_us},\"max_us\":{max_us},\"mean_us\":{mean_us}{extra}}}"
56    );
57    Ok(())
58}
59
60async fn produce_ack(
61    producer: &Producer,
62    topic: &str,
63    value: &Bytes,
64    n: u64,
65) -> partitionline::Result<Vec<u64>> {
66    let mut samples = Vec::with_capacity(usize::try_from(n).unwrap_or(0));
67    for _ in 0..n {
68        let start = Instant::now();
69        let _md = producer
70            .send(ProduceRecord::to(topic).value(value.clone()))
71            .await?;
72        samples.push(u64::try_from(start.elapsed().as_micros()).unwrap_or(u64::MAX));
73    }
74    Ok(samples)
75}
76
77async fn fetch_rpc(consumer: &mut Consumer, count: u64) -> partitionline::Result<(Vec<u64>, u64)> {
78    let mut samples = Vec::new();
79    let mut got = 0u64;
80    let mut empty = 0u32;
81    while got < count {
82        let start = Instant::now();
83        let recs = consumer.fetch().await?;
84        let elapsed = u64::try_from(start.elapsed().as_micros()).unwrap_or(u64::MAX);
85        if recs.is_empty() {
86            empty += 1;
87            if empty > 600 {
88                return Err(partitionline::Error::Timeout);
89            }
90            continue;
91        }
92        empty = 0;
93        samples.push(elapsed);
94        got += recs.len() as u64;
95    }
96    Ok((samples, got))
97}
98
99#[tokio::main]
100async fn main() -> partitionline::Result<()> {
101    let bootstrap = env_or("KAFKA_BOOTSTRAP", "127.0.0.1:9092");
102    let topic = env_or("KAFKA_TOPIC", "pllat");
103    let payload = env_parse("PAYLOAD_BYTES", 100usize);
104    let warmup = env_parse("WARMUP", 1_000u64);
105    let count = env_parse("COUNT", 10_000u64);
106    let linger_ms = env_parse("LINGER_MS", 0u64);
107    let acks = env_parse("ACKS", 1i16);
108    let max_wait_ms = env_parse("MAX_WAIT_MS", 100i32);
109    let max_bytes = env_parse("MAX_BYTES", 4_096i32);
110    let min_bytes = env_parse("MIN_BYTES", 1i32);
111    let mode = env_or("MODE", "both");
112
113    let mut pcfg = ProducerConfig::bootstrap([bootstrap.clone()]);
114    pcfg.linger = Duration::from_millis(linger_ms);
115    pcfg.acks = acks;
116    pcfg.batch_records = 1;
117    pcfg.batch_bytes = payload.saturating_add(256).max(1);
118    pcfg.connections = 1;
119    pcfg.max_in_flight = 1;
120    let producer = Producer::new(pcfg).await?;
121    let value = Bytes::from(vec![b'x'; payload]);
122
123    if warmup > 0 {
124        let _ = produce_ack(&producer, &topic, &value, warmup).await?;
125    }
126
127    if mode == "produce" || mode == "both" {
128        let samples = produce_ack(&producer, &topic, &value, count).await?;
129        print_latency(
130            "produce_ack",
131            samples,
132            &format!(
133                ",\"payload_bytes\":{payload},\"acks\":{acks},\"linger_ms\":{linger_ms},\"client\":\"partitionline\""
134            ),
135        )?;
136    }
137    producer.close().await?;
138
139    if mode == "fetch" || mode == "both" {
140        if mode == "fetch" {
141            let producer = Producer::new({
142                let mut cfg = ProducerConfig::bootstrap([bootstrap.clone()]);
143                cfg.linger = Duration::ZERO;
144                cfg.acks = acks;
145                cfg.batch_records = 1;
146                cfg.connections = 1;
147                cfg.max_in_flight = 1;
148                cfg
149            })
150            .await?;
151            let _ = produce_ack(&producer, &topic, &value, count).await?;
152            producer.close().await?;
153        }
154        let mut ccfg = ConsumerConfig::bootstrap([bootstrap]);
155        ccfg.max_wait_ms = max_wait_ms;
156        ccfg.max_bytes = max_bytes;
157        ccfg.min_bytes = min_bytes;
158        let mut consumer = Consumer::new(ccfg).await?;
159        consumer.assign(&topic, 0, 0).await?;
160        let (samples, got) = fetch_rpc(&mut consumer, count).await?;
161        print_latency(
162            "fetch_rpc",
163            samples,
164            &format!(
165                ",\"consumed\":{got},\"max_wait_ms\":{max_wait_ms},\"max_bytes\":{max_bytes},\"min_bytes\":{min_bytes},\"client\":\"partitionline\""
166            ),
167        )?;
168    }
169    Ok(())
170}
bench_latency/bench_latency.rs

bench_latency/
bench_latency.rs
