From 05be749369d11e18ecbd7bf0d4d82b0054ae9288 Mon Sep 17 00:00:00 2001 From: Floyd Finnegan Date: Tue, 31 Aug 2021 13:03:01 +0200 Subject: [PATCH 1/3] Allow shared reqwest client configuration --- CHANGELOG.md | 2 +- src/goose.rs | 56 ++++++++++++++++++++++++++++------------------ src/lib.rs | 62 +++++++++++++++++++++++++++++++++++++-------------- src/worker.rs | 3 +++ 4 files changed, 83 insertions(+), 40 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 26b11765..78da269d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,7 @@ - update `GooseTaskSet::set_wait_time()` to accept `std::time::Duration` instead of `usize` allowing more granularity (API change) - use request name when displaying errors to avoid having a large volume of distinct error for the same endpoint when using path parameters - updated `tungstenite` dependency to [`0.15`](https://github.com/snapview/tungstenite-rs/blob/master/CHANGELOG.md) - + - 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 96b1fb35..638288fd 100644 --- a/src/goose.rs +++ b/src/goose.rs @@ -303,9 +303,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 { @@ -857,18 +854,13 @@ pub struct GooseUser { impl GooseUser { /// Create a new user state. pub fn new( + client: Client, task_sets_index: usize, base_url: Url, configuration: &GooseConfiguration, load_test_hash: u64, ) -> Result { 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(), @@ -891,8 +883,12 @@ impl GooseUser { } /// Create a new single-use user. - pub fn single(base_url: Url, configuration: &GooseConfiguration) -> Result { - let mut single_user = GooseUser::new(0, base_url, configuration, 0)?; + pub fn single( + client: Client, + base_url: Url, + configuration: &GooseConfiguration, + ) -> Result { + let mut single_user = GooseUser::new(client, 0, base_url, 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 @@ -2606,10 +2602,11 @@ mod tests { const EMPTY_ARGS: Vec<&str> = vec![]; fn setup_user(server: &MockServer) -> Result { + 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] @@ -2797,9 +2794,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, &configuration, 0).unwrap(); + let user = GooseUser::new(client.clone(), 0, base_url, &configuration, 0).unwrap(); assert_eq!(user.task_sets_index, 0); assert_eq!(user.weighted_users_index, usize::max_value()); @@ -2826,7 +2824,7 @@ mod tests { Some("http://www.example.com/".to_string()), ) .unwrap(); - let user2 = GooseUser::new(0, base_url, &configuration, 0).unwrap(); + let user2 = GooseUser::new(client.clone(), 0, base_url, &configuration, 0).unwrap(); // Confirm the URLs are correctly built using the task_set_host. let url = user2.build_url("/foo").unwrap(); @@ -2885,7 +2883,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, &configuration, 0).unwrap(); + let user = GooseUser::new(client, 0, base_url, &configuration, 0).unwrap(); // Confirm the URLs are correctly built using the default_host that includes a path. let url = user.build_url("foo").unwrap(); @@ -2987,9 +2985,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()); @@ -3011,10 +3014,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); @@ -3042,9 +3049,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 933fb99a..98ba0c25 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}; @@ -490,6 +491,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)>; @@ -782,6 +786,8 @@ pub struct GooseAttack { started: Option, /// All metrics merged together. metrics: GooseMetrics, + /// Shared request client + client: Client, } /// Goose's internal global state. impl GooseAttack { @@ -794,21 +800,7 @@ impl GooseAttack { /// let mut goose_attack = GooseAttack::initialize(); /// ``` pub fn initialize() -> Result { - 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 @@ -828,6 +820,13 @@ impl GooseAttack { pub fn initialize_with_config( configuration: GooseConfiguration, ) -> Result { + 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, @@ -842,9 +841,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. @@ -1131,6 +1156,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.configuration, @@ -1589,7 +1615,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; } @@ -1615,7 +1642,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 b03971d7..4c8814ef 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.config, From 66e6123bda8c731766684241e0f54402d39ff254 Mon Sep 17 00:00:00 2001 From: Floyd Finnegan Date: Sun, 5 Sep 2021 17:51:05 +0200 Subject: [PATCH 2/3] Disable cookies in reqwest::Client and handle cookies in GooseUser so we can share the client across all running GooseUser --- src/goose.rs | 19 ++++++++++++++++--- src/lib.rs | 15 +++++++-------- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/src/goose.rs b/src/goose.rs index 638288fd..90c117e1 100644 --- a/src/goose.rs +++ b/src/goose.rs @@ -287,7 +287,8 @@ use downcast_rs::{impl_downcast, Downcast}; use http::method::Method; -use reqwest::{header, Client, ClientBuilder, RequestBuilder, Response}; +use reqwest::cookie::{CookieStore, Jar}; +use reqwest::{cookie, header, Client, ClientBuilder, RequestBuilder, Response}; use serde::{Deserialize, Serialize}; use std::hash::{Hash, Hasher}; use std::sync::Arc; @@ -850,6 +851,8 @@ pub struct GooseUser { pub(crate) task_name: Option, session_data: Option>, + + cookie_store: Jar, } impl GooseUser { /// Create a new user state. @@ -879,6 +882,7 @@ impl GooseUser { slept: 0, task_name: None, session_data: None, + cookie_store: cookie::Jar::default(), }) } @@ -1533,8 +1537,7 @@ impl GooseUser { }; let started = Instant::now(); - let request = request_builder.build()?; - + let mut request = request_builder.build()?; // String version of request path. let path = match Url::parse(&request.url().to_string()) { Ok(u) => u.path().to_string(), @@ -1578,6 +1581,8 @@ impl GooseUser { self.weighted_users_index, ); + self.add_cookie_header(&mut request); + // Make the actual request. let response = self.client.execute(request).await; request_metric.set_response_time(started.elapsed().as_millis()); @@ -1638,6 +1643,14 @@ impl GooseUser { Ok(GooseResponse::new(request_metric, response)) } + fn add_cookie_header(&self, request: &mut reqwest::Request) { + if request.headers().get(header::COOKIE).is_none() { + if let Some(header) = self.cookie_store.cookies(request.url()) { + request.headers_mut().insert(header::COOKIE, header); + } + } + } + /// Tracks the time it takes for the current GooseUser to loop through all GooseTasks /// if Coordinated Omission Mitigation is enabled. pub(crate) async fn update_request_cadence(&mut self, thread_number: usize) { diff --git a/src/lib.rs b/src/lib.rs index 98ba0c25..07e8ca01 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -459,7 +459,7 @@ use lazy_static::lazy_static; use nng::Socket; use rand::seq::SliceRandom; use rand::thread_rng; -use reqwest::Client; +use reqwest::{Client, ClientBuilder}; use std::collections::hash_map::DefaultHasher; use std::collections::BTreeMap; use std::hash::{Hash, Hasher}; @@ -822,7 +822,7 @@ impl GooseAttack { ) -> Result { let client = Client::builder() .user_agent(APP_USER_AGENT) - .cookie_store(true) + .cookie_store(false) // Enable gzip unless `--no-gzip` flag is enabled. .gzip(!configuration.no_gzip) .build()?; @@ -855,19 +855,18 @@ impl GooseAttack { /// /// #[tokio::main] /// async fn main() -> Result<(), GooseError> { - /// let client = Client::builder() - /// .build()?; + /// let client_builder = Client::builder(); /// /// GooseAttack::initialize()? /// .set_scheduler(GooseScheduler::Random) - /// .set_client(client); + /// .set_client(client_builder)?; /// /// Ok(()) /// } /// ``` - pub fn set_client(mut self, client: Client) -> Self { - self.client = client; - self + pub fn set_client(mut self, client_builder: ClientBuilder) -> Result { + self.client = client_builder.cookie_store(false).build()?; + Ok(self) } /// Define the order [`GooseTaskSet`](./goose/struct.GooseTaskSet.html)s are From 0df6e15779f05f700380946f560fa2fa2d721a14 Mon Sep 17 00:00:00 2001 From: Floyd Finnegan Date: Mon, 6 Sep 2021 10:31:43 +0200 Subject: [PATCH 3/3] Add add_cookie_str and add_default_header to GooseUser API. Remove set_client_builder from GooseUser API. Rename set_client to set_client_builder on GooseAttack --- src/goose.rs | 174 +++++++++++++++------------------------------------ src/lib.rs | 29 +++++++-- 2 files changed, 74 insertions(+), 129 deletions(-) diff --git a/src/goose.rs b/src/goose.rs index 90c117e1..5b906752 100644 --- a/src/goose.rs +++ b/src/goose.rs @@ -286,9 +286,11 @@ //! limitations under the License. use downcast_rs::{impl_downcast, Downcast}; +use http::header::{Entry, IntoHeaderName}; use http::method::Method; +use http::HeaderValue; use reqwest::cookie::{CookieStore, Jar}; -use reqwest::{cookie, header, Client, ClientBuilder, RequestBuilder, Response}; +use reqwest::{cookie, header, Client, RequestBuilder, Response}; use serde::{Deserialize, Serialize}; use std::hash::{Hash, Hasher}; use std::sync::Arc; @@ -853,6 +855,8 @@ pub struct GooseUser { session_data: Option>, cookie_store: Jar, + + headers: header::HeaderMap, } impl GooseUser { /// Create a new user state. @@ -883,6 +887,7 @@ impl GooseUser { task_name: None, session_data: None, cookie_store: cookie::Jar::default(), + headers: header::HeaderMap::default(), }) } @@ -1581,8 +1586,12 @@ impl GooseUser { self.weighted_users_index, ); + // Set cookies self.add_cookie_header(&mut request); + // Set default headers + self.add_default_headers(&mut request); + // Make the actual request. let response = self.client.execute(request).await; request_metric.set_response_time(started.elapsed().as_millis()); @@ -1643,6 +1652,7 @@ impl GooseUser { Ok(GooseResponse::new(request_metric, response)) } + // Add cookies from the cookie store. fn add_cookie_header(&self, request: &mut reqwest::Request) { if request.headers().get(header::COOKIE).is_none() { if let Some(header) = self.cookie_store.cookies(request.url()) { @@ -1651,6 +1661,17 @@ impl GooseUser { } } + // insert default headers in the request headers + // without overwriting already appended headers. + fn add_default_headers(&self, request: &mut reqwest::Request) { + let headers = request.headers_mut(); + for (key, value) in &self.headers { + if let Entry::Vacant(entry) = headers.entry(key) { + entry.insert(value.clone()); + } + } + } + /// Tracks the time it takes for the current GooseUser to loop through all GooseTasks /// if Coordinated Omission Mitigation is enabled. pub(crate) async fn update_request_cadence(&mut self, thread_number: usize) { @@ -2050,152 +2071,55 @@ impl GooseUser { Ok(()) } - /// Manually build a - /// [`reqwest::Client`](https://docs.rs/reqwest/*/reqwest/struct.Client.html). - /// - /// By default, Goose configures two options when building a - /// [`reqwest::Client`](https://docs.rs/reqwest/*/reqwest/struct.Client.html). The first - /// configures Goose to report itself as the - /// [`user_agent`](https://docs.rs/reqwest/*/reqwest/struct.ClientBuilder.html#method.user_agent) - /// requesting web pages (ie `goose/0.11.2`). The second option configures - /// [`reqwest`](https://docs.rs/reqwest/) to - /// [store cookies](https://docs.rs/reqwest/*/reqwest/struct.ClientBuilder.html#method.cookie_store), - /// which is generally necessary if you aim to simulate logged in users. - /// - /// # Default configuration: - /// - /// ```rust - /// use reqwest::Client; - /// - /// static APP_USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION")); - /// - /// let builder = Client::builder() - /// .user_agent(APP_USER_AGENT) - /// .cookie_store(true) - /// .gzip(true); - /// ``` - /// - /// Alternatively, you can use this function to manually build a - /// [`reqwest::Client`](https://docs.rs/reqwest/*/reqwest/struct.Client.html). - /// with custom configuration. Available options are found in the - /// [`reqwest::ClientBuilder`](https://docs.rs/reqwest/*/reqwest/struct.ClientBuilder.html) - /// documentation. - /// - /// When manually building a - /// [`reqwest::Client`](https://docs.rs/reqwest/*/reqwest/struct.Client.html), - /// there are a few things to be aware of: - /// - Manually building a client in [`test_start`](../struct.GooseAttack.html#method.test_start) - /// will only affect requests made during test setup; - /// - Manually building a client in [`test_stop`](../struct.GooseAttack.html#method.test_stop) - /// will only affect requests made during test teardown; - /// - A manually built client is specific to a single Goose thread -- if you are - /// generating a large load test with many users, each will need to manually build their - /// own client (typically you'd do this in a Task that is registered with - /// [`GooseTask::set_on_start()`] in each Task Set requiring a custom client; - /// - Manually building a client will completely replace the automatically built client - /// with a brand new one, so any configuration, cookies or headers set in the previously - /// built client will be gone; - /// - You must include all desired configuration, as you are completely replacing Goose - /// defaults. For example, if you want Goose clients to store cookies, you will have to - /// include - /// [`.cookie_store(true)`](https://docs.rs/reqwest/*/reqwest/struct.ClientBuilder.html#method.cookie_store). - /// - /// In the following example, the Goose client is configured with a different user agent, - /// sets a default header on every request, stores cookies, and supports gzip compression. + /// Custom cookies can be manually set to each [`GooseUser`](./goose/struct.GooseUser.html). + /// Those cookies will automatically been added to each request. /// /// ## Example /// ```rust /// use goose::prelude::*; /// - /// task!(setup_custom_client).set_on_start(); - /// - /// async fn setup_custom_client(user: &mut GooseUser) -> GooseTaskResult { - /// use reqwest::{Client, header}; + /// task!(custom_cookie).set_on_start(); /// - /// // Build a custom HeaderMap to include with all requests made by this client. - /// let mut headers = header::HeaderMap::new(); - /// headers.insert("X-Custom-Header", header::HeaderValue::from_str("custom value").unwrap()); - /// - /// // Build a custom client. - /// let builder = Client::builder() - /// .default_headers(headers) - /// .user_agent("custom user agent") - /// .cookie_store(true) - /// .gzip(true); + /// async fn custom_cookie(user: &mut GooseUser) -> GooseTaskResult { + /// // Prepare the contents of a custom cookie. + /// let cookie = "my-custom-cookie=custom-value"; /// - /// // Assign the custom client to this GooseUser. - /// user.set_client_builder(builder).await?; + /// // Pre-load one or more cookies into the GooseUser. + /// user.add_cookie_str( + /// cookie, + /// ); /// /// Ok(()) /// } /// ``` + pub fn add_cookie_str(&mut self, cookie: &str) { + self.cookie_store.add_cookie_str(cookie, &self.base_url); + } + + /// Default headers can be manually set to each [`GooseUser`](./goose/struct.GooseUser.html). + /// The headers will automatically been added to each request made by the [`GooseUser`](./goose/struct.GooseUser.html). /// - /// # Alternative Compression Algorithms - /// Reqwest also supports - /// [`brotli`](https://docs.rs/reqwest/*/reqwest/struct.ClientBuilder.html#method.brotli) and - /// [`deflate`](https://docs.rs/reqwest/*/reqwest/struct.ClientBuilder.html#method.deflate) compression. - /// - /// To enable either, you must enable the features in your load test's `Cargo.toml`, for example: - /// ```text - /// reqwest = { version = "^0.11.4", default-features = false, features = [ - /// "brotli", - /// "cookies", - /// "deflate", - /// "gzip", - /// "json", - /// ] } - /// ``` - /// - /// Once enabled, you can add `.brotli(true)` and/or `.deflate(true)` to your custom - /// [`reqwest::Client::builder()`], following the documentation above. - /// - /// # Custom Cookies - /// Custom cookies can also be manually set when building a custom [`reqwest::Client`]. This requires - /// loading the [`GooseUser::base_url`] being load tested in order to properly build the cookie. Then - /// a custom [`reqwest::cookie::Jar`] is created and the custom cookie is added with - /// [`reqwest::cookie::Jar::add_cookie_str`]. Finally, the new cookie jar must be specified as the - /// [`reqwest::ClientBuilder::cookie_provider`] for the custom client. - /// + /// Default headers doesn't overwrite already appended headers. /// ## Example /// ```rust - /// use reqwest::{cookie::Jar, Client}; - /// use std::sync::Arc; - /// /// use goose::prelude::*; + /// use reqwest::header; /// - /// task!(custom_cookie_with_custom_client).set_on_start(); + /// task!(add_default_header).set_on_start(); /// - /// async fn custom_cookie_with_custom_client(user: &mut GooseUser) -> GooseTaskResult { - /// // Prepare the contents of a custom cookie. - /// let cookie = "my-custom-cookie=custom-value"; - /// - /// // Pre-load one or more cookies into a custom cookie jar to use with this client. - /// let jar = Jar::default(); - /// jar.add_cookie_str( - /// cookie, - /// &user.base_url, - /// ); - /// - /// // Build a custom client. - /// let builder = Client::builder() - /// .user_agent("example-loadtest") - /// .cookie_store(true) - /// .cookie_provider(Arc::new(jar)) - /// .gzip(true); + /// async fn add_default_header(user: &mut GooseUser) -> GooseTaskResult { /// - /// // Assign the custom client to this GooseUser. - /// user.set_client_builder(builder).await?; + /// // Build a custom HeaderMap to include with all requests made by this client. + /// user.add_default_header("X-Custom-Header", header::HeaderValue::from_str("custom value").unwrap()); /// /// Ok(()) /// } /// ``` - pub async fn set_client_builder( - &mut self, - builder: ClientBuilder, - ) -> Result<(), GooseTaskError> { - self.client = builder.build()?; - - Ok(()) + pub fn add_default_header(&mut self, key: K, val: HeaderValue) + where + K: IntoHeaderName, + { + self.headers.insert(key, val); } /// Some websites use multiple domains to serve traffic, redirecting depending on diff --git a/src/lib.rs b/src/lib.rs index 07e8ca01..c49e8aa4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -845,26 +845,47 @@ impl GooseAttack { }) } - /// Define a reqwest::Client shared among all the + /// Define a reqwest::ClientBuilder used to build a reqwest Client shared among all the /// [`GooseUser`](./goose/struct.GooseUser.html)s running. /// + /// # Alternative Compression Algorithms + /// Reqwest also supports + /// [`brotli`](https://docs.rs/reqwest/*/reqwest/struct.ClientBuilder.html#method.brotli) and + /// [`deflate`](https://docs.rs/reqwest/*/reqwest/struct.ClientBuilder.html#method.deflate) compression. + /// + /// To enable either, you must enable the features in your load test's `Cargo.toml`, for example: + /// ```text + /// reqwest = { version = "^0.11.4", default-features = false, features = [ + /// "brotli", + /// "cookies", + /// "deflate", + /// "gzip", + /// "json", + /// ] } + /// ``` + /// + /// Once enabled, you can add `.brotli(true)` and/or `.deflate(true)` to your custom + /// [`reqwest::Client::builder()`], following the documentation above. + /// /// # Example /// ```rust /// use goose::prelude::*; /// use reqwest::Client; + /// static APP_USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION")); /// /// #[tokio::main] /// async fn main() -> Result<(), GooseError> { - /// let client_builder = Client::builder(); + /// let client_builder = Client::builder() + /// .user_agent(APP_USER_AGENT); /// /// GooseAttack::initialize()? /// .set_scheduler(GooseScheduler::Random) - /// .set_client(client_builder)?; + /// .set_client_builder(client_builder)?; /// /// Ok(()) /// } /// ``` - pub fn set_client(mut self, client_builder: ClientBuilder) -> Result { + pub fn set_client_builder(mut self, client_builder: ClientBuilder) -> Result { self.client = client_builder.cookie_store(false).build()?; Ok(self) }