LISTEN wake-up (next_notification) and RowStream::rows_affected

.github/workflows/ci.yml

@@ -74,6 +74,9 @@ jobs:
       - uses: Swatinem/rust-cache@v2
 
       - name: Test
+        env:
+          # severed_connection proxies to this port (defaults to 5432).
+          TEST_PG_PORT: "5433"
         run: cargo test --all
 
       - name: PostgreSQL logs

src/client.rs

@@ -192,26 +192,57 @@ impl Client {
     /// so `LISTEN` could be issued but nothing could ever be observed.
     ///
     /// The queue is unbounded and is only drained by the caller, so a client
-    /// that issues `LISTEN` and never reads will grow it without limit. Poll it
-    /// as part of the same loop that does the work:
+    /// that issues `LISTEN` and never reads will grow it without limit. The
+    /// connection fills it as notifications arrive, without any request from
+    /// the client; to wait for the next one rather than poll, use
+    /// [`Client::next_notification`].
+    #[inline]
+    pub fn notifications(&self) -> &may::queue::mpsc::Queue<Notification> {
+        self.inner.sender.notifications().queue()
+    }
+
+    /// Wait up to `timeout` for the next notification.
+    ///
+    /// Returns as soon as one is queued (or immediately if one already is),
+    /// parking the coroutine (or thread) in the meantime, so a `LISTEN`
+    /// consumer needs no polling interval. Returns `None` when `timeout`
+    /// passes with nothing delivered, or when the connection has closed and
+    /// nothing is left in the queue.
     ///
     /// ```no_run
     /// # fn main() -> Result<(), may_postgres::Error> {
+    /// # use std::time::Duration;
     /// let client = may_postgres::connect("host=localhost user=postgres")?;
     /// client.batch_execute("LISTEN my_channel")?;
     ///
-    /// loop {
-    ///     // Any round trip lets the connection surface pending notifications.
-    ///     client.batch_execute("SELECT 1")?;
-    ///     while let Some(notification) = client.notifications().pop() {
-    ///         println!("{}: {}", notification.channel(), notification.payload());
+    /// while !client.is_closed() {
+    ///     if let Some(n) = client.next_notification(Duration::from_secs(1)) {
+    ///         println!("{}: {}", n.channel(), n.payload());
     ///     }
     /// }
+    /// # Ok(())
     /// # }
     /// ```
-    #[inline]
-    pub fn notifications(&self) -> &may::queue::mpsc::Queue<Notification> {
-        self.inner.sender.notifications()
+    pub fn next_notification(&self, timeout: std::time::Duration) -> Option<Notification> {
+        let notifications = self.inner.sender.notifications();
+        let deadline = std::time::Instant::now() + timeout;
+        loop {
+            if let Some(n) = notifications.queue().pop() {
+                return Some(n);
+            }
+            if self.is_closed() {
+                // Pass the wake-up on, so every waiter on this connection
+                // returns rather than one of them.
+                notifications.ready().post();
+                return notifications.queue().pop();
+            }
+            let left = deadline.saturating_duration_since(std::time::Instant::now());
+            if left.is_zero() || !notifications.ready().wait_timeout(left) {
+                return notifications.queue().pop();
+            }
+            // A permit can outnumber queued notifications when the queue was
+            // also drained with `notifications().pop()`; loop and re-check.
+        }
     }
 
     pub(crate) fn new(sender: Connection, process_id: i32, secret_key: i32) -> Client {

src/connection.rs

@@ -5,7 +5,7 @@ use may::coroutine::JoinHandle;
 use may::go;
 use may::io::{WaitIo, WaitIoWaker};
 use may::queue::mpsc::Queue;
-use may::sync::spsc;
+use may::sync::{spsc, Semphore};
 use postgres_protocol::message::backend::Message;
 use postgres_protocol::message::frontend;
 
@@ -62,6 +62,41 @@ pub struct Response {
     tx: Arc<spsc::Sender<ResponseResult>>,
 }
 
+/// Asynchronous notifications delivered by the server: a queue the connection
+/// coroutine fills and the client drains, plus a count of pushes a consumer can
+/// park on instead of polling the queue.
+pub(crate) struct Notifications {
+    queue: Queue<Notification>,
+    ready: Semphore,
+}
+
+impl Notifications {
+    fn new() -> Self {
+        Notifications {
+            queue: Queue::new(),
+            ready: Semphore::new(0),
+        }
+    }
+
+    fn push(&self, notification: Notification) {
+        self.queue.push(notification);
+        self.ready.post();
+    }
+
+    /// Wake a consumer without a notification (the connection closed).
+    fn wake(&self) {
+        self.ready.post();
+    }
+
+    pub(crate) fn queue(&self) -> &Queue<Notification> {
+        &self.queue
+    }
+
+    pub(crate) fn ready(&self) -> &Semphore {
+        &self.ready
+    }
+}
+
 /// A connection to a PostgreSQL database.
 pub(crate) struct Connection {
     /// Held only to keep the io coroutine owned; dropping it DETACHES. The
@@ -72,7 +107,7 @@ pub(crate) struct Connection {
     req_queue: Arc<Queue<Request>>,
     /// Asynchronous notifications delivered by the server, filled by the
     /// connection coroutine and drained by the client.
-    notifications: Arc<Queue<Notification>>,
+    notifications: Arc<Notifications>,
     waker: Arc<WaitIoWaker>,
     /// Cooperative shutdown: `drop` sets this and wakes the coroutine; the
     /// connection loop polls it and exits on its own stack, closing its own
@@ -176,7 +211,7 @@ fn decode_messages(
     read_buf: &mut BytesMut,
     rsp_queue: &mut VecDeque<Response>,
     parameters: &mut HashMap<String, String>,
-    notifications: &Queue<Notification>,
+    notifications: &Notifications,
 ) -> Result<(), Error> {
     use crate::codec::PostgresCodec;
 
@@ -320,7 +355,7 @@ fn connection_loop(
     req_queue: &Queue<Request>,
     rsp_queue: &mut VecDeque<Response>,
     mut params: HashMap<String, String>,
-    notifications: Arc<Queue<Notification>>,
+    notifications: Arc<Notifications>,
     shutdown: &AtomicBool,
 ) -> Result<(), Error> {
     let mut read_buf = BytesMut::with_capacity(IO_BUF_SIZE);
@@ -373,8 +408,9 @@ impl Connection {
 
         let req_queue = Arc::new(Queue::new());
         let req_queue_dup = req_queue.clone();
-        let notifications = Arc::new(Queue::new());
+        let notifications = Arc::new(Notifications::new());
         let notifications_dup = notifications.clone();
+        let notifications_exit = notifications.clone();
         let shutdown = Arc::new(AtomicBool::new(false));
         let shutdown_dup = shutdown.clone();
         let dead = Arc::new(AtomicBool::new(false));
@@ -430,6 +466,9 @@ impl Connection {
                     let _g = drain_lock_dup.lock();
                     fail_queued_requests(&req_queue_dup);
                 }
+                // A consumer parked waiting for a notification would otherwise
+                // wait out its timeout on a connection that will never deliver.
+                notifications_exit.wake();
                 // Send the protocol Terminate and close the socket, on this
                 // coroutine's own stack. This is what releases the server backend
                 // the moment the client goes away, instead of at process exit.
@@ -487,7 +526,7 @@ impl Connection {
     ///
     /// Populated by `LISTEN`; see [`crate::Client::notifications`].
     #[inline]
-    pub fn notifications(&self) -> &Queue<Notification> {
+    pub fn notifications(&self) -> &Notifications {
         &self.notifications
     }
 

src/query.rs

@@ -17,6 +17,7 @@ pub fn query(
     Ok(RowStream {
         statement,
         responses,
+        rows_affected: None,
     })
 }
 
@@ -32,6 +33,7 @@ pub fn query_portal(client: &Client, portal: &Portal, max_rows: i32) -> Result<R
     Ok(RowStream {
         statement: portal.statement().clone(),
         responses,
+        rows_affected: None,
     })
 }
 
@@ -123,6 +125,20 @@ pub fn encode_bind(
 pub struct RowStream {
     statement: Statement,
     responses: Responses,
+    rows_affected: Option<u64>,
+}
+
+impl RowStream {
+    /// The number of rows the statement affected, from its `CommandComplete`
+    /// tag (`INSERT 0 3`, `UPDATE 5`, `SELECT 2`, ...).
+    ///
+    /// `None` until the stream has been read to its end, and for a portal
+    /// that was suspended before completing. This lets a statement sent with
+    /// `query_raw` (for example to pipeline it with others) report its count
+    /// the way `execute` does, without a second round trip.
+    pub fn rows_affected(&self) -> Option<u64> {
+        self.rows_affected
+    }
 }
 
 impl Iterator for RowStream {
@@ -135,9 +151,19 @@ impl Iterator for RowStream {
                 Message::DataRow(body) => {
                     return Some(Ok(o_try!(Row::new(self.statement.clone(), body))))
                 }
-                Message::EmptyQueryResponse
-                | Message::CommandComplete(_)
-                | Message::PortalSuspended => return None,
+                Message::CommandComplete(body) => {
+                    self.rows_affected = body
+                        .tag()
+                        .ok()
+                        .and_then(|tag| tag.rsplit(' ').next())
+                        .and_then(|n| n.parse().ok());
+                    return None;
+                }
+                Message::EmptyQueryResponse => {
+                    self.rows_affected = Some(0);
+                    return None;
+                }
+                Message::PortalSuspended => return None,
                 Message::ErrorResponse(body) => return Some(Err(Error::db(body))),
                 Message::BindComplete => {}
                 _ => return Some(Err(Error::unexpected_message())),

tests/test/main.rs

@@ -649,6 +649,98 @@ fn a_client_that_never_listens_receives_nothing() {
     );
 }
 
+#[test]
+fn next_notification_wakes_a_waiting_client_without_polling() {
+    let client = connect("user=postgres");
+    client.batch_execute("LISTEN test_wakeup").unwrap();
+
+    // Raise the notification from another coroutine after the waiter is parked;
+    // the waiter issues no request in the meantime, so only the connection's
+    // own read can deliver it.
+    let notify = may::go!(|| {
+        may::coroutine::sleep(std::time::Duration::from_millis(200));
+        let notifier = connect("user=postgres");
+        notifier
+            .batch_execute("NOTIFY test_wakeup, 'woken'")
+            .unwrap();
+    });
+
+    let started = std::time::Instant::now();
+    let n = client
+        .next_notification(std::time::Duration::from_secs(10))
+        .expect("the notification should wake the waiter");
+    let waited = started.elapsed();
+    notify.join().unwrap();
+
+    assert_eq!(n.payload(), "woken");
+    assert!(
+        waited >= std::time::Duration::from_millis(150),
+        "returned before the notification was raised: {waited:?}"
+    );
+    assert!(
+        waited < std::time::Duration::from_secs(5),
+        "the waiter must wake on arrival, not at its timeout: {waited:?}"
+    );
+}
+
+#[test]
+fn next_notification_returns_queued_notifications_first_and_times_out_empty() {
+    let client = connect("user=postgres");
+    client.batch_execute("LISTEN test_queued").unwrap();
+    let notifier = connect("user=postgres");
+    notifier
+        .batch_execute("NOTIFY test_queued, 'a'; NOTIFY test_queued, 'b';")
+        .unwrap();
+
+    let a = client
+        .next_notification(std::time::Duration::from_secs(5))
+        .expect("first");
+    let b = client
+        .next_notification(std::time::Duration::from_secs(5))
+        .expect("second");
+    assert_eq!((a.payload(), b.payload()), ("a", "b"));
+
+    let started = std::time::Instant::now();
+    assert!(client
+        .next_notification(std::time::Duration::from_millis(100))
+        .is_none());
+    let waited = started.elapsed();
+    assert!(
+        waited >= std::time::Duration::from_millis(90)
+            && waited < std::time::Duration::from_secs(2),
+        "an empty wait lasts about its timeout: {waited:?}"
+    );
+}
+
+#[test]
+fn next_notification_returns_when_the_connection_closes() {
+    let client = connect("user=postgres");
+    client.batch_execute("LISTEN test_closed").unwrap();
+    let pid: i32 = client
+        .query_one("SELECT pg_backend_pid()", &[])
+        .unwrap()
+        .get(0);
+
+    let killer = may::go!(move || {
+        may::coroutine::sleep(std::time::Duration::from_millis(200));
+        let admin = connect("user=postgres");
+        admin
+            .query_one("SELECT pg_terminate_backend($1)", &[&pid])
+            .unwrap();
+    });
+
+    let started = std::time::Instant::now();
+    assert!(client
+        .next_notification(std::time::Duration::from_secs(10))
+        .is_none());
+    killer.join().unwrap();
+    assert!(
+        started.elapsed() < std::time::Duration::from_secs(5),
+        "a closed connection must release the waiter, not hold it to its timeout"
+    );
+    assert!(client.is_closed());
+}
+
 /// Poll for a notification, giving the connection round trips to surface it.
 ///
 /// Notifications arrive asynchronously: the server sends them when it chooses,
@@ -665,6 +757,54 @@ fn wait_for_notification(client: &Client) -> Option<may_postgres::Notification>
     None
 }
 
+#[test]
+fn row_stream_reports_rows_affected() {
+    let client = connect("user=postgres");
+    client
+        .batch_execute(
+            "CREATE TEMPORARY TABLE rows_affected_test (id INT PRIMARY KEY, n INT);
+             INSERT INTO rows_affected_test VALUES (1, 0), (2, 0), (3, 0);",
+        )
+        .unwrap();
+
+    // A write without RETURNING: no rows, but a count, as `execute` reports.
+    let stmt = client
+        .prepare("UPDATE rows_affected_test SET n = n + 1 WHERE id <= $1")
+        .unwrap();
+    let mut stream = client.query_raw(&stmt, &[&2i32]).unwrap();
+    assert_eq!(
+        stream.rows_affected(),
+        None,
+        "unknown until the stream ends"
+    );
+    assert!(stream.next().is_none());
+    assert_eq!(stream.rows_affected(), Some(2));
+
+    // With RETURNING: the rows and the count.
+    let stmt = client
+        .prepare("UPDATE rows_affected_test SET n = n + 1 RETURNING id")
+        .unwrap();
+    let mut stream = client.query_raw(&stmt, &[]).unwrap();
+    let rows = stream.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
+    assert_eq!(rows.len(), 3);
+    assert_eq!(stream.rows_affected(), Some(3));
+
+    // Pipelined: both statements are sent before either is read, and each
+    // stream reports its own count.
+    let insert = client
+        .prepare("INSERT INTO rows_affected_test VALUES ($1, 0)")
+        .unwrap();
+    let delete = client
+        .prepare("DELETE FROM rows_affected_test WHERE id <= $1")
+        .unwrap();
+    let mut a = client.query_raw(&insert, &[&10i32]).unwrap();
+    let mut b = client.query_raw(&delete, &[&3i32]).unwrap();
+    assert!(a.next().is_none());
+    assert!(b.next().is_none());
+    assert_eq!(a.rows_affected(), Some(1));
+    assert_eq!(b.rows_affected(), Some(3));
+}
+
 #[test]
 fn query_portal() {
     let mut client = connect("user=postgres");