From afd37db4f5e728afdbf86fb99a60471d01620d82 Mon Sep 17 00:00:00 2001 From: Charles Sibbald Date: Tue, 27 Oct 2026 18:57:40 +0200 Subject: [PATCH 1/3] feat(query): expose the affected-row count on RowStream `execute` reports how many rows a statement affected, but a statement sent with `query_raw` returns a `RowStream`, which dropped the `CommandComplete` tag. A caller that pipelines several statements with `query_raw` (send them all, then read the streams in order) therefore had no way to learn a write's row count without a separate round trip. `RowStream::rows_affected()` returns the count parsed from the tag (`UPDATE 5`, `INSERT 0 3`, `SELECT 2`) once the stream has been read to its end, `Some(0)` for an empty query, and `None` before the end or for a suspended portal. tokio-postgres has the same method. Test: counts for an UPDATE without RETURNING, with RETURNING, and for two pipelined statements read in order. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Y3TSRmEZAYU6RBUXhkn7WF --- src/query.rs | 32 ++++++++++++++++++++++++++++--- tests/test/main.rs | 48 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 3 deletions(-) diff --git a/src/query.rs b/src/query.rs index e84b5c6..cebfe8c 100644 --- a/src/query.rs +++ b/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, +} + +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 { + 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())), diff --git a/tests/test/main.rs b/tests/test/main.rs index 5617208..93c40c8 100644 --- a/tests/test/main.rs +++ b/tests/test/main.rs @@ -665,6 +665,54 @@ fn wait_for_notification(client: &Client) -> Option 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::, _>>().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"); From 3597ac0e02be14b774a930382624c97e08a960d2 Mon Sep 17 00:00:00 2001 From: Charles Sibbald Date: Tue, 27 Oct 2026 18:57:40 +0200 Subject: [PATCH 2/3] feat(client): wake LISTEN consumers when a notification arrives Notifications are queued by the connection coroutine as they arrive, but the queue had no wake-up, so a consumer had to poll it: either burning CPU on a short interval or adding up to the interval in latency. The queue now carries a semaphore that the connection posts once per notification, and once more when the connection closes. `Client::next_notification(timeout)` returns a queued notification at once, otherwise parks the coroutine (or thread) until one arrives, the timeout passes, or the connection closes. `Client::notifications()` is unchanged; mixing direct pops with waits only costs a spurious wake-up. Tests: a parked waiter wakes on a notification raised by another session with no request of its own; queued notifications come first and an empty wait lasts its timeout; terminating the backend releases a waiter well before its timeout. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Y3TSRmEZAYU6RBUXhkn7WF --- src/client.rs | 51 ++++++++++++++++++++----- src/connection.rs | 51 ++++++++++++++++++++++--- tests/test/main.rs | 92 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 178 insertions(+), 16 deletions(-) diff --git a/src/client.rs b/src/client.rs index a6f0d9d..27dd691 100644 --- a/src/client.rs +++ b/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 { + 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 { - self.inner.sender.notifications() + pub fn next_notification(&self, timeout: std::time::Duration) -> Option { + 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 { diff --git a/src/connection.rs b/src/connection.rs index b7676a4..4e7bff1 100644 --- a/src/connection.rs +++ b/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>, } +/// 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, + 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 { + &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>, /// Asynchronous notifications delivered by the server, filled by the /// connection coroutine and drained by the client. - notifications: Arc>, + notifications: Arc, waker: Arc, /// 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, parameters: &mut HashMap, - notifications: &Queue, + notifications: &Notifications, ) -> Result<(), Error> { use crate::codec::PostgresCodec; @@ -320,7 +355,7 @@ fn connection_loop( req_queue: &Queue, rsp_queue: &mut VecDeque, mut params: HashMap, - notifications: Arc>, + notifications: Arc, 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 { + pub fn notifications(&self) -> &Notifications { &self.notifications } diff --git a/tests/test/main.rs b/tests/test/main.rs index 93c40c8..9e12501 100644 --- a/tests/test/main.rs +++ b/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, From 6ca08d171fde1e547707d1bf476823b9d57277ff Mon Sep 17 00:00:00 2001 From: Charles Sibbald Date: Tue, 27 Oct 2026 19:01:39 +0200 Subject: [PATCH 3/3] ci: point severed_connection tests at the CI postgres port The test postgres is published on 127.0.0.1:5433, but tests/severed_connection.rs reads TEST_PG_PORT and defaults to 5432, so both of its tests failed on every CI run (master included). Set TEST_PG_PORT=5433 for the test step. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Y3TSRmEZAYU6RBUXhkn7WF --- .github/workflows/ci.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index caea076..b613820 100644 --- a/.github/workflows/ci.yml +++ b/.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