diff --git a/src/lib.rs b/src/lib.rs index 6d8936a0..04b0e2c6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -527,6 +527,24 @@ impl InnerConnection { } } + fn read_message_with_notification_nonblocking(&mut self) + -> std::io::Result> { + debug_assert!(!self.desynchronized); + loop { + match try_desync!(self, self.stream.read_message_nonblocking()) { + Some(NoticeResponse { fields }) => { + if let Ok(err) = DbError::new_raw(fields) { + self.notice_handler.handle_notice(err); + } + } + Some(ParameterStatus { parameter, value }) => { + self.parameters.insert(parameter, value); + } + val => return Ok(val), + } + } + } + fn read_message(&mut self) -> std_io::Result { loop { match try!(self.read_message_with_notification()) { diff --git a/src/message.rs b/src/message.rs index 18d0cf9a..b454268e 100644 --- a/src/message.rs +++ b/src/message.rs @@ -6,7 +6,7 @@ use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; use types::Oid; use util; -use priv_io::ReadTimeout; +use priv_io::StreamOptions; use self::BackendMessage::*; use self::FrontendMessage::*; @@ -287,10 +287,12 @@ pub trait ReadMessage { fn read_message_timeout(&mut self, timeout: Duration) -> io::Result>; + fn read_message_nonblocking(&mut self) -> io::Result>; + fn finish_read_message(&mut self, ident: u8) -> io::Result; } -impl ReadMessage for R { +impl ReadMessage for R { fn read_message(&mut self) -> io::Result { let ident = try!(self.read_u8()); self.finish_read_message(ident) @@ -314,6 +316,24 @@ impl ReadMessage for R { } } + fn read_message_nonblocking(&mut self) -> io::Result> { + try!(self.set_nonblocking(true)); + let ident = self.read_u8(); + try!(self.set_nonblocking(false)); + + 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 { + Ok(None) + } else { + Err(e) + } + } + } + } + fn finish_read_message(&mut self, ident: u8) -> io::Result { // subtract size of length value let len = try!(self.read_u32::()) - mem::size_of::() as u32; diff --git a/src/notification.rs b/src/notification.rs index bb07742e..b05d7e0a 100644 --- a/src/notification.rs +++ b/src/notification.rs @@ -42,8 +42,7 @@ impl<'conn> Notifications<'conn> { /// # Note /// /// This iterator may start returning `Some` after previously returning - /// `None` if more notifications are received. However, those notifications - /// will not be registered until the connection is used in some way. + /// `None` if more notifications are received. pub fn iter<'a>(&'a self) -> Iter<'a> { Iter { conn: self.conn } } @@ -72,7 +71,7 @@ impl<'conn> Notifications<'conn> { } impl<'a, 'conn> IntoIterator for &'a Notifications<'conn> { - type Item = Notification; + type Item = Result; type IntoIter = Iter<'a>; fn into_iter(self) -> Iter<'a> { @@ -92,10 +91,27 @@ pub struct Iter<'a> { } impl<'a> Iterator for Iter<'a> { - type Item = Notification; + type Item = Result; - fn next(&mut self) -> Option { - self.conn.conn.borrow_mut().notifications.pop_front() + fn next(&mut self) -> Option> { + let mut conn = self.conn.conn.borrow_mut(); + + if let Some(notification) = conn.notifications.pop_front() { + return Some(Ok(notification)); + } + + match conn.read_message_with_notification_nonblocking() { + Ok(Some(NotificationResponse { pid, channel, payload })) => { + Some(Ok(Notification { + pid: pid, + channel: channel, + payload: payload, + })) + } + Ok(None) => None, + Err(err) => Some(Err(Error::Io(err))), + _ => unreachable!(), + } } } diff --git a/src/priv_io.rs b/src/priv_io.rs index 03590c56..251760ff 100644 --- a/src/priv_io.rs +++ b/src/priv_io.rs @@ -22,11 +22,12 @@ use message::FrontendMessage::SslRequest; const DEFAULT_PORT: u16 = 5432; #[doc(hidden)] -pub trait ReadTimeout { +pub trait StreamOptions { fn set_read_timeout(&self, timeout: Option) -> io::Result<()>; + fn set_nonblocking(&self, nonblock: bool) -> io::Result<()>; } -impl ReadTimeout for BufStream> { +impl StreamOptions for BufStream> { fn set_read_timeout(&self, timeout: Option) -> io::Result<()> { match self.get_ref().get_ref().0 { InternalStream::Tcp(ref s) => { @@ -36,6 +37,14 @@ impl ReadTimeout for BufStream> { InternalStream::Unix(ref s) => s.set_read_timeout(timeout), } } + + fn set_nonblocking(&self, nonblock: bool) -> io::Result<()> { + match self.get_ref().get_ref().0 { + InternalStream::Tcp(ref s) => s.set_nonblocking(nonblock), + #[cfg(feature = "unix_socket")] + InternalStream::Unix(ref s) => s.set_nonblocking(nonblock), + } + } } /// A connection to the Postgres server. diff --git a/tests/test.rs b/tests/test.rs index 4b380f3e..e7a61b1c 100644 --- a/tests/test.rs +++ b/tests/test.rs @@ -575,12 +575,12 @@ fn test_notification_iterator_some() { pid: 0, channel: "test_notification_iterator_one_channel".to_string(), payload: "hello".to_string() - }, it.next().unwrap()); + }, it.next().unwrap().unwrap()); check_notification(Notification { pid: 0, channel: "test_notification_iterator_one_channel2".to_string(), payload: "world".to_string() - }, it.next().unwrap()); + }, it.next().unwrap().unwrap()); assert!(it.next().is_none()); or_panic!(conn.execute("NOTIFY test_notification_iterator_one_channel, '!'", &[])); @@ -588,7 +588,7 @@ fn test_notification_iterator_some() { pid: 0, channel: "test_notification_iterator_one_channel".to_string(), payload: "!".to_string() - }, it.next().unwrap()); + }, it.next().unwrap().unwrap()); assert!(it.next().is_none()); }