Apply read and write timeouts when connections are opened

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<Duration>) -> 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<Duration>) -> 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)
+        })
     }
 }
 

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<Duration>,
+    write_timeout: Option<Duration>,
 }
 
 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<Self::Connection, Self::Error> {
         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<Duration>) -> 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<Duration>) -> 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) {

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,16 +47,48 @@ 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 || {
         for stream in listener.incoming() {
             thread::spawn(move || serve(stream.unwrap()));
         }
     });
+    port
+}
+
+fn assert_times_out(client: &memcache::Client) {
+    let start = Instant::now();
+    assert!(client.get::<String>("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()