save request results into sqlite database

Cargo.toml

@@ -69,6 +69,7 @@ pin-project-lite = "0.2.14"
 http-body-util = "0.1.2"
 hyper-util = { version = "0.1.6", features = ["tokio"] }
 tokio-vsock = { version = "0.5.0", optional = true }
+rusqlite = { version = "0.32.0", features = ["bundled"] }
 
 [target.'cfg(unix)'.dependencies]
 rlimit = "0.10.1"

src/db.rs

@@ -0,0 +1,62 @@
+use rusqlite::Connection;
+
+use crate::client::RequestResult;
+
+fn create_db(conn: &Connection) -> Result<usize, rusqlite::Error> {
+    conn.execute(
+        "CREATE TABLE oha (
+            start REAL NOT NULL,
+            start_latency_correction REAL,
+            end REAL NOT NULL,
+            duration REAL NOT NULL,
+            status INTEGER NOT NULL,
+            len_bytes INTEGER NOT NULL
+        )",
+        (),
+    )
+}
+
+pub fn store(
+    db_url: &str,
+    start: std::time::Instant,
+    request_records: &[RequestResult],
+) -> Result<usize, rusqlite::Error> {
+    let mut conn = Connection::open(db_url)?;
+    create_db(&conn)?;
+
+    let t = conn.transaction()?;
+    let affected_rows =
+        request_records
+            .iter()
+            .map(|req| {
+                t.execute(
+          "INSERT INTO oha (start, start_latency_correction, end, duration, status, len_bytes) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
+        ((req.start - start).as_secs_f64(), req.start_latency_correction.map(|d| (d-start).as_secs_f64()), (req.end - start).as_secs_f64(), req.duration().as_secs_f64(), req.status.as_u16() as i64, req.len_bytes)
+        ).unwrap_or(0)
+            })
+            .sum();
+    t.commit()?;
+
+    Ok(affected_rows)
+}
+
+#[cfg(test)]
+mod test_db {
+    use super::*;
+
+    #[test]
+    fn test_store() {
+        let start = std::time::Instant::now();
+        let test_val = RequestResult {
+            status: hyper::StatusCode::OK,
+            len_bytes: 100,
+            start_latency_correction: None,
+            start: std::time::Instant::now(),
+            connection_time: None,
+            end: std::time::Instant::now(),
+        };
+        let test_vec = vec![test_val.clone(), test_val.clone()];
+        let result = store(":memory:", start, &test_vec);
+        assert_eq!(result.unwrap(), 2);
+    }
+}

src/main.rs

@@ -17,6 +17,7 @@ use url::Url;
 use url_generator::UrlGenerator;
 
 mod client;
+mod db;
 mod histogram;
 mod monitor;
 mod printer;
@@ -188,6 +189,11 @@ Note: If qps is specified, burst will be ignored",
         long = "stats-success-breakdown"
     )]
     stats_success_breakdown: bool,
+    #[clap(
+        help = "Write succeeded requests to sqlite database url E.G test.db",
+        long = "db-url"
+    )]
+    db_url: Option<String>,
 }
 
 /// An entry specified by `connect-to` to override DNS resolution and default
@@ -654,6 +660,11 @@ async fn main() -> anyhow::Result<()> {
         opts.stats_success_breakdown,
     )?;
 
+    if let Some(db_url) = opts.db_url {
+        eprintln!("Storing results to {db_url}");
+        let _ = db::store(&db_url, start, res.success());
+    }
+
     Ok(())
 }