feat(client): add a Connection Pool
This adds a connection pool to the Client that is used by default. It accepts any other NetworkConnector, and simply acts as a NetworkConnector itself. Other Pools can exist by simply providing a custom NetworkConnector. This Pool is only used by default if you also use the default connector, which is `HttpConnector`. If you wish to use the Pool with a custom connector, you'll need to create the Pool with your custom connector, and then pass that pool to the Client::with_connector. This also adds a method to `NetworkStream`, `close`, which can be used to know when the Stream should be put down, because a server requested that the connection close instead of be kept alive. Closes #363 Closes #41
This commit is contained in:
		| @@ -46,15 +46,17 @@ use status::StatusClass::Redirection; | ||||
| use {Url, HttpResult}; | ||||
| use HttpError::HttpUriError; | ||||
|  | ||||
| pub use self::pool::Pool; | ||||
| pub use self::request::Request; | ||||
| pub use self::response::Response; | ||||
|  | ||||
| pub mod pool; | ||||
| pub mod request; | ||||
| pub mod response; | ||||
|  | ||||
| /// A Client to use additional features with Requests. | ||||
| /// | ||||
| /// Clients can handle things such as: redirect policy. | ||||
| /// Clients can handle things such as: redirect policy, connection pooling. | ||||
| pub struct Client { | ||||
|     connector: Connector, | ||||
|     redirect_policy: RedirectPolicy, | ||||
| @@ -64,7 +66,12 @@ impl Client { | ||||
|  | ||||
|     /// Create a new Client. | ||||
|     pub fn new() -> Client { | ||||
|         Client::with_connector(HttpConnector(None)) | ||||
|         Client::with_pool_config(Default::default()) | ||||
|     } | ||||
|  | ||||
|     /// Create a new Client with a configured Pool Config. | ||||
|     pub fn with_pool_config(config: pool::Config) -> Client { | ||||
|         Client::with_connector(Pool::new(config)) | ||||
|     } | ||||
|  | ||||
|     /// Create a new client with a specific connector. | ||||
| @@ -78,7 +85,10 @@ impl Client { | ||||
|  | ||||
|     /// Set the SSL verifier callback for use with OpenSSL. | ||||
|     pub fn set_ssl_verifier(&mut self, verifier: ContextVerifier) { | ||||
|         self.connector = with_connector(HttpConnector(Some(verifier))); | ||||
|         self.connector = with_connector(Pool::with_connector( | ||||
|             Default::default(), | ||||
|             HttpConnector(Some(verifier)) | ||||
|         )); | ||||
|     } | ||||
|  | ||||
|     /// Set the RedirectPolicy. | ||||
|   | ||||
							
								
								
									
										227
									
								
								src/client/pool.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										227
									
								
								src/client/pool.rs
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,227 @@ | ||||
| //! Client Connection Pooling | ||||
| use std::borrow::ToOwned; | ||||
| use std::collections::HashMap; | ||||
| use std::io::{self, Read, Write}; | ||||
| use std::net::{SocketAddr, Shutdown}; | ||||
| use std::sync::{Arc, Mutex}; | ||||
|  | ||||
| use net::{NetworkConnector, NetworkStream, HttpConnector}; | ||||
|  | ||||
| /// The `NetworkConnector` that behaves as a connection pool used by hyper's `Client`. | ||||
| pub struct Pool<C: NetworkConnector> { | ||||
|     connector: C, | ||||
|     inner: Arc<Mutex<PoolImpl<<C as NetworkConnector>::Stream>>> | ||||
| } | ||||
|  | ||||
| /// Config options for the `Pool`. | ||||
| #[derive(Debug)] | ||||
| pub struct Config { | ||||
|     /// The maximum idle connections *per host*. | ||||
|     pub max_idle: usize, | ||||
| } | ||||
|  | ||||
| impl Default for Config { | ||||
|     #[inline] | ||||
|     fn default() -> Config { | ||||
|         Config { | ||||
|             max_idle: 5, | ||||
|         } | ||||
|     } | ||||
| } | ||||
|  | ||||
| #[derive(Debug)] | ||||
| struct PoolImpl<S> { | ||||
|     conns: HashMap<Key, Vec<S>>, | ||||
|     config: Config, | ||||
| } | ||||
|  | ||||
| type Key = (String, u16, Scheme); | ||||
|  | ||||
| fn key<T: Into<Scheme>>(host: &str, port: u16, scheme: T) -> Key { | ||||
|     (host.to_owned(), port, scheme.into()) | ||||
| } | ||||
|  | ||||
| #[derive(Clone, PartialEq, Eq, Debug, Hash)] | ||||
| enum Scheme { | ||||
|     Http, | ||||
|     Https, | ||||
|     Other(String) | ||||
| } | ||||
|  | ||||
| impl<'a> From<&'a str> for Scheme { | ||||
|     fn from(s: &'a str) -> Scheme { | ||||
|         match s { | ||||
|             "http" => Scheme::Http, | ||||
|             "https" => Scheme::Https, | ||||
|             s => Scheme::Other(String::from(s)) | ||||
|         } | ||||
|     } | ||||
| } | ||||
|  | ||||
| impl Pool<HttpConnector> { | ||||
|     /// Creates a `Pool` with an `HttpConnector`. | ||||
|     #[inline] | ||||
|     pub fn new(config: Config) -> Pool<HttpConnector> { | ||||
|         Pool::with_connector(config, HttpConnector(None)) | ||||
|     } | ||||
| } | ||||
|  | ||||
| impl<C: NetworkConnector> Pool<C> { | ||||
|     /// Creates a `Pool` with a specified `NetworkConnector`. | ||||
|     #[inline] | ||||
|     pub fn with_connector(config: Config, connector: C) -> Pool<C> { | ||||
|         Pool { | ||||
|             connector: connector, | ||||
|             inner: Arc::new(Mutex::new(PoolImpl { | ||||
|                 conns: HashMap::new(), | ||||
|                 config: config, | ||||
|             })) | ||||
|         } | ||||
|     } | ||||
|  | ||||
|     /// Clear all idle connections from the Pool, closing them. | ||||
|     #[inline] | ||||
|     pub fn clear_idle(&mut self) { | ||||
|         self.inner.lock().unwrap().conns.clear(); | ||||
|     } | ||||
| } | ||||
|  | ||||
| impl<S> PoolImpl<S> { | ||||
|     fn reuse(&mut self, key: Key, conn: S) { | ||||
|         trace!("reuse {:?}", key); | ||||
|         let conns = self.conns.entry(key).or_insert(vec![]); | ||||
|         if conns.len() < self.config.max_idle { | ||||
|             conns.push(conn); | ||||
|         } | ||||
|     } | ||||
| } | ||||
|  | ||||
| impl<C: NetworkConnector<Stream=S>, S: NetworkStream + Send> NetworkConnector for Pool<C> { | ||||
|     type Stream = PooledStream<S>; | ||||
|     fn connect(&mut self, host: &str, port: u16, scheme: &str) -> io::Result<PooledStream<S>> { | ||||
|         let key = key(host, port, scheme); | ||||
|         let mut locked = self.inner.lock().unwrap(); | ||||
|         let mut should_remove = false; | ||||
|         let conn = match locked.conns.get_mut(&key) { | ||||
|             Some(ref mut vec) => { | ||||
|                 should_remove = vec.len() == 1; | ||||
|                 vec.pop().unwrap() | ||||
|             } | ||||
|             _ => try!(self.connector.connect(host, port, scheme)) | ||||
|         }; | ||||
|         if should_remove { | ||||
|             locked.conns.remove(&key); | ||||
|         } | ||||
|         Ok(PooledStream { | ||||
|             inner: Some((key, conn)), | ||||
|             is_closed: false, | ||||
|             is_drained: false, | ||||
|             pool: self.inner.clone() | ||||
|         }) | ||||
|     } | ||||
| } | ||||
|  | ||||
| /// A Stream that will try to be returned to the Pool when dropped. | ||||
| pub struct PooledStream<S> { | ||||
|     inner: Option<(Key, S)>, | ||||
|     is_closed: bool, | ||||
|     is_drained: bool, | ||||
|     pool: Arc<Mutex<PoolImpl<S>>> | ||||
| } | ||||
|  | ||||
| impl<S: NetworkStream> Read for PooledStream<S> { | ||||
|     #[inline] | ||||
|     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { | ||||
|         match self.inner.as_mut().unwrap().1.read(buf) { | ||||
|             Ok(0) => { | ||||
|                 self.is_drained = true; | ||||
|                 Ok(0) | ||||
|             } | ||||
|             r => r | ||||
|         } | ||||
|     } | ||||
| } | ||||
|  | ||||
| impl<S: NetworkStream> Write for PooledStream<S> { | ||||
|     #[inline] | ||||
|     fn write(&mut self, buf: &[u8]) -> io::Result<usize> { | ||||
|         self.inner.as_mut().unwrap().1.write(buf) | ||||
|     } | ||||
|  | ||||
|     #[inline] | ||||
|     fn flush(&mut self) -> io::Result<()> { | ||||
|         self.inner.as_mut().unwrap().1.flush() | ||||
|     } | ||||
| } | ||||
|  | ||||
| impl<S: NetworkStream> NetworkStream for PooledStream<S> { | ||||
|     #[inline] | ||||
|     fn peer_addr(&mut self) -> io::Result<SocketAddr> { | ||||
|         self.inner.as_mut().unwrap().1.peer_addr() | ||||
|     } | ||||
|  | ||||
|     #[inline] | ||||
|     fn close(&mut self, how: Shutdown) -> io::Result<()> { | ||||
|         self.is_closed = true; | ||||
|         self.inner.as_mut().unwrap().1.close(how) | ||||
|     } | ||||
| } | ||||
|  | ||||
| impl<S> Drop for PooledStream<S> { | ||||
|     fn drop(&mut self) { | ||||
|         trace!("PooledStream.drop, is_closed={}, is_drained={}", self.is_closed, self.is_drained); | ||||
|         if !self.is_closed && self.is_drained { | ||||
|             self.inner.take().map(|(key, conn)| { | ||||
|                 if let Ok(mut pool) = self.pool.lock() { | ||||
|                     pool.reuse(key, conn); | ||||
|                 } | ||||
|                 // else poisoned, give up | ||||
|             }); | ||||
|         } | ||||
|     } | ||||
| } | ||||
|  | ||||
| #[cfg(test)] | ||||
| mod tests { | ||||
|     use std::net::Shutdown; | ||||
|     use mock::MockConnector; | ||||
|     use net::{NetworkConnector, NetworkStream}; | ||||
|  | ||||
|     use super::{Pool, key}; | ||||
|  | ||||
|     macro_rules! mocked { | ||||
|         () => ({ | ||||
|             Pool::with_connector(Default::default(), MockConnector) | ||||
|         }) | ||||
|     } | ||||
|  | ||||
|     #[test] | ||||
|     fn test_connect_and_drop() { | ||||
|         let mut pool = mocked!(); | ||||
|         let key = key("127.0.0.1", 3000, "http"); | ||||
|         pool.connect("127.0.0.1", 3000, "http").unwrap().is_drained = true; | ||||
|         { | ||||
|             let locked = pool.inner.lock().unwrap(); | ||||
|             assert_eq!(locked.conns.len(), 1); | ||||
|             assert_eq!(locked.conns.get(&key).unwrap().len(), 1); | ||||
|         } | ||||
|         pool.connect("127.0.0.1", 3000, "http").unwrap().is_drained = true; //reused | ||||
|         { | ||||
|             let locked = pool.inner.lock().unwrap(); | ||||
|             assert_eq!(locked.conns.len(), 1); | ||||
|             assert_eq!(locked.conns.get(&key).unwrap().len(), 1); | ||||
|         } | ||||
|     } | ||||
|  | ||||
|     #[test] | ||||
|     fn test_closed() { | ||||
|         let mut pool = mocked!(); | ||||
|         let mut stream = pool.connect("127.0.0.1", 3000, "http").unwrap(); | ||||
|         stream.close(Shutdown::Both).unwrap(); | ||||
|         drop(stream); | ||||
|         let locked = pool.inner.lock().unwrap(); | ||||
|         assert_eq!(locked.conns.len(), 0); | ||||
|     } | ||||
|  | ||||
|  | ||||
| } | ||||
| @@ -1,6 +1,7 @@ | ||||
| //! Client Requests | ||||
| use std::marker::PhantomData; | ||||
| use std::io::{self, Write, BufWriter}; | ||||
| use std::net::Shutdown; | ||||
|  | ||||
| use url::Url; | ||||
|  | ||||
| @@ -8,7 +9,7 @@ use method::{self, Method}; | ||||
| use header::Headers; | ||||
| use header::{self, Host}; | ||||
| use net::{NetworkStream, NetworkConnector, HttpConnector, Fresh, Streaming}; | ||||
| use http::{HttpWriter, LINE_ENDING}; | ||||
| use http::{self, HttpWriter, LINE_ENDING}; | ||||
| use http::HttpWriter::{ThroughWriter, ChunkedWriter, SizedWriter, EmptyWriter}; | ||||
| use version; | ||||
| use HttpResult; | ||||
| @@ -154,7 +155,10 @@ impl Request<Streaming> { | ||||
|     /// | ||||
|     /// Consumes the Request. | ||||
|     pub fn send(self) -> HttpResult<Response> { | ||||
|         let raw = try!(self.body.end()).into_inner().unwrap(); // end() already flushes | ||||
|         let mut raw = try!(self.body.end()).into_inner().unwrap(); // end() already flushes | ||||
|         if !http::should_keep_alive(self.version, &self.headers) { | ||||
|             try!(raw.close(Shutdown::Write)); | ||||
|         } | ||||
|         Response::new(raw) | ||||
|     } | ||||
| } | ||||
|   | ||||
| @@ -1,6 +1,7 @@ | ||||
| //! Client Responses | ||||
| use std::io::{self, Read}; | ||||
| use std::marker::PhantomData; | ||||
| use std::net::Shutdown; | ||||
|  | ||||
| use buffer::BufReader; | ||||
| use header; | ||||
| @@ -42,6 +43,10 @@ impl Response { | ||||
|         debug!("version={:?}, status={:?}", head.version, status); | ||||
|         debug!("headers={:?}", headers); | ||||
|  | ||||
|         if !http::should_keep_alive(head.version, &headers) { | ||||
|             try!(stream.get_mut().close(Shutdown::Write)); | ||||
|         } | ||||
|  | ||||
|         let body = if headers.has::<TransferEncoding>() { | ||||
|             match headers.get::<TransferEncoding>() { | ||||
|                 Some(&TransferEncoding(ref codings)) => { | ||||
|   | ||||
		Reference in New Issue
	
	Block a user