fix: replica lag evalution in load balancer (#1122)

- fix: the primary/replica evaluation order was inverted, so the replica
ban never triggered
- fix: `prefer_primary` incorrectly excluded primary when all replicas
were banned @stewart-glabs 🙏
- feat: show the reason why pools are unbanned (helps debugging)

Author: levkk

pgdog/src/admin/ban.rs

@@ -53,7 +53,7 @@ impl Command for Ban {
                     }
 
                     if self.unban {
-                        ban.unban(false);
+                        ban.unban(false, pool::lb::UnbanReason::Manual);
                     } else {
                         ban.ban(pool::Error::ManualBan, Duration::MAX);
                     }

pgdog/src/backend/pool/lb/ban.rs

@@ -1,6 +1,6 @@
 use super::*;
 use parking_lot::RwLock;
-use std::time::Instant;
+use std::{fmt::Display, time::Instant};
 
 use tracing::{error, warn};
 
@@ -11,6 +11,23 @@ pub struct Ban {
     pool: Pool,
 }
 
+#[derive(Debug, Copy, Clone)]
+pub enum UnbanReason {
+    AllTargetsBanned,
+    Expired,
+    Manual,
+}
+
+impl Display for UnbanReason {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::AllTargetsBanned => write!(f, "all targets banned"),
+            Self::Expired => write!(f, "expired"),
+            Self::Manual => write!(f, "manual"),
+        }
+    }
+}
+
 impl Ban {
     /// Create new ban handler.
     pub(super) fn new(pool: &Pool) -> Self {
@@ -48,15 +65,25 @@ impl Ban {
     }
 
     /// Unban the database.
-    pub fn unban(&self, manual_check: bool) {
+    ///
+    /// FIXME(lev): `reason` seems like it should be
+    /// used as an operand but it's only used for logging.
+    /// We should unify methods and provide one public interface to this.
+    ///
+    pub fn unban(&self, manual_check: bool, reason: UnbanReason) {
         let mut guard = self.inner.upgradable_read();
         if let Some(ref ban) = guard.ban {
+            let mut unbanned = false;
             if ban.error != Error::ManualBan || !manual_check {
                 guard.with_upgraded(|guard| {
                     guard.ban = None;
                 });
+                unbanned = true;
+            }
+
+            if unbanned {
+                warn!("resuming read queries: {} [{}]", reason, self.pool.addr());
             }
-            warn!("resuming read queries [{}]", self.pool.addr());
         }
     }
 
@@ -114,7 +141,11 @@ impl Ban {
         };
         drop(guard);
         if unbanned {
-            warn!("resuming read queries [{}]", self.pool.addr());
+            warn!(
+                "resuming read queries: {} [{}]",
+                UnbanReason::Expired,
+                self.pool.addr()
+            );
         }
         unbanned
     }
@@ -185,7 +216,7 @@ mod tests {
         let pool = Pool::new_test();
         let ban = Ban::new(&pool);
         ban.ban(Error::ServerError, Duration::from_secs(1));
-        ban.unban(false);
+        ban.unban(false, UnbanReason::Expired);
         assert!(!ban.banned());
         assert!(ban.error().is_none());
     }
@@ -195,7 +226,7 @@ mod tests {
         let pool = Pool::new_test();
         let ban = Ban::new(&pool);
         ban.ban(Error::ManualBan, Duration::from_secs(1));
-        ban.unban(true);
+        ban.unban(true, UnbanReason::Expired);
         assert!(ban.banned());
         assert_eq!(ban.error(), Some(Error::ManualBan));
     }
@@ -205,7 +236,7 @@ mod tests {
         let pool = Pool::new_test();
         let ban = Ban::new(&pool);
         ban.ban(Error::ServerError, Duration::from_secs(1));
-        ban.unban(true);
+        ban.unban(true, UnbanReason::Manual);
         assert!(!ban.banned());
         assert!(ban.error().is_none());
     }
@@ -303,7 +334,7 @@ mod tests {
 
         for _ in 0..100 {
             ban.ban(Error::ServerError, Duration::from_secs(1));
-            ban.unban(false);
+            ban.unban(false, UnbanReason::Expired);
         }
 
         h1.join().unwrap();
@@ -374,7 +405,7 @@ mod tests {
         for _ in 0..10 {
             assert!(ban.ban(Error::ServerError, Duration::from_secs(1)));
             assert!(ban.banned());
-            ban.unban(false);
+            ban.unban(false, UnbanReason::AllTargetsBanned);
             assert!(!ban.banned());
         }
     }

pgdog/src/backend/pool/lb/mod.rs

@@ -25,6 +25,7 @@ pub mod monitor;
 pub mod target_health;
 
 use ban::Ban;
+pub use ban::UnbanReason;
 use monitor::*;
 pub use target_health::*;
 
@@ -349,12 +350,16 @@ impl LoadBalancer {
         let primary_reads = match self.rw_split {
             IncludePrimary => true,
             IncludePrimaryIfReplicaBanned => candidates.iter().any(|target| target.ban.banned()),
-            // Under PreferPrimary, default queries already route to the primary as writes;
-            // a read only reaches here when it explicitly opted into replicas, so honor it
-            // like ExcludePrimary and fall back to the primary only if there are no replicas.
-            ExcludePrimary | PreferPrimary => !candidates
+            // we read from the primary if we have no replicas
+            ExcludePrimary => !candidates
                 .iter()
                 .any(|target| matches!(target.role(), Role::Replica | Role::Auto)),
+            // 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()
+            }),
         };
 
         if !primary_reads {
@@ -427,7 +432,9 @@ impl LoadBalancer {
             }
         }
 
-        candidates.iter().for_each(|target| target.ban.unban(true));
+        candidates
+            .iter()
+            .for_each(|target| target.ban.unban(true, UnbanReason::AllTargetsBanned));
 
         Err(Error::AllReplicasDown)
     }

pgdog/src/backend/pool/lb/monitor.rs

@@ -128,7 +128,7 @@ impl Monitor {
         // intent.
         if targets.len() == unavailable {
             targets.iter().for_each(|target| {
-                target.ban.unban(true);
+                target.ban.unban(true, UnbanReason::AllTargetsBanned);
             });
         } else {
             for (i, reason) in ban_targets {

pgdog/src/backend/pool/lb/test.rs

@@ -79,7 +79,7 @@ async fn test_replica_manual_unban() {
     assert!(ban.banned());
 
     // Manually unban
-    ban.unban(false);
+    ban.unban(false, UnbanReason::Manual);
 
     assert!(!ban.banned());
 
@@ -634,6 +634,41 @@ async fn test_read_write_split_with_banned_replicas() {
     replicas.shutdown();
 }
 
+#[tokio::test]
+async fn test_prefer_primary_with_banned_replicas_falls_back_to_primary() {
+    let primary_config = create_test_pool_config("127.0.0.1", 5432);
+    let primary_pool = Pool::new(&primary_config);
+    primary_pool.launch();
+
+    let replica_configs = [create_test_pool_config("localhost", 5432)];
+
+    let replicas = LoadBalancer::new(
+        &Some(primary_pool),
+        &replica_configs,
+        LoadBalancingStrategy::Random,
+        ReadWriteSplit::PreferPrimary,
+    );
+    replicas.launch();
+
+    let replica_ban = &replicas.targets[0].ban;
+    replica_ban.ban(Error::ServerError, Duration::from_millis(1000));
+
+    let request = Request::default();
+
+    let mut used_pool_ids = HashSet::new();
+    for _ in 0..10 {
+        let conn = replicas.get(&request).await.unwrap();
+        used_pool_ids.insert(conn.pool.id());
+    }
+
+    assert_eq!(used_pool_ids.len(), 1);
+
+    let primary_id = replicas.primary().unwrap().id();
+    assert!(used_pool_ids.contains(&primary_id));
+
+    replicas.shutdown();
+}
+
 #[tokio::test]
 async fn test_read_write_split_exclude_primary_with_round_robin() {
     let primary_config = create_test_pool_config("127.0.0.1", 5432);