fix: keep client connected when LISTEN/NOTIFY run with pub/sub disabled

pgdog/src/backend/error.rs

@@ -144,12 +144,14 @@ impl Error {
 
     /// It's not a supported feature so we should tell the client so.
     pub(crate) fn is_unsupported(&self) -> bool {
+        use crate::backend::pool::Error as PoolError;
         matches!(
             self,
             Self::Aggregate(AggregateError::UnsupportedAggregation { .. })
                 | Self::MultiShard(MultiShardError::Aggregate(
                     AggregateError::UnsupportedAggregation { .. }
                 ))
+                | Self::Pool(PoolError::PubSubDisabled)
         )
     }
 

pgdog/src/frontend/client/query_engine/pub_sub.rs

@@ -1,4 +1,6 @@
 use super::*;
+use crate::backend::Error as BackendError;
+use crate::backend::pool::Error as PoolError;
 use crate::frontend::client::query_engine::fake::FakeResponse;
 use crate::net::ProtocolMessage;
 
@@ -10,7 +12,10 @@ impl QueryEngine {
         channel: &str,
         shard: Shard,
     ) -> Result<(), Error> {
-        self.backend.listen(channel, shard).await?;
+        let result = self.backend.listen(channel, shard).await;
+        if self.pub_sub_error(context, result).await? {
+            return Ok(());
+        }
         self.fake_command_response(context, client_messages, &FakeResponse::command("LISTEN"))
             .await?;
 
@@ -25,13 +30,20 @@ impl QueryEngine {
         payload: &str,
         shard: &Shard,
     ) -> Result<(), Error> {
-        if context.in_transaction() {
+        let result = if !context.in_transaction() {
+            // Send immediately if not in transaction
+            self.backend.notify(channel, payload, shard.clone()).await
+        } else if config().pub_sub_enabled() {
             // Buffer the NOTIFY command if we're in a transaction
             self.notify_buffer
                 .add(channel.to_string(), payload.to_string(), shard.clone());
+            Ok(())
         } else {
-            // Send immediately if not in transaction
-            self.backend.notify(channel, payload, shard.clone()).await?;
+            // Fail now: the buffer is flushed after COMMIT has already succeeded.
+            Err(BackendError::Pool(PoolError::PubSubDisabled))
+        };
+        if self.pub_sub_error(context, result).await? {
+            return Ok(());
         }
         self.fake_command_response(context, client_messages, &FakeResponse::command("NOTIFY"))
             .await?;
@@ -58,4 +70,26 @@ impl QueryEngine {
         }
         Ok(())
     }
+
+    /// Report an unsupported pub/sub request to the client as an `ERROR`,
+    /// keeping the connection open. Returns `true` if an error was sent.
+    async fn pub_sub_error(
+        &mut self,
+        context: &mut QueryEngineContext<'_>,
+        result: Result<(), BackendError>,
+    ) -> Result<bool, Error> {
+        match result {
+            Ok(()) => Ok(false),
+            Err(err) if err.is_unsupported() => {
+                if let Some(transaction) = context.transaction.as_mut() {
+                    transaction.set_error();
+                }
+                let mut error = ErrorResponse::unsupported(&err);
+                error.hint = Some("set pub_sub_channel_size to enable LISTEN/NOTIFY".into());
+                self.send_error_response(context, error).await?;
+                Ok(true)
+            }
+            Err(err) => Err(err.into()),
+        }
+    }
 }

pgdog/src/frontend/client/query_engine/test/pub_sub.rs

@@ -1,12 +1,79 @@
 use std::io::ErrorKind;
+use std::time::Duration;
+
+use tokio::time::timeout;
 
 use crate::{
     expect_message,
-    net::{BindComplete, CommandComplete, Parameters, ParseComplete},
+    net::{BindComplete, CommandComplete, ErrorResponse, Parameters, ParseComplete, ReadyForQuery},
 };
 
 use super::prelude::*;
 
+#[tokio::test]
+async fn pub_sub_disabled_keeps_connection_usable() {
+    let mut client = SpawnedClient::new_sharded(Parameters::default()).await;
+
+    timeout(Duration::from_secs(5), async {
+        for extended in [false, true] {
+            for (begin, query) in [
+                (false, "LISTEN pgdog_disabled"),
+                (false, "NOTIFY pgdog_disabled"),
+                (true, "LISTEN pgdog_disabled"),
+                (true, "NOTIFY pgdog_disabled"),
+            ] {
+                if begin {
+                    client.send(Query::new("BEGIN")).await;
+                    client.read_until('Z').await;
+                }
+
+                if extended {
+                    client.send(Parse::new_anonymous(query)).await;
+                    client.send(Bind::new_statement("")).await;
+                    client.send(Execute::new()).await;
+                    client.send(Sync).await;
+                } else {
+                    client.send(Query::new(query)).await;
+                }
+                let error = expect_message!(client.read().await, ErrorResponse);
+                assert_eq!(error.severity, "ERROR", "{query}");
+                assert_eq!(error.code, "0A000", "{query}");
+                assert_eq!(error.message, "pub/sub disabled", "{query}");
+                assert!(
+                    error
+                        .hint
+                        .as_deref()
+                        .is_some_and(|hint| hint.contains("pub_sub_channel_size")),
+                    "{query}"
+                );
+                assert_eq!(
+                    expect_message!(client.read().await, ReadyForQuery).status,
+                    if begin { 'E' } else { 'I' },
+                    "{query}"
+                );
+
+                if begin {
+                    client.send(Query::new("ROLLBACK")).await;
+                    assert_eq!(
+                        expect_message!(client.read().await, CommandComplete).command(),
+                        "ROLLBACK"
+                    );
+                    expect_message!(client.read().await, ReadyForQuery);
+                }
+            }
+        }
+
+        client.send(Query::new("SELECT 1")).await;
+        let messages = client.read_until('Z').await;
+        assert!(
+            messages.iter().any(|message| message.code() == 'D'),
+            "connection still serves queries"
+        );
+    })
+    .await
+    .expect("pub/sub disabled error or subsequent query hung");
+}
+
 #[tokio::test]
 async fn extended_unlisten_flush_does_not_emit_ready_for_query() {
     let mut client = TestClient::new_sharded(Parameters::default()).await;