wrk_api_bench/
wrk.rs

1use std::{
2    collections::HashMap,
3    fs::{self, File},
4    io::{BufReader, BufWriter},
5    ops::Sub,
6    path::{Path, PathBuf},
7    process::Command,
8    time::Duration,
9};
10
11use chrono::{DateTime, Duration as ChronoDuration, NaiveDateTime, Utc};
12use getset::{Getters, MutGetters, Setters};
13use serde::{Deserialize, Serialize};
14use tempfile::NamedTempFile;
15use url::Url;
16
17use crate::{
18    benchmark::{Benchmark, BenchmarkBuilder},
19    error::WrkError,
20    result::{Deviation, WrkResult, WrkResultBuilder},
21    Gnuplot, LuaScript, Result,
22};
23
24const DATE_FORMAT: &str = "%Y-%m-%d-%H:%M:%S-%z";
25
26#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
27pub enum HistoryPeriod {
28    Last,
29    Hour,
30    Day,
31    Week,
32    Month,
33    Forever,
34}
35
36impl Default for HistoryPeriod {
37    fn default() -> Self {
38        HistoryPeriod::Last
39    }
40}
41
42impl HistoryPeriod {
43    pub fn last_valid_datapoint(&self) -> DateTime<Utc> {
44        let now = Utc::now();
45        match self {
46            Self::Last => now,
47            Self::Hour => now.sub(ChronoDuration::hours(1)),
48            Self::Day => now.sub(ChronoDuration::days(1)),
49            Self::Week => now.sub(ChronoDuration::weeks(1)),
50            Self::Month => now.sub(ChronoDuration::weeks(4)),
51            Self::Forever => DateTime::from_utc(NaiveDateTime::from_timestamp(1, 0), Utc),
52        }
53    }
54}
55
56pub type Benchmarks = Vec<WrkResult>;
57pub type Headers = HashMap<String, String>;
58
59/// Wrapper around Wrk enabling to run benchmarks, record historical data and plot graphs.
60#[derive(Debug, Clone, Serialize, Deserialize, Getters, Setters, MutGetters, Builder)]
61pub struct Wrk {
62    /// Url of the service to benchmark against. Use the full URL of the request.
63    /// IE: http://localhost:1234/some/uri.
64    #[getset(get = "pub", set = "pub", get_mut = "pub")]
65    url: String,
66    /// Wrk timeout in seconds
67    #[serde(skip)]
68    #[builder(default = "1")]
69    #[getset(get = "pub", set = "pub", get_mut = "pub")]
70    timeout: u8,
71    /// Set of benchmarks for the current instance.
72    #[builder(default)]
73    #[getset(get = "pub", set = "pub", get_mut = "pub")]
74    benchmarks: Benchmarks,
75    /// Historical benchmarks data, indexed by dates.
76    #[builder(default)]
77    #[getset(get = "pub", set = "pub", get_mut = "pub")]
78    benchmarks_history: Benchmarks,
79    /// Directory on disk where to store and read the historical benchmark data.
80    #[builder(default = "Path::new(\".\").join(\".wrk-api-bench\")")]
81    #[getset(get = "pub", set = "pub", get_mut = "pub")]
82    history_dir: PathBuf,
83    /// User defined LUA script to run through wrk.
84    /// **NOTE: This script MUST not override the wrk function `done()` as it already
85    /// overriden by this crate to allow wrk to spit out a parsable JSON output.
86    #[builder(default)]
87    #[getset(get = "pub", set = "pub", get_mut = "pub")]
88    user_script: Option<PathBuf>,
89    /// Header to add to the wrk request.
90    #[builder(default)]
91    #[getset(get = "pub", set = "pub", get_mut = "pub")]
92    headers: Headers,
93    /// Method for the wrk request.
94    #[builder(default = "String::from(\"GET\")")]
95    #[getset(get = "pub", set = "pub", get_mut = "pub")]
96    method: String,
97    /// Body for the wrk request.
98    #[builder(default)]
99    #[getset(get = "pub", set = "pub", get_mut = "pub")]
100    body: String,
101    /// Max percentage of errors vs total request to conside a benchmark healthy.
102    #[builder(default = "2")]
103    #[getset(get = "pub", set = "pub", get_mut = "pub")]
104    max_error_percentage: u8,
105    /// Current benchmark date and time.
106    #[serde(skip)]
107    #[builder(default)]
108    #[getset(get = "pub", set = "pub", get_mut = "pub")]
109    benchmark_date: Option<DateTime<Utc>>,
110}
111
112impl Wrk {
113    fn wrk_args(&self, benchmark: &Benchmark, url: &Url, lua_script: &Path) -> Result<Vec<String>> {
114        Ok(vec![
115            "-t".to_string(),
116            benchmark.threads().to_string(),
117            "-c".to_string(),
118            benchmark.connections().to_string(),
119            "-d".to_string(),
120            format!("{}s", benchmark.duration().as_secs()),
121            "--timeout".to_string(),
122            format!("{}s", self.timeout()),
123            "-s".to_string(),
124            lua_script.to_string_lossy().to_string(),
125            url.to_string(),
126        ])
127    }
128
129    fn wrk_result(&self, wrk_json: &str) -> WrkResult {
130        match serde_json::from_str::<WrkResult>(wrk_json) {
131            Ok(mut run) => {
132                let error_percentage = run.errors() / 100.0 * run.requests();
133                if error_percentage < *self.max_error_percentage() as f64 {
134                    *run.success_mut() = true;
135                } else {
136                    error!(
137                        "Errors percentage is {}%, which is more than {}%",
138                        error_percentage, self.max_error_percentage
139                    );
140                }
141                run
142            }
143            Err(e) => {
144                error!("Wrk JSON result deserialize failed: {}", e);
145                WrkResult::fail(e.to_string())
146            }
147        }
148    }
149
150    pub fn bench(&mut self, benchmarks: &Vec<Benchmark>) -> Result<()> {
151        if !self.history_dir().exists() {
152            fs::create_dir(self.history_dir()).unwrap_or_else(|e| {
153                error!(
154                    "Unable to create storage dir {}: {}. Statistics calculation could be impaired",
155                    self.history_dir().display(),
156                    e
157                );
158            });
159        }
160        let date = Utc::now();
161        *self.benchmark_date_mut() = Some(date);
162        let url = Url::parse(self.url())?;
163        let mut script_file = NamedTempFile::new()?;
164        LuaScript::render(
165            &mut script_file,
166            self.user_script().as_ref(),
167            url.path(),
168            self.method(),
169            self.headers(),
170            self.body(),
171        )?;
172        for benchmark in benchmarks {
173            let mut run = match Command::new("wrk")
174                .args(self.wrk_args(benchmark, &url, script_file.path())?)
175                .output()
176            {
177                Ok(wrk) => {
178                    let output = String::from_utf8_lossy(&wrk.stdout);
179                    let error = String::from_utf8_lossy(&wrk.stderr);
180                    if wrk.status.success() {
181                        debug!("Wrk execution succeded:\n{}", output);
182                        let wrk_json = output
183                            .split("JSON")
184                            .nth(1)
185                            .ok_or_else(|| WrkError::Lua("Wrk returned empty JSON".to_string()))?;
186                        self.wrk_result(wrk_json)
187                    } else {
188                        error!("Wrk execution failed.\nOutput: {}\nError: {}", output, error);
189                        WrkResult::fail(error.to_string())
190                    }
191                }
192                Err(e) => {
193                    error!("Wrk execution failed: {}", e);
194                    WrkResult::fail(e.to_string())
195                }
196            };
197            *run.date_mut() = date;
198            *run.benchmark_mut() = benchmark.clone();
199            self.benchmarks_mut().push(run);
200        }
201        script_file.keep()?;
202        self.dump(date)?;
203        Ok(())
204    }
205
206    pub fn bench_exponential(&mut self, duration: Option<Duration>) -> Result<()> {
207        self.bench(&BenchmarkBuilder::exponential(duration))?;
208        Ok(())
209    }
210
211    fn dump(&self, date: DateTime<Utc>) -> Result<()> {
212        let filename = format!("result.{}.json", date.format(DATE_FORMAT));
213        let file = File::create(self.history_dir().join(&filename))?;
214        let writer = BufWriter::new(file);
215        println!("Writing current benchmark to {}", filename);
216        serde_json::to_writer(writer, &self.benchmarks())?;
217        Ok(())
218    }
219
220    fn load(&mut self, period: HistoryPeriod, best: bool) -> Result<()> {
221        if !self.history_dir().exists() {
222            fs::create_dir(self.history_dir())?;
223        }
224        let mut paths: Vec<_> = fs::read_dir(self.history_dir())?.map(|r| r.unwrap()).collect();
225        paths.sort_by_key(|dir| {
226            let metadata = fs::metadata(dir.path()).unwrap();
227            metadata.modified().unwrap()
228        });
229        let mut history = Benchmarks::new();
230        if period == HistoryPeriod::Last {
231            let file = File::open(paths.pop().unwrap().path())?;
232            let mut reader = BufReader::new(file);
233            history = serde_json::from_reader(&mut reader)?;
234            let benchmark = history.pop().unwrap();
235            if let Some(benchmark_date) = self.benchmark_date() {
236                if benchmark_date == benchmark.date() && !paths.is_empty() {
237                    let file = File::open(paths.pop().unwrap().path())?;
238                    let mut reader = BufReader::new(file);
239                    history = serde_json::from_reader(&mut reader)?;
240                    if best {
241                        let best = self.best_benchmark(&history)?;
242                        history = vec![best];
243                    }
244                } else {
245                    return Err(WrkError::History(
246                        "Unable to load history with a single measurement".to_string(),
247                    ));
248                }
249            }
250        } else {
251            for path in paths {
252                if let Some(date_str) = path.file_name().to_string_lossy().split('.').nth(1) {
253                    let date = DateTime::parse_from_str(date_str, DATE_FORMAT)?;
254                    if date >= period.last_valid_datapoint() {
255                        let file = File::open(path.path())?;
256                        let mut reader = BufReader::new(file);
257                        let mut benchmarks: Vec<_> = serde_json::from_reader(&mut reader)?;
258                        benchmarks.retain(|x| !self.benchmarks_history().contains(x));
259                        if best {
260                            let best = self.best_benchmark(&benchmarks)?;
261                            history.push(best);
262                        } else {
263                            history.append(&mut benchmarks);
264                        }
265                    }
266                }
267            }
268        }
269        *self.benchmarks_history_mut() = history;
270        Ok(())
271    }
272
273    fn best_benchmark(&self, benchmarks: &Benchmarks) -> Result<WrkResult> {
274        let best = benchmarks.iter().filter(|v| *v.success()).max_by(|a, b| {
275            (*a.requests_sec() as i64)
276                .cmp(&(*b.requests_sec() as i64))
277                .then((*a.successes() as i64).cmp(&(*b.successes() as i64)))
278                .then((*a.requests() as i64).cmp(&(*b.requests() as i64)))
279                .then((*a.requests() as i64).cmp(&(*b.requests() as i64)))
280                .then((*a.transfer_mb() as i64).cmp(&(*b.transfer_mb() as i64)))
281        });
282        best.cloned().ok_or_else(|| {
283            WrkError::Stats(format!(
284                "Unable to calculate best in a set of {} elements",
285                benchmarks.len()
286            ))
287        })
288    }
289
290    fn best(&self) -> Result<WrkResult> {
291        self.best_benchmark(self.benchmarks())
292    }
293
294    fn historical_best(&self) -> Result<WrkResult> {
295        self.best_benchmark(self.benchmarks_history())
296    }
297
298    pub fn all_benchmarks(&self) -> Benchmarks {
299        let mut history = self.benchmarks_history().clone();
300        history.append(&mut self.benchmarks().clone());
301        history
302    }
303
304    pub fn deviation(&mut self, period: HistoryPeriod) -> Result<Deviation> {
305        self.load(period, false)?;
306        let new = self.best()?;
307        let old = self.historical_best()?;
308        Ok(Deviation::new(new, old))
309    }
310
311    pub fn plot(&self, title: &str, output: &Path, benchmarks: &Benchmarks) -> Result<()> {
312        Gnuplot::new(title, output).plot(benchmarks)
313    }
314}
315
316#[cfg(test)]
317mod tests {
318    use std::{net::SocketAddr, thread, time::Duration};
319
320    use super::*;
321    use crate::benchmark::BenchmarkBuilder;
322    use axum::{
323        http::StatusCode,
324        response::IntoResponse,
325        routing::{get, post},
326        Json, Router,
327    };
328    use http::Request;
329    use hyper::Body;
330
331    async fn server() {
332        let app = Router::new().route("/", get(|| async { "Hello, world!" }));
333
334        let addr = SocketAddr::from(([127, 0, 0, 1], 13734));
335        println!("server listening on {}", addr);
336        axum::Server::bind(&addr).serve(app.into_make_service()).await.unwrap();
337    }
338
339    #[tokio::test]
340    async fn benchmark() {
341        tokio::spawn(server());
342        tokio::time::sleep(Duration::from_secs(1)).await;
343        let client = hyper::Client::new();
344
345        let response = client
346            .request(
347                Request::builder()
348                    .uri("http://127.0.0.1:13734/")
349                    .body(Body::empty())
350                    .unwrap(),
351            )
352            .await
353            .unwrap();
354        let body = hyper::body::to_bytes(response.into_body()).await.unwrap();
355
356        let mut wrk = WrkBuilder::default()
357            .url("http://127.0.0.1:13734".to_string())
358            .build()
359            .unwrap();
360        // wrk.bench_exponential(Some(Duration::from_secs(30))).unwrap();
361        wrk.bench(&vec![BenchmarkBuilder::default()
362            .duration(Duration::from_secs(5))
363            .build()
364            .unwrap()])
365            .unwrap();
366        // println!("{}", wrk.deviation(HistoryPeriod::Hour).unwrap());
367        // wrk.load(HistoryPeriod::Day, false).unwrap();
368        // wrk.plot("Wrk Weeeeeee", Path::new("./some.png"), &wrk.all_benchmarks())
369        // .unwrap();
370    }
371}
wrk_api_bench/wrk.rs

wrk_api_bench/
wrk.rs
