+ behaviours

cchexcode · Sep 16, 2023 · febfa81 · febfa81
1 parent 51fe585
commit febfa81
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 25 deletions.
diff --git a/res/example.yaml b/res/example.yaml
@@ -17,9 +17,11 @@ campaigns:
             per_page: [1000]
             from: [1262304000]
             to: [1262307600]
-        # behaviours:
-        #   - match: ^(200)$
-        #     mark: !success
-        #   - match: .*
-        #     mark: !error
-        #     sleep: 1000
+        behaviours:
+          ok:
+            - match: ^(200)$
+              mark: !success
+            - match: .*
+              mark: !error
+              sleep: 1000
+          error: !backoff 1000
diff --git a/src/config.rs b/src/config.rs
@@ -28,7 +28,7 @@ pub struct Phase {
     pub timeout_ms: u64,
     pub report: Report,
     pub spec: Spec,
-    // pub behaviours: Vec<Behaviour>,
+    pub behaviours: Behaviours,
 }
 
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
@@ -52,20 +52,33 @@ pub enum Spec {
     },
 }
 
-// #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
-// #[serde(rename_all = "snake_case")]
-// pub struct Behaviour {
-//     #[serde(rename = "match")]
-//     pub match_: String,
-//     pub mark: Mark,
-// }
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub struct Behaviours {
+    pub ok: Vec<Behaviour>,
+    pub error: ErrorBehaviour,
+}
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum ErrorBehaviour {
+    Backoff(u64),
+}
 
-// #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
-// #[serde(rename_all = "snake_case")]
-// pub enum Mark {
-//     Success,
-//     Error,
-// }
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub struct Behaviour {
+    #[serde(rename = "match")]
+    pub match_: String,
+    pub mark: Mark,
+}
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum Mark {
+    Success,
+    Error,
+}
 
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 #[serde(rename_all = "snake_case")]

diff --git a/src/engine.rs b/src/engine.rs
@@ -2,9 +2,12 @@ use {
     crate::config::{
         Campaign,
         End,
+        ErrorBehaviour,
+        Mark,
         Spec,
     },
     anyhow::Result,
+    fancy_regex::Regex,
     itertools::Itertools,
     reqwest::{
         header::{
@@ -129,20 +132,33 @@ impl Engine {
                 },
             };
 
+            let mut behaviours = Vec::<(Regex, &Mark)>::new();
+            for behav in &phase.behaviours.ok {
+                behaviours.push((Regex::new(&behav.match_).unwrap(), &behav.mark));
+            }
+
             for msg in status_rx.iter() {
                 let stats = &mut thread_stats.get_mut(&msg.0).unwrap();
                 match msg.1 {
                     | ThreadEvent::Success { status_code } => {
                         stats.count += 1;
-                        if let StatusCode::OK = status_code {
-                            stats.success += 1;
-                        } else {
-                            stats.error += 1;
-                        };
+
+                        let s_code = status_code.to_string();
+                        for b in &behaviours {
+                            if b.0.is_match(&s_code).unwrap() {
+                                match b.1 {
+                                    | Mark::Success => stats.success += 1,
+                                    | Mark::Error => stats.error += 1,
+                                }
+                            }
+                        }
                     },
                     | ThreadEvent::Error {} => {
                         stats.count += 1;
                         stats.client_error += 1;
+                        match phase.behaviours.error {
+                            | ErrorBehaviour::Backoff(v) => std::thread::sleep(Duration::from_millis(v)),
+                        }
                     },
                 };