From 734d79bdbc612f60b81f53a2d5a5007474fd0211 Mon Sep 17 00:00:00 2001
From: Floyd Finnegan <floyd.finnegan@protonmail.com>
Date: Tue, 31 Aug 2021 13:03:01 +0200
Subject: [PATCH] Allow shared reqwest client configuration

---
 CHANGELOG.md  |  3 ++-
 src/goose.rs  | 56 ++++++++++++++++++++++++++++------------------
 src/lib.rs    | 62 +++++++++++++++++++++++++++++++++++++--------------
 src/worker.rs |  3 +++
 4 files changed, 84 insertions(+), 40 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 23d058dad..beb1c0d78 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,7 +4,8 @@
  - Add the possibility to attach custom session data `GooseUserData` to each `GooseUser` (API change)
  - Change `GooseTask` signature to take a mutable reference of `GooseUser` (API change)
  - Remove `Clone` trait from `GooseUser` and `GooseAttack`
- 
+ - Use a shared reqwest client among all `GooseUser` 
+ - 
 ## 0.13.3 August 25, 2021
  - document GooseConfiguration fields that were only documented as gumpdrop parameters (in order to generate new lines in the help output) so now they're also documented in the code
  - fix panic when `--no-task-metrics` is enabled and metrics are printed; add tests to prevent further regressions
diff --git a/src/goose.rs b/src/goose.rs
index bf92775c1..1de1d3eb9 100644
--- a/src/goose.rs
+++ b/src/goose.rs
@@ -301,9 +301,6 @@ use crate::metrics::{
 };
 use crate::{GooseConfiguration, GooseError, WeightedGooseTasks};
 
