use rustis::{Result, client::Client, commands::StringCommands};
use std::{sync::Arc, time::Instant};
fn env_usize(key: &str, default: usize) -> usize {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
async fn run(clients: &[Client], tasks: usize, reqs: usize, keys: Arc<Vec<String>>) -> f64 {
let start = Instant::now();
let mut handles = Vec::with_capacity(tasks);
for t in 0..tasks {
let client = clients[t % clients.len()].clone();
let keys = keys.clone();
handles.push(tokio::spawn(async move {
for i in 0..reqs {
let _: String = client.get(&keys[i % keys.len()]).await.unwrap();
}
}));
}
for h in handles {
let _ = h.await;
}
let elapsed = start.elapsed().as_secs_f64();
(tasks * reqs) as f64 / elapsed
}
#[tokio::main]
async fn main() -> Result<()> {
let redis_host = std::env::var("REDIS_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
let tasks = env_usize("PROBE_TASKS", 256);
let reqs = env_usize("PROBE_REQS", 2000);
let keys: Arc<Vec<String>> = Arc::new((0..100).map(|i| format!("key{i}")).collect());
let seed = Client::connect(redis_host.clone()).await?;
let data: Vec<_> = (0..100)
.map(|i| (format!("key{i}"), format!("value{i}")))
.collect();
let _: () = seed.mset(data).await?;
println!("{tasks} tasks × {reqs} GETs, round-robin over N connections\n");
for &n in &[1usize, 2, 4, 8] {
let mut clients = Vec::with_capacity(n);
for _ in 0..n {
clients.push(Client::connect(redis_host.clone()).await?);
}
let _ = run(&clients, tasks, reqs / 4, keys.clone()).await;
let mut best = 0.0f64;
for _ in 0..3 {
let ops = run(&clients, tasks, reqs, keys.clone()).await;
if ops > best {
best = ops;
}
}
println!("{n:>2} conn {:>10.0} ops/s", best);
}
Ok(())
}