Add setting for controlling connection cleanup (#611)

* Add setting for controlling connection cleanup

* early return

* test early return

* var name

Author: levkk

pgdog-config/src/general.rs

@@ -4,6 +4,8 @@ use std::net::Ipv4Addr;
 use std::path::PathBuf;
 use std::time::Duration;
 
+use crate::pooling::ConnectionRecovery;
+
 use super::auth::{AuthType, PassthoughAuth};
 use super::database::{LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy};
 use super::networking::TlsVerifyMode;
@@ -167,6 +169,9 @@ pub struct General {
     /// Stats averaging period (in milliseconds).
     #[serde(default = "General::stats_period")]
     pub stats_period: u64,
+    /// Connection cleanup algorithm.
+    #[serde(default = "General::connection_recovery")]
+    pub connection_recovery: ConnectionRecovery,
 }
 
 impl Default for General {
@@ -227,6 +232,7 @@ impl Default for General {
             expanded_explain: Self::expanded_explain(),
             server_lifetime: Self::server_lifetime(),
             stats_period: Self::stats_period(),
+            connection_recovery: Self::connection_recovery(),
         }
     }
 }
@@ -509,6 +515,10 @@ impl General {
         )
     }
 
+    pub fn connection_recovery() -> ConnectionRecovery {
+        Self::env_enum_or_default("PGDOG_CONNECTION_RECOVERY")
+    }
+
     fn stats_period() -> u64 {
         Self::env_or_default("PGDOG_STATS_PERIOD", 15_000)
     }

pgdog-config/src/pooling.rs

@@ -71,3 +71,38 @@ impl FromStr for PoolerMode {
         }
     }
 }
+
+#[derive(Serialize, Deserialize, Debug, Clone, Copy, Default, PartialEq, Eq, Ord, PartialOrd)]
+#[serde(rename_all = "snake_case")]
+pub enum ConnectionRecovery {
+    #[default]
+    Recover,
+    RollbackOnly,
+    Drop,
+}
+
+impl ConnectionRecovery {
+    pub fn can_recover(&self) -> bool {
+        matches!(self, ConnectionRecovery::Recover)
+    }
+
+    pub fn can_rollback(&self) -> bool {
+        matches!(
+            self,
+            ConnectionRecovery::Recover | ConnectionRecovery::RollbackOnly
+        )
+    }
+}
+
+impl FromStr for ConnectionRecovery {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "recover" => Ok(Self::Recover),
+            "rollbackonly" => Ok(Self::RollbackOnly),
+            "drop" => Ok(Self::Drop),
+            _ => Err(format!("Invalid pooler mode: {}", s)),
+        }
+    }
+}

pgdog/src/backend/pool/cluster.rs

@@ -16,7 +16,8 @@ use crate::{
         Schema, ShardedTables,
     },
     config::{
-        General, MultiTenant, PoolerMode, ReadWriteSplit, ReadWriteStrategy, ShardedTable, User,
+        ConnectionRecovery, General, MultiTenant, PoolerMode, ReadWriteSplit, ReadWriteStrategy,
+        ShardedTable, User,
     },
     net::{messages::BackendKeyData, Query},
 };
@@ -59,6 +60,7 @@ pub struct Cluster {
     expanded_explain: bool,
     pub_sub_channel_size: usize,
     query_parser_enabled: bool,
+    connection_recovery: ConnectionRecovery,
 }
 
 /// Sharding configuration from the cluster.
@@ -107,6 +109,7 @@ pub struct ClusterConfig<'a> {
     pub expanded_explain: bool,
     pub pub_sub_channel_size: usize,
     pub query_parser_enabled: bool,
+    pub connection_recovery: ConnectionRecovery,
 }
 
 impl<'a> ClusterConfig<'a> {
@@ -146,6 +149,7 @@ impl<'a> ClusterConfig<'a> {
             expanded_explain: general.expanded_explain,
             pub_sub_channel_size: general.pub_sub_channel_size,
             query_parser_enabled: general.query_parser_enabled,
+            connection_recovery: general.connection_recovery,
         }
     }
 }
@@ -176,6 +180,7 @@ impl Cluster {
             expanded_explain,
             pub_sub_channel_size,
             query_parser_enabled,
+            connection_recovery,
         } = config;
 
         Self {
@@ -207,6 +212,7 @@ impl Cluster {
             expanded_explain,
             pub_sub_channel_size,
             query_parser_enabled,
+            connection_recovery,
         }
     }
 
@@ -304,6 +310,10 @@ impl Cluster {
         &self.prepared_statements
     }
 
+    pub fn connection_recovery(&self) -> &ConnectionRecovery {
+        &self.connection_recovery
+    }
+
     pub fn dry_run(&self) -> bool {
         self.dry_run
     }

pgdog/src/backend/pool/config.rs

@@ -2,6 +2,7 @@
 
 use std::time::Duration;
 
+use pgdog_config::pooling::ConnectionRecovery;
 use serde::{Deserialize, Serialize};
 
 use crate::config::{Database, General, PoolerMode, User};
@@ -59,6 +60,8 @@ pub struct Config {
     pub prepared_statements_limit: usize,
     /// Stats averaging period.
     pub stats_period: Duration,
+    /// Recovery algo.
+    pub connection_recovery: ConnectionRecovery,
 }
 
 impl Config {
@@ -186,6 +189,7 @@ impl Config {
             prepared_statements_limit: general.prepared_statements_limit,
             stats_period: Duration::from_millis(general.stats_period),
             bannable: !is_only_replica,
+            connection_recovery: general.connection_recovery,
             ..Default::default()
         }
     }
@@ -219,6 +223,7 @@ impl Default for Config {
             prepared_statements_limit: usize::MAX,
             stats_period: Duration::from_millis(15_000),
             dns_ttl: Duration::from_millis(60_000),
+            connection_recovery: ConnectionRecovery::Recover,
         }
     }
 }