Skip to content

Commit

Permalink
Allow shared reqwest client configuration
Browse files Browse the repository at this point in the history
  • Loading branch information
finnefloyd committed Sep 1, 2021
1 parent 1f81d47 commit 5a3104c
Show file tree
Hide file tree
Showing 4 changed files with 84 additions and 40 deletions.
3 changes: 2 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
56 changes: 34 additions & 22 deletions src/goose.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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,
Expand All @@ -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(),
Expand All @@ -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
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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());

Expand All @@ -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);

Expand Down Expand Up @@ -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());

Expand Down
62 changes: 45 additions & 17 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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)>;

Expand Down Expand Up @@ -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 {
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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.
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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;
}
Expand All @@ -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;
}
Expand Down
3 changes: 3 additions & 0 deletions src/worker.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 5a3104c

Please sign in to comment.