From db5fbe787de65e10af4b38cfe9ce2166b25e6d71 Mon Sep 17 00:00:00 2001 From: 4thel00z <22024133+4thel00z@users.noreply.github.com> Date: Wed, 29 Jul 2026 21:29:57 +0200 Subject: [PATCH 1/2] fix: scope pub/sub channels to their pool The channel registry is a process-global map keyed by channel name alone, shared by every listener via `CHANNELS.clone()`. Each listener owns a connection to one database, and PostgreSQL scopes NOTIFY to a database, so that key is missing a dimension. The visible consequence is in `listen()`: when the name is already present it returns a subscriber on the existing channel and never sends `LISTEN` to its own pool. So the first pool to register a name owns the only real subscription; later subscribers from other pools attach to it and receive that pool's notifications, while notifications raised in their own database reach nobody. Key the registry on the pool instead: database, user and shard, as configured in pgdog. Two databases using the same channel name now get their own channel and their own `LISTEN`, and so do two users of the same database -- they have separate pools and separate backend connections, so sharing one entry left whichever of them lost the race silently dependent on the other. The identity deliberately comes from `Shard`'s identifier and number rather than the server address, so it survives a replica being promoted underneath the listener: `init_pub_sub` rebuilds against the new primary, which has a different host but the same pgdog-side identity. Nesting the map (pool -> channel -> state) rather than flattening the key into one struct keeps the scoping structural. The restart path had the mirror of the same bug -- it re-subscribed to every key in the shared map, issuing `LISTEN` for channels owned by other pools, which it can never receive and which channel teardown does not reach, since the `UNLISTEN` only runs against the connection the notification arrived on. It now iterates one pool's channels by construction rather than filtering, and the receive path no longer builds a key per notification. `SHOW LISTENERS` and the `pub_sub_*` metrics gain database, user and shard: the channel name is no longer a unique row or series identity. --- pgdog/src/admin/show_listeners.rs | 23 ++- pgdog/src/backend/pool/shard/mod.rs | 4 +- pgdog/src/backend/pub_sub/listener.rs | 202 ++++++++++++++++++++++++-- pgdog/src/stats/listeners.rs | 9 +- 4 files changed, 216 insertions(+), 22 deletions(-) diff --git a/pgdog/src/admin/show_listeners.rs b/pgdog/src/admin/show_listeners.rs index e19f10bf5..26819c3e1 100644 --- a/pgdog/src/admin/show_listeners.rs +++ b/pgdog/src/admin/show_listeners.rs @@ -22,6 +22,9 @@ impl Command for ShowListeners { let mut messages = vec![ RowDescription::new(&[ + Field::text("database"), + Field::text("user"), + Field::numeric("shard"), Field::text("channel"), Field::numeric("listeners"), Field::numeric("received"), @@ -30,10 +33,13 @@ impl Command for ShowListeners { .message()?, ]; - for (channel, stats) in channels { + for (key, stats) in channels { let mut data_row = DataRow::new(); data_row - .add(channel.as_str()) + .add(key.pool.database.as_str()) + .add(key.pool.user.as_str()) + .add(key.pool.shard as i64) + .add(key.channel.as_str()) .add(stats.listeners as i64) .add(stats.recv as i64) .add(stats.dropped as i64); @@ -67,6 +73,17 @@ mod tests { .map(|field| field.name.as_str()) .collect(); - assert_eq!(columns, ["channel", "listeners", "received", "dropped"]); + assert_eq!( + columns, + [ + "database", + "user", + "shard", + "channel", + "listeners", + "received", + "dropped" + ] + ); } } diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index 81b3db133..0373d0ba5 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -287,7 +287,9 @@ impl Shard { // This is useful if we promoted a primary // from a replica. let primary = self.lb.primary().cloned(); - let pub_sub = primary.as_ref().map(PubSubListener::new); + let pub_sub = primary + .as_ref() + .map(|primary| PubSubListener::new(primary, self.identifier(), self.number())); // Launch the new listener first! if let Some(ref pub_sub) = pub_sub { diff --git a/pgdog/src/backend/pub_sub/listener.rs b/pgdog/src/backend/pub_sub/listener.rs index c157d786b..98554d45a 100644 --- a/pgdog/src/backend/pub_sub/listener.rs +++ b/pgdog/src/backend/pub_sub/listener.rs @@ -22,7 +22,7 @@ use tracing::{debug, error, info}; use super::{Stats, StatsSnapshot, channel_size}; use crate::{ - backend::{self, ConnectReason, DisconnectReason, Pool, pool::Error}, + backend::{self, ConnectReason, DisconnectReason, Pool, databases::User, pool::Error}, config::config, net::{ FromBytes, FrontendPid, NotificationResponse, Parameter, Parameters, Protocol, @@ -50,16 +50,59 @@ impl From for ProtocolMessage { } } -type Channels = Arc>>; +/// Pool a set of channels belongs to. `NOTIFY` is scoped to a database, so +/// channels have to be scoped the same way. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct PoolKey { + /// Database name, as configured in pgdog. + pub database: String, + /// User, as configured in pgdog. + pub user: String, + /// Shard number. + pub shard: usize, +} + +impl PoolKey { + fn new(identifier: &User, shard: usize) -> Self { + Self { + database: identifier.database.clone(), + user: identifier.user.clone(), + shard, + } + } + + /// Key for one of this pool's channels. + fn channel(&self, channel: &str) -> ChannelKey { + ChannelKey { + pool: self.clone(), + channel: channel.to_owned(), + } + } +} + +/// One channel on one pool. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct ChannelKey { + /// Pool the channel belongs to. + pub pool: PoolKey, + /// Channel name used by `LISTEN`/`NOTIFY`. + pub channel: String, +} + +type Channels = Arc>>>; static CHANNELS: Lazy = Lazy::new(|| Arc::new(Mutex::new(HashMap::new()))); /// Get stats for all channels. -pub fn stats() -> HashMap { +pub fn stats() -> HashMap { CHANNELS .lock() .iter() - .map(|(name, channel)| (name.to_string(), channel.stats.get())) + .flat_map(|(pool, channels)| { + channels + .iter() + .map(|(channel, state)| (pool.channel(channel), state.stats.get())) + }) .collect() } @@ -160,6 +203,7 @@ struct Comms { pub struct PubSubListener { id: FrontendPid, pool: Pool, + pool_key: PoolKey, tx: mpsc::Sender, channels: Channels, comms: Arc, @@ -167,7 +211,11 @@ pub struct PubSubListener { impl PubSubListener { /// Create new listener on the server connection. - pub fn new(pool: &Pool) -> Self { + /// + /// `identifier` and `shard` scope the channels this listener owns. They are + /// the pgdog-side identity of the pool, so they survive a primary being + /// promoted underneath us, which the server address would not. + pub fn new(pool: &Pool, identifier: &User, shard: usize) -> Self { let (tx, mut rx) = mpsc::channel(channel_size()); let pool = pool.clone(); @@ -176,6 +224,7 @@ impl PubSubListener { let listener = Self { id: FrontendPid::new(), pool: pool.clone(), + pool_key: PoolKey::new(identifier, shard), tx, channels, comms: Arc::new(Comms { @@ -186,6 +235,7 @@ impl PubSubListener { let id = listener.id; let channels = listener.channels.clone(); + let pool_key = listener.pool_key.clone(); let pool = listener.pool.clone(); let comms = listener.comms.clone(); tasks::spawn("pub sub", async move { @@ -206,7 +256,7 @@ impl PubSubListener { rx.close(); // Drain remaining messages. } - result = Self::run(id, &pool, &mut rx, channels.clone()) => { + result = Self::run(id, &pool, &pool_key, &mut rx, channels.clone()) => { if let Err(err) = result { error!("pub/sub error: {} [{}]", err, pool.addr()); // Don't reconnect for another connect attempt delay @@ -242,8 +292,9 @@ impl PubSubListener { pub async fn listen(&self, channel_name: &str) -> Result { let listener = { let mut guard = self.channels.lock(); + let channels = guard.entry(self.pool_key.clone()).or_default(); - if let Some(channel) = guard.get(channel_name) { + if let Some(channel) = channels.get(channel_name) { return Ok(Listener::new(channel)); } @@ -256,7 +307,7 @@ impl PubSubListener { }; let listener = Listener::new(&channel); - guard.insert(channel_name.to_string(), channel); + channels.insert(channel_name.to_string(), channel); listener }; @@ -284,6 +335,7 @@ impl PubSubListener { async fn run( id: FrontendPid, pool: &Pool, + pool_key: &PoolKey, rx: &mut mpsc::Receiver, channels: Channels, ) -> Result<(), backend::Error> { @@ -302,13 +354,18 @@ impl PubSubListener { ) .await?; - // Re-listen on all channels when re-starting the task. + // Re-listen on this pool's channels when re-starting the task. // We don't lose LISTEN commands. let resub = channels .lock() - .keys() - .map(|channel| Request::Subscribe(channel.to_string()).into()) - .collect::>(); + .get(pool_key) + .map(|channels| { + channels + .keys() + .map(|channel| Request::Subscribe(channel.clone()).into()) + .collect::>() + }) + .unwrap_or_default(); if !resub.is_empty() { server.send(&resub.into()).await?; @@ -323,7 +380,11 @@ impl PubSubListener { if message.code() == 'A' { let notification = NotificationResponse::from_bytes(message.to_bytes())?; let mut unsub = None; - if let Some(channel) = channels.lock().get(notification.channel()) { + if let Some(channel) = channels + .lock() + .get(pool_key) + .and_then(|channels| channels.get(notification.channel())) + { match channel.tx.send(notification) { Ok(_) => (), Err(err) => unsub = Some(err.0.channel().to_string()), @@ -331,7 +392,9 @@ impl PubSubListener { } if let Some(unsub) = unsub { - channels.lock().remove(&unsub); + if let Some(channels) = channels.lock().get_mut(pool_key) { + channels.remove(&unsub); + } server.send(&vec![Request::Unsubscribe(unsub).into()].into()).await?; } } @@ -367,15 +430,37 @@ mod test { use super::{test_support::TestChannel, *}; + fn test_user(user: &str, database: &str) -> User { + User { + user: user.into(), + database: database.into(), + } + } + fn test_pub_sub_listener() -> (PubSubListener, mpsc::Receiver) { + test_pub_sub_listener_on( + Arc::new(Mutex::new(HashMap::new())), + &test_user("pgdog", "pgdog"), + 0, + ) + } + + /// A listener for `identifier`/`shard`, sharing `channels` with any other + /// listener built from the same map. + fn test_pub_sub_listener_on( + channels: Channels, + identifier: &User, + shard: usize, + ) -> (PubSubListener, mpsc::Receiver) { let (tx, rx) = mpsc::channel(4); ( PubSubListener { id: FrontendPid::new(), pool: Pool::new_test(), + pool_key: PoolKey::new(identifier, shard), tx, - channels: Arc::new(Mutex::new(HashMap::new())), + channels, comms: Arc::new(Comms { start: Notify::new(), shutdown: CancellationToken::new(), @@ -385,6 +470,16 @@ mod test { ) } + /// Assert a Subscribe request is already queued. Deliberately non-blocking: + /// the sender is alive, so an `await` here would hang rather than fail if a + /// regression stopped the request being sent. + fn expect_subscribe_now(rx: &mut mpsc::Receiver, expected: &str) { + match rx.try_recv() { + Ok(Request::Subscribe(channel)) => assert_eq!(channel, expected), + other => panic!("expected subscribe request for {expected}, got {other:?}"), + } + } + fn assert_snapshot(snapshot: StatsSnapshot, recv: u64, dropped: u64, listeners: u64) { assert_eq!(snapshot.recv, recv); assert_eq!(snapshot.dropped, dropped); @@ -479,13 +574,88 @@ mod test { let stats = pub_sub .channels .lock() - .get("events") + .get(&pub_sub.pool_key) + .and_then(|channels| channels.get("events")) .expect("events channel") .stats .get(); assert_snapshot(stats, 0, 0, 0); } + /// Two listeners sharing the registry must not share a channel just because + /// they were handed the same name. Each has to be sent its own LISTEN, + /// otherwise the one that misses out never receives its own database's + /// notifications. + async fn assert_channels_not_shared( + left: &User, + left_shard: usize, + right: &User, + right_shard: usize, + ) { + let channels: Channels = Arc::new(Mutex::new(HashMap::new())); + let (first, mut first_rx) = test_pub_sub_listener_on(channels.clone(), left, left_shard); + let (second, mut second_rx) = + test_pub_sub_listener_on(channels.clone(), right, right_shard); + + let _first = first.listen("events").await.expect("first listen"); + let _second = second.listen("events").await.expect("second listen"); + + expect_subscribe_now(&mut first_rx, "events"); + expect_subscribe_now(&mut second_rx, "events"); + + let guard = channels.lock(); + assert_eq!(guard.len(), 2, "each pool needs its own channel set"); + for pool_key in [&first.pool_key, &second.pool_key] { + let channel = guard + .get(pool_key) + .and_then(|channels| channels.get("events")) + .expect("channel for pool"); + assert_snapshot(channel.stats.get(), 0, 0, 1); + } + } + + #[tokio::test] + async fn channels_are_not_shared_between_databases() { + assert_channels_not_shared( + &test_user("pgdog", "first_database"), + 0, + &test_user("pgdog", "second_database"), + 0, + ) + .await; + } + + #[tokio::test] + async fn channels_are_not_shared_between_users() { + assert_channels_not_shared( + &test_user("alice", "pgdog"), + 0, + &test_user("bob", "pgdog"), + 0, + ) + .await; + } + + #[tokio::test] + async fn channels_are_not_shared_between_shards() { + let user = test_user("pgdog", "pgdog"); + assert_channels_not_shared(&user, 0, &user, 1).await; + } + + #[test] + fn pool_key_is_built_from_the_pgdog_side_identity() { + let key = PoolKey::new(&test_user("alice", "shop"), 2); + + assert_eq!(key.database, "shop"); + assert_eq!(key.user, "alice"); + assert_eq!(key.shard, 2); + + // Every component discriminates. + assert_ne!(key, PoolKey::new(&test_user("alice", "other"), 2)); + assert_ne!(key, PoolKey::new(&test_user("bob", "shop"), 2)); + assert_ne!(key, PoolKey::new(&test_user("alice", "shop"), 3)); + } + #[tokio::test] async fn notify_queues_notify_request() { let (pub_sub, mut rx) = test_pub_sub_listener(); diff --git a/pgdog/src/stats/listeners.rs b/pgdog/src/stats/listeners.rs index 5bc619e54..89ea5a1b0 100644 --- a/pgdog/src/stats/listeners.rs +++ b/pgdog/src/stats/listeners.rs @@ -13,8 +13,13 @@ impl Listeners { let mut received = vec![]; let mut dropped = vec![]; - for (channel, stats) in stats { - let labels = vec![("channel".into(), channel)]; + for (key, stats) in stats { + let labels = vec![ + ("database".into(), key.pool.database), + ("user".into(), key.pool.user), + ("shard".into(), key.pool.shard.to_string()), + ("channel".into(), key.channel), + ]; listeners.push(Measurement { labels: labels.clone(), From fdac36ff25b425c19d78fff90341713675a62d47 Mon Sep 17 00:00:00 2001 From: 4thel00z <4thel00z@gmail.com> Date: Fri, 31 Jul 2026 11:36:09 +0200 Subject: [PATCH 2/2] fix: drop empty pool entries, make registry keys crate-private MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups: channel removal now drops the pool's registry entry once its channel map is empty, so churned pools don't accumulate empty maps. stats(), PoolKey and ChannelKey are pub(crate) — their only consumers are SHOW LISTENERS and the metrics endpoint. --- pgdog/src/backend/pub_sub/listener.rs | 63 ++++++++++++++++++++++----- 1 file changed, 52 insertions(+), 11 deletions(-) diff --git a/pgdog/src/backend/pub_sub/listener.rs b/pgdog/src/backend/pub_sub/listener.rs index 98554d45a..3c3369055 100644 --- a/pgdog/src/backend/pub_sub/listener.rs +++ b/pgdog/src/backend/pub_sub/listener.rs @@ -53,13 +53,13 @@ impl From for ProtocolMessage { /// Pool a set of channels belongs to. `NOTIFY` is scoped to a database, so /// channels have to be scoped the same way. #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -pub struct PoolKey { +pub(crate) struct PoolKey { /// Database name, as configured in pgdog. - pub database: String, + pub(crate) database: String, /// User, as configured in pgdog. - pub user: String, + pub(crate) user: String, /// Shard number. - pub shard: usize, + pub(crate) shard: usize, } impl PoolKey { @@ -82,11 +82,11 @@ impl PoolKey { /// One channel on one pool. #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -pub struct ChannelKey { +pub(crate) struct ChannelKey { /// Pool the channel belongs to. - pub pool: PoolKey, + pub(crate) pool: PoolKey, /// Channel name used by `LISTEN`/`NOTIFY`. - pub channel: String, + pub(crate) channel: String, } type Channels = Arc>>>; @@ -94,7 +94,7 @@ type Channels = Arc>>>; static CHANNELS: Lazy = Lazy::new(|| Arc::new(Mutex::new(HashMap::new()))); /// Get stats for all channels. -pub fn stats() -> HashMap { +pub(crate) fn stats() -> HashMap { CHANNELS .lock() .iter() @@ -106,6 +106,18 @@ pub fn stats() -> HashMap { .collect() } +/// Remove a channel from its pool's set, dropping the pool's entry once no +/// channels remain, so pools that churn don't accumulate empty maps. +fn remove_channel(channels: &Channels, pool_key: &PoolKey, channel: &str) { + let mut guard = channels.lock(); + if let Some(pool_channels) = guard.get_mut(pool_key) { + pool_channels.remove(channel); + if pool_channels.is_empty() { + guard.remove(pool_key); + } + } +} + #[derive(Debug)] struct Channel { tx: broadcast::Sender, @@ -392,9 +404,7 @@ impl PubSubListener { } if let Some(unsub) = unsub { - if let Some(channels) = channels.lock().get_mut(pool_key) { - channels.remove(&unsub); - } + remove_channel(&channels, pool_key, &unsub); server.send(&vec![Request::Unsubscribe(unsub).into()].into()).await?; } } @@ -656,6 +666,37 @@ mod test { assert_ne!(key, PoolKey::new(&test_user("alice", "shop"), 3)); } + #[tokio::test] + async fn removing_the_last_channel_drops_the_pool_entry() { + let (pub_sub, mut rx) = test_pub_sub_listener(); + + let _listener = pub_sub.listen("events").await.expect("listen"); + expect_subscribe(&mut rx, "events").await; + + remove_channel(&pub_sub.channels, &pub_sub.pool_key, "events"); + assert!( + pub_sub.channels.lock().is_empty(), + "empty pool entries must not accumulate" + ); + } + + #[tokio::test] + async fn removing_one_channel_keeps_the_pool_entry_for_the_rest() { + let (pub_sub, mut rx) = test_pub_sub_listener(); + + let _events = pub_sub.listen("events").await.expect("listen events"); + let _jobs = pub_sub.listen("jobs").await.expect("listen jobs"); + expect_subscribe(&mut rx, "events").await; + expect_subscribe(&mut rx, "jobs").await; + + remove_channel(&pub_sub.channels, &pub_sub.pool_key, "events"); + + let guard = pub_sub.channels.lock(); + let channels = guard.get(&pub_sub.pool_key).expect("pool entry remains"); + assert!(channels.contains_key("jobs")); + assert!(!channels.contains_key("events")); + } + #[tokio::test] async fn notify_queues_notify_request() { let (pub_sub, mut rx) = test_pub_sub_listener();