From a27abe638fa4f465319f09d1dbf42144abe4aa53 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Oko=C5=84ski?= Date: Fri, 21 Aug 2026 09:41:04 +0200 Subject: [PATCH 1/2] fix(pool): fail closed automatic-primary evidence --- pgdog/src/backend/pool/lb/mod.rs | 89 ++++--- pgdog/src/backend/pool/lb/test.rs | 90 ++++++- pgdog/src/backend/pool/lsn_monitor.rs | 236 ++++++++++++++++-- pgdog/src/backend/pool/pool_impl.rs | 35 +++ pgdog/src/backend/pool/shard/monitor.rs | 54 ++-- pgdog/src/backend/pool/shard/role_detector.rs | 3 +- pgdog/src/backend/server.rs | 78 ++++++ 7 files changed, 498 insertions(+), 87 deletions(-) diff --git a/pgdog/src/backend/pool/lb/mod.rs b/pgdog/src/backend/pool/lb/mod.rs index a8fb839bb..9f1339272 100644 --- a/pgdog/src/backend/pool/lb/mod.rs +++ b/pgdog/src/backend/pool/lb/mod.rs @@ -8,6 +8,7 @@ use std::{ time::{Duration, SystemTime}, }; +use parking_lot::Mutex; use rand::seq::SliceRandom; use tokio::sync::Notify; use tracing::warn; @@ -79,6 +80,18 @@ impl Target { pub(super) fn health(&self) -> &TargetHealth { &self.pool.inner().health } + + fn is_qualified_primary(&self) -> bool { + if self.role() != Role::Primary { + return false; + } + if self.pool.addr().configured_role != Role::Auto { + return true; + } + + let stats = self.pool.lsn_stats(); + stats.valid() && !stats.replica + } } /// Load balancer. @@ -96,6 +109,8 @@ pub struct LoadBalancer { pub(super) maintenance: Arc, /// Role detection waiter. pub(super) role_detection: Arc, + /// Automatic-role election lock. + election: Arc>, /// Read/write split. pub(super) rw_split: ReadWriteSplit, } @@ -147,6 +162,7 @@ impl LoadBalancer { lb_strategy, maintenance: Arc::new(Notify::new()), role_detection: Arc::new(Notify::new()), + election: Arc::new(Mutex::new(())), rw_split, } } @@ -170,12 +186,16 @@ impl LoadBalancer { /// Detect database roles from pg_is_in_recovery() and /// return new primary (if any), and replicas. pub fn redetect_roles(&self) -> bool { - let mut promoted = false; + let _election = self.election.lock(); + let previous_primary = self + .primary_target() + .filter(|target| target.pool.addr().configured_role == Role::Auto) + .map(|target| target.pool.id()); let mut targets = self .targets - .clone() - .into_iter() + .iter() + .filter(|target| target.pool.addr().configured_role == Role::Auto) .map(|target| (target.pool.lsn_stats(), target)) .collect::>(); @@ -192,11 +212,13 @@ impl LoadBalancer { let primary = targets .iter() .position(|target| !target.0.replica && target.0.valid()); + let current_primary = primary.map(|index| targets[index].1.pool.id()); + let primary_changed = previous_primary != current_primary; if let Some(primary) = primary { - promoted = targets[primary].1.set_role(Role::Primary); + targets[primary].1.set_role(Role::Primary); - if promoted { + if primary_changed { warn!("new primary chosen: {}", targets[primary].1.pool.addr()); } @@ -208,18 +230,17 @@ impl LoadBalancer { .for_each(|(_, target)| { target.1.set_role(Role::Replica); }); - } else if targets.iter().all(|target| target.0.valid()) { - // All targets are replicas until we get a primary. + } else { targets.iter().for_each(|target| { target.1.set_role(Role::Replica); }); } - if promoted { - self.role_detection.notify_one(); + if current_primary.is_some() || primary_changed { + self.role_detection.notify_waiters(); } - promoted + primary_changed } /// Launch replica pools and start the monitor. @@ -323,37 +344,39 @@ impl LoadBalancer { result } - /// Block until automatic role detection elects a primary. - /// - /// Static replica-only configurations return immediately. In automatic - /// mode, callers wait until a primary is elected or checkout times out. - async fn wait_primary(&self) -> Result<(), Error> { - if self.primary_target().is_none() && self.role_detection_enabled() { - if safe_timeout(self.checkout_timeout, self.role_detection.notified()) - .await - .is_err() - { - return Err(Error::CheckoutTimeout); - }; - // Chain the wakeup so any other waiter that arrived after us - // also gets released without needing another promotion event. - self.role_detection.notify_one(); + fn qualified_primary_target(&self) -> Option<&Target> { + self.targets + .iter() + .rev() + .find(|target| target.is_qualified_primary()) + } + + async fn wait_primary_target(&self) -> Result<&Target, Error> { + if !self.role_detection_enabled() { + return self.qualified_primary_target().ok_or(Error::NoPrimary); } - Ok(()) + loop { + let notified = self.role_detection.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + + if let Some(target) = self.qualified_primary_target() { + return Ok(target); + } + + notified.await; + } } pub(super) async fn get_primary(&self, request: &Request) -> Result { - self.get_primary_internal(request).await + safe_timeout(self.checkout_timeout, self.get_primary_internal(request)) + .await + .map_err(|_| Error::CheckoutTimeout)? } async fn get_primary_internal(&self, request: &Request) -> Result { - self.wait_primary().await?; - self.primary_target() - .ok_or(Error::NoPrimary)? - .pool - .get(request) - .await + self.wait_primary_target().await?.pool.get(request).await } async fn get_internal(&self, request: &Request) -> Result { diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs index 92fec24de..2464ef48a 100644 --- a/pgdog/src/backend/pool/lb/test.rs +++ b/pgdog/src/backend/pool/lb/test.rs @@ -72,7 +72,7 @@ fn set_lsn_stats(target: &Target, replica: bool, lsn: i64) { ..Default::default() } .into(); - *target.pool.inner().lsn_stats.write() = stats; + target.pool.publish_lsn_stats(stats); } impl LoadBalancer { @@ -1503,7 +1503,10 @@ async fn test_auto_mode_waits_for_primary_election() { Default::default(), ); - assert_eq!(lb.wait_primary().await, Err(Error::CheckoutTimeout)); + assert_eq!( + lb.get_primary(&Request::default()).await.unwrap_err(), + Error::CheckoutTimeout + ); } #[tokio::test] @@ -1520,11 +1523,11 @@ async fn test_auto_mode_primary_election_releases_writes() { tokio::spawn(async move { sleep(Duration::from_millis(10)).await; - election.targets[0].set_role(Role::Primary); - election.role_detection.notify_one(); + set_lsn_stats(&election.targets[0], false, 100); + election.redetect_roles(); }); - assert_eq!(lb.wait_primary().await, Ok(())); + assert!(lb.wait_primary_target().await.is_ok()); assert!(lb.primary().is_some()); } @@ -1539,13 +1542,88 @@ async fn test_static_replica_only_does_not_wait_for_primary() { Default::default(), ); - assert_eq!(lb.wait_primary().await, Ok(())); assert!(matches!( lb.get_primary(&Request::default()).await, Err(Error::NoPrimary) )); } +#[test] +fn test_automatic_primary_requires_valid_writer_evidence() { + let lb = LoadBalancer::new( + &None, + &[create_auto_test_pool_config("127.0.0.1", 5432)], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + Default::default(), + ); + let target = &lb.targets[0]; + + target.set_role(Role::Primary); + assert!(!target.is_qualified_primary()); + + set_lsn_stats(target, false, 100); + assert!(target.is_qualified_primary()); + + target.pool.revoke_automatic_primary_evidence(); + assert!(!target.is_qualified_primary()); +} + +#[test] +fn test_redetect_roles_demotes_primary_after_evidence_revocation() { + let lb = LoadBalancer::new( + &None, + &[create_auto_test_pool_config("127.0.0.1", 5432)], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + Default::default(), + ); + let target = &lb.targets[0]; + + set_lsn_stats(target, false, 100); + assert!(lb.redetect_roles()); + assert_eq!(target.role(), Role::Primary); + + target.pool.revoke_automatic_primary_evidence(); + assert!(lb.redetect_roles()); + assert_eq!(target.role(), Role::Replica); +} + +#[test] +fn test_concurrent_role_detection_keeps_one_qualified_primary() { + let lb = LoadBalancer::new( + &None, + &[ + create_auto_test_pool_config("127.0.0.1", 5432), + create_auto_test_pool_config("localhost", 5432), + ], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + Default::default(), + ); + set_lsn_stats(&lb.targets[0], false, 100); + set_lsn_stats(&lb.targets[1], false, 200); + + std::thread::scope(|scope| { + for _ in 0..8 { + let lb = lb.clone(); + scope.spawn(move || { + for _ in 0..100 { + lb.redetect_roles(); + } + }); + } + }); + + let primaries = lb + .targets + .iter() + .filter(|target| target.is_qualified_primary()) + .collect::>(); + assert_eq!(primaries.len(), 1); + assert_eq!(primaries[0].pool.id(), lb.targets[1].pool.id()); +} + #[tokio::test] async fn test_can_move_conns_to_different_addresses() { let pool_config1 = create_test_pool_config("127.0.0.1", 5432); diff --git a/pgdog/src/backend/pool/lsn_monitor.rs b/pgdog/src/backend/pool/lsn_monitor.rs index 1166ba9f3..564c9531b 100644 --- a/pgdog/src/backend/pool/lsn_monitor.rs +++ b/pgdog/src/backend/pool/lsn_monitor.rs @@ -97,16 +97,18 @@ impl LsnStats { } impl LsnStats { - fn from_row(value: DataRow, aurora: bool) -> Self { - StatsLsnStats { - replica: value.get(0, Format::Text).unwrap_or_default(), - lsn: value.get(1, Format::Text).unwrap_or_default(), - offset_bytes: value.get(2, Format::Text).unwrap_or_default(), - timestamp: value.get(3, Format::Text).unwrap_or_default(), - fetched: SystemTime::now(), - aurora, - } - .into() + fn from_row(value: DataRow, aurora: bool) -> Option { + Some( + StatsLsnStats { + replica: value.get(0, Format::Text)?, + lsn: value.get(1, Format::Text)?, + offset_bytes: value.get(2, Format::Text)?, + timestamp: value.get(3, Format::Text)?, + fetched: SystemTime::now(), + aurora, + } + .into(), + ) } } @@ -126,12 +128,24 @@ impl LsnMonitor { async fn run_query(&self, conn: &mut Server, query: &str) -> Option { match safe_timeout(self.pool.config().lsn_check_timeout, conn.fetch_all(query)).await { - Ok(Ok(rows)) => rows.into_iter().next(), + Ok(Ok(rows)) => match rows.into_iter().next() { + Some(row) => Some(row), + None => { + self.revoke_automatic_primary_evidence(); + error!( + "lsn monitor query returned zero rows [{}]", + self.pool.addr() + ); + None + } + }, Ok(Err(err)) => { + self.revoke_automatic_primary_evidence(); error!("lsn monitor query error: {} [{}]", err, self.pool.addr()); None } Err(_) => { + self.revoke_automatic_primary_evidence(); error!("lsn monitor query timeout [{}]", self.pool.addr()); None } @@ -145,12 +159,21 @@ impl LsnMonitor { ) .await { - Ok(Ok(_)) => { + Ok(Ok(rows)) if !rows.is_empty() => { debug!("aurora detected [{}]", self.pool.addr()); Some(true) } + Ok(Ok(_)) => { + self.revoke_automatic_primary_evidence(); + error!( + "lsn monitor aurora detection returned zero rows [{}]", + self.pool.addr() + ); + None + } Ok(Err(crate::backend::Error::ExecutionError(_))) => Some(false), Ok(Err(err)) => { + self.revoke_automatic_primary_evidence(); error!( "lsn monitor aurora detection error: {} [{}]", err, @@ -159,6 +182,7 @@ impl LsnMonitor { None } Err(_) => { + self.revoke_automatic_primary_evidence(); error!( "lsn monitor aurora detection timeout [{}]", self.pool.addr() @@ -198,8 +222,12 @@ impl LsnMonitor { async fn run_check(&self, mut aurora_detected: Option) -> Result, Error> { let mut conn = match self.get_connection().await { Ok(conn) => conn, - Err(Error::Offline) => return Err(Error::Offline), + Err(Error::Offline) => { + self.revoke_automatic_primary_evidence(); + return Err(Error::Offline); + } Err(err) => { + self.revoke_automatic_primary_evidence(); error!("lsn monitor checkout error: {} [{}]", err, self.pool.addr()); return Err(err); } @@ -218,22 +246,29 @@ impl LsnMonitor { if let Some(row) = self.run_query(&mut conn, query).await { drop(conn); - let stats = LsnStats::from_row(row, aurora); - { - let mut guard = self.pool.inner().lsn_stats.write(); - // Notify that the role changed and the shard monitor - // should immediately resynchronize. - if stats.replica != guard.replica { - self.pool.inner().lsn_role_change.notify_one(); - } - (*guard) = stats; - } - trace!("lsn monitor stats updated [{}]", self.pool.addr()); + self.update_stats(row, aurora); } Ok(aurora_detected) } + fn update_stats(&self, row: DataRow, aurora: bool) { + if let Some(stats) = LsnStats::from_row(row, aurora) { + self.pool.publish_lsn_stats(stats); + trace!("lsn monitor stats updated [{}]", self.pool.addr()); + } else { + self.revoke_automatic_primary_evidence(); + error!( + "lsn monitor returned malformed stats row [{}]", + self.pool.addr() + ); + } + } + + fn revoke_automatic_primary_evidence(&self) { + self.pool.revoke_automatic_primary_evidence(); + } + async fn get_connection(&self) -> Result { match self.pool.get(&Request::default()).await { Ok(conn) => Ok(LsnConnection::Guard(conn)), @@ -273,10 +308,16 @@ impl DerefMut for LsnConnection { #[cfg(test)] mod test { + use std::time::{Duration, SystemTime}; + use super::*; + use crate::{ + backend::pool::{Address, Config, PoolConfig, lb::LoadBalancer}, + config::{LoadBalancingStrategy, ReadWriteSplit, Role}, + }; use pgdog_postgres_types::TimestampTz; - use pgdog_stats::Lsn; + use pgdog_stats::{Lsn, LsnStats as StatsLsnStats}; use tokio::time::timeout; // A launched pool against the local Postgres. The default `lsn_check_delay` @@ -289,6 +330,151 @@ mod test { LsnMonitor { pool } } + fn lsn_row(role: &str) -> DataRow { + let mut row = DataRow::new(); + row.add(role) + .add("0/64") + .add(100_i64) + .add("2026-07-01 13:33:10.000000+00"); + row + } + + fn automatic_primary_monitor() -> (LoadBalancer, LsnMonitor) { + let mut config = PoolConfig { + address: Address::new_test(), + config: Config::default(), + }; + config.address.configured_role = Role::Auto; + config.config.inner.role_detection = true; + config.config.inner.lsn_check_timeout = Duration::from_millis(10); + let lb = LoadBalancer::new( + &None, + &[config], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + Default::default(), + ); + let monitor = LsnMonitor { + pool: lb.targets[0].pool.clone(), + }; + publish_writer_and_elect(&lb, &monitor); + (lb, monitor) + } + + fn publish_writer_and_elect(lb: &LoadBalancer, monitor: &LsnMonitor) { + monitor.pool.publish_lsn_stats( + StatsLsnStats { + replica: false, + lsn: Lsn::from_i64(100), + offset_bytes: 100, + fetched: SystemTime::now(), + ..Default::default() + } + .into(), + ); + assert!(lb.redetect_roles()); + assert_eq!(lb.targets[0].role(), Role::Primary); + } + + fn assert_writer_evidence_revoked(lb: &LoadBalancer, monitor: &LsnMonitor) { + assert!(!monitor.pool.lsn_stats().valid()); + assert!(lb.redetect_roles()); + assert_eq!(lb.targets[0].role(), Role::Replica); + } + + #[tokio::test] + async fn test_malformed_row_revokes_writer_evidence() { + let (lb, monitor) = automatic_primary_monitor(); + let notified = monitor.pool.inner().lsn_role_change.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + + monitor.update_stats(lsn_row("invalid"), false); + + assert!(timeout(Duration::from_millis(10), notified).await.is_ok()); + assert_writer_evidence_revoked(&lb, &monitor); + } + + #[tokio::test] + async fn test_lsn_query_failures_revoke_writer_evidence() { + let (lb, monitor) = automatic_primary_monitor(); + let servers = [ + crate::backend::server::test::automatic_role_error_server().await, + crate::backend::server::test::automatic_role_empty_server().await, + crate::backend::server::test::automatic_role_server(None).await, + ]; + + for mut server in servers { + assert!(monitor.run_query(&mut server, LSN_QUERY).await.is_none()); + assert_writer_evidence_revoked(&lb, &monitor); + publish_writer_and_elect(&lb, &monitor); + } + } + + #[tokio::test] + async fn test_aurora_detection_failures_revoke_writer_evidence() { + let (lb, monitor) = automatic_primary_monitor(); + let servers = [ + crate::backend::server::test::automatic_role_disconnect_server().await, + crate::backend::server::test::automatic_role_empty_server().await, + crate::backend::server::test::automatic_role_server(None).await, + ]; + + for mut server in servers { + assert_eq!(monitor.detect_aurora(&mut server).await, None); + assert_writer_evidence_revoked(&lb, &monitor); + publish_writer_and_elect(&lb, &monitor); + } + } + + #[tokio::test] + async fn test_offline_monitor_checkout_revokes_writer_evidence() { + let (lb, monitor) = automatic_primary_monitor(); + + assert_eq!(monitor.run_check(None).await, Err(Error::Offline)); + assert_writer_evidence_revoked(&lb, &monitor); + } + + #[tokio::test] + async fn test_evidence_revocation_notifies_only_on_valid_to_invalid_transition() { + let (_lb, monitor) = automatic_primary_monitor(); + assert!( + timeout( + Duration::from_millis(10), + monitor.pool.inner().lsn_role_change.notified() + ) + .await + .is_ok() + ); + let first = monitor.pool.inner().lsn_role_change.notified(); + tokio::pin!(first); + first.as_mut().enable(); + + monitor.pool.revoke_automatic_primary_evidence(); + assert!(timeout(Duration::from_millis(10), first).await.is_ok()); + + let second = monitor.pool.inner().lsn_role_change.notified(); + tokio::pin!(second); + second.as_mut().enable(); + monitor.pool.revoke_automatic_primary_evidence(); + assert!(timeout(Duration::from_millis(10), second).await.is_err()); + } + + #[test] + fn test_lsn_stats_from_row_requires_all_fields() { + let stats = LsnStats::from_row(lsn_row("f"), false).unwrap(); + assert!(!stats.replica); + assert_eq!(stats.lsn, Lsn::from_i64(100)); + assert_eq!(stats.offset_bytes, 100); + + assert!(LsnStats::from_row(lsn_row("invalid"), false).is_none()); + assert!(LsnStats::from_row(DataRow::new(), false).is_none()); + + let mut null_role = lsn_row("f"); + null_role.insert(0, "", true); + assert!(LsnStats::from_row(null_role, false).is_none()); + } + #[tokio::test] async fn test_run_check_detects_non_aurora() { let monitor = monitor(); diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index ff2f72234..f7112de8a 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -104,6 +104,41 @@ impl Pool { self.inner.health.healthy() } + pub(super) fn revoke_automatic_primary_evidence(&self) { + if !self.config().role_detection { + return; + } + + let was_valid = { + let mut stats = self.inner.lsn_stats.write(); + let was_valid = stats.valid(); + if was_valid { + *stats = Default::default(); + } + was_valid + }; + + if was_valid { + self.inner.lsn_role_change.notify_one(); + } + } + + pub(super) fn publish_lsn_stats(&self, stats: LsnStats) { + let notify = { + let mut current = self.inner.lsn_stats.write(); + let current_valid = current.valid(); + let stats_valid = stats.valid(); + let notify = current_valid != stats_valid + || (current_valid && stats_valid && current.replica != stats.replica); + *current = stats; + notify + }; + + if notify { + self.inner.lsn_role_change.notify_one(); + } + } + /// Launch the maintenance loop, bringing the pool online. pub fn launch(&self) { let mut guard = self.lock(); diff --git a/pgdog/src/backend/pool/shard/monitor.rs b/pgdog/src/backend/pool/shard/monitor.rs index af3df05e0..da074802c 100644 --- a/pgdog/src/backend/pool/shard/monitor.rs +++ b/pgdog/src/backend/pool/shard/monitor.rs @@ -111,20 +111,19 @@ fn update_replica_lag(pools: &[Pool]) { let primary = pools .iter() .map(|pool| (pool, pool.lsn_stats())) - .find(|(_, stats)| !stats.replica); - - // There is a primary. If not, replica lag cannot be calculated. - if let Some((primary_pool, primary_stats)) = primary { - for replica_pool in pools { - let replica_stats = replica_pool.lsn_stats(); - if !replica_stats.replica { - continue; + .find(|(_, stats)| stats.valid() && !stats.replica); + + for pool in pools { + let stats = pool.lsn_stats(); + let lag = match &primary { + Some((primary_pool, primary_stats)) + if pool.id() != primary_pool.id() && stats.valid() && stats.replica => + { + calculate_replica_lag(primary_stats, &stats) } - - let lag = calculate_replica_lag(&primary_stats, &replica_stats); - replica_pool.lock().replica_lag = lag; - } - primary_pool.lock().replica_lag = ReplicaLag::default(); + _ => ReplicaLag::default(), + }; + pool.lock().replica_lag = lag; } } @@ -248,7 +247,7 @@ mod test { } #[test] - fn test_update_replica_lag_assigns_primary_minus_replica_to_replica_pool() { + fn test_update_replica_lag_clears_lag_from_invalid_evidence() { let primary = Pool::new(&PoolConfig { address: Address::new_test(), config: Config::default(), @@ -264,15 +263,23 @@ mod test { set_pool_lsn_stats(&primary, false, 200, "2026-07-01 13:33:10.000000+00"); set_pool_lsn_stats(&replica, true, 100, "2026-07-01 13:33:00.000000+00"); - update_replica_lag(&[replica.clone(), primary.clone()]); + update_replica_lag(&[primary.clone(), replica.clone()]); + assert_eq!(replica.replica_lag().bytes, 100); + assert_eq!(replica.replica_lag().duration.as_secs(), 10); - let replica_lag = replica.replica_lag(); - assert_eq!(replica_lag.bytes, 100); - assert_eq!(replica_lag.duration.as_secs(), 10); + *primary.inner().lsn_stats.write() = LsnStats::default(); + update_replica_lag(&[primary.clone(), replica.clone()]); + assert_eq!(replica.replica_lag().bytes, 0); + assert_eq!(replica.replica_lag().duration, Duration::default()); - let primary_lag = primary.replica_lag(); - assert_eq!(primary_lag.bytes, 0); - assert_eq!(primary_lag.duration, Duration::default()); + set_pool_lsn_stats(&primary, false, 200, "2026-07-01 13:33:10.000000+00"); + update_replica_lag(&[primary.clone(), replica.clone()]); + assert_eq!(replica.replica_lag().bytes, 100); + + *replica.inner().lsn_stats.write() = LsnStats::default(); + update_replica_lag(&[primary, replica.clone()]); + assert_eq!(replica.replica_lag().bytes, 0); + assert_eq!(replica.replica_lag().duration, Duration::default()); } // The shard monitor reacts to an `lsn_role_change` notification by @@ -282,7 +289,10 @@ mod test { async fn test_monitor_updates_roles_on_failover() { crate::logger(); - let primary = Some(&pool_config(Address::new_test())); + let primary = Some(&pool_config(Address { + configured_role: Role::Auto, + ..Address::new_test() + })); let replicas = [pool_config(Address { configured_role: Role::Auto, ..Address::new_test() diff --git a/pgdog/src/backend/pool/shard/role_detector.rs b/pgdog/src/backend/pool/shard/role_detector.rs index 6fcc8eadc..14a1438e1 100644 --- a/pgdog/src/backend/pool/shard/role_detector.rs +++ b/pgdog/src/backend/pool/shard/role_detector.rs @@ -104,7 +104,7 @@ mod test { } #[test] - fn test_changed_returns_false_when_lsn_stats_invalid() { + fn test_changed_revokes_primary_when_lsn_stats_invalid() { let primary = Some(create_test_pool_config("127.0.0.1", 5432, true)); let replicas = [create_test_pool_config("localhost", 5432, true)]; let shard = create_test_shard(primary.as_ref(), &replicas); @@ -112,6 +112,7 @@ mod test { let mut detector = RoleDetector::new(&shard); assert!(detector.enabled()); + assert!(detector.changed()); assert!(!detector.changed()); } diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index ad8a2e87e..abf28aca6 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -1408,6 +1408,84 @@ pub mod test { } } + pub(crate) async fn automatic_role_server(value: Option<&'static str>) -> Server { + automatic_role_server_response(value.map(AutomaticRoleResponse::Value)).await + } + + pub(crate) async fn automatic_role_error_server() -> Server { + automatic_role_server_response(Some(AutomaticRoleResponse::Error)).await + } + + pub(crate) async fn automatic_role_empty_server() -> Server { + automatic_role_server_response(Some(AutomaticRoleResponse::Empty)).await + } + + pub(crate) async fn automatic_role_disconnect_server() -> Server { + automatic_role_server_response(Some(AutomaticRoleResponse::Disconnect)).await + } + + enum AutomaticRoleResponse { + Value(&'static str), + Error, + Empty, + Disconnect, + } + + async fn automatic_role_server_response(response: Option) -> Server { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let connect = tokio::net::TcpStream::connect(listener.local_addr().unwrap()); + let (client, accepted) = tokio::join!(connect, listener.accept()); + let mut socket = accepted.unwrap().0; + tokio::spawn(async move { + let code = socket.read_u8().await.unwrap(); + let len = socket.read_i32().await.unwrap(); + assert_eq!(code, b'Q'); + let mut payload = vec![0; len as usize - 4]; + socket.read_exact(&mut payload).await.unwrap(); + + let Some(response) = response else { + std::future::pending::<()>().await; + return; + }; + + let mut bytes = BytesMut::new(); + match response { + AutomaticRoleResponse::Value(value) => { + let mut row = DataRow::new(); + row.add(value); + bytes.extend_from_slice( + &RowDescription::new(&[Field::bool("pg_is_in_recovery")]).to_bytes(), + ); + bytes.extend_from_slice(&row.to_bytes()); + bytes.extend_from_slice(&CommandComplete::new("SELECT 1").to_bytes()); + } + AutomaticRoleResponse::Error => { + bytes.extend_from_slice( + &ErrorResponse { + code: "XX000".into(), + message: "backend role query failed".into(), + ..Default::default() + } + .to_bytes(), + ); + } + AutomaticRoleResponse::Empty => { + bytes.extend_from_slice( + &RowDescription::new(&[Field::bool("pg_is_in_recovery")]).to_bytes(), + ); + bytes.extend_from_slice(&CommandComplete::new("SELECT 0").to_bytes()); + } + AutomaticRoleResponse::Disconnect => return, + } + bytes.extend_from_slice(&ReadyForQuery::idle().to_bytes()); + socket.write_all(&bytes).await.unwrap(); + }); + + let mut server = Server::default(); + server.stream = Some(Stream::plain(client.unwrap(), 4096)); + server + } + pub(crate) async fn test_server() -> Server { Server::connect( &Address::new_test(), From 5bc0f705137c04aa4191146b6d96e8fbc44e3073 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Oko=C5=84ski?= Date: Fri, 21 Aug 2026 09:49:30 +0200 Subject: [PATCH 2/2] fix(pool): verify automatic-primary checkouts --- pgdog/src/backend/pool/lb/mod.rs | 135 +++++++++++++++-- pgdog/src/backend/pool/lb/test.rs | 137 +++++++++++++++++- pgdog/src/backend/pool/shard/role_detector.rs | 11 +- pgdog/src/backend/server.rs | 68 +++++++++ 4 files changed, 335 insertions(+), 16 deletions(-) diff --git a/pgdog/src/backend/pool/lb/mod.rs b/pgdog/src/backend/pool/lb/mod.rs index 9f1339272..83e0f3ded 100644 --- a/pgdog/src/backend/pool/lb/mod.rs +++ b/pgdog/src/backend/pool/lb/mod.rs @@ -20,7 +20,7 @@ use crate::{ }; use super::{Error, Guard, Oids, Pool, PoolConfig, PoolRole, Request}; -use crate::util::safe_timeout; +use crate::{state::State, util::safe_timeout}; pub mod ban; pub mod monitor; @@ -92,6 +92,10 @@ impl Target { let stats = self.pool.lsn_stats(); stats.valid() && !stats.replica } + + fn is_automatic_primary(&self) -> bool { + self.role() == Role::Primary && self.pool.addr().configured_role == Role::Auto + } } /// Load balancer. @@ -344,16 +348,15 @@ impl LoadBalancer { result } - fn qualified_primary_target(&self) -> Option<&Target> { - self.targets - .iter() - .rev() - .find(|target| target.is_qualified_primary()) + fn qualified_primary_target(&self, excluded_pools: &[u64]) -> Option<&Target> { + self.targets.iter().rev().find(|target| { + !excluded_pools.contains(&target.pool.id()) && target.is_qualified_primary() + }) } async fn wait_primary_target(&self) -> Result<&Target, Error> { if !self.role_detection_enabled() { - return self.qualified_primary_target().ok_or(Error::NoPrimary); + return self.qualified_primary_target(&[]).ok_or(Error::NoPrimary); } loop { @@ -361,7 +364,7 @@ impl LoadBalancer { tokio::pin!(notified); notified.as_mut().enable(); - if let Some(target) = self.qualified_primary_target() { + if let Some(target) = self.qualified_primary_target(&[]) { return Ok(target); } @@ -376,7 +379,97 @@ impl LoadBalancer { } async fn get_primary_internal(&self, request: &Request) -> Result { - self.wait_primary_target().await?.pool.get(request).await + use smallvec::SmallVec; + + let first = self.wait_primary_target().await?; + let mut attempted: SmallVec<[u64; 8]> = SmallVec::new(); + + while attempted.len() < self.targets.len() { + let target = if attempted.is_empty() { + first + } else { + self.qualified_primary_target(&attempted) + .ok_or(Error::NoPrimary)? + }; + attempted.push(target.pool.id()); + + match self.checkout_target(target, request, true).await { + Err(Error::NoPrimary) => continue, + result => return result, + } + } + + Err(Error::NoPrimary) + } + + async fn checkout_target( + &self, + target: &Target, + request: &Request, + primary_required: bool, + ) -> Result { + let automatic_primary_before = target.is_automatic_primary(); + if (primary_required || automatic_primary_before) && !target.is_qualified_primary() { + return Err(Error::NoPrimary); + } + + let guard = target.pool.get(request).await?; + if primary_required || automatic_primary_before || target.is_automatic_primary() { + self.check_automatic_primary_guard(target, guard).await + } else { + Ok(guard) + } + } + + async fn check_automatic_primary_guard( + &self, + target: &Target, + guard: Guard, + ) -> Result { + let mut guard = guard; + if !target.is_qualified_primary() { + guard.stats_mut().state(State::ForceClose); + return Err(Error::NoPrimary); + } + + if target.pool.addr().configured_role == Role::Auto { + match guard + .check_automatic_primary_backend(target.pool.config().lsn_check_timeout) + .await + { + Ok(true) => {} + Ok(false) => { + warn!( + "automatic primary checkout rejected: backend {} is in recovery [{}]", + guard.id(), + guard.addr(), + ); + self.reject_automatic_primary(target, &mut guard); + return Err(Error::NoPrimary); + } + Err(err) => { + self.reject_automatic_primary(target, &mut guard); + return Err(if matches!(err, crate::backend::Error::ReadTimeout) { + Error::CheckoutTimeout + } else { + Error::ServerError + }); + } + } + } + + if !target.is_qualified_primary() { + guard.stats_mut().state(State::ForceClose); + return Err(Error::NoPrimary); + } + + Ok(guard) + } + + fn reject_automatic_primary(&self, target: &Target, guard: &mut Guard) { + guard.stats_mut().state(State::ForceClose); + target.pool.revoke_automatic_primary_evidence(); + self.redetect_roles(); } async fn get_internal(&self, request: &Request) -> Result { @@ -461,14 +554,26 @@ impl LoadBalancer { // Only ban a candidate pool if there are more than one // and we have alternates. let bannable = candidates.len() > 1; + let mut automatic_primary_rejected = false; + let mut automatic_primary_error = None; for target in &candidates { if target.ban.banned() { continue; } - match target.pool.get(request).await { + let automatic_primary = target.is_automatic_primary(); + match self.checkout_target(target, request, false).await { Ok(conn) => return Ok(conn), - Err(Error::Offline) => { + Err(Error::Offline) => continue, + Err(Error::NoPrimary) => { + automatic_primary_rejected = true; + continue; + } + Err(err) + if matches!(err, Error::CheckoutTimeout | Error::ServerError) + && automatic_primary => + { + automatic_primary_error.get_or_insert(err); continue; } Err(err) => { @@ -483,7 +588,13 @@ impl LoadBalancer { .iter() .for_each(|target| target.ban.unban(true, UnbanReason::AllTargetsBanned)); - Err(Error::AllReplicasDown) + Err(if automatic_primary_rejected { + Error::NoPrimary + } else if let Some(err) = automatic_primary_error { + err + } else { + Error::AllReplicasDown + }) } /// Shutdown replica pools. diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs index 2464ef48a..3f510004c 100644 --- a/pgdog/src/backend/pool/lb/test.rs +++ b/pgdog/src/backend/pool/lb/test.rs @@ -1,9 +1,12 @@ use std::collections::HashSet; use std::time::Duration; -use tokio::time::sleep; +use tokio::time::{Instant, sleep, timeout}; -use crate::backend::pool::{Address, Config, Error, PoolConfig, Request}; use crate::backend::replication::publisher::Lsn; +use crate::backend::{ + Server, + pool::{Address, Config, Error, Pool, PoolConfig, Request}, +}; use crate::config::{LoadBalancingStrategy, Role}; use itertools::*; use pgdog_stats::{LsnStats as StatsLsnStats, ReplicaLag}; @@ -75,6 +78,15 @@ fn set_lsn_stats(target: &Target, replica: bool, lsn: i64) { target.pool.publish_lsn_stats(stats); } +fn install_idle_server(pool: &Pool, mut server: Server) { + server.stats_mut().healthcheck(); + let mut inner = pool.lock(); + inner.online = true; + inner + .put(Box::new(server), Instant::now()) + .expect("test server should check in"); +} + impl LoadBalancer { /// Replica pools handle. pub fn pools(&self) -> Vec<&Pool> { @@ -1624,6 +1636,127 @@ fn test_concurrent_role_detection_keeps_one_qualified_primary() { assert_eq!(primaries[0].pool.id(), lb.targets[1].pool.id()); } +#[tokio::test] +async fn test_primary_checkout_retries_after_standby_automatic_backend() { + let configs = [ + create_auto_test_pool_config("127.0.0.1", 5432), + create_auto_test_pool_config("localhost", 5432), + ]; + let lb = LoadBalancer::new( + &None, + &configs, + LoadBalancingStrategy::RoundRobin, + ReadWriteSplit::IncludePrimary, + Default::default(), + ); + set_lsn_stats(&lb.targets[0], false, 100); + set_lsn_stats(&lb.targets[1], false, 200); + assert!(lb.redetect_roles()); + + let stale = crate::backend::server::test::automatic_role_server(Some("t")).await; + let writer = crate::backend::server::test::automatic_role_server(Some("f")).await; + let writer_id = writer.id(); + install_idle_server(&lb.targets[1].pool, stale); + install_idle_server(&lb.targets[0].pool, writer); + + let guard = lb.get_primary(&Request::default()).await.unwrap(); + assert_eq!(guard.id(), writer_id); + assert_eq!(lb.targets[1].role(), Role::Replica); + assert!(!lb.targets[1].pool.lsn_stats().valid()); + assert_eq!(lb.targets[1].pool.lock().force_close, 1); +} + +#[tokio::test] +async fn test_read_checkout_falls_through_after_standby_automatic_primary() { + let mut primary = create_auto_test_pool_config("primary", 5432); + primary.config.inner.lsn_check_timeout = Duration::from_millis(50); + let replica = create_test_pool_config("replica", 5432); + let lb = LoadBalancer::new( + &None, + &[primary, replica], + LoadBalancingStrategy::RoundRobin, + ReadWriteSplit::IncludePrimary, + Default::default(), + ); + set_lsn_stats(&lb.targets[0], false, 100); + assert!(lb.redetect_roles()); + + install_idle_server( + &lb.targets[0].pool, + crate::backend::server::test::automatic_role_server(Some("t")).await, + ); + let replica_server = crate::backend::server::test::live_server().await; + let replica_id = replica_server.id(); + install_idle_server(&lb.targets[1].pool, replica_server); + + let guard = lb.get(&Request::default()).await.unwrap(); + assert_eq!(guard.id(), replica_id); + assert_eq!(lb.targets[0].role(), Role::Replica); +} + +async fn assert_automatic_primary_probe_error(server: Server, expected: Error) { + let mut config = create_auto_test_pool_config("127.0.0.1", 5432); + config.config.inner.lsn_check_timeout = Duration::from_millis(10); + let lb = LoadBalancer::new( + &None, + &[config], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + Default::default(), + ); + set_lsn_stats(&lb.targets[0], false, 100); + assert!(lb.redetect_roles()); + install_idle_server(&lb.targets[0].pool, server); + + assert_eq!( + lb.get_primary(&Request::default()).await.unwrap_err(), + expected + ); + assert!(!lb.targets[0].pool.lsn_stats().valid()); + assert_eq!(lb.targets[0].role(), Role::Replica); + assert_eq!(lb.targets[0].pool.lock().force_close, 1); +} + +#[tokio::test] +async fn test_automatic_primary_probe_errors_are_distinct_and_fail_closed() { + assert_automatic_primary_probe_error( + crate::backend::server::test::automatic_role_error_server().await, + Error::ServerError, + ) + .await; + assert_automatic_primary_probe_error( + crate::backend::server::test::automatic_role_server(None).await, + Error::CheckoutTimeout, + ) + .await; +} + +#[tokio::test] +async fn test_static_primary_checkout_does_not_probe_backend_role() { + let mut config = create_test_pool_config("127.0.0.1", 5432); + config.address.configured_role = Role::Primary; + let primary = Pool::new(&config); + let server = crate::backend::server::test::live_server().await; + let server_id = server.id(); + install_idle_server(&primary, server); + let lb = LoadBalancer::new( + &Some(primary), + &[], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + Default::default(), + ); + + let result = timeout( + Duration::from_millis(50), + lb.get_primary(&Request::default()), + ) + .await + .expect("static primary checkout must not issue a role probe") + .unwrap(); + assert_eq!(result.id(), server_id); +} + #[tokio::test] async fn test_can_move_conns_to_different_addresses() { let pool_config1 = create_test_pool_config("127.0.0.1", 5432); diff --git a/pgdog/src/backend/pool/shard/role_detector.rs b/pgdog/src/backend/pool/shard/role_detector.rs index 14a1438e1..bfa48906b 100644 --- a/pgdog/src/backend/pool/shard/role_detector.rs +++ b/pgdog/src/backend/pool/shard/role_detector.rs @@ -2,17 +2,20 @@ use super::Shard; pub(super) struct RoleDetector { enabled: bool, + primary_id: Option, shard: Shard, } impl RoleDetector { /// Create new role change detector. pub(super) fn new(shard: &Shard) -> Self { + let primary_id = shard.lb.primary().map(|pool| pool.id()); Self { enabled: shard .pools() .iter() .all(|pool| pool.config().role_detection), + primary_id, shard: shard.clone(), } } @@ -20,11 +23,14 @@ impl RoleDetector { /// Detect role change in the shard. pub(super) fn changed(&mut self) -> bool { if self.enabled() { - let changed = self.shard.redetect_roles(); + let lb_changed = self.shard.redetect_roles(); + let primary_id = self.shard.lb.primary().map(|pool| pool.id()); + let changed = lb_changed || primary_id != self.primary_id; if changed { // Re-initialize pub/sub channel. self.shard.init_pub_sub(); } + self.primary_id = primary_id; changed } else { false @@ -132,7 +138,7 @@ mod test { } #[test] - fn test_changed_returns_true_on_failover() { + fn test_changed_returns_true_after_external_failover_detection() { let primary = Some(create_test_pool_config("127.0.0.1", 5432, true)); let replicas = [create_test_pool_config("localhost", 5432, true)]; let shard = create_test_shard(primary.as_ref(), &replicas); @@ -148,6 +154,7 @@ mod test { set_lsn_stats(&shard, 0, false, 300); set_lsn_stats(&shard, 1, true, 200); + assert!(shard.redetect_roles()); assert!(detector.changed()); } diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index abf28aca6..699b468b8 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -32,6 +32,7 @@ use crate::{ }, }, stats::memory::MemoryUsage, + util::safe_timeout, }; use crate::{ config::{PoolerMode, TlsVerifyMode, config}, @@ -874,6 +875,23 @@ impl Server { &self.params } + pub(super) async fn check_automatic_primary_backend( + &mut self, + timeout: Duration, + ) -> Result { + let replica: bool = safe_timeout( + timeout, + self.fetch_all::("SELECT pg_is_in_recovery()"), + ) + .await + .unwrap_or(Err(Error::ReadTimeout)) + .and_then(|rows| rows.into_iter().next().ok_or(Error::DecoderRowError)) + .and_then(|row| row.get(0, Format::Text).ok_or(Error::DecoderRowError)) + .inspect_err(|_| self.stats.state(State::ForceClose))?; + + Ok(!replica) + } + /// Execute a batch of queries and return all results. pub async fn execute_batch( &mut self, @@ -1424,6 +1442,21 @@ pub mod test { automatic_role_server_response(Some(AutomaticRoleResponse::Disconnect)).await } + pub(crate) async fn live_server() -> Server { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let connect = tokio::net::TcpStream::connect(listener.local_addr().unwrap()); + let (client, accepted) = tokio::join!(connect, listener.accept()); + let socket = accepted.unwrap().0; + tokio::spawn(async move { + let _socket = socket; + std::future::pending::<()>().await; + }); + + let mut server = Server::default(); + server.stream = Some(Stream::plain(client.unwrap(), 4096)); + server + } + enum AutomaticRoleResponse { Value(&'static str), Error, @@ -1486,6 +1519,41 @@ pub mod test { server } + async fn assert_automatic_primary_check_fails_closed(mut server: Server) { + let result = server + .check_automatic_primary_backend(Duration::from_millis(25)) + .await; + assert!(result.is_err()); + assert_eq!(server.stats().get_state(), State::ForceClose); + } + + #[tokio::test] + async fn test_automatic_primary_backend_check() { + let timeout = Duration::from_millis(50); + let mut writer = automatic_role_server(Some("f")).await; + let mut standby = automatic_role_server(Some("t")).await; + + assert!( + writer + .check_automatic_primary_backend(timeout) + .await + .unwrap() + ); + assert!( + !standby + .check_automatic_primary_backend(timeout) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_automatic_primary_backend_check_fails_closed() { + assert_automatic_primary_check_fails_closed(automatic_role_error_server().await).await; + assert_automatic_primary_check_fails_closed(automatic_role_empty_server().await).await; + assert_automatic_primary_check_fails_closed(automatic_role_server(None).await).await; + } + pub(crate) async fn test_server() -> Server { Server::connect( &Address::new_test(),