wrk_api_bench/
lua.rs

1use std::{
2    env,
3    fs::{self, File},
4    io::{BufReader, Read, Write},
5    path::{Path, PathBuf},
6};
7
8use assert_cmd::prelude::OutputOkExt;
9use getset::{Getters, MutGetters, Setters};
10use rslua::lexer::Lexer;
11use tempfile::NamedTempFile;
12
13use crate::{Headers, Result, WrkError};
14
15const LUA_DEFAULT_DONE_FUNCTION: &str = r#"
16-- The done() function is called at the end of wrk execution
17-- and allows us to produce a well formed JSON output, prefixed
18-- by the string "JSON" which allows us to parse the wrk output
19-- easily.
20done = function(summary, latency, requests)
21    local errors = summary.errors.connect
22        + summary.errors.read
23        + summary.errors.write
24        + summary.errors.status
25        + summary.errors.timeout
26    io.write("JSON")
27    io.write(string.format(
28        [[{
29    "requests": %.2f,
30    "errors": %.2f,
31    "successes": %.2f,
32    "requests_sec": %.2f,
33    "avg_latency_ms": %.2f,
34    "min_latency_ms": %.2f,
35    "max_latency_ms": %.2f,
36    "stdev_latency_ms": %.2f,
37    "transfer_mb": %.2f,
38    "errors_connect": %.2f,
39    "errors_read": %.2f,
40    "errors_write": %.2f,
41    "errors_status": %.2f,
42    "errors_timeout": %.2f
43}
44]],
45        summary.requests,
46        errors,
47        summary.requests - errors,
48        summary.requests / (summary.duration / 1000000),
49        (latency.mean / 1000),
50        (latency.min / 1000),
51        (latency.max / 1000),
52        (latency.stdev / 1000),
53        (summary.bytes / 1048576),
54        summary.errors.connect,
55        summary.errors.read,
56        summary.errors.write,
57        summary.errors.status,
58        summary.errors.timeout
59    ))
60end
61"#;
62
63#[derive(Debug)]
64pub struct LuaScript {}
65
66impl LuaScript {
67    fn lua_script_from_config(&mut self, uri: &str, method: &str, headers: &Headers, body: &str) -> Result<String> {
68        let request = format!(
69            r#"
70-- The request() function is called by wrk on all requests
71-- and allow us to configure things like headers, method, body, etc..
72request = function()
73    wrk.method = "{}"
74    wrk.body = "{}"
75    {}
76    return wrk.format("{}", "{}")
77end
78        "#,
79            method,
80            body,
81            self.lua_headers(headers)?,
82            method,
83            uri
84        );
85        let buffer = request + LUA_DEFAULT_DONE_FUNCTION;
86        Ok(buffer)
87    }
88
89    fn lua_script_from_user(&mut self, lua_script: &Path) -> Result<String> {
90        let file = File::open(lua_script)?;
91        let mut reader = BufReader::new(file);
92        let mut buffer = String::new();
93        reader.read_to_string(&mut buffer)?;
94        let mut lexer = Lexer::new();
95        let tokens = lexer.run(&buffer).map_err(|e| WrkError::Lua(format!("{:?}", e)))?;
96        let buffer = buffer + LUA_DEFAULT_DONE_FUNCTION;
97        Ok(buffer)
98    }
99
100    fn lua_headers(&self, headers: &Headers) -> Result<String> {
101        let mut result = String::new();
102        for (k, v) in headers {
103            result += &format!(r#"wrk.headers["{}"] = "{}"\n"#, k, v);
104        }
105        Ok(result)
106    }
107
108    pub fn render(
109        script_file: &mut NamedTempFile,
110        user_script: Option<&PathBuf>,
111        uri: &str,
112        method: &str,
113        headers: &Headers,
114        body: &str,
115    ) -> Result<()> {
116        let mut this = Self {};
117        let script = match user_script {
118            Some(lua_script) => {
119                if !lua_script.exists() {
120                    error!(
121                        "Wrk Lua file {} not found in {}",
122                        env::current_dir().expect("unable to get current directory").display(),
123                        lua_script.display()
124                    );
125                    return Err(WrkError::Lua("Wrk Lua file not found".to_string()));
126                } else {
127                    this.lua_script_from_user(&lua_script)?
128                }
129            }
130            None => this.lua_script_from_config(uri, method, headers, body)?,
131        };
132        script_file.write_all(script.as_bytes())?;
133        Ok(())
134    }
135}
wrk_api_bench/lua.rs

wrk_api_bench/
lua.rs
