use futures::*;
use std::time::Duration;
use std::sync::{Arc, RwLock};
use rand::Rng;
use log::*;
use tokio::time::sleep as t_sleep;
const RAMP_UP_SLEEP: Duration = Duration::from_millis(10);
pub struct PerfTool<IN, OUT>
where
IN: Send + 'static,
OUT: Future + Send + 'static,
OUT::Output: Send + 'static,
{
wps: u64,
ramp_up: Duration,
test_duration: Duration,
init: Box<dyn Fn() -> IN>,
worker: Box<dyn Fn(IN) -> OUT>,
num_workers: Arc<RwLock<u64>>,
}
impl<IN, OUT> PerfTool<IN, OUT>
where
IN: Send + 'static,
OUT: Future + Send + 'static,
OUT::Output: Send + 'static,
{
pub fn new(wps: u64,
ramp_up: Duration,
test_duration: Duration,
init_task: impl Fn() -> IN + 'static,
task: impl Fn(IN) -> OUT + 'static,
) -> PerfTool<IN, OUT> {
PerfTool{
wps,
ramp_up,
test_duration,
init: Box::new(init_task),
worker: Box::new(task),
num_workers: Arc::new(RwLock::new(0)),
}
}
pub async fn start(&self) {
debug!("starting ramp up");
let ramp_up_steps = (self.ramp_up.as_millis() / RAMP_UP_SLEEP.as_millis()) as u64;
let ramp_up_speed: f64 = self.wps as f64 / ramp_up_steps as f64;
let mut ramp_handles = self.spawn_loop(ramp_up_steps, |i| {
(ramp_up_speed * (i + 1) as f64).ceil() as u64
}).await;
debug!("ramp up done");
debug!("starting test with constant load");
let test_steps = (self.test_duration.as_millis() / RAMP_UP_SLEEP.as_millis()) as u64;
let handles = self.spawn_loop(test_steps, |_| self.wps).await;
ramp_handles.extend(handles);
debug!("test done, waiting for workers to finish");
for h in ramp_handles {
let _ = h.await;
}
debug!("all workers done")
}
async fn spawn_loop(&self, steps: u64, expected_num_workers: impl Fn(u64) -> u64) -> Vec<tokio::task::JoinHandle<()>> {
let mut handles = vec![];
for i in 0..steps {
t_sleep(RAMP_UP_SLEEP).await;
let expected_num_workers = expected_num_workers(i);
let actual_num_workers = *self.num_workers.read().unwrap();
if actual_num_workers >= expected_num_workers {
continue
}
debug!("actual #workers - {}, expected #workers - {}", actual_num_workers, expected_num_workers);
for _ in actual_num_workers..expected_num_workers {
let input = (self.init)();
let future = (self.worker)(input);
let handle = self.spawn_worker(future);
handles.push(handle);
}
}
handles
}
fn spawn_worker(&self, task: OUT) -> tokio::task::JoinHandle<()> {
let num_workers = Arc::clone(&self.num_workers);
tokio::spawn(async move {
*num_workers.write().unwrap() += 1;
task.await;
*num_workers.write().unwrap() -= 1;
})
}
}
#[cfg(test)]
mod tests {
use super::*;
struct WorkerInput<'a> {
id: u8,
name: &'a str,
}
impl<'a> WorkerInput<'a> {
fn new(n: &str) -> WorkerInput {
WorkerInput{
id: rand::thread_rng().gen(),
name: n,
}
}
}
async fn test_worker(input: WorkerInput<'_>) {
info!("hello, {}, my name is {}", input.id, input.name);
t_sleep(Duration::from_secs(2)).await;
info!("I'm done, {}", input.id);
}
#[tokio::test]
async fn test_simple() {
std::env::set_var("RUST_LOG", "debug");
env_logger::init();
let wps = 10;
let ramp_up = Duration::from_secs(5);
let test_time = Duration::from_secs(10);
let test_data = vec!["JORA", "KOLEA", "FEDEA", "VASEA", "VANEA", "PETEA"];
let init = move || {
let n = test_data.len();
let idx: usize = rand::thread_rng().gen();
WorkerInput::new(test_data[idx % n])
};
let tool = PerfTool::new(wps, ramp_up, test_time,init, test_worker);
let num_workers = Arc::clone(&tool.num_workers);
let checker = tokio::spawn(async move {
info!("test waiting for ramp up to finish");
t_sleep(ramp_up).await;
info!("test starts checkng actual wps");
for _ in 0..test_time.as_secs() {
t_sleep(Duration::from_secs(1)).await;
let actual_wps = *num_workers.read().unwrap();
if actual_wps + 2 <= wps {
return Err(format!("expected wps {}, actual wps {}", wps, actual_wps));
}
}
Ok(())
});
tool.start().await;
let result = checker.await;
match result {
Ok(Ok(())) => {},
Ok(Err(msg)) => assert!(false, msg),
Err(msg) => assert!(false, "{}", msg)
}
}
}