Add a notification API to the blocking client

This mirrors the implementation in the old 0.15 release, but is quite a
bit simpler now that we're built on the nonblocking API!

Author: sfackler

postgres/Cargo.toml

@@ -35,7 +35,7 @@ fallible-iterator = "0.2"
 futures = "0.3"
 tokio-postgres = { version = "0.5.3", path = "../tokio-postgres" }
 
-tokio = { version = "0.2", features = ["rt-core"] }
+tokio = { version = "0.2", features = ["rt-core", "time"] }
 log = "0.4"
 
 [dev-dependencies]

postgres/src/client.rs

@@ -1,7 +1,7 @@
 use crate::connection::Connection;
 use crate::{
-    CancelToken, Config, CopyInWriter, CopyOutReader, RowIter, Statement, ToStatement, Transaction,
-    TransactionBuilder,
+    CancelToken, Config, CopyInWriter, CopyOutReader, Notifications, RowIter, Statement,
+    ToStatement, Transaction, TransactionBuilder,
 };
 use tokio_postgres::tls::{MakeTlsConnect, TlsConnect};
 use tokio_postgres::types::{ToSql, Type};
@@ -471,6 +471,13 @@ impl Client {
         TransactionBuilder::new(self.connection.as_ref(), self.client.build_transaction())
     }
 
+    /// Returns a structure providing access to asynchronous notifications.
+    ///
+    /// Use the `LISTEN` command to register this connection for notifications.
+    pub fn notifications(&mut self) -> Notifications<'_> {
+        Notifications::new(self.connection.as_ref())
+    }
+
     /// Constructs a cancellation token that can later be used to request
     /// cancellation of a query running on this connection.
     ///