-/// By default Goose sets the following User-Agent header when making requests.
-static APP_USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));
-
 /// `task!(foo)` expands to `GooseTask::new(foo)`, but also does some boxing to work around a limitation in the compiler.
 #[macro_export]
 macro_rules! task {
@@ -866,6 +863,7 @@ pub struct GooseUser {
 impl GooseUser {
     /// Create a new user state.
     pub fn new(
+        client: Client,
         task_sets_index: usize,
         base_url: Url,
         min_wait: usize,
@@ -874,12 +872,6 @@ impl GooseUser {
         load_test_hash: u64,
     ) -> Result<Self, GooseError> {
         trace!("new GooseUser");
-        let client = Client::builder()
-            .user_agent(APP_USER_AGENT)
-            .cookie_store(true)
-            // Enable gzip unless `--no-gzip` flag is enabled.
-            .gzip(!configuration.no_gzip)
-            .build()?;
 
         Ok(GooseUser {
             started: Instant::now(),
@@ -904,8 +896,12 @@ impl GooseUser {
     }
 
     /// Create a new single-use user.
-    pub fn single(base_url: Url, configuration: &GooseConfiguration) -> Result<Self, GooseError> {
-        let mut single_user = GooseUser::new(0, base_url, 0, 0, configuration, 0)?;
+    pub fn single(
+        client: Client,
+        base_url: Url,
+        configuration: &GooseConfiguration,
+    ) -> Result<Self, GooseError> {
+        let mut single_user = GooseUser::new(client, 0, base_url, 0, 0, configuration, 0)?;
         // Only one user, so index is 0.
         single_user.weighted_users_index = 0;
         // Do not throttle [`test_start`](../struct.GooseAttack.html#method.test_start) (setup) and
@@ -2594,10 +2590,11 @@ mod tests {
     const EMPTY_ARGS: Vec<&str> = vec![];
 
     fn setup_user(server: &MockServer) -> Result<GooseUser, GooseError> {
+        let client = Client::builder().build().unwrap();
         let mut configuration = GooseConfiguration::parse_args_default(&EMPTY_ARGS).unwrap();
         configuration.co_mitigation = Some(GooseCoordinatedOmissionMitigation::Average);
         let base_url = get_base_url(Some(server.url("/")), None, None).unwrap();
-        GooseUser::single(base_url, &configuration)
+        GooseUser::single(client, base_url, &configuration)
     }
 
     #[test]
@@ -2784,9 +2781,10 @@ mod tests {
     #[tokio::test]
     async fn goose_user() {
         const HOST: &str = "http://example.com/";
+        let client = Client::builder().build().unwrap();
         let configuration = GooseConfiguration::parse_args_default(&EMPTY_ARGS).unwrap();
         let base_url = get_base_url(Some(HOST.to_string()), None, None).unwrap();
-        let user = GooseUser::new(0, base_url, 0, 0, &configuration, 0).unwrap();
+        let user = GooseUser::new(client.clone(), 0, base_url, 0, 0, &configuration, 0).unwrap();
         assert_eq!(user.task_sets_index, 0);
         assert_eq!(user.min_wait, 0);
         assert_eq!(user.max_wait, 0);
@@ -2815,7 +2813,7 @@ mod tests {
             Some("http://www.example.com/".to_string()),
         )
         .unwrap();
-        let user2 = GooseUser::new(0, base_url, 1, 3, &configuration, 0).unwrap();
+        let user2 = GooseUser::new(client.clone(), 0, base_url, 1, 3, &configuration, 0).unwrap();
         assert_eq!(user2.min_wait, 1);
         assert_eq!(user2.max_wait, 3);
 
@@ -2876,7 +2874,7 @@ mod tests {
         // Confirm Goose can build a base_url that includes a path.
         const HOST_WITH_PATH: &str = "http://example.com/with/path/";
         let base_url = get_base_url(Some(HOST_WITH_PATH.to_string()), None, None).unwrap();
-        let user = GooseUser::new(0, base_url, 0, 0, &configuration, 0).unwrap();
+        let user = GooseUser::new(client, 0, base_url, 0, 0, &configuration, 0).unwrap();
 
         // Confirm the URLs are correctly built using the default_host that includes a path.
         let url = user.build_url("foo").unwrap();
@@ -2978,9 +2976,14 @@ mod tests {
             data: "foo".to_owned(),
         };
 
+        let client = Client::builder().build().unwrap();
         let configuration = GooseConfiguration::parse_args_default(&EMPTY_ARGS).unwrap();
-        let mut user =
-            GooseUser::single("http://localhost:8080".parse().unwrap(), &configuration).unwrap();
+        let mut user = GooseUser::single(
+            client,
+            "http://localhost:8080".parse().unwrap(),
+            &configuration,
+        )
+        .unwrap();
 
         user.set_session_data(session_data.clone());
 
@@ -3002,10 +3005,14 @@ mod tests {
         let session_data = CustomSessionData {
             data: "foo".to_owned(),
         };
-
+        let client = Client::builder().build().unwrap();
         let configuration = GooseConfiguration::parse_args_default(&EMPTY_ARGS).unwrap();
-        let mut user =
-            GooseUser::single("http://localhost:8080".parse().unwrap(), &configuration).unwrap();
+        let mut user = GooseUser::single(
+            client,
+            "http://localhost:8080".parse().unwrap(),
+            &configuration,
+        )
+        .unwrap();
 
         user.set_session_data(session_data);
 
@@ -3033,9 +3040,14 @@ mod tests {
             data: "foo".to_owned(),
         };
 
+        let client = Client::builder().build().unwrap();
         let configuration = GooseConfiguration::parse_args_default(&EMPTY_ARGS).unwrap();
-        let mut user =
-            GooseUser::single("http://localhost:8080".parse().unwrap(), &configuration).unwrap();
+        let mut user = GooseUser::single(
+            client,
+            "http://localhost:8080".parse().unwrap(),
+            &configuration,
+        )
+        .unwrap();
 
         user.set_session_data(session_data.clone());
 
diff --git a/src/lib.rs b/src/lib.rs
index fce61b5ad..8dfd80c52 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -459,6 +459,7 @@ use lazy_static::lazy_static;
 use nng::Socket;
 use rand::seq::SliceRandom;
 use rand::thread_rng;
+use reqwest::Client;
 use std::collections::hash_map::DefaultHasher;
 use std::collections::BTreeMap;
 use std::hash::{Hash, Hasher};
@@ -489,6 +490,9 @@ lazy_static! {
     static ref WORKER_ID: AtomicUsize = AtomicUsize::new(0);
 }
 
+/// By default Goose sets the following User-Agent header when making requests.
+pub static APP_USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));
+
 /// Internal representation of a weighted task list.
 type WeightedGooseTasks = Vec<(usize, String)>;
 
@@ -781,6 +785,8 @@ pub struct GooseAttack {
     started: Option<time::Instant>,
     /// All metrics merged together.
     metrics: GooseMetrics,
+    /// Shared request client
+    client: Client,
 }
 /// Goose's internal global state.
 impl GooseAttack {
@@ -793,21 +799,7 @@ impl GooseAttack {
     /// let mut goose_attack = GooseAttack::initialize();
     /// ```
     pub fn initialize() -> Result<GooseAttack, GooseError> {
-        Ok(GooseAttack {
-            test_start_task: None,
-            test_stop_task: None,
-            task_sets: Vec::new(),
-            weighted_users: Vec::new(),
-            weighted_gaggle_users: Vec::new(),
-            defaults: GooseDefaults::default(),
-            configuration: GooseConfiguration::parse_args_default_or_exit(),
-            run_time: 0,
-            attack_mode: AttackMode::Undefined,
-            attack_phase: AttackPhase::Idle,
-            scheduler: GooseScheduler::RoundRobin,
-            started: None,
-            metrics: GooseMetrics::default(),
-        })
+        Self::initialize_with_config(GooseConfiguration::parse_args_default_or_exit())
     }
 
     /// Initialize a [`GooseAttack`](./struct.GooseAttack.html) with an already loaded
@@ -827,6 +819,13 @@ impl GooseAttack {
     pub fn initialize_with_config(
         configuration: GooseConfiguration,
     ) -> Result<GooseAttack, GooseError> {
+        let client = Client::builder()
+            .user_agent(APP_USER_AGENT)
+            .cookie_store(true)
+            // Enable gzip unless `--no-gzip` flag is enabled.
+            .gzip(!configuration.no_gzip)
+            .build()?;
+
         Ok(GooseAttack {
             test_start_task: None,
             test_stop_task: None,
@@ -841,9 +840,35 @@ impl GooseAttack {
             scheduler: GooseScheduler::RoundRobin,
             started: None,
             metrics: GooseMetrics::default(),
+            client,
         })
     }
 
+    /// Define a reqwest::Client shared among all the
+    /// [`GooseUser`](./goose/struct.GooseUser.html)s running.
+    ///
+    /// # Example
+    /// ```rust
+    /// use goose::prelude::*;
+    /// use reqwest::Client;
+    ///
+    /// #[tokio::main]
+    /// async fn main() -> Result<(), GooseError> {
+    ///     let client = Client::builder()
+    ///         .build()?;
+    ///
+    ///     GooseAttack::initialize()?
+    ///         .set_scheduler(GooseScheduler::Random)
+    ///         .set_client(client);
+    ///
+    ///     Ok(())
+    /// }
+    /// ```
+    pub fn set_client(mut self, client: Client) -> Self {
+        self.client = client;
+        self
+    }
+
     /// Define the order [`GooseTaskSet`](./goose/struct.GooseTaskSet.html)s are
     /// allocated to new [`GooseUser`](./goose/struct.GooseUser.html)s as they are
     /// launched.
@@ -1130,6 +1155,7 @@ impl GooseAttack {
                     self.defaults.host.clone(),
                 )?;
                 weighted_users.push(GooseUser::new(
+                    self.client.clone(),
                     self.task_sets[*task_sets_index].task_sets_index,
                     base_url,
                     self.task_sets[*task_sets_index].min_wait,
@@ -1592,7 +1618,8 @@ impl GooseAttack {
                         None,
                         self.defaults.host.clone(),
                     )?;
-                    let mut user = GooseUser::single(base_url, &self.configuration)?;
+                    let mut user =
+                        GooseUser::single(self.client.clone(), base_url, &self.configuration)?;
                     let function = &t.function;
                     let _ = function(&mut user).await;
                 }
@@ -1618,7 +1645,8 @@ impl GooseAttack {
                         None,
                         self.defaults.host.clone(),
                     )?;
-                    let mut user = GooseUser::single(base_url, &self.configuration)?;
+                    let mut user =
+                        GooseUser::single(self.client.clone(), base_url, &self.configuration)?;
                     let function = &t.function;
                     let _ = function(&mut user).await;
                 }
diff --git a/src/worker.rs b/src/worker.rs
index 2848e5549..371220c97 100644
--- a/src/worker.rs
+++ b/src/worker.rs
@@ -1,5 +1,6 @@
 use gumdrop::Options;
 use nng::*;
+use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use std::io::BufWriter;
 use std::sync::atomic::Ordering;
@@ -136,7 +137,9 @@ pub(crate) async fn worker_main(goose_attack: &GooseAttack) -> GooseAttack {
         if worker_id == 0 {
             worker_id = initializer.worker_id;
         }
+        let client = Client::builder().build().unwrap();
         let user = GooseUser::new(
+            client,
             initializer.task_sets_index,
             Url::parse(&initializer.base_url).unwrap(),
             initializer.min_wait,