Move Agent construction to new AgentBuilder.

In the process, rename set_foo methods to just foo, since methods on the
builder will always be setters.

Adds a new() method on ConnectionPool so it can be constructed directly
with the desired limits. Removes the setter methods on ConnectionPool
for those limits. This means that connection limits can only be set when
an Agent is built.

There were two tests that verify Send and Sync implementations, one for
Agent and one for Request. This PR moves the Request test to request.rs,
and changes both tests to more directly verify the traits. There may be
another way to do this, I'm not sure.
This commit is contained in:
Jacob Hoffman-Andrews
2020-10-17 23:07:53 -07:00
committed by Martin Algesten
parent 0f083a4436
commit 30162bf3bb
9 changed files with 211 additions and 262 deletions

View File

@@ -61,7 +61,7 @@ fn get_and_write(agent: &ureq::Agent, url: &String) -> Result<()> {
} }
fn get_many(urls: Vec<String>, simultaneous_fetches: usize) -> Result<()> { fn get_many(urls: Vec<String>, simultaneous_fetches: usize) -> Result<()> {
let agent = ureq::Agent::default().build(); let agent = ureq::Agent::default();
let pool = rayon::ThreadPoolBuilder::new() let pool = rayon::ThreadPoolBuilder::new()
.num_threads(simultaneous_fetches) .num_threads(simultaneous_fetches)
.build()?; .build()?;

View File

@@ -13,6 +13,25 @@ use crate::proxy::Proxy;
use crate::request::Request; use crate::request::Request;
use crate::resolve::ArcResolver; use crate::resolve::ArcResolver;
#[derive(Debug, Default)]
pub struct AgentBuilder {
headers: Vec<Header>,
proxy: Option<Proxy>,
max_idle_connections: usize,
max_idle_connections_per_host: usize,
/// Cookies saved between requests.
/// Invariant: All cookies must have a nonempty domain and path.
#[cfg(feature = "cookie")]
jar: CookieStore,
resolver: ArcResolver,
}
impl Default for Agent {
fn default() -> Self {
AgentBuilder::new().build()
}
}
/// Agents keep state between requests. /// Agents keep state between requests.
/// ///
/// By default, no state, such as cookies, is kept between requests. /// By default, no state, such as cookies, is kept between requests.
@@ -41,7 +60,7 @@ use crate::resolve::ArcResolver;
/// println!("Secret is: {}", secret.unwrap().into_string().unwrap()); /// println!("Secret is: {}", secret.unwrap().into_string().unwrap());
/// } /// }
/// ``` /// ```
#[derive(Debug, Default, Clone)] #[derive(Debug, Clone)]
pub struct Agent { pub struct Agent {
/// Copied into each request of this agent. /// Copied into each request of this agent.
pub(crate) headers: Vec<Header>, pub(crate) headers: Vec<Header>,
@@ -65,98 +84,12 @@ pub(crate) struct AgentState {
} }
impl AgentState { impl AgentState {
fn new() -> Self { pub(crate) fn pool(&mut self) -> &mut ConnectionPool {
Self::default()
}
pub fn pool(&mut self) -> &mut ConnectionPool {
&mut self.pool &mut self.pool
} }
} }
impl Agent { impl Agent {
/// Creates a new agent. Typically you'd use [`ureq::agent()`](fn.agent.html) to
/// do this.
///
/// ```
/// let agent = ureq::Agent::new()
/// .set("X-My-Header", "Foo") // present on all requests from this agent
/// .build();
///
/// agent.get("/foo");
/// ```
pub fn new() -> Agent {
Default::default()
}
/// Create a new agent after treating it as a builder.
/// This actually clones the internal state to a new one and instantiates
/// a new connection pool that is reused between connects.
pub fn build(&self) -> Self {
Agent {
headers: self.headers.clone(),
state: Arc::new(Mutex::new(AgentState::new())),
}
}
/// Set a header field that will be present in all requests using the agent.
///
/// ```
/// let agent = ureq::agent()
/// .set("X-API-Key", "foobar")
/// .set("Accept", "text/plain")
/// .build();
///
/// let r = agent
/// .get("/my-page")
/// .call();
///
/// if let Ok(resp) = r {
/// println!("yay got {}", resp.into_string().unwrap());
/// } else {
/// println!("Oh no error!");
/// }
/// ```
pub fn set(&mut self, header: &str, value: &str) -> &mut Agent {
header::add_header(&mut self.headers, Header::new(header, value));
self
}
/// Basic auth that will be present in all requests using the agent.
///
/// ```
/// let agent = ureq::agent()
/// .auth("martin", "rubbermashgum")
/// .build();
///
/// let r = agent
/// .get("/my_page")
/// .call();
/// println!("{:?}", r);
/// ```
pub fn auth(&mut self, user: &str, pass: &str) -> &mut Agent {
let pass = basic_auth(user, pass);
self.auth_kind("Basic", &pass)
}
/// Auth of other kinds such as `Digest`, `Token` etc, that will be present
/// in all requests using the agent.
///
/// ```
/// // sets a header "Authorization: token secret"
/// let agent = ureq::agent()
/// .auth_kind("token", "secret")
/// .build();
///
/// let r = agent
/// .get("/my_page")
/// .call();
/// ```
pub fn auth_kind(&mut self, kind: &str, pass: &str) -> &mut Agent {
let value = format!("{} {}", kind, pass);
self.set("Authorization", &value);
self
}
/// Request by providing the HTTP verb such as `GET`, `POST`... /// Request by providing the HTTP verb such as `GET`, `POST`...
/// ///
/// ``` /// ```
@@ -171,73 +104,6 @@ impl Agent {
Request::new(&self, method.into(), path.into()) Request::new(&self, method.into(), path.into())
} }
/// Sets the maximum number of connections allowed in the connection pool.
/// By default, this is set to 100. Setting this to zero would disable
/// connection pooling.
///
/// ```
/// let agent = ureq::agent();
/// agent.set_max_pool_connections(200);
/// ```
pub fn set_max_pool_connections(&self, max_connections: usize) {
let mut state = self.state.lock().unwrap();
state.pool.set_max_idle_connections(max_connections);
}
/// Sets the maximum number of connections per host to keep in the
/// connection pool. By default, this is set to 1. Setting this to zero
/// would disable connection pooling.
///
/// ```
/// let agent = ureq::agent();
/// agent.set_max_pool_connections_per_host(10);
/// ```
pub fn set_max_pool_connections_per_host(&self, max_connections: usize) {
let mut state = self.state.lock().unwrap();
state
.pool
.set_max_idle_connections_per_host(max_connections);
}
/// Configures a custom resolver to be used by this agent. By default,
/// address-resolution is done by std::net::ToSocketAddrs. This allows you
/// to override that resolution with your own alternative. Useful for
/// testing and special-cases like DNS-based load balancing.
///
/// A `Fn(&str) -> io::Result<Vec<SocketAddr>>` is a valid resolver,
/// passing a closure is a simple way to override. Note that you might need
/// explicit type `&str` on the closure argument for type inference to
/// succeed.
/// ```
/// use std::net::ToSocketAddrs;
///
/// let mut agent = ureq::agent();
/// agent.set_resolver(|addr: &str| match addr {
/// "example.com" => Ok(vec![([127,0,0,1], 8096).into()]),
/// addr => addr.to_socket_addrs().map(Iterator::collect),
/// });
/// ```
pub fn set_resolver(&mut self, resolver: impl crate::Resolver + 'static) -> &mut Self {
self.state.lock().unwrap().resolver = resolver.into();
self
}
/// Set the proxy server to use for all connections from this Agent.
///
/// Example:
/// ```
/// let proxy = ureq::Proxy::new("user:password@cool.proxy:9090").unwrap();
/// let agent = ureq::agent()
/// .set_proxy(proxy)
/// .build();
/// ```
pub fn set_proxy(&mut self, proxy: Proxy) -> &mut Agent {
let mut state = self.state.lock().unwrap();
state.proxy = Some(proxy);
drop(state);
self
}
/// Gets a cookie in this agent by name. Cookies are available /// Gets a cookie in this agent by name. Cookies are available
/// either by setting it in the agent, or by making requests /// either by setting it in the agent, or by making requests
/// that `Set-Cookie` in the agent. /// that `Set-Cookie` in the agent.
@@ -345,6 +211,156 @@ impl Agent {
} }
} }
impl AgentBuilder {
pub fn new() -> AgentBuilder {
AgentBuilder {
max_idle_connections: 100,
max_idle_connections_per_host: 1,
..Default::default()
}
}
/// Create a new agent.
// Note: This could take &self as the first argument, allowing one
// AgentBuilder to be used multiple times, except CookieStore does
// not implement clone, so we have to give ownership to the newly
// built Agent.
pub fn build(self) -> Agent {
Agent {
headers: self.headers.clone(),
state: Arc::new(Mutex::new(AgentState {
pool: ConnectionPool::new(
self.max_idle_connections,
self.max_idle_connections_per_host,
),
proxy: self.proxy.clone(),
#[cfg(feature = "cookie")]
jar: self.jar,
resolver: self.resolver,
})),
}
}
/// Set a header field that will be present in all requests using the agent.
///
/// ```
/// let agent = ureq::AgentBuilder::new()
/// .set("X-API-Key", "foobar")
/// .set("Accept", "text/plain")
/// .build();
///
/// let r = agent
/// .get("/my-page")
/// .call();
///
/// if let Ok(resp) = r {
/// println!("yay got {}", resp.into_string().unwrap());
/// } else {
/// println!("Oh no error!");
/// }
/// ```
pub fn set(mut self, header: &str, value: &str) -> Self {
header::add_header(&mut self.headers, Header::new(header, value));
self
}
/// Basic auth that will be present in all requests using the agent.
///
/// ```
/// let agent = ureq::AgentBuilder::new()
/// .auth("martin", "rubbermashgum")
/// .build();
///
/// let r = agent
/// .get("/my_page")
/// .call();
/// println!("{:?}", r);
/// ```
pub fn auth(self, user: &str, pass: &str) -> Self {
let pass = basic_auth(user, pass);
self.auth_kind("Basic", &pass)
}
/// Auth of other kinds such as `Digest`, `Token` etc, that will be present
/// in all requests using the agent.
///
/// ```
/// // sets a header "Authorization: token secret"
/// let agent = ureq::AgentBuilder::new()
/// .auth_kind("token", "secret")
/// .build();
///
/// let r = agent
/// .get("/my_page")
/// .call();
/// ```
pub fn auth_kind(self, kind: &str, pass: &str) -> Self {
let value = format!("{} {}", kind, pass);
self.set("Authorization", &value)
}
/// Sets the maximum number of connections allowed in the connection pool.
/// By default, this is set to 100. Setting this to zero would disable
/// connection pooling.
///
/// ```
/// let agent = ureq::AgentBuilder::new().max_pool_connections(200).build();
/// ```
pub fn max_pool_connections(mut self, max: usize) -> Self {
self.max_idle_connections = max;
self
}
/// Sets the maximum number of connections per host to keep in the
/// connection pool. By default, this is set to 1. Setting this to zero
/// would disable connection pooling.
///
/// ```
/// let agent = ureq::AgentBuilder::new().max_pool_connections_per_host(200).build();
/// ```
pub fn max_pool_connections_per_host(mut self, max: usize) -> Self {
self.max_idle_connections_per_host = max;
self
}
/// Configures a custom resolver to be used by this agent. By default,
/// address-resolution is done by std::net::ToSocketAddrs. This allows you
/// to override that resolution with your own alternative. Useful for
/// testing and special-cases like DNS-based load balancing.
///
/// A `Fn(&str) -> io::Result<Vec<SocketAddr>>` is a valid resolver,
/// passing a closure is a simple way to override. Note that you might need
/// explicit type `&str` on the closure argument for type inference to
/// succeed.
/// ```
/// use std::net::ToSocketAddrs;
///
/// let mut agent = ureq::AgentBuilder::new()
/// .resolver(|addr: &str| match addr {
/// "example.com" => Ok(vec![([127,0,0,1], 8096).into()]),
/// addr => addr.to_socket_addrs().map(Iterator::collect),
/// })
/// .build();
/// ```
pub fn resolver(mut self, resolver: impl crate::Resolver + 'static) -> Self {
self.resolver = resolver.into();
self
}
/// Set the proxy server to use for all connections from this Agent.
///
/// Example:
/// ```
/// let proxy = ureq::Proxy::new("user:password@cool.proxy:9090").unwrap();
/// let agent = ureq::AgentBuilder::new()
/// .proxy(proxy)
/// .build();
/// ```
pub fn proxy(mut self, proxy: Proxy) -> Self {
self.proxy = Some(proxy);
self
}
}
pub(crate) fn basic_auth(user: &str, pass: &str) -> String { pub(crate) fn basic_auth(user: &str, pass: &str) -> String {
let safe = match user.find(':') { let safe = match user.find(':') {
Some(idx) => &user[..idx], Some(idx) => &user[..idx],
@@ -356,16 +372,13 @@ pub(crate) fn basic_auth(user: &str, pass: &str) -> String {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use std::thread;
///////////////////// AGENT TESTS ////////////////////////////// ///////////////////// AGENT TESTS //////////////////////////////
#[test] #[test]
fn agent_implements_send() { fn agent_implements_send_and_sync() {
let mut agent = Agent::new(); let _agent: Box<dyn Send> = Box::new(AgentBuilder::new().build());
thread::spawn(move || { let _agent: Box<dyn Sync> = Box::new(AgentBuilder::new().build());
agent.set("Foo", "Bar");
});
} }
#[test] #[test]
@@ -395,15 +408,4 @@ mod tests {
let mut buf = vec![]; let mut buf = vec![];
reader.read_to_end(&mut buf).unwrap(); reader.read_to_end(&mut buf).unwrap();
} }
//////////////////// REQUEST TESTS /////////////////////////////
#[test]
fn request_implements_send() {
let agent = Agent::new();
let mut request = Request::new(&agent, "GET".to_string(), "/foo".to_string());
thread::spawn(move || {
request.set("Foo", "Bar");
});
}
} }

View File

@@ -120,6 +120,7 @@ mod serde_macros;
mod test; mod test;
pub use crate::agent::Agent; pub use crate::agent::Agent;
pub use crate::agent::AgentBuilder;
pub use crate::error::Error; pub use crate::error::Error;
pub use crate::header::Header; pub use crate::header::Header;
pub use crate::proxy::Proxy; pub use crate::proxy::Proxy;
@@ -135,7 +136,7 @@ pub use serde_json::{to_value as serde_to_value, Map as SerdeMap, Value as Serde
/// Agents are used to keep state between requests. /// Agents are used to keep state between requests.
pub fn agent() -> Agent { pub fn agent() -> Agent {
Agent::new().build() Agent::default()
} }
/// Make a request setting the HTTP method via a string. /// Make a request setting the HTTP method via a string.
@@ -144,7 +145,7 @@ pub fn agent() -> Agent {
/// ureq::request("GET", "http://example.com").call().unwrap(); /// ureq::request("GET", "http://example.com").call().unwrap();
/// ``` /// ```
pub fn request(method: &str, path: &str) -> Request { pub fn request(method: &str, path: &str) -> Request {
Agent::new().request(method, path) agent().request(method, path)
} }
/// Make a GET request. /// Make a GET request.

View File

@@ -73,16 +73,13 @@ impl Default for ConnectionPool {
} }
impl ConnectionPool { impl ConnectionPool {
pub fn set_max_idle_connections(&mut self, max_connections: usize) { pub(crate) fn new(max_idle_connections: usize, max_idle_connections_per_host: usize) -> Self {
if self.max_idle_connections == max_connections { ConnectionPool {
return; recycle: Default::default(),
lru: Default::default(),
max_idle_connections,
max_idle_connections_per_host,
} }
// Remove any extra connections if the number was decreased.
while self.lru.len() > max_connections {
self.remove_oldest();
}
self.max_idle_connections = max_connections;
} }
/// Return true if either of the max_* settings is 0, meaning we should do no work. /// Return true if either of the max_* settings is 0, meaning we should do no work.
@@ -90,30 +87,6 @@ impl ConnectionPool {
self.max_idle_connections == 0 || self.max_idle_connections_per_host == 0 self.max_idle_connections == 0 || self.max_idle_connections_per_host == 0
} }
pub fn set_max_idle_connections_per_host(&mut self, max_connections: usize) {
if self.max_idle_connections_per_host == max_connections {
return;
}
if max_connections == 0 {
// Clear the connection pool, caching is disabled.
self.lru.clear();
self.recycle.clear();
return;
}
// Remove any extra streams if the number was decreased.
for (key, val) in self.recycle.iter_mut() {
while val.len() > max_connections {
// Remove the oldest entry
val.pop_front();
remove_first_match(&mut self.lru, key)
.expect("invariant failed: key in recycle but not in lru");
}
}
self.max_idle_connections_per_host = max_connections;
}
/// How the unit::connect tries to get a pooled connection. /// How the unit::connect tries to get a pooled connection.
pub fn try_get_connection(&mut self, url: &Url, proxy: &Option<Proxy>) -> Option<Stream> { pub fn try_get_connection(&mut self, url: &Url, proxy: &Option<Proxy>) -> Option<Stream> {
let key = PoolKey::new(url, proxy); let key = PoolKey::new(url, proxy);
@@ -287,50 +260,6 @@ fn pool_per_host_connections_limit() {
assert_eq!(pool.len(), 0); assert_eq!(pool.len(), 0);
} }
#[test]
fn pool_update_connection_limit() {
let mut pool = ConnectionPool::default();
pool.set_max_idle_connections(50);
let hostnames = (0..pool.max_idle_connections).map(|i| format!("{}.example", i));
let poolkeys = hostnames.map(|hostname| PoolKey {
scheme: "https".to_string(),
hostname,
port: Some(999),
proxy: None,
});
for key in poolkeys.clone() {
pool.add(key, Stream::Cursor(std::io::Cursor::new(vec![])));
}
assert_eq!(pool.len(), 50);
pool.set_max_idle_connections(25);
assert_eq!(pool.len(), 25);
}
#[test]
fn pool_update_per_host_connection_limit() {
let mut pool = ConnectionPool::default();
pool.set_max_idle_connections(50);
pool.set_max_idle_connections_per_host(50);
let poolkey = PoolKey {
scheme: "https".to_string(),
hostname: "example.com".to_string(),
port: Some(999),
proxy: None,
};
for _ in 0..50 {
pool.add(
poolkey.clone(),
Stream::Cursor(std::io::Cursor::new(vec![])),
);
}
assert_eq!(pool.len(), 50);
pool.set_max_idle_connections_per_host(25);
assert_eq!(pool.len(), 25);
}
#[test] #[test]
fn pool_checks_proxy() { fn pool_checks_proxy() {
// Test inserting different poolkeys with same address but different proxies. // Test inserting different poolkeys with same address but different proxies.

View File

@@ -690,3 +690,17 @@ fn no_hostname() {
); );
assert!(req.get_host().is_err()); assert!(req.get_host().is_err());
} }
#[test]
fn request_implements_send_and_sync() {
let _request: Box<dyn Send> = Box::new(Request::new(
&Agent::default(),
"GET".to_string(),
"https://example.com/".to_string(),
));
let _request: Box<dyn Sync> = Box::new(Request::new(
&Agent::default(),
"GET".to_string(),
"https://example.com/".to_string(),
));
}

View File

@@ -10,7 +10,9 @@ use super::super::*;
#[test] #[test]
fn agent_reuse_headers() { fn agent_reuse_headers() {
let agent = agent().set("Authorization", "Foo 12345").build(); let agent = AgentBuilder::new()
.set("Authorization", "Foo 12345")
.build();
test::set_handler("/agent_reuse_headers", |unit| { test::set_handler("/agent_reuse_headers", |unit| {
assert!(unit.has("Authorization")); assert!(unit.has("Authorization"));
@@ -44,7 +46,7 @@ fn idle_timeout_handler(mut stream: TcpStream) -> io::Result<()> {
fn connection_reuse() { fn connection_reuse() {
let testserver = TestServer::new(idle_timeout_handler); let testserver = TestServer::new(idle_timeout_handler);
let url = format!("http://localhost:{}", testserver.port); let url = format!("http://localhost:{}", testserver.port);
let agent = Agent::default().build(); let agent = Agent::default();
let resp = agent.get(&url).call().unwrap(); let resp = agent.get(&url).call().unwrap();
// use up the connection so it gets returned to the pool // use up the connection so it gets returned to the pool
@@ -87,8 +89,9 @@ fn custom_resolver() {
buf buf
}); });
crate::agent() AgentBuilder::new()
.set_resolver(move |_: &str| Ok(vec![local_addr])) .resolver(move |_: &str| Ok(vec![local_addr]))
.build()
.get("http://cool.server/") .get("http://cool.server/")
.call() .call()
.ok(); .ok();
@@ -144,7 +147,7 @@ fn cookie_and_redirect(mut stream: TcpStream) -> io::Result<()> {
fn test_cookies_on_redirect() -> Result<(), Error> { fn test_cookies_on_redirect() -> Result<(), Error> {
let testserver = TestServer::new(cookie_and_redirect); let testserver = TestServer::new(cookie_and_redirect);
let url = format!("http://localhost:{}/first", testserver.port); let url = format!("http://localhost:{}/first", testserver.port);
let agent = Agent::default().build(); let agent = Agent::default();
agent.post(&url).call()?; agent.post(&url).call()?;
assert!(agent.cookie("first").is_some()); assert!(agent.cookie("first").is_some());
assert!(agent.cookie("second").is_some()); assert!(agent.cookie("second").is_some());
@@ -168,7 +171,7 @@ fn dirty_streams_not_returned() -> Result<(), Error> {
Ok(()) Ok(())
}); });
let url = format!("http://localhost:{}/", testserver.port); let url = format!("http://localhost:{}/", testserver.port);
let agent = Agent::default().build(); let agent = Agent::default();
let resp = agent.get(&url).call()?; let resp = agent.get(&url).call()?;
let resp_str = resp.into_string()?; let resp_str = resp.into_string()?;
assert_eq!(resp_str, "corgidachsund"); assert_eq!(resp_str, "corgidachsund");

View File

@@ -55,7 +55,7 @@ fn url_auth_overridden() {
); );
test::make_response(200, "OK", vec![], vec![]) test::make_response(200, "OK", vec![], vec![])
}); });
let agent = agent().auth("martin", "rubbermashgum").build(); let agent = AgentBuilder::new().auth("martin", "rubbermashgum").build();
let resp = agent let resp = agent
.get("test://Aladdin:OpenSesame@host/url_auth_overridden") .get("test://Aladdin:OpenSesame@host/url_auth_overridden")
.call() .call()

View File

@@ -25,7 +25,7 @@ fn dribble_body_respond(mut stream: TcpStream, contents: &[u8]) -> io::Result<()
} }
fn get_and_expect_timeout(url: String) { fn get_and_expect_timeout(url: String) {
let agent = Agent::default().build(); let agent = Agent::default();
let timeout = Duration::from_millis(500); let timeout = Duration::from_millis(500);
let resp = agent.get(&url).timeout(timeout).call().unwrap(); let resp = agent.get(&url).timeout(timeout).call().unwrap();
@@ -86,7 +86,7 @@ fn overall_timeout_reading_json() {
}); });
let url = format!("http://localhost:{}/", server.port); let url = format!("http://localhost:{}/", server.port);
let agent = Agent::default().build(); let agent = Agent::default();
let timeout = Duration::from_millis(500); let timeout = Duration::from_millis(500);
let resp = agent.get(&url).timeout(timeout).call().unwrap(); let resp = agent.get(&url).timeout(timeout).call().unwrap();

View File

@@ -102,7 +102,7 @@ m0Wqhhi8/24Sy934t5Txgkfoltg8ahkx934WjP6WWRnSAu+cf+vW
#[cfg(feature = "tls")] #[cfg(feature = "tls")]
#[test] #[test]
fn tls_client_certificate() { fn tls_client_certificate() {
let agent = ureq::Agent::default().build(); let agent = ureq::Agent::default();
let mut tls_config = rustls::ClientConfig::new(); let mut tls_config = rustls::ClientConfig::new();