bench_produce/
bench_produce.rs

1//! Locked produce throughput example.
2
3use std::time::{Duration, Instant};
4
5use bytes::Bytes;
6use partitionline::{Compression, ProduceRecord, Producer, ProducerConfig, TlsConfig};
7
8#[tokio::main]
9async fn main() -> partitionline::Result<()> {
10    let bootstrap = std::env::var("KAFKA_BOOTSTRAP").unwrap_or_else(|_| "127.0.0.1:9092".into());
11    let topic = std::env::var("KAFKA_TOPIC").unwrap_or_else(|_| "partitionline".into());
12    let payload = std::env::var("PAYLOAD_BYTES")
13        .ok()
14        .and_then(|s| s.parse().ok())
15        .unwrap_or(100usize);
16    let warmup = Duration::from_secs(
17        std::env::var("WARMUP_SECS")
18            .ok()
19            .and_then(|s| s.parse().ok())
20            .unwrap_or(2),
21    );
22    let measure = Duration::from_secs(
23        std::env::var("MEASURE_SECS")
24            .ok()
25            .and_then(|s| s.parse().ok())
26            .unwrap_or(5),
27    );
28    let linger_ms = std::env::var("LINGER_MS")
29        .ok()
30        .and_then(|s| s.parse().ok())
31        .unwrap_or(5u64);
32    let acks = std::env::var("ACKS")
33        .ok()
34        .and_then(|s| s.parse().ok())
35        .unwrap_or(1i16);
36
37    let mut cfg = ProducerConfig::bootstrap([bootstrap]);
38    cfg.linger = Duration::from_millis(linger_ms);
39    cfg.batch_records = 32_768;
40    cfg.batch_bytes = 1_000_000;
41    cfg.acks = acks;
42    cfg.connections = std::env::var("CONNECTIONS")
43        .ok()
44        .and_then(|s| s.parse().ok())
45        .unwrap_or(8);
46    cfg.max_in_flight = std::env::var("MAX_IN_FLIGHT")
47        .ok()
48        .and_then(|s| s.parse().ok())
49        .unwrap_or(16);
50    let compression =
51        Compression::from_name(&std::env::var("COMPRESSION").unwrap_or_else(|_| "none".into()))?;
52    cfg.compression = compression;
53    let idempotent = std::env::var("IDEMPOTENT").ok().as_deref() == Some("1");
54    if idempotent {
55        cfg.enable_idempotence = true;
56    }
57    let tls_on = if let Ok(ca_path) = std::env::var("TLS_CA_PEM") {
58        let mut tls = TlsConfig {
59            ca_pem: Some(tokio::fs::read(&ca_path).await.map_err(|e| {
60                partitionline::Error::protocol(format!("read TLS_CA_PEM {ca_path}: {e}"))
61            })?),
62            ..TlsConfig::default()
63        };
64        if let Ok(name) = std::env::var("TLS_SERVER_NAME") {
65            if !name.is_empty() {
66                tls.server_name = Some(name);
67            }
68        }
69        cfg.tls = Some(tls);
70        true
71    } else {
72        false
73    };
74    let mut scram_on = false;
75    let mut scram512_on = false;
76    let mut oauth_on = false;
77    let mech = std::env::var("SASL_MECHANISM").unwrap_or_else(|_| "PLAIN".into());
78    if mech == "OAUTHBEARER" {
79        let principal = std::env::var("SASL_OAUTH_PRINCIPAL").unwrap_or_else(|_| "alice".into());
80        cfg.sasl_oauthbearer = Some(principal);
81        oauth_on = true;
82    } else if let (Ok(user), Ok(pass)) = (
83        std::env::var("SASL_USERNAME"),
84        std::env::var("SASL_PASSWORD"),
85    ) {
86        match mech.as_str() {
87            "SCRAM-SHA-256" => {
88                cfg.sasl_scram = Some((user, pass));
89                scram_on = true;
90            }
91            "SCRAM-SHA-512" => {
92                cfg.sasl_scram_sha512 = Some((user, pass));
93                scram512_on = true;
94            }
95            "PLAIN" => cfg.sasl_plain = Some((user, pass)),
96            other => {
97                return Err(partitionline::Error::protocol(format!(
98                    "unknown SASL_MECHANISM {other}"
99                )));
100            }
101        }
102    }
103    let acks_out = if idempotent { -1 } else { acks };
104    let producer = Producer::new(cfg).await?;
105    let topic: std::sync::Arc<str> = topic.into();
106    let value = Bytes::from(vec![b'x'; payload]);
107    let count: Option<u64> = std::env::var("COUNT").ok().and_then(|s| s.parse().ok());
108
109    async fn send_one(
110        producer: &Producer,
111        topic: &std::sync::Arc<str>,
112        value: &Bytes,
113        spins: &mut u32,
114    ) -> partitionline::Result<bool> {
115        match producer.try_send(ProduceRecord::to(topic.clone()).value(value.clone())) {
116            Ok(()) => {
117                *spins = 0;
118                Ok(true)
119            }
120            Err(partitionline::Error::QueueFull) => {
121                *spins += 1;
122                if *spins > 32 {
123                    tokio::task::yield_now().await;
124                    *spins = 0;
125                }
126                Ok(false)
127            }
128            Err(e) => Err(e),
129        }
130    }
131
132    async fn drive(
133        producer: &Producer,
134        topic: &std::sync::Arc<str>,
135        value: &Bytes,
136        dur: Duration,
137    ) -> partitionline::Result<u64> {
138        if dur.is_zero() {
139            producer.flush().await?;
140            return Ok(0);
141        }
142        let deadline = Instant::now() + dur;
143        let mut sent = 0u64;
144        let mut spins = 0u32;
145        loop {
146            for _ in 0..1024 {
147                if send_one(producer, topic, value, &mut spins).await? {
148                    sent += 1;
149                }
150            }
151            if Instant::now() >= deadline {
152                break;
153            }
154        }
155        producer.flush().await?;
156        Ok(sent)
157    }
158
159    async fn drive_count(
160        producer: &Producer,
161        topic: &std::sync::Arc<str>,
162        value: &Bytes,
163        n: u64,
164    ) -> partitionline::Result<u64> {
165        let mut sent = 0u64;
166        let mut spins = 0u32;
167        while sent < n {
168            if send_one(producer, topic, value, &mut spins).await? {
169                sent += 1;
170            }
171        }
172        producer.flush().await?;
173        Ok(sent)
174    }
175
176    if let Some(n) = count {
177        let _ = drive(&producer, &topic, &value, warmup).await?;
178        let start = Instant::now();
179        let acked = drive_count(&producer, &topic, &value, n).await?;
180        let elapsed = start.elapsed().as_secs_f64();
181        let rec_s = acked as f64 / elapsed;
182        println!(
183            "{{\"acked\":{acked},\"elapsed_s\":{elapsed:.6},\"acked_rec_s\":{rec_s:.3},\"payload_bytes\":{payload},\"acks\":{acks_out},\"linger_ms\":{linger_ms},\"compression\":\"{}\",\"idempotent\":{},\"tls\":{},\"scram\":{},\"scram512\":{},\"oauthbearer\":{}}}",
184            compression.as_str(),
185            idempotent,
186            tls_on,
187            scram_on,
188            scram512_on,
189            oauth_on
190        );
191    } else {
192        let _ = drive(&producer, &topic, &value, warmup).await?;
193        let start = Instant::now();
194        let acked = drive(&producer, &topic, &value, measure).await?;
195        let elapsed = start.elapsed().as_secs_f64();
196        let rec_s = acked as f64 / elapsed;
197        println!(
198            "{{\"acked\":{acked},\"elapsed_s\":{elapsed:.6},\"acked_rec_s\":{rec_s:.3},\"payload_bytes\":{payload},\"acks\":{acks_out},\"linger_ms\":{linger_ms},\"compression\":\"{}\",\"idempotent\":{},\"tls\":{},\"scram\":{},\"scram512\":{},\"oauthbearer\":{}}}",
199            compression.as_str(),
200            idempotent,
201            tls_on,
202            scram_on,
203            scram512_on,
204            oauth_on
205        );
206    }
207    producer.close().await?;
208    Ok(())
209}
bench_produce/bench_produce.rs

bench_produce/
bench_produce.rs
