sqlx-postgres: Update PgListener

Author: joeydewaal

sqlx-postgres/src/connection/mod.rs

@@ -50,7 +50,7 @@ pub struct PgConnectionInner {
 
     chan: UnboundedSender<IoRequest>,
 
-    notifications: UnboundedReceiver<Notification>,
+    pub(crate) notifications: UnboundedReceiver<Notification>,
 
     // process id of this backend
     // used to send cancel requests
@@ -74,9 +74,6 @@ pub struct PgConnectionInner {
     cache_type_oid: HashMap<UStr, Oid>,
     cache_elem_type_to_array: HashMap<Oid, Oid>,
 
-    // number of ReadyForQuery messages that we are currently expecting
-    pub(crate) pending_ready_for_query_count: usize,
-
     // current transaction status
     transaction_status: TransactionStatus,
     pub(crate) transaction_depth: usize,
@@ -101,7 +98,7 @@ impl PgConnection {
     ///
     /// Used for rolling back transactions and releasing advisory locks.
     #[inline(always)]
-    pub(crate) fn queue_simple_query(&mut self, query: &str) -> Result<Pipe, Error> {
+    pub(crate) fn queue_simple_query(&self, query: &str) -> Result<Pipe, Error> {
         self.pipe(|buf| buf.write_msg(Query(query)))
     }
 
@@ -132,7 +129,6 @@ impl PgConnection {
                 cache_elem_type_to_array: HashMap::new(),
                 transaction_depth: 0,
                 stream,
-                pending_ready_for_query_count: 0,
                 transaction_status: TransactionStatus::Idle,
             }),
         }

sqlx-postgres/src/connection/stream.rs

@@ -2,16 +2,14 @@ use std::collections::BTreeMap;
 use std::ops::{ControlFlow, Deref, DerefMut};
 use std::str::FromStr;
 
-use futures_channel::mpsc::UnboundedSender;
-use futures_util::SinkExt;
 use log::Level;
 use sqlx_core::bytes::Buf;
 
 use crate::connection::tls::MaybeUpgradeTls;
 use crate::error::Error;
 use crate::message::{
-    BackendMessage, BackendMessageFormat, EncodeMessage, FrontendMessage, Notice, Notification,
-    ParameterStatus, ReceivedMessage,
+    BackendMessage, BackendMessageFormat, EncodeMessage, FrontendMessage, Notice, ParameterStatus,
+    ReceivedMessage,
 };
 use crate::net::{self, BufferedSocket, Socket};
 use crate::{PgConnectOptions, PgDatabaseError, PgSeverity};
@@ -30,11 +28,6 @@ pub struct PgStream {
     // function call as well as the syscall.
     inner: BufferedSocket<Box<dyn Socket>>,
 
-    // buffer of unreceived notification messages from `PUBLISH`
-    // this is set when creating a PgListener and only written to if that listener is
-    // re-used for query execution in-between receiving messages
-    pub(crate) notifications: Option<UnboundedSender<Notification>>,
-
     pub(crate) parameter_statuses: BTreeMap<String, String>,
 
     pub(crate) server_version_num: Option<u32>,
@@ -55,7 +48,6 @@ impl PgStream {
 
         Ok(Self {
             inner: BufferedSocket::new(socket),
-            notifications: None,
             parameter_statuses: BTreeMap::default(),
             server_version_num: None,
         })
@@ -134,15 +126,6 @@ impl PgStream {
                     return Err(message.decode::<PgDatabaseError>()?.into());
                 }
 
-                BackendMessageFormat::NotificationResponse => {
-                    if let Some(buffer) = &mut self.notifications {
-                        let notification: Notification = message.decode()?;
-                        let _ = buffer.send(notification).await;
-
-                        continue;
-                    }
-                }
-
                 BackendMessageFormat::ParameterStatus => {
                     // informs the frontend about the current (initial)
                     // setting of backend parameters

sqlx-postgres/src/listener.rs

@@ -1,20 +1,17 @@
 use std::fmt::{self, Debug};
-use std::io;
 use std::str::from_utf8;
 
-use futures_channel::mpsc;
 use futures_core::future::BoxFuture;
 use futures_core::stream::{BoxStream, Stream};
 use futures_util::{FutureExt, StreamExt, TryFutureExt, TryStreamExt};
 use sqlx_core::acquire::Acquire;
 use sqlx_core::transaction::Transaction;
 use sqlx_core::Either;
-use tracing::Instrument;
 
 use crate::describe::Describe;
 use crate::error::Error;
 use crate::executor::{Execute, Executor};
-use crate::message::{BackendMessageFormat, Notification};
+use crate::message::Notification;
 use crate::pool::PoolOptions;
 use crate::pool::{Pool, PoolConnection};
 use crate::{PgConnection, PgQueryResult, PgRow, PgStatement, PgTypeInfo, Postgres};
@@ -28,8 +25,6 @@ use crate::{PgConnection, PgQueryResult, PgRow, PgStatement, PgTypeInfo, Postgre
 pub struct PgListener {
     pool: Pool<Postgres>,
     connection: Option<PoolConnection<Postgres>>,
-    buffer_rx: mpsc::UnboundedReceiver<Notification>,
-    buffer_tx: Option<mpsc::UnboundedSender<Notification>>,
     channels: Vec<String>,
     ignore_close_event: bool,
     eager_reconnect: bool,
@@ -58,17 +53,11 @@ impl PgListener {
 
     pub async fn connect_with(pool: &Pool<Postgres>) -> Result<Self, Error> {
         // Pull out an initial connection
-        let mut connection = pool.acquire().await?;
-
-        // Setup a notification buffer
-        let (sender, receiver) = mpsc::unbounded();
-        connection.inner.stream.notifications = Some(sender);
+        let connection = pool.acquire().await?;
 
         Ok(Self {
             pool: pool.clone(),
             connection: Some(connection),
-            buffer_rx: receiver,
-            buffer_tx: None,
             channels: Vec::new(),
             ignore_close_event: false,
             eager_reconnect: true,
@@ -173,7 +162,6 @@ impl PgListener {
     async fn connect_if_needed(&mut self) -> Result<(), Error> {
         if self.connection.is_none() {
             let mut connection = self.pool.acquire().await?;
-            connection.inner.stream.notifications = self.buffer_tx.take();
 
             connection
                 .execute(&*build_listen_all_query(&self.channels))
@@ -263,67 +251,37 @@ impl PgListener {
         // Fetch our `CloseEvent` listener, if applicable.
         let mut close_event = (!self.ignore_close_event).then(|| self.pool.close_event());
 
-        loop {
-            let next_message = self.connection().await?.inner.stream.recv_unchecked();
-
-            let res = if let Some(ref mut close_event) = close_event {
-                // cancels the wait and returns `Err(PoolClosed)` if the pool is closed
-                // before `next_message` returns, or if the pool was already closed
-                close_event.do_until(next_message).await?
-            } else {
-                next_message.await
-            };
-
-            let message = match res {
-                Ok(message) => message,
-
-                // The connection is dead, ensure that it is dropped,
-                // update self state, and loop to try again.
-                Err(Error::Io(err))
-                    if matches!(
-                        err.kind(),
-                        io::ErrorKind::ConnectionAborted |
-                        io::ErrorKind::UnexpectedEof |
-                        // see ERRORS section in tcp(7) man page (https://man7.org/linux/man-pages/man7/tcp.7.html)
-                        io::ErrorKind::TimedOut |
-                        io::ErrorKind::BrokenPipe
-                    ) =>
-                {
-                    if let Some(mut conn) = self.connection.take() {
-                        self.buffer_tx = conn.inner.stream.notifications.take();
-                        // Close the connection in a background task, so we can continue.
-                        conn.close_on_drop();
-                    }
-
-                    if self.eager_reconnect {
-                        self.connect_if_needed().await?;
-                    }
-
-                    // lost connection
-                    return Ok(None);
-                }
-
-                // Forward other errors
-                Err(error) => {
-                    return Err(error);
-                }
-            };
+        let next_message = self.connection().await?.inner.notifications.next();
 
-            match message.format {
-                // We've received an async notification, return it.
-                BackendMessageFormat::NotificationResponse => {
-                    return Ok(Some(PgNotification(message.decode()?)));
+        let res = if let Some(ref mut close_event) = close_event {
+            // cancels the wait and returns `Err(PoolClosed)` if the pool is closed
+            // before `next_message` returns, or if the pool was already closed
+            close_event.do_until(next_message).await?
+        } else {
+            next_message.await
+        };
+
+        let message = match res {
+            Some(message) => message,
+
+            // The connection is dead, ensure that it is dropped,
+            // update self state, and loop to try again.
+            None => {
+                if let Some(mut conn) = self.connection.take() {
+                    // Close the connection in a background task, so we can continue.
+                    conn.close_on_drop();
                 }
 
-                // Mark the connection as ready for another query
-                BackendMessageFormat::ReadyForQuery => {
-                    self.connection().await?.inner.pending_ready_for_query_count -= 1;
+                if self.eager_reconnect {
+                    self.connect_if_needed().await?;
                 }
 
-                // Ignore unexpected messages
-                _ => {}
+                // lost connection
+                return Ok(None);
             }
-        }
+        };
+
+        Ok(Some(PgNotification(message)))
     }
 
     /// Receives the next notification that already exists in the connection buffer, if any.
@@ -332,7 +290,7 @@ impl PgListener {
     ///
     /// This is helpful if you want to retrieve all buffered notifications and process them in batches.
     pub fn next_buffered(&mut self) -> Option<PgNotification> {
-        if let Ok(Some(notification)) = self.buffer_rx.try_next() {
+        if let Ok(Some(notification)) = self.connection.as_mut()?.inner.notifications.try_next() {
             Some(PgNotification(notification))
         } else {
             None
@@ -356,18 +314,8 @@ impl PgListener {
 
 impl Drop for PgListener {
     fn drop(&mut self) {
-        if let Some(mut conn) = self.connection.take() {
-            let fut = async move {
-                let _ = conn.execute("UNLISTEN *").await;
-
-                // inline the drop handler from `PoolConnection` so it doesn't try to spawn another task
-                // otherwise, it may trigger a panic if this task is dropped because the runtime is going away:
-                // https://github.com/launchbadge/sqlx/issues/1389
-                conn.return_to_pool().await;
-            };
-
-            // Unregister any listeners before returning the connection to the pool.
-            crate::rt::spawn(fut.in_current_span());
+        if let Some(conn) = self.connection.take() {
+            let _ = conn.queue_simple_query("UNLISTEN *");
         }
     }
 }

tests/postgres/postgres.rs

@@ -1126,9 +1126,6 @@ async fn test_listener_try_recv_buffered() -> anyhow::Result<()> {
         txn.commit().await?;
     }
 
-    // Still no notifications buffered, since we haven't awaited the listener yet.
-    assert!(listener.next_buffered().is_none());
-
     // Activate connection.
     sqlx::query!("SELECT 1 AS one")
         .fetch_all(&mut listener)