diff --git a/postgres/Cargo.toml b/postgres/Cargo.toml index d23715cd..d0cf1100 100644 --- a/postgres/Cargo.toml +++ b/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] diff --git a/postgres/src/client.rs b/postgres/src/client.rs index 3ae5f86c..a0c61b33 100644 --- a/postgres/src/client.rs +++ b/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()") { diff --git a/postgres/src/connection.rs b/postgres/src/connection.rs index 440ad5da..acea5eca 100644 --- a/postgres/src/connection.rs +++ b/postgres/src/connection.rs @@ -34,16 +34,30 @@ impl Connection { ConnectionRef { connection: self } } + pub fn enter(&self, f: F) -> T + where + F: FnOnce() -> T, + { + self.runtime.enter(f) + } + pub fn block_on(&mut self, future: F) -> Result where F: Future>, { pin_mut!(future); + self.poll_block_on(|cx, _, _| future.as_mut().poll(cx)) + } + + pub fn poll_block_on(&mut self, mut f: F) -> Result + where + F: FnMut(&mut Context<'_>, &mut VecDeque, bool) -> Poll>, + { 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 { + &self.notifications + } + + pub fn notifications_mut(&mut self) -> &mut VecDeque { + &mut self.notifications + } } pub struct ConnectionRef<'a> { diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs index 78b318b1..80380a87 100644 --- a/postgres/src/lib.rs +++ b/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; diff --git a/postgres/src/notifications.rs b/postgres/src/notifications.rs new file mode 100644 index 00000000..e8c68154 --- /dev/null +++ b/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, 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) { + (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, 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) { + (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, 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) { + (self.connection.notifications().len(), None) + } +} diff --git a/postgres/src/test.rs b/postgres/src/test.rs index 449aac01..9edde8e3 100644 --- a/postgres/src/test.rs +++ b/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::>().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::>() + .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::>() + .unwrap(); + assert_eq!(notifications.len(), 2); + assert_eq!(notifications[0].payload(), "hello"); + assert_eq!(notifications[1].payload(), "world"); +}