lb: default auto mode pools to replica

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<Guard, Error> {
         self.get_primary_internal(request).await
     }
 
     async fn get_primary_internal(&self, request: &Request) -> Result<Guard, Error> {
-        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() {

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);

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);
+    }
 }