Read message with timeout backend support

Author: sfackler

Cargo.toml

@@ -29,11 +29,12 @@ byteorder = "0.3"
 log = "0.3"
 phf = "0.7"
 rustc-serialize = "0.3"
+net2 = { version = "0.2", features = ["nightly"] }
 chrono = { version = "0.2.14", optional = true }
 openssl = { version = "0.6.4", optional = true }
 serde = { version = "0.3", optional = true }
 time = { version = "0.1.14", optional = true }
-unix_socket = { version = ">= 0.3, < 0.5", optional = true }
+unix_socket = { version = ">= 0.3, < 0.5", optional = true, features = ["socket_timeout"] }
 uuid = { version = "0.1", optional = true }
 
 [dev-dependencies]

src/lib.rs

@@ -52,6 +52,7 @@ extern crate phf;
 extern crate rustc_serialize as serialize;
 #[cfg(feature = "unix_socket")]
 extern crate unix_socket;
+extern crate net2;
 
 use bufstream::BufStream;
 use md5::Md5;

src/message.rs

@@ -1,10 +1,12 @@
 use std::io;
 use std::io::prelude::*;
 use std::mem;
+use std::time::Duration;
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 
 use types::Oid;
 use util;
+use priv_io::ReadTimeout;
 
 use self::BackendMessage::*;
 use self::FrontendMessage::*;
@@ -282,12 +284,39 @@ impl<R: BufRead> ReadCStr for R {
 #[doc(hidden)]
 pub trait ReadMessage {
     fn read_message(&mut self) -> io::Result<BackendMessage>;
+
+    fn read_message_timeout(&mut self, timeout: Duration)
+                            -> io::Result<Option<BackendMessage>>;
+
+    fn finish_read_message(&mut self, ident: u8) -> io::Result<BackendMessage>;
 }
 
-impl<R: BufRead> ReadMessage for R {
+impl<R: BufRead + ReadTimeout> ReadMessage for R {
     fn read_message(&mut self) -> io::Result<BackendMessage> {
         let ident = try!(self.read_u8());
+        self.finish_read_message(ident)
+    }
+
+    fn read_message_timeout(&mut self, timeout: Duration)
+                            -> io::Result<Option<BackendMessage>> {
+        try!(self.set_read_timeout(Some(timeout)));
+        let ident = self.read_u8();
+        try!(self.set_read_timeout(None));
+
+        match ident {
+            Ok(ident) => self.finish_read_message(ident).map(Some),
+            Err(e) => {
+                let e: io::Error = e.into();
+                if e.kind() == io::ErrorKind::WouldBlock || e.kind() == io::ErrorKind::TimedOut {
+                    Ok(None)
+                } else {
+                    Err(e)
+                }
+            }
+        }
+    }
 
+    fn finish_read_message(&mut self, ident: u8) -> io::Result<BackendMessage> {
         // subtract size of length value
         let len = try!(self.read_u32::<BigEndian>()) - mem::size_of::<u32>() as u32;
         let mut rdr = self.by_ref().take(len as u64);

src/priv_io.rs

@@ -1,7 +1,10 @@
 use byteorder::ReadBytesExt;
+use net2::TcpStreamExt;
 use std::io;
 use std::io::prelude::*;
 use std::net::TcpStream;
+use std::time::Duration;
+use bufstream::BufStream;
 #[cfg(feature = "unix_socket")]
 use unix_socket::UnixStream;
 #[cfg(unix)]
@@ -17,6 +20,21 @@ use message::FrontendMessage::SslRequest;
 
 const DEFAULT_PORT: u16 = 5432;
 
+#[doc(hidden)]
+pub trait ReadTimeout {
+    fn set_read_timeout(&self, timeout: Option<Duration>) -> io::Result<()>;
+}
+
+impl ReadTimeout for BufStream<Box<StreamWrapper>> {
+    fn set_read_timeout(&self, timeout: Option<Duration>) -> io::Result<()> {
+        match self.get_ref().get_ref().0 {
+            InternalStream::Tcp(ref s) => <TcpStream as TcpStreamExt>::set_read_timeout(s, timeout),
+            #[cfg(feature = "unix_socket")]
+            InternalStream::Unix(ref s) => s.set_read_timeout(timeout),
+        }
+    }
+}
+
 /// A connection to the Postgres server.
 ///
 /// It implements `Read`, `Write` and `StreamWrapper`, as well as `AsRawFd` on