Minor cleanup of TargetHealth (#1353)

Fixed visibility, removed dead code, replaced cloned value with an
accessor method since we always have access to the original

Author: sgrif

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

@@ -28,7 +28,7 @@ pub mod target_health;
 use ban::Ban;
 pub use ban::UnbanReason;
 use monitor::*;
-pub use target_health::*;
+pub(crate) use target_health::*;
 
 #[cfg(test)]
 mod test;
@@ -39,7 +39,6 @@ pub struct Target {
     pub pool: Pool,
     pub ban: Ban,
     role: PoolRole,
-    pub health: TargetHealth,
     /// Smooth weighted round-robin current weight tracker.
     current_weight: Arc<AtomicI64>,
 }
@@ -54,7 +53,6 @@ impl Target {
         Self {
             ban,
             role: PoolRole::new(role),
-            health: pool.inner().health.clone(),
             pool,
             current_weight: Arc::new(AtomicI64::new(0)),
         }
@@ -77,6 +75,10 @@ impl Target {
 
         lb && pool
     }
+
+    pub(super) fn health(&self) -> &TargetHealth {
+        &self.pool.inner().health
+    }
 }
 
 /// Load balancer.

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

@@ -89,7 +89,7 @@ impl Monitor {
         let targets = &self.replicas.targets;
 
         for (i, target) in targets.iter().enumerate() {
-            let healthy = target.health.healthy();
+            let healthy = target.health().healthy();
             let replica_lag_bad = target
                 .pool
                 .replica_lag()

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

@@ -6,25 +6,22 @@ use std::sync::{
 };
 
 #[derive(Clone, Debug)]
-pub struct TargetHealth {
-    #[allow(dead_code)]
-    pub(super) id: u64,
+pub(crate) struct TargetHealth {
     pub(super) healthy: Arc<AtomicBool>,
 }
 
 impl TargetHealth {
-    pub fn new(id: u64) -> Self {
+    pub(crate) fn new() -> Self {
         Self {
-            id,
             healthy: Arc::new(AtomicBool::new(true)),
         }
     }
 
-    pub fn toggle(&self, healthy: bool) {
+    pub(crate) fn toggle(&self, healthy: bool) {
         self.healthy.swap(healthy, Ordering::SeqCst);
     }
 
-    pub fn healthy(&self) -> bool {
+    pub(crate) fn healthy(&self) -> bool {
         self.healthy.load(Ordering::Relaxed)
     }
 }

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

@@ -277,7 +277,7 @@ async fn test_unban_if_expired_checks_pool_health() {
     ban.ban(Error::ServerError, Duration::from_millis(50));
     assert!(ban.banned());
 
-    pool.inner().health.toggle(false);
+    pool.inner().health().toggle(false);
 
     sleep(Duration::from_millis(60)).await;
 
@@ -814,7 +814,7 @@ async fn test_monitor_shuts_down_on_notify() {
 async fn test_monitor_bans_unhealthy_target() {
     let replicas = setup_test_replicas();
 
-    replicas.targets[0].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
 
     sleep(Duration::from_millis(400)).await;
 
@@ -851,7 +851,7 @@ async fn test_monitor_does_not_ban_single_target() {
     );
     replicas.launch();
 
-    replicas.targets[0].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
 
     sleep(Duration::from_millis(400)).await;
 
@@ -864,8 +864,8 @@ async fn test_monitor_does_not_ban_single_target() {
 async fn test_monitor_unbans_all_when_all_unhealthy() {
     let replicas = setup_test_replicas();
 
-    replicas.targets[0].health.toggle(false);
-    replicas.targets[1].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
+    replicas.targets[1].health().toggle(false);
 
     sleep(Duration::from_millis(400)).await;
 
@@ -924,7 +924,7 @@ async fn test_monitor_does_not_ban_with_zero_ban_timeout() {
     );
     replicas.launch();
 
-    replicas.targets[0].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
 
     sleep(Duration::from_millis(400)).await;
 
@@ -942,9 +942,9 @@ async fn test_monitor_health_state_race() {
 
     let toggle_task = spawn(async move {
         for _ in 0..50 {
-            target.health.toggle(false);
+            target.health().toggle(false);
             sleep(Duration::from_micros(100)).await;
-            target.health.toggle(true);
+            target.health().toggle(true);
             sleep(Duration::from_micros(100)).await;
         }
     });
@@ -954,7 +954,7 @@ async fn test_monitor_health_state_race() {
     toggle_task.await.unwrap();
 
     let banned = replicas.targets[0].ban.banned();
-    let healthy = replicas.targets[0].health.healthy();
+    let healthy = replicas.targets[0].health().healthy();
 
     assert!(
         !banned || !healthy,
@@ -1471,7 +1471,7 @@ async fn test_monitor_unbans_all_when_second_target_becomes_unhealthy_after_firs
     let replicas = setup_test_replicas();
 
     // First target becomes unhealthy
-    replicas.targets[0].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
 
     // Wait for monitor to ban the first target
     sleep(Duration::from_millis(400)).await;
@@ -1486,7 +1486,7 @@ async fn test_monitor_unbans_all_when_second_target_becomes_unhealthy_after_firs
     );
 
     // Now second target becomes unhealthy (first is already banned)
-    replicas.targets[1].health.toggle(false);
+    replicas.targets[1].health().toggle(false);
 
     // Wait for monitor to process - should unban all since all are unhealthy
     sleep(Duration::from_millis(400)).await;
@@ -1739,7 +1739,7 @@ fn test_ban_check_does_not_clear_expired_ban_when_healthy_with_bad_lag() {
     let replicas = setup_test_replicas_no_launch();
 
     // Target is healthy (default)
-    assert!(replicas.targets[0].health.healthy());
+    assert!(replicas.targets[0].health().healthy());
 
     // Ban with short timeout
     replicas.targets[0]
@@ -1776,7 +1776,7 @@ fn test_ban_check_does_not_clear_expired_ban_when_unhealthy_with_bad_lag() {
     let replicas = setup_test_replicas_no_launch();
 
     // Set target as unhealthy
-    replicas.targets[0].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
 
     // Ban with short timeout
     replicas.targets[0]
@@ -1813,7 +1813,7 @@ fn test_ban_check_bans_unhealthy_replica_with_bad_lag() {
     let replicas = setup_test_replicas_no_launch();
 
     // Set target as unhealthy
-    replicas.targets[0].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
 
     // Set replica lag on the pool
     replicas.targets[0].pool.lock().replica_lag = ReplicaLag {
@@ -1842,7 +1842,7 @@ fn test_ban_check_bans_healthy_replica_with_bad_lag() {
     let replicas = setup_test_replicas_no_launch();
 
     // Target stays healthy (default)
-    assert!(replicas.targets[0].health.healthy());
+    assert!(replicas.targets[0].health().healthy());
 
     // Set replica lag on the pool
     replicas.targets[0].pool.lock().replica_lag = ReplicaLag {
@@ -1874,7 +1874,7 @@ fn test_ban_check_bans_with_pool_unhealthy_reason() {
     let replicas = setup_test_replicas_no_launch();
 
     // Set target as unhealthy
-    replicas.targets[0].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
 
     // No replica lag set (defaults to zero)
 
@@ -1908,7 +1908,7 @@ fn test_ban_check_does_not_ban_single_target() {
     // Don't launch - we're unit testing ban_check
 
     // Set target as unhealthy
-    replicas.targets[0].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
 
     let monitor = Monitor::new_test(&replicas);
     let threshold = ReplicaLag {
@@ -1973,7 +1973,7 @@ fn test_ban_check_does_not_ban_with_zero_ban_timeout() {
     );
 
     // Set target as unhealthy
-    replicas.targets[0].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
 
     let monitor = Monitor::new_test(&replicas);
     let threshold = ReplicaLag {
@@ -2002,8 +2002,8 @@ fn test_ban_check_unbans_all_when_all_unhealthy() {
         .ban(Error::ServerError, Duration::from_secs(60));
 
     // Set both as unhealthy
-    replicas.targets[0].health.toggle(false);
-    replicas.targets[1].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
+    replicas.targets[1].health().toggle(false);
 
     assert!(replicas.targets[0].ban.banned());
     assert!(replicas.targets[1].ban.banned());
@@ -2043,8 +2043,8 @@ fn test_ban_check_unbans_all_when_all_healthy_but_banned() {
         .ban
         .ban(Error::ConnectTimeout, Duration::from_secs(60));
 
-    assert!(replicas.targets[0].health.healthy());
-    assert!(replicas.targets[1].health.healthy());
+    assert!(replicas.targets[0].health().healthy());
+    assert!(replicas.targets[1].health().healthy());
     assert!(replicas.targets[0].ban.banned());
     assert!(replicas.targets[1].ban.banned());
 
@@ -2126,7 +2126,7 @@ fn test_ban_check_default_threshold_does_not_ban_healthy_replica_with_high_lag()
     let replicas = setup_test_replicas_no_launch();
 
     // Target is healthy (default)
-    assert!(replicas.targets[0].health.healthy());
+    assert!(replicas.targets[0].health().healthy());
 
     // Set very high replica lag on the pool
     replicas.targets[0].pool.lock().replica_lag = ReplicaLag {
@@ -2160,7 +2160,7 @@ fn test_ban_check_default_threshold_bans_unhealthy_with_pool_unhealthy_reason()
     };
 
     // Set target as unhealthy
-    replicas.targets[0].health.toggle(false);
+    replicas.targets[0].health().toggle(false);
 
     let monitor = Monitor::new_test(&replicas);
     // Use default config thresholds (MAX values)

pgdog/src/backend/pool/pool_impl.rs

@@ -76,7 +76,7 @@ impl Pool {
                 inner: Mutex::new(Inner::new(config.config, id)),
                 id,
                 config: config.config,
-                health: TargetHealth::new(id),
+                health: TargetHealth::new(),
                 params: OnceCell::new(),
                 lsn_stats: RwLock::new(LsnStats::default()),
                 lsn_role_change: Notify::new(),