ends time & request count

cchexcode · Sep 16, 2023 · 7d028bb · 7d028bb
1 parent febfa81
commit 7d028bb
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 26 deletions.
diff --git a/res/example.yaml b/res/example.yaml
@@ -3,20 +3,26 @@ version: "0.0"
 campaigns:
   main:
     phases:
-      - target: ""
+      - target: "https://solar-monitoring-dev.nexteraanalytics.com/v1/telemetry/sites"
         threads: 32
-        ends: !requests 2000
+        # ends: !requests 2000
+        ends: !seconds 60
         timeout_ms: 1000
         report:
           interval: !seconds 4
         spec: !get
           header:
-            x-api-key: [""]
+            x-api-key:
+              - !env "API_KEY"
           query:
-            page: [0]
-            per_page: [1000]
-            from: [1262304000]
-            to: [1262307600]
+            page:
+              - !static 0
+            per_page:
+              - !static 1000
+            from:
+              - !static 1262304000
+            to:
+              - !static 1262307600
         behaviours:
           ok:
             - match: ^(200)$

diff --git a/src/config.rs b/src/config.rs
@@ -47,11 +47,18 @@ pub enum Interval {
 #[serde(rename_all = "snake_case")]
 pub enum Spec {
     Get {
-        header: HashMap<String, Vec<String>>,
-        query: HashMap<String, Vec<String>>,
+        header: HashMap<String, Vec<ValueParser>>,
+        query: HashMap<String, Vec<ValueParser>>,
     },
 }
 
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum ValueParser {
+    Static(String),
+    Env(String),
+}
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 #[serde(rename_all = "snake_case")]
 pub struct Behaviours {
@@ -84,6 +91,7 @@ pub enum Mark {
 #[serde(rename_all = "snake_case")]
 pub enum End {
     Requests(usize),
+    Seconds(u64),
 }
 
 #[cfg(test)]

diff --git a/src/engine.rs b/src/engine.rs
@@ -5,6 +5,7 @@ use {
         ErrorBehaviour,
         Mark,
         Spec,
+        ValueParser,
     },
     anyhow::Result,
     fancy_regex::Regex,
@@ -50,8 +51,8 @@ impl Engine {
         for phase in &campaign.phases {
             let phase_start = std::time::Instant::now();
             let (tasks_tx, tasks_rx) =
-                flume::unbounded::<(Method, String, HeaderMap, Vec<(String, String)>, Duration)>();
-            let (status_tx, status_rx) = flume::unbounded::<(usize, ThreadEvent)>();
+                flume::bounded::<(Method, String, HeaderMap, Vec<(String, String)>, Duration)>(phase.threads * 2);
+            let (status_tx, status_rx) = flume::bounded::<(usize, ThreadEvent)>(phase.threads * 2);
 
             let mut threads = Vec::<JoinHandle<_>>::with_capacity(phase.threads);
             let mut thread_stats = BTreeMap::<usize, ThreadStats>::new();
@@ -97,23 +98,48 @@ impl Engine {
 
             match &phase.spec {
                 | Spec::Get { header, query } => {
+                    let header_map = HeaderMap::from_iter(
+                        header
+                            .iter()
+                            .map(|v| {
+                                (
+                                    v.0.parse().unwrap(),
+                                    v.1.into_iter()
+                                        .map(|v| {
+                                            match v {
+                                                | ValueParser::Static(v) => v.to_owned(),
+                                                | ValueParser::Env(v) => std::env::var(v).unwrap(),
+                                            }
+                                        })
+                                        .join(",")
+                                        .parse()
+                                        .unwrap(),
+                                )
+                            })
+                            .collect::<Vec<(HeaderName, HeaderValue)>>(),
+                    );
+                    let query_map = query
+                        .iter()
+                        .map(|v| {
+                            (
+                                v.0.clone(),
+                                v.1.into_iter()
+                                    .map(|v| {
+                                        match v {
+                                            | ValueParser::Static(v) => v.to_owned(),
+                                            | ValueParser::Env(v) => std::env::var(v).unwrap(),
+                                        }
+                                    })
+                                    .join(","),
+                            )
+                        })
+                        .collect::<Vec<_>>();
+
+                    let target = phase.target.clone();
+                    let timeout_ms = phase.timeout_ms;
+
                     match phase.ends {
                         | End::Requests(v) => {
-                            let header_map = HeaderMap::from_iter(
-                                header
-                                    .iter()
-                                    .map(|v| (v.0.parse().unwrap(), v.1.into_iter().join(",").parse().unwrap()))
-                                    .collect::<Vec<(HeaderName, HeaderValue)>>(),
-                            );
-                            let query_map = query
-                                .iter()
-                                .map(|v| (v.0.clone(), v.1.into_iter().join(",")))
-                                .collect::<Vec<_>>();
-
-                            let target = phase.target.clone();
-                            let timeout_ms = phase.timeout_ms;
-
-                            // producer thread
                             spawn(move || {
                                 for _ in 0..v {
                                     tasks_tx
@@ -128,6 +154,22 @@ impl Engine {
                                 }
                             });
                         },
+                        | End::Seconds(v) => {
+                            let now = std::time::Instant::now();
+                            spawn(move || {
+                                while now.elapsed().as_secs() < v {
+                                    tasks_tx
+                                        .send((
+                                            Method::GET,
+                                            target.clone(),
+                                            header_map.clone(),
+                                            query_map.clone(),
+                                            Duration::from_millis(timeout_ms),
+                                        ))
+                                        .unwrap();
+                                }
+                            });
+                        },
                     }
                 },
             };