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,