From 23288315ced5c96a9193aaf2089a4b0fdea2760f Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Thu, 22 Oct 2026 20:48:39 +0530 Subject: [PATCH 1/2] fix: keep client connected when LISTEN/NOTIFY run with pub/sub disabled --- pgdog/src/backend/error.rs | 2 + .../frontend/client/query_engine/pub_sub.rs | 42 +++++++++++-- .../client/query_engine/test/pub_sub.rs | 60 ++++++++++++++++++- 3 files changed, 99 insertions(+), 5 deletions(-) diff --git a/pgdog/src/backend/error.rs b/pgdog/src/backend/error.rs index 5696f3051..a05248e90 100644 --- a/pgdog/src/backend/error.rs +++ b/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) ) } diff --git a/pgdog/src/frontend/client/query_engine/pub_sub.rs b/pgdog/src/frontend/client/query_engine/pub_sub.rs index 8cb149160..933e7d28e 100644 --- a/pgdog/src/frontend/client/query_engine/pub_sub.rs +++ b/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 { + 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()), + } + } } diff --git a/pgdog/src/frontend/client/query_engine/test/pub_sub.rs b/pgdog/src/frontend/client/query_engine/test/pub_sub.rs index 7dc72d706..7d7d49fe5 100644 --- a/pgdog/src/frontend/client/query_engine/test/pub_sub.rs +++ b/pgdog/src/frontend/client/query_engine/test/pub_sub.rs @@ -1,12 +1,70 @@ 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 (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; + } + + 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; From 7f8c9185263664a1399817e3f7dec04bb03c51bd Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Wed, 28 Oct 2026 11:23:20 +0530 Subject: [PATCH 2/2] extended proto test --- .../client/query_engine/test/pub_sub.rs | 75 +++++++++++-------- 1 file changed, 42 insertions(+), 33 deletions(-) diff --git a/pgdog/src/frontend/client/query_engine/test/pub_sub.rs b/pgdog/src/frontend/client/query_engine/test/pub_sub.rs index 7d7d49fe5..08ffa3f46 100644 --- a/pgdog/src/frontend/client/query_engine/test/pub_sub.rs +++ b/pgdog/src/frontend/client/query_engine/test/pub_sub.rs @@ -15,42 +15,51 @@ async fn pub_sub_disabled_keeps_connection_usable() { let mut client = SpawnedClient::new_sharded(Parameters::default()).await; timeout(Duration::from_secs(5), async { - 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; - } - - 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}" - ); + 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 begin { - client.send(Query::new("ROLLBACK")).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, CommandComplete).command(), - "ROLLBACK" + expect_message!(client.read().await, ReadyForQuery).status, + if begin { 'E' } else { 'I' }, + "{query}" ); - expect_message!(client.read().await, ReadyForQuery); + + if begin { + client.send(Query::new("ROLLBACK")).await; + assert_eq!( + expect_message!(client.read().await, CommandComplete).command(), + "ROLLBACK" + ); + expect_message!(client.read().await, ReadyForQuery); + } } }