From 60a6b20a04e3632233f58d16e66424d55aa9660d Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Fri, 7 Aug 2026 13:55:06 +0530 Subject: [PATCH] lb: default auto mode pools to replica --- pgdog/src/backend/pool/lb/mod.rs | 59 +++++++------- pgdog/src/backend/pool/lb/test.rs | 119 +++++++++++++++++++++------- pgdog/src/backend/pool/shard/mod.rs | 34 +++++++- 3 files changed, 149 insertions(+), 63 deletions(-) diff --git a/pgdog/src/backend/pool/lb/mod.rs b/pgdog/src/backend/pool/lb/mod.rs index 7bf0b71a3..ddf95c0b5 100644 --- a/pgdog/src/backend/pool/lb/mod.rs +++ b/pgdog/src/backend/pool/lb/mod.rs @@ -122,10 +122,11 @@ impl LoadBalancer { let mut targets: Vec<_> = addrs .iter() .map(|config| { - Target::new( - Pool::with_oid_mapping(config, Arc::clone(&oids)), - config.address.configured_role, - ) + let role = match config.address.configured_role { + Role::Auto => Role::Replica, + role => role, + }; + Target::new(Pool::with_oid_mapping(config, Arc::clone(&oids)), role) }) .collect(); @@ -168,7 +169,6 @@ impl LoadBalancer { /// return new primary (if any), and replicas. pub fn redetect_roles(&self) -> bool { let mut promoted = false; - let roles_detected_before = self.roles_detected(); let mut targets = self .targets @@ -213,7 +213,7 @@ impl LoadBalancer { }); } - if promoted || (!roles_detected_before && self.roles_detected()) { + if promoted { self.role_detection.notify_one(); } @@ -285,14 +285,19 @@ impl LoadBalancer { } /// True if the LB has any target that can serve replica reads. - /// - /// An `Auto` target counts as a potential replica until role detection - /// converges, so callers may briefly route reads to a target that turns - /// out to be the primary. pub fn has_replicas(&self) -> bool { self.targets .iter() - .any(|target| matches!(target.role(), Role::Replica | Role::Auto)) + .any(|target| target.role() == Role::Replica) + } + + /// True if target roles are detected automatically. + pub fn role_detection_enabled(&self) -> bool { + !self.targets.is_empty() + && self + .targets + .iter() + .all(|target| target.pool.config().role_detection) } /// Cancel a query if one is running. @@ -320,12 +325,12 @@ impl LoadBalancer { result } - /// Block until role detection has assigned every `Auto` target to - /// `Primary` or `Replica`. The wakeup is driven by `pick_primary`, so - /// if no primary is ever elected (e.g. LSN stats never populate), - /// callers will block until their `checkout_timeout` fires. - async fn wait_roles_detected(&self) -> Result<(), Error> { - if !self.roles_detected() { + /// 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() @@ -340,20 +345,12 @@ impl LoadBalancer { Ok(()) } - /// True once no target is still in the `Auto` state. - pub fn roles_detected(&self) -> bool { - !self - .targets - .iter() - .any(|target| target.role() == Role::Auto) - } - pub(super) async fn get_primary(&self, request: &Request) -> Result { self.get_primary_internal(request).await } async fn get_primary_internal(&self, request: &Request) -> Result { - self.wait_roles_detected().await?; + self.wait_primary().await?; self.primary_target() .ok_or(Error::NoPrimary)? .pool @@ -380,17 +377,17 @@ impl LoadBalancer { // we read from the primary if we have no replicas ExcludePrimary => !candidates .iter() - .any(|target| matches!(target.role(), Role::Replica | Role::Auto)), + .any(|target| target.role() == Role::Replica), // PreferPrimary makes all queries writes. If a query lands here, // it's because of pgdog.role=replica. Let it use the primary only if // no replicas are available. - PreferPrimary => !candidates.iter().any(|target| { - matches!(target.role(), Role::Replica | Role::Auto) && !target.ban.banned() - }), + PreferPrimary => !candidates + .iter() + .any(|target| target.role() == Role::Replica && !target.ban.banned()), }; if !primary_reads { - candidates.retain(|target| matches!(target.role(), Role::Replica | Role::Auto)); + candidates.retain(|target| target.role() == Role::Replica); } if candidates.is_empty() { diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs index 35bbf99dd..50fd6c1b4 100644 --- a/pgdog/src/backend/pool/lb/test.rs +++ b/pgdog/src/backend/pool/lb/test.rs @@ -32,6 +32,13 @@ fn create_test_pool_config(host: &str, port: u16) -> PoolConfig { } } +fn create_auto_test_pool_config(host: &str, port: u16) -> PoolConfig { + let mut config = create_test_pool_config(host, port); + config.address.configured_role = Role::Auto; + config.config.inner.role_detection = true; + config +} + fn setup_test_replicas() -> LoadBalancer { let pool_config1 = create_test_pool_config("127.0.0.1", 5432); let pool_config2 = create_test_pool_config("localhost", 5432); @@ -1260,11 +1267,8 @@ async fn test_move_conns_to_with_added_replica_matches_by_address() { #[tokio::test] async fn test_redetect_roles_marks_added_auto_target_replica_when_primary_unchanged() { - let mut primary_config = create_test_pool_config("127.0.0.1", 5432); - primary_config.address.configured_role = Role::Auto; - - let mut existing_replica_config = create_test_pool_config("localhost", 5432); - existing_replica_config.address.configured_role = Role::Auto; + let primary_config = create_auto_test_pool_config("127.0.0.1", 5432); + let existing_replica_config = create_auto_test_pool_config("localhost", 5432); let old_primary = Pool::new(&primary_config); let lb_old = LoadBalancer::new( @@ -1278,10 +1282,9 @@ async fn test_redetect_roles_marks_added_auto_target_replica_when_primary_unchan set_lsn_stats(&lb_old.targets[0], true, 100); set_lsn_stats(&lb_old.targets[1], false, 200); assert!(!lb_old.redetect_roles()); - assert!(lb_old.roles_detected()); + assert!(lb_old.role_detection_enabled()); - let mut added_replica_config = create_test_pool_config("localhost", 5433); - added_replica_config.address.configured_role = Role::Auto; + let added_replica_config = create_auto_test_pool_config("localhost", 5433); let new_primary = Pool::new(&primary_config); let lb_new = LoadBalancer::new( @@ -1299,7 +1302,7 @@ async fn test_redetect_roles_marks_added_auto_target_replica_when_primary_unchan .iter() .find(|target| target.pool.addr().port == 5433) .expect("newly added target should exist"); - assert_eq!(added_target.role(), Role::Auto); + assert_eq!(added_target.role(), Role::Replica); set_lsn_stats(&lb_new.targets[0], true, 100); set_lsn_stats(&lb_new.targets[1], true, 90); @@ -1310,16 +1313,13 @@ async fn test_redetect_roles_marks_added_auto_target_replica_when_primary_unchan "primary did not change, so role detector reports no promotion" ); assert_eq!(added_target.role(), Role::Replica); - assert!(lb_new.roles_detected()); + assert!(lb_new.role_detection_enabled()); } #[tokio::test] -async fn test_redetect_roles_leaves_auto_targets_pending_when_stats_are_invalid() { - let mut config1 = create_test_pool_config("127.0.0.1", 5432); - config1.address.configured_role = Role::Auto; - - let mut config2 = create_test_pool_config("localhost", 5432); - config2.address.configured_role = Role::Auto; +async fn test_auto_targets_remain_replicas_when_stats_are_invalid() { + let config1 = create_auto_test_pool_config("127.0.0.1", 5432); + let config2 = create_auto_test_pool_config("localhost", 5432); let lb = LoadBalancer::new( &None, @@ -1329,26 +1329,30 @@ async fn test_redetect_roles_leaves_auto_targets_pending_when_stats_are_invalid( Default::default(), ); - assert!(lb.targets.iter().all(|target| target.role() == Role::Auto)); - assert!(!lb.roles_detected()); + assert!( + lb.targets + .iter() + .all(|target| target.role() == Role::Replica) + ); + assert!(lb.role_detection_enabled()); assert!( !lb.redetect_roles(), "no valid primary was discovered, so no promotion is reported" ); - assert!(lb.targets.iter().all(|target| target.role() == Role::Auto)); - assert!(!lb.roles_detected()); + assert!( + lb.targets + .iter() + .all(|target| target.role() == Role::Replica) + ); assert!(lb.has_replicas()); } #[tokio::test] async fn test_redetect_roles_marks_auto_targets_replicas_when_all_valid_targets_are_replicas() { - let mut config1 = create_test_pool_config("127.0.0.1", 5432); - config1.address.configured_role = Role::Auto; - - let mut config2 = create_test_pool_config("localhost", 5432); - config2.address.configured_role = Role::Auto; + let config1 = create_auto_test_pool_config("127.0.0.1", 5432); + let config2 = create_auto_test_pool_config("localhost", 5432); let lb = LoadBalancer::new( &None, @@ -1361,8 +1365,12 @@ async fn test_redetect_roles_marks_auto_targets_replicas_when_all_valid_targets_ set_lsn_stats(&lb.targets[0], true, 100); set_lsn_stats(&lb.targets[1], true, 90); - assert!(lb.targets.iter().all(|target| target.role() == Role::Auto)); - assert!(!lb.roles_detected()); + assert!( + lb.targets + .iter() + .all(|target| target.role() == Role::Replica) + ); + assert!(lb.role_detection_enabled()); assert!( !lb.redetect_roles(), @@ -1374,10 +1382,65 @@ async fn test_redetect_roles_marks_auto_targets_replicas_when_all_valid_targets_ .iter() .all(|target| target.role() == Role::Replica) ); - assert!(lb.roles_detected()); assert!(lb.has_replicas()); } +#[tokio::test] +async fn test_auto_mode_waits_for_primary_election() { + let mut config = create_auto_test_pool_config("127.0.0.1", 5432); + config.config.inner.checkout_timeout = Duration::from_millis(10); + + let lb = LoadBalancer::new( + &None, + &[config], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + Default::default(), + ); + + assert_eq!(lb.wait_primary().await, Err(Error::CheckoutTimeout)); +} + +#[tokio::test] +async fn test_auto_mode_primary_election_releases_writes() { + let config = create_auto_test_pool_config("127.0.0.1", 5432); + let lb = LoadBalancer::new( + &None, + &[config], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + Default::default(), + ); + let election = lb.clone(); + + tokio::spawn(async move { + sleep(Duration::from_millis(10)).await; + election.targets[0].set_role(Role::Primary); + election.role_detection.notify_one(); + }); + + assert_eq!(lb.wait_primary().await, Ok(())); + assert!(lb.primary().is_some()); +} + +#[tokio::test] +async fn test_static_replica_only_does_not_wait_for_primary() { + let config = create_test_pool_config("127.0.0.1", 5432); + let lb = LoadBalancer::new( + &None, + &[config], + LoadBalancingStrategy::Random, + ReadWriteSplit::IncludePrimary, + Default::default(), + ); + + assert_eq!(lb.wait_primary().await, Ok(())); + assert!(matches!( + lb.get_primary(&Request::default()).await, + Err(Error::NoPrimary) + )); +} + #[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/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index 79c812e62..1dab2e3af 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -196,10 +196,7 @@ impl Shard { /// Returns true if the shard has a primary database. pub fn has_primary(&self) -> bool { - match self.lb.primary() { - Some(_) => true, - None => !self.lb.roles_detected(), // Assume there is a primary, until proven otherwise. - } + self.lb.primary().is_some() || self.lb.role_detection_enabled() } /// Returns true if the shard has any replica databases. @@ -465,4 +462,33 @@ mod test { assert_eq!(ids.len(), 2); } + + #[test] + fn test_auto_mode_is_read_ready_while_primary_election_is_pending() { + let replicas = &[PoolConfig { + address: Address { + configured_role: Role::Auto, + ..Address::new_test() + }, + config: super::super::Config { + inner: pgdog_stats::Config { + role_detection: true, + ..Default::default() + }, + }, + }]; + + let shard = Shard::new(ShardConfig { + replicas, + identifier: Arc::new(User { + user: "pgdog".into(), + database: "pgdog".into(), + }), + ..Default::default() + }); + + assert!(shard.has_primary()); + assert!(shard.has_replicas()); + assert_eq!(shard.lb.targets[0].role(), Role::Replica); + } }