@@ -490,7 +497,7 @@ impl Client {
     /// thread::spawn(move || {
     ///     // Abort the query after 5s.
     ///     thread::sleep(Duration::from_secs(5));
-    ///     cancel_token.cancel_query(NoTls);
+    ///     let _ = cancel_token.cancel_query(NoTls);
     /// });
     ///
     /// match client.simple_query("SELECT long_running_query()") {

postgres/src/connection.rs

@@ -34,16 +34,30 @@ impl Connection {
         ConnectionRef { connection: self }
     }
 
+    pub fn enter<F, T>(&self, f: F) -> T
+    where
+        F: FnOnce() -> T,
+    {
+        self.runtime.enter(f)
+    }
+
     pub fn block_on<F, T>(&mut self, future: F) -> Result<T, Error>
     where
         F: Future<Output = Result<T, Error>>,
     {
         pin_mut!(future);
+        self.poll_block_on(|cx, _, _| future.as_mut().poll(cx))
+    }
+
+    pub fn poll_block_on<F, T>(&mut self, mut f: F) -> Result<T, Error>
+    where
+        F: FnMut(&mut Context<'_>, &mut VecDeque<Notification>, bool) -> Poll<Result<T, Error>>,
+    {
         let connection = &mut self.connection;
         let notifications = &mut self.notifications;
         self.runtime.block_on({
             future::poll_fn(|cx| {
-                loop {
+                let done = loop {
                     match connection.as_mut().poll_next(cx) {
                         Poll::Ready(Some(Ok(AsyncMessage::Notification(notification)))) => {
                             notifications.push_back(notification);
@@ -53,14 +67,23 @@ impl Connection {
                         }
                         Poll::Ready(Some(Ok(_))) => {}
                         Poll::Ready(Some(Err(e))) => return Poll::Ready(Err(e)),
-                        Poll::Ready(None) | Poll::Pending => break,
+                        Poll::Ready(None) => break true,
+                        Poll::Pending => break false,
                     }
-                }
+                };
 
-                future.as_mut().poll(cx)
+                f(cx, notifications, done)
             })
         })
     }
+
+    pub fn notifications(&self) -> &VecDeque<Notification> {
+        &self.notifications
+    }
+
+    pub fn notifications_mut(&mut self) -> &mut VecDeque<Notification> {
+        &mut self.notifications
+    }
 }
 
 pub struct ConnectionRef<'a> {

postgres/src/lib.rs

@@ -77,6 +77,8 @@ pub use crate::copy_out_reader::CopyOutReader;
 #[doc(no_inline)]
 pub use crate::error::Error;
 pub use crate::generic_client::GenericClient;
+#[doc(inline)]
+pub use crate::notifications::Notifications;
 #[doc(no_inline)]
 pub use crate::row::{Row, SimpleQueryRow};
 pub use crate::row_iter::RowIter;
@@ -94,6 +96,7 @@ mod copy_in_writer;
 mod copy_out_reader;
 mod generic_client;
 mod lazy_pin;
+pub mod notifications;
 mod row_iter;
 mod transaction;
 mod transaction_builder;

postgres/src/notifications.rs

@@ -0,0 +1,161 @@
+//! Asynchronous notifications.
+
+use crate::connection::ConnectionRef;
+use crate::{Error, Notification};
+use fallible_iterator::FallibleIterator;
+use futures::{ready, FutureExt};
+use std::task::Poll;
+use std::time::Duration;
+use tokio::time::{self, Delay, Instant};
+
+/// Notifications from a PostgreSQL backend.
+pub struct Notifications<'a> {
+    connection: ConnectionRef<'a>,
+}
+
+impl<'a> Notifications<'a> {
+    pub(crate) fn new(connection: ConnectionRef<'a>) -> Notifications<'a> {
+        Notifications { connection }
+    }
+
+    /// Returns the number of already buffered pending notifications.
+    pub fn len(&self) -> usize {
+        self.connection.notifications().len()
+    }
+
+    /// Determines if there are any already buffered pending notifications.
+    pub fn is_empty(&self) -> bool {
+        self.connection.notifications().is_empty()
+    }
+
+    /// Returns a nonblocking iterator over notifications.
+    ///
+    /// If there are no already buffered pending notifications, this iterator will poll the connection but will not
+    /// block waiting on notifications over the network. A return value of `None` either indicates that there are no
+    /// pending notifications or that the server has disconnected.
+    ///
+    /// # Note
+    ///
+    /// This iterator may start returning `Some` after previously returning `None` if more notifications are received.
+    pub fn iter(&mut self) -> Iter<'_> {
+        Iter {
+            connection: self.connection.as_ref(),
+        }
+    }
+
+    /// Returns a blocking iterator over notifications.
+    ///
+    /// If there are no already buffered pending notifications, this iterator will block indefinitely waiting on the
+    /// PostgreSQL backend server to send one. It will only return `None` if the server has disconnected.
+    pub fn blocking_iter(&mut self) -> BlockingIter<'_> {
+        BlockingIter {
+            connection: self.connection.as_ref(),
+        }
+    }
+
+    /// Returns an iterator over notifications which blocks a limited amount of time.
+    ///
+    /// If there are no already buffered pending notifications, this iterator will block waiting on the PostgreSQL
+    /// backend server to send one up to the provided timeout. A return value of `None` either indicates that there are
+    /// no pending notifications or that the server has disconnected.
+    ///
+    /// # Note
+    ///
+    /// This iterator may start returning `Some` after previously returning `None` if more notifications are received.
+    pub fn timeout_iter(&mut self, timeout: Duration) -> TimeoutIter<'_> {
+        TimeoutIter {
+            delay: self.connection.enter(|| time::delay_for(timeout)),
+            timeout,
+            connection: self.connection.as_ref(),
+        }
+    }
+}
+
+/// A nonblocking iterator over pending notifications.
+pub struct Iter<'a> {
+    connection: ConnectionRef<'a>,
+}
+
+impl<'a> FallibleIterator for Iter<'a> {
+    type Item = Notification;
+    type Error = Error;
+
+    fn next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
+        if let Some(notification) = self.connection.notifications_mut().pop_front() {
+            return Ok(Some(notification));
+        }
+
+        self.connection
+            .poll_block_on(|_, notifications, _| Poll::Ready(Ok(notifications.pop_front())))
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.connection.notifications().len(), None)
+    }
+}
+
+/// A blocking iterator over pending notifications.
+pub struct BlockingIter<'a> {
+    connection: ConnectionRef<'a>,
+}
+
+impl<'a> FallibleIterator for BlockingIter<'a> {
+    type Item = Notification;
+    type Error = Error;
+
+    fn next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
+        if let Some(notification) = self.connection.notifications_mut().pop_front() {
+            return Ok(Some(notification));
+        }
+
+        self.connection
+            .poll_block_on(|_, notifications, done| match notifications.pop_front() {
+                Some(notification) => Poll::Ready(Ok(Some(notification))),
+                None if done => Poll::Ready(Ok(None)),
+                None => Poll::Pending,
+            })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.connection.notifications().len(), None)
+    }
+}
+
+/// A time-limited blocking iterator over pending notifications.
+pub struct TimeoutIter<'a> {
+    connection: ConnectionRef<'a>,
+    delay: Delay,
+    timeout: Duration,
+}
+
+impl<'a> FallibleIterator for TimeoutIter<'a> {
+    type Item = Notification;
+    type Error = Error;
+
+    fn next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
+        if let Some(notification) = self.connection.notifications_mut().pop_front() {
+            self.delay.reset(Instant::now() + self.timeout);
+            return Ok(Some(notification));
+        }
+
+        let delay = &mut self.delay;
+        let timeout = self.timeout;
+        self.connection.poll_block_on(|cx, notifications, done| {
+            match notifications.pop_front() {
+                Some(notification) => {
+                    delay.reset(Instant::now() + timeout);
+                    return Poll::Ready(Ok(Some(notification)));
+                }
+                None if done => return Poll::Ready(Ok(None)),
+                None => {}
+            }
+
+            ready!(delay.poll_unpin(cx));
+            Poll::Ready(Ok(None))
+        })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.connection.notifications().len(), None)
+    }
+}

