bench_fetch/
bench_fetch.rs

1//! Locked fetch throughput example.
2
3use std::time::Instant;
4
5use partitionline::{Consumer, ConsumerConfig, TlsConfig};
6
7#[tokio::main]
8async fn main() -> partitionline::Result<()> {
9    let bootstrap = std::env::var("KAFKA_BOOTSTRAP").unwrap_or_else(|_| "127.0.0.1:9092".into());
10    let topic = std::env::var("KAFKA_TOPIC").unwrap_or_else(|_| "plbench".into());
11    let count: u64 = std::env::var("COUNT")
12        .ok()
13        .and_then(|s| s.parse().ok())
14        .unwrap_or(8_000_000);
15    let max_wait_ms = std::env::var("MAX_WAIT_MS")
16        .ok()
17        .and_then(|s| s.parse().ok())
18        .unwrap_or(100i32);
19    let max_bytes = std::env::var("MAX_BYTES")
20        .ok()
21        .and_then(|s| s.parse().ok())
22        .unwrap_or(16_777_216i32);
23    let min_bytes = std::env::var("MIN_BYTES")
24        .ok()
25        .and_then(|s| s.parse().ok())
26        .unwrap_or(1i32);
27
28    let mut cfg = ConsumerConfig::bootstrap([bootstrap]);
29    cfg.max_wait_ms = max_wait_ms;
30    cfg.max_bytes = max_bytes;
31    cfg.min_bytes = min_bytes;
32    if let Ok(ca_path) = std::env::var("TLS_CA_PEM") {
33        let mut tls = TlsConfig {
34            ca_pem: Some(tokio::fs::read(&ca_path).await.map_err(|e| {
35                partitionline::Error::protocol(format!("read TLS_CA_PEM {ca_path}: {e}"))
36            })?),
37            ..TlsConfig::default()
38        };
39        if let Ok(name) = std::env::var("TLS_SERVER_NAME") {
40            if !name.is_empty() {
41                tls.server_name = Some(name);
42            }
43        }
44        cfg.tls = Some(tls);
45    }
46    let mech = std::env::var("SASL_MECHANISM").unwrap_or_else(|_| "PLAIN".into());
47    if mech == "OAUTHBEARER" {
48        let principal = std::env::var("SASL_OAUTH_PRINCIPAL").unwrap_or_else(|_| "alice".into());
49        cfg.sasl_oauthbearer = Some(principal);
50    } else if let (Ok(user), Ok(pass)) = (
51        std::env::var("SASL_USERNAME"),
52        std::env::var("SASL_PASSWORD"),
53    ) {
54        match mech.as_str() {
55            "SCRAM-SHA-256" => cfg.sasl_scram = Some((user, pass)),
56            "SCRAM-SHA-512" => cfg.sasl_scram_sha512 = Some((user, pass)),
57            "PLAIN" => cfg.sasl_plain = Some((user, pass)),
58            other => {
59                return Err(partitionline::Error::protocol(format!(
60                    "unknown SASL_MECHANISM {other}"
61                )));
62            }
63        }
64    }
65
66    let mut consumer = Consumer::new(cfg).await?;
67    consumer.assign_topic(&topic, 0).await?;
68    let assigned = consumer.assignment().len();
69    let start = Instant::now();
70    let mut got = 0u64;
71    let mut empty = 0u32;
72    while got < count {
73        let recs = consumer.fetch().await?;
74        if recs.is_empty() {
75            empty += 1;
76            if empty > 600 {
77                return Err(partitionline::Error::Timeout);
78            }
79            continue;
80        }
81        empty = 0;
82        got += recs.len() as u64;
83    }
84    let elapsed = start.elapsed().as_secs_f64();
85    let rec_s = got as f64 / elapsed.max(1e-9);
86    println!(
87        "{{\"consumed\":{got},\"elapsed_s\":{elapsed:.6},\"consumed_rec_s\":{rec_s:.3},\"partitions\":{assigned},\"max_wait_ms\":{max_wait_ms},\"max_bytes\":{max_bytes}}}"
88    );
89    Ok(())
90}
bench_fetch/bench_fetch.rs

bench_fetch/
bench_fetch.rs
