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..3c3369055 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,19 +50,74 @@ 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(crate) struct PoolKey { + /// Database name, as configured in pgdog. + pub(crate) database: String, + /// User, as configured in pgdog. + pub(crate) user: String, + /// Shard number. + pub(crate) 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(crate) struct ChannelKey { + /// Pool the channel belongs to. + pub(crate) pool: PoolKey, + /// Channel name used by `LISTEN`/`NOTIFY`. + pub(crate) 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(crate) 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() } +/// 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, @@ -160,6 +215,7 @@ struct Comms { pub struct PubSubListener { id: FrontendPid, pool: Pool, + pool_key: PoolKey, tx: mpsc::Sender, channels: Channels, comms: Arc, @@ -167,7 +223,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 +236,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 +247,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 +268,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 +304,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 +319,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 +347,7 @@ impl PubSubListener { async fn run( id: FrontendPid, pool: &Pool, + pool_key: &PoolKey, rx: &mut mpsc::Receiver, channels: Channels, ) -> Result<(), backend::Error> { @@ -302,13 +366,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 +392,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 +404,7 @@ impl PubSubListener { } if let Some(unsub) = unsub { - channels.lock().remove(&unsub); + remove_channel(&channels, pool_key, &unsub); server.send(&vec![Request::Unsubscribe(unsub).into()].into()).await?; } } @@ -367,15 +440,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 +480,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 +584,119 @@ 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 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(); 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(),