From bb837bd8722ef16f085ea76e473f0ae3fbae4d9f Mon Sep 17 00:00:00 2001 From: Steven Fackler Date: Sun, 27 Dec 2015 10:13:42 -0700 Subject: [PATCH] Make the default notifications iterator read nonblocking It is always super confusing as to when a notification that's been sent to the client will actually show up in the old version of this iterator, so it's best to have it see if there's anything waiting in the TCP buffer. Closes #149 --- src/lib.rs | 18 ++++++++++++++++++ src/message.rs | 24 ++++++++++++++++++++++-- src/notification.rs | 28 ++++++++++++++++++++++------ src/priv_io.rs | 13 +++++++++++-- tests/test.rs | 6 +++--- 5 files changed, 76 insertions(+), 13 deletions(-) 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()); }