postgres/src/test.rs

@@ -309,3 +309,93 @@ fn cancel_query() {
 
     cancel_thread.join().unwrap();
 }
+
+#[test]
+fn notifications_iter() {
+    let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
+
+    client
+        .batch_execute(
+            "\
+        LISTEN notifications_iter;
+        NOTIFY notifications_iter, 'hello';
+        NOTIFY notifications_iter, 'world';
+    ",
+        )
+        .unwrap();
+
+    let notifications = client.notifications().iter().collect::<Vec<_>>().unwrap();
+    assert_eq!(notifications.len(), 2);
+    assert_eq!(notifications[0].payload(), "hello");
+    assert_eq!(notifications[1].payload(), "world");
+}
+
+#[test]
+fn notifications_blocking_iter() {
+    let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
+
+    client
+        .batch_execute(
+            "\
+        LISTEN notifications_blocking_iter;
+        NOTIFY notifications_blocking_iter, 'hello';
+    ",
+        )
+        .unwrap();
+
+    thread::spawn(|| {
+        let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
+
+        thread::sleep(Duration::from_secs(1));
+        client
+            .batch_execute("NOTIFY notifications_blocking_iter, 'world'")
+            .unwrap();
+    });
+
+    let notifications = client
+        .notifications()
+        .blocking_iter()
+        .take(2)
+        .collect::<Vec<_>>()
+        .unwrap();
+    assert_eq!(notifications.len(), 2);
+    assert_eq!(notifications[0].payload(), "hello");
+    assert_eq!(notifications[1].payload(), "world");
+}
+
+#[test]
+fn notifications_timeout_iter() {
+    let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
+
+    client
+        .batch_execute(
+            "\
+        LISTEN notifications_timeout_iter;
+        NOTIFY notifications_timeout_iter, 'hello';
+    ",
+        )
+        .unwrap();
+
+    thread::spawn(|| {
+        let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
+
+        thread::sleep(Duration::from_secs(1));
+        client
+            .batch_execute("NOTIFY notifications_timeout_iter, 'world'")
+            .unwrap();
+
+        thread::sleep(Duration::from_secs(10));
+        client
+            .batch_execute("NOTIFY notifications_timeout_iter, '!'")
+            .unwrap();
+    });
+
+    let notifications = client
+        .notifications()
+        .timeout_iter(Duration::from_secs(2))
+        .collect::<Vec<_>>()
+        .unwrap();
+    assert_eq!(notifications.len(), 2);
+    assert_eq!(notifications[0].payload(), "hello");
+    assert_eq!(notifications[1].payload(), "world");
+}