diff --git a/src/client.rs b/src/client.rs index 5ba24a3..6715f43 100644 --- a/src/client.rs +++ b/src/client.rs @@ -7,7 +7,7 @@ use url::Url; use crate::connection::{ConnectionManager, with_connection}; use crate::error::{ClientError, MemcacheError}; -use crate::protocol::{Protocol, ProtocolTrait}; +use crate::protocol::ProtocolTrait; use crate::stream::Stream; use crate::value::{FromMemcacheValueExt, ToMemcacheValue}; use r2d2::Pool; @@ -156,11 +156,7 @@ impl Client { /// ``` pub fn set_read_timeout(&self, timeout: Option) -> Result<(), MemcacheError> { for conn in self.connections.iter() { - let mut conn = conn.get()?; - match **conn { - Protocol::Ascii(ref mut protocol) => protocol.stream().set_read_timeout(timeout)?, - Protocol::Binary(ref mut protocol) => protocol.stream.set_read_timeout(timeout)?, - } + conn.get()?.set_read_timeout(timeout)?; } Ok(()) } @@ -175,11 +171,7 @@ impl Client { /// ``` pub fn set_write_timeout(&self, timeout: Option) -> Result<(), MemcacheError> { for conn in self.connections.iter() { - let mut conn = conn.get()?; - match **conn { - Protocol::Ascii(ref mut protocol) => protocol.stream().set_write_timeout(timeout)?, - Protocol::Binary(ref mut protocol) => protocol.stream.set_write_timeout(timeout)?, - } + conn.get()?.set_write_timeout(timeout)?; } Ok(()) } @@ -574,22 +566,23 @@ impl ClientBuilder { builder = builder.connection_timeout(timeout); } - let connection = builder - .build(ConnectionManager::new(url)) - .map_err(|e| MemcacheError::PoolError(e))?; + let mut manager = ConnectionManager::new(url); + if let Some(timeout) = self.read_timeout { + manager = manager.with_read_timeout(timeout); + } + if let Some(timeout) = self.write_timeout { + manager = manager.with_write_timeout(timeout); + } + + let connection = builder.build(manager).map_err(|e| MemcacheError::PoolError(e))?; connections.push(connection); } - let client = Client { + Ok(Client { connections, hash_function: self.hash_function, - }; - - client.set_read_timeout(self.read_timeout)?; - client.set_write_timeout(self.write_timeout)?; - - Ok(client) + }) } } diff --git a/src/connection.rs b/src/connection.rs index bad6f7d..0b04ab6 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -38,12 +38,30 @@ impl Deref for Connection { /// Memcache connection manager implementing rd2d Pool ManageConnection pub struct ConnectionManager { url: Url, + read_timeout: Option, + write_timeout: Option, } impl ConnectionManager { /// Initialize connection manager with given Url pub fn new(url: Url) -> Self { - Self { url } + Self { + url, + read_timeout: None, + write_timeout: None, + } + } + + /// Set the socket read timeout for every connection opened by this manager. + pub fn with_read_timeout(mut self, timeout: Duration) -> Self { + self.read_timeout = Some(timeout); + self + } + + /// Set the socket write timeout for every connection opened by this manager. + pub fn with_write_timeout(mut self, timeout: Duration) -> Self { + self.write_timeout = Some(timeout); + self } } @@ -54,6 +72,12 @@ impl ManageConnection for ConnectionManager { fn connect(&self) -> Result { let url = &self.url; let mut connection = Connection::connect(url)?; + if self.read_timeout.is_some() { + connection.set_read_timeout(self.read_timeout)?; + } + if self.write_timeout.is_some() { + connection.set_write_timeout(self.write_timeout)?; + } if url.has_authority() && !url.username().is_empty() && url.password().is_some() { let username = url.username(); let password = url.password().unwrap(); @@ -239,6 +263,20 @@ impl Connection { self.url.to_string() } + pub(crate) fn set_read_timeout(&mut self, timeout: Option) -> Result<(), MemcacheError> { + match self.protocol { + Protocol::Ascii(ref mut protocol) => protocol.stream().set_read_timeout(timeout), + Protocol::Binary(ref mut protocol) => protocol.stream.set_read_timeout(timeout), + } + } + + pub(crate) fn set_write_timeout(&mut self, timeout: Option) -> Result<(), MemcacheError> { + match self.protocol { + Protocol::Ascii(ref mut protocol) => protocol.stream().set_write_timeout(timeout), + Protocol::Binary(ref mut protocol) => protocol.stream.set_write_timeout(timeout), + } + } + /// Flag the connection so the pool drops it if `err` may have left unread /// data on the stream, otherwise later commands would read stale responses. fn mark_broken_on(&mut self, err: &MemcacheError) { diff --git a/tests/test_broken_connection.rs b/tests/test_broken_connection.rs index 5e6c840..717925d 100644 --- a/tests/test_broken_connection.rs +++ b/tests/test_broken_connection.rs @@ -1,7 +1,7 @@ use std::io::{Read, Write}; use std::net::{TcpListener, TcpStream}; use std::thread; -use std::time::Duration; +use std::time::{Duration, Instant}; use byteorder::{BigEndian, ByteOrder, WriteBytesExt}; @@ -47,8 +47,7 @@ fn serve(mut stream: TcpStream) { } } -#[test] -fn test_connection_dropped_after_read_timeout() { +fn start_server() -> u16 { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let port = listener.local_addr().unwrap().port(); thread::spawn(move || { @@ -56,7 +55,40 @@ fn test_connection_dropped_after_read_timeout() { thread::spawn(move || serve(stream.unwrap())); } }); + port +} + +fn assert_times_out(client: &memcache::Client) { + let start = Instant::now(); + assert!(client.get::("slow").is_err()); + assert!(start.elapsed() < Duration::from_millis(400)); +} +#[test] +fn test_url_timeout_survives_connect() { + let port = start_server(); + let client = memcache::connect(format!("memcache://127.0.0.1:{}?timeout=0.1", port)).unwrap(); + assert_times_out(&client); +} + +#[test] +fn test_builder_timeout_applies_to_new_connections() { + let port = start_server(); + let client = memcache::Client::builder() + .add_server(format!("memcache://127.0.0.1:{}", port)) + .unwrap() + .with_min_idle_conns(0) + .with_read_timeout(Duration::from_millis(100)) + .build() + .unwrap(); + + assert_times_out(&client); + assert_times_out(&client); +} + +#[test] +fn test_connection_dropped_after_read_timeout() { + let port = start_server(); let client = memcache::Client::builder() .add_server(format!("memcache://127.0.0.1:{}", port)) .unwrap()