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 5, 2021
1 parent f8e1954 commit 05be749
Show file tree
Hide file tree
Showing 4 changed files with 83 additions and 40 deletions.
2 changes: 1 addition & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
56 changes: 34 additions & 22 deletions src/goose.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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<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 @@ -891,8 +883,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, 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, 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 @@ -2606,10 +2602,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 @@ -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());

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

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

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

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

Expand Down Expand Up @@ -782,6 +786,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 @@ -794,21 +800,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 @@ -828,6 +820,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 @@ -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.
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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;
}
Expand All @@ -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;
}
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.config,
Expand Down

0 comments on commit 05be749

Please sign in to comment.