From a53db2a238886e0cb8df9d64560dbdd8faa6109b Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Mon, 3 Aug 2026 21:29:25 -0700 Subject: [PATCH 1/5] feat: role-specific pool configs --- pgdog-config/src/database.rs | 6 ++ pgdog-config/src/lib.rs | 2 + pgdog-config/src/role_config.rs | 52 ++++++++++++ pgdog-config/src/users.rs | 15 ++++ pgdog-stats/src/pool.rs | 36 +++++++- pgdog/src/backend/pool/config.rs | 125 ++++++++++++++-------------- pgdog/src/backend/pool/inner.rs | 64 +++++++++++++- pgdog/src/backend/pool/lb/mod.rs | 27 ++++-- pgdog/src/backend/pool/mod.rs | 2 + pgdog/src/backend/pool/monitor.rs | 6 +- pgdog/src/backend/pool/pool_impl.rs | 6 ++ pgdog/src/backend/pool/role.rs | 37 ++++++++ 12 files changed, 298 insertions(+), 80 deletions(-) create mode 100644 pgdog-config/src/role_config.rs create mode 100644 pgdog/src/backend/pool/role.rs diff --git a/pgdog-config/src/database.rs b/pgdog-config/src/database.rs index bb31809f0..09ed66e7c 100644 --- a/pgdog-config/src/database.rs +++ b/pgdog-config/src/database.rs @@ -6,6 +6,8 @@ use std::{ str::FromStr, }; +use crate::RoleConfig; + use super::pooling::PoolerMode; /// How aggressive the query parser should be in determining read vs. write queries. @@ -207,6 +209,10 @@ pub struct Database { /// Used for weighted load balancing. #[serde(default = "Database::lb_weight")] pub lb_weight: u8, + + /// Role-specific settings. + #[serde(flatten, default)] + pub role_config: RoleConfig, } impl Database { diff --git a/pgdog-config/src/lib.rs b/pgdog-config/src/lib.rs index 142713fec..b70ecb55d 100644 --- a/pgdog-config/src/lib.rs +++ b/pgdog-config/src/lib.rs @@ -12,6 +12,7 @@ pub mod overrides; pub mod pooling; pub mod replication; pub mod rewrite; +pub mod role_config; pub mod sharding; pub mod system_catalogs; #[cfg(test)] @@ -37,6 +38,7 @@ pub use overrides::Overrides; pub use pooling::{PoolerMode, PreparedStatements}; pub use replication::*; pub use rewrite::{Rewrite, RewriteMode}; +pub use role_config::RoleConfig; pub use sharding::*; pub use system_catalogs::system_catalogs; pub use users::{Admin, Plugin, ServerAuth, User, Users}; diff --git a/pgdog-config/src/role_config.rs b/pgdog-config/src/role_config.rs new file mode 100644 index 000000000..f34d013b5 --- /dev/null +++ b/pgdog-config/src/role_config.rs @@ -0,0 +1,52 @@ +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// Database-role specific configuration. +#[derive( + Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema, PartialOrd, Eq, Ord, +)] +#[serde(rename_all = "snake_case", deny_unknown_fields)] +pub struct RoleConfig { + /// Maximum pool size for the primary database in the cluster. + pub pool_size_primary: Option, + /// Maximum pool size for the replica databases in the cluster. + pub pool_size_replica: Option, + /// Minimum pool size for the primary database in the cluster. + pub min_pool_size_primary: Option, + /// Minimum pool size for the replica databases in the cluster. + pub min_pool_size_replica: Option, + /// Idle connection timeout for the primary database in the cluster. + pub idle_timeout_primary: Option, + /// Idle connection timeout for the replica databases in the cluster. + pub idle_timeout_replica: Option, +} + +#[cfg(test)] +mod test { + use super::super::Users; + + #[test] + fn test_role_config_with_users() { + let config = r#" +[[users]] +name = "pgdog" +database = "prod" +pool_size_primary = 10 +min_pool_size_primary = 2 +idle_timeout_primary = 500 +idle_timeout_replica = 1_000 + "#; + + let users: Users = toml::from_str(&config).unwrap(); + assert_eq!(users.users[0].role_config.pool_size_primary, Some(10)); + assert_eq!( + users.users[0].role_config.min_pool_size_primary, + Some(2) + ); + assert_eq!( + users.users[0].role_config.idle_timeout_primary, + Some(500) + ); + assert_eq!(users.users[0].role_config.idle_timeout_replica, Some(1_000)); + } +} diff --git a/pgdog-config/src/users.rs b/pgdog-config/src/users.rs index 9bf8a52f0..b24719cea 100644 --- a/pgdog-config/src/users.rs +++ b/pgdog-config/src/users.rs @@ -7,6 +7,7 @@ use tracing::warn; use super::core::Config; use super::pooling::PoolerMode; +use crate::RoleConfig; use crate::util::random_string; use schemars::JsonSchema; @@ -390,6 +391,20 @@ pub struct User { /// Maximum random adjustment applied to `server_lifetime` per backend connection (milliseconds). /// Overrides the database-level and general-level `server_lifetime_jitter` setting for this user. pub server_lifetime_jitter: Option, + + /// Role-specific configuration settings. + /// + /// These are flattened, so they appear directly in the [[users]] entry, e.g., + /// + /// ```toml + /// [[users]] + /// name = "pgdog" + /// database = "pgdog" + /// pool_size_primary = 10 + /// ``` + /// + #[serde(flatten, default)] + pub role_config: RoleConfig, } impl User { diff --git a/pgdog-stats/src/pool.rs b/pgdog-stats/src/pool.rs index 16f4112a1..8b2382fc2 100644 --- a/pgdog-stats/src/pool.rs +++ b/pgdog-stats/src/pool.rs @@ -3,7 +3,7 @@ use std::{ time::Duration, }; -use pgdog_config::{PoolerMode, PreparedStatements, pooling::ConnectionRecovery}; +use pgdog_config::{PoolerMode, PreparedStatements, Role, pooling::ConnectionRecovery}; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; @@ -275,14 +275,26 @@ pub struct State { pub struct Config { /// Minimum connections that should be in the pool. pub min: usize, + /// Minimum connections that should be in the pool if it's a primary. + pub min_primary: Option, + /// Minimum connections that should be in the pool if it's a replica. + pub min_replica: Option, /// Maximum connections allowed in the pool. pub max: usize, + /// Maximum connection allowed in the pool if its a primary. + pub max_primary: Option, + /// Maximum connections allowed in the pool if its a replica. + pub max_replica: Option, /// How long to wait for a connection before giving up. pub checkout_timeout: Duration, // ms /// Interval duration of DNS cache refresh. pub dns_ttl: Duration, // ms /// Close connections that have been idle for longer than this. pub idle_timeout: Duration, // ms + /// Close primary connections that have been idle for longer than this. + pub idle_timeout_primary: Option, // ms + /// Close replica connections that have been idle for longer than this. + pub idle_timeout_replica: Option, // ms /// How long to wait for connections to be created. pub connect_timeout: Duration, // ms /// How many times to attempt a connection before returning an error. @@ -348,13 +360,35 @@ pub struct Config { pub prepared_statements_level: PreparedStatements, } +pub struct RoleSpecificConfig { + pub value: T, + pub value_primary: Option, + pub value_replica: Option, +} + +impl RoleSpecificConfig { + /// Get value given role. + pub fn value(&self, role: Role) -> T { + match role { + Role::Auto | Role::Replica => self.value_replica.unwrap_or(self.value), + Role::Primary => self.value_primary.unwrap_or(self.value), + } + } +} + impl Default for Config { fn default() -> Self { Self { min: 1, + min_primary: None, + min_replica: None, max: 10, + max_primary: None, + max_replica: None, checkout_timeout: Duration::from_millis(5_000), idle_timeout: Duration::from_millis(60_000), + idle_timeout_primary: None, + idle_timeout_replica: None, connect_timeout: Duration::from_millis(5_000), connect_attempts: 1, connect_attempt_delay: Duration::from_millis(10), diff --git a/pgdog/src/backend/pool/config.rs b/pgdog/src/backend/pool/config.rs index 26a8363d1..c1fcfb0bd 100644 --- a/pgdog/src/backend/pool/config.rs +++ b/pgdog/src/backend/pool/config.rs @@ -31,70 +31,6 @@ impl DerefMut for Config { } impl Config { - /// Connect timeout duration. - pub fn connect_timeout(&self) -> Duration { - self.connect_timeout - } - - /// Checkout timeout duration. - pub fn checkout_timeout(&self) -> Duration { - self.checkout_timeout - } - - /// DNS TTL duration. - pub fn dns_ttl(&self) -> Duration { - self.dns_ttl - } - - /// Idle timeout duration. - pub fn idle_timeout(&self) -> Duration { - self.idle_timeout - } - - /// Max age duration. - pub fn max_age(&self) -> Duration { - self.max_age - } - - /// Healthcheck timeout. - pub fn healthcheck_timeout(&self) -> Duration { - self.healthcheck_timeout - } - - /// How long to wait between healtchecks. - pub fn healthcheck_interval(&self) -> Duration { - self.healthcheck_interval - } - - /// Idle healtcheck interval. - pub fn idle_healthcheck_interval(&self) -> Duration { - self.idle_healthcheck_interval - } - - /// Idle healtcheck delay. - pub fn idle_healthcheck_delay(&self) -> Duration { - self.idle_healthcheck_delay - } - - /// Ban timeout. - pub fn ban_timeout(&self) -> Duration { - self.ban_timeout - } - - /// Rollback timeout. - pub fn rollback_timeout(&self) -> Duration { - self.rollback_timeout - } - - /// Read timeout. - pub fn read_timeout(&self) -> Duration { - self.read_timeout - } - - pub fn query_timeout(&self) -> Duration { - self.query_timeout - } - /// Create from database/user configuration. pub fn new(general: &General, database: &Database, user: &User, is_only_replica: bool) -> Self { Self { @@ -102,9 +38,29 @@ impl Config { min: user .min_pool_size .unwrap_or(database.min_pool_size.unwrap_or(general.min_pool_size)), + min_primary: if let Some(user_setting) = user.role_config.min_pool_size_primary { + Some(user_setting) + } else { + database.role_config.min_pool_size_primary + }, + min_replica: if let Some(user_setting) = user.role_config.min_pool_size_replica { + Some(user_setting) + } else { + database.role_config.min_pool_size_replica + }, max: user .pool_size .unwrap_or(database.pool_size.unwrap_or(general.default_pool_size)), + max_primary: if let Some(user_setting) = user.role_config.pool_size_primary { + Some(user_setting) + } else { + database.role_config.pool_size_primary + }, + max_replica: if let Some(user_setting) = user.role_config.pool_size_replica { + Some(user_setting) + } else { + database.role_config.pool_size_replica + }, max_age: Duration::from_millis( user.server_lifetime .unwrap_or(database.server_lifetime.unwrap_or(general.server_lifetime)), @@ -143,6 +99,26 @@ impl Config { user.idle_timeout .unwrap_or(database.idle_timeout.unwrap_or(general.idle_timeout)), ), + idle_timeout_primary: if let Some(user_setting) = + user.role_config.idle_timeout_primary + { + Some(Duration::from_millis(user_setting)) + } else { + database + .role_config + .idle_timeout_primary + .map(Duration::from_millis) + }, + idle_timeout_replica: if let Some(user_setting) = + user.role_config.idle_timeout_replica + { + Some(Duration::from_millis(user_setting)) + } else { + database + .role_config + .idle_timeout_replica + .map(Duration::from_millis) + }, read_only: user .read_only .unwrap_or(database.read_only.unwrap_or_default()), @@ -231,6 +207,27 @@ mod test { assert!(config.read_only); } + #[test] + fn test_role_specific_pool_config_user_takes_precedence_over_database() { + let general = General::default(); + let mut user = User::default(); + user.role_config.min_pool_size_primary = Some(2); + user.role_config.idle_timeout_primary = Some(20); + + let mut database = Database::default(); + database.role_config.min_pool_size_primary = Some(3); + database.role_config.min_pool_size_replica = Some(4); + database.role_config.idle_timeout_primary = Some(30); + database.role_config.idle_timeout_replica = Some(40); + + let config = Config::new(&general, &database, &user, false); + + assert_eq!(config.min_primary, Some(2)); + assert_eq!(config.min_replica, Some(4)); + assert_eq!(config.idle_timeout_primary, Some(Duration::from_millis(20))); + assert_eq!(config.idle_timeout_replica, Some(Duration::from_millis(40))); + } + #[test] fn test_jitter_falls_through_general_to_database_to_user() { let general = General { diff --git a/pgdog/src/backend/pool/inner.rs b/pgdog/src/backend/pool/inner.rs index a44004f8b..1d3fbbfdd 100644 --- a/pgdog/src/backend/pool/inner.rs +++ b/pgdog/src/backend/pool/inner.rs @@ -3,11 +3,14 @@ use std::cmp::max; use std::collections::VecDeque; use std::fmt::Display; +use std::time::Duration; use crate::backend::{ConnectReason, DisconnectReason}; use crate::backend::{Server, stats::Counts as BackendCounts}; use crate::net::messages::{BackendKeyData, BackendPid, FrontendPid}; +use pgdog_config::Role; +use pgdog_stats::RoleSpecificConfig; use tokio::time::Instant; use super::{Config, Error, Pool, Request, Stats, Taken, Waiter, lsn_monitor::ReplicaLag}; @@ -49,6 +52,8 @@ pub(super) struct Inner { /// Bumped each time Vault credentials rotate. Connections stamped with /// an older generation are closed on check-in rather than reused. pub(super) credentials_generation: u64, + /// Pool role. + pub(super) role: Role, } impl std::fmt::Debug for Inner { @@ -82,6 +87,7 @@ impl Inner { id, replica_lag: ReplicaLag::default(), credentials_generation: 0, + role: Role::Auto, } } /// Total number of connections managed by the pool. @@ -151,13 +157,34 @@ impl Inner { /// Minimum number of connections the pool should keep open. #[inline] pub(super) fn min(&self) -> usize { - self.config.min + RoleSpecificConfig { + value: self.config.min, + value_primary: self.config.min_primary, + value_replica: self.config.min_replica, + } + .value(self.role) } /// Maximum number of connections in the pool. #[inline] pub(super) fn max(&self) -> usize { - self.config.max + RoleSpecificConfig { + value: self.config.max, + value_primary: self.config.max_primary, + value_replica: self.config.max_replica, + } + .value(self.role) + } + + /// Close connections that have been idle for longer than this. + #[inline] + pub(super) fn idle_timeout(&self) -> Duration { + RoleSpecificConfig { + value: self.config.idle_timeout, + value_primary: self.config.idle_timeout_primary, + value_replica: self.config.idle_timeout_replica, + } + .value(self.role) } /// The pool should create more connections now. @@ -216,7 +243,7 @@ impl Inner { #[inline] pub(crate) fn close_idle(&mut self, now: Instant) -> usize { let (mut remove, mut removed) = (self.can_remove(), 0); - let idle_timeout = self.config.idle_timeout; + let idle_timeout = self.idle_timeout(); self.idle_connections.retain_mut(|c| { let idle_for = c.idle_for(now); @@ -451,6 +478,14 @@ impl Inner { let _ = waiter.tx.send(Err(err)); } } + + #[inline] + pub(super) fn set_role(&mut self, role: Role) -> bool { + let changed = role != self.role; + self.role = role; + + changed + } } /// Result of connection check into the pool. @@ -661,6 +696,29 @@ mod test { )); } + #[test] + fn test_role_specific_min_and_idle_timeout() { + let mut inner = Inner::default(); + inner.config.min = 1; + inner.config.min_primary = Some(2); + inner.config.min_replica = Some(3); + inner.config.idle_timeout = Duration::from_millis(10); + inner.config.idle_timeout_primary = Some(Duration::from_millis(20)); + inner.config.idle_timeout_replica = Some(Duration::from_millis(30)); + + // Auto targets use replica settings until role detection completes. + assert_eq!(inner.min(), 3); + assert_eq!(inner.idle_timeout(), Duration::from_millis(30)); + + inner.set_role(Role::Primary); + assert_eq!(inner.min(), 2); + assert_eq!(inner.idle_timeout(), Duration::from_millis(20)); + + inner.set_role(Role::Replica); + assert_eq!(inner.min(), 3); + assert_eq!(inner.idle_timeout(), Duration::from_millis(30)); + } + #[test] fn test_should_not_create_at_max() { let mut inner = Inner { diff --git a/pgdog/src/backend/pool/lb/mod.rs b/pgdog/src/backend/pool/lb/mod.rs index 91e39a198..a4d4c0805 100644 --- a/pgdog/src/backend/pool/lb/mod.rs +++ b/pgdog/src/backend/pool/lb/mod.rs @@ -3,7 +3,7 @@ use std::{ sync::{ Arc, - atomic::{AtomicI64, AtomicU8, AtomicUsize, Ordering}, + atomic::{AtomicI64, AtomicUsize, Ordering}, }, time::{Duration, SystemTime}, }; @@ -18,7 +18,7 @@ use crate::{ net::Parameters, }; -use super::{Error, Guard, Pool, PoolConfig, Request}; +use super::{Error, Guard, Pool, PoolConfig, PoolRole, Request}; use crate::util::safe_timeout; pub mod ban; @@ -38,7 +38,7 @@ mod test; pub struct Target { pub pool: Pool, pub ban: Ban, - role: Arc, + role: PoolRole, pub health: TargetHealth, /// Smooth weighted round-robin current weight tracker. current_weight: Arc, @@ -47,9 +47,13 @@ pub struct Target { impl Target { pub(super) fn new(pool: Pool, role: Role) -> Self { let ban = Ban::new(&pool); + + // Set pool to last known role. + pool.set_role(role); + Self { ban, - role: Arc::new(AtomicU8::new(role.into())), + role: PoolRole::new(role), health: pool.inner().health.clone(), pool, current_weight: Arc::new(AtomicI64::new(0)), @@ -58,15 +62,20 @@ impl Target { /// Get role. pub(super) fn role(&self) -> Role { - let role = self.role.load(Ordering::Relaxed); - role.try_into().expect("valid role") + self.role.role() } /// Set role. pub(super) fn set_role(&self, role: Role) -> bool { - let value = u8::from(role); - let old = self.role.swap(value, Ordering::Relaxed); - value != old + let lb = self.role.set_role(role); + let pool = self.pool.set_role(role); + + debug_assert_eq!( + lb, pool, + "pool and lb role must agree: lb={lb}, pool={pool}" + ); + + lb && pool } } diff --git a/pgdog/src/backend/pool/mod.rs b/pgdog/src/backend/pool/mod.rs index 5550a0e68..dc85de745 100644 --- a/pgdog/src/backend/pool/mod.rs +++ b/pgdog/src/backend/pool/mod.rs @@ -19,6 +19,7 @@ pub mod monitor; pub mod password; pub mod pool_impl; pub mod request; +pub mod role; pub mod shard; pub mod state; pub mod stats; @@ -40,6 +41,7 @@ pub use monitor::Monitor; pub use password::Password; pub use pool_impl::Pool; pub use request::Request; +pub(crate) use role::PoolRole; pub use shard::Shard; pub use state::State; pub use stats::Stats; diff --git a/pgdog/src/backend/pool/monitor.rs b/pgdog/src/backend/pool/monitor.rs index c6d36bb04..fcbaeb871 100644 --- a/pgdog/src/backend/pool/monitor.rs +++ b/pgdog/src/backend/pool/monitor.rs @@ -100,8 +100,8 @@ impl Monitor { let lock = pool.lock(); let config = lock.config(); ( - config.idle_healthcheck_delay(), - config.idle_healthcheck_interval(), + config.idle_healthcheck_delay, + config.idle_healthcheck_interval, config.replication_mode, ) }; @@ -256,7 +256,7 @@ impl Monitor { /// /// Runs regularly and ensures the pool triggers health checks on idle connections. async fn healthchecks(pool: Pool) { - let mut tick = safe_interval(pool.lock().config().idle_healthcheck_interval()); + let mut tick = safe_interval(pool.lock().config().idle_healthcheck_interval); let comms = pool.comms(); debug!("health checks running [{}]", pool.addr()); diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index 1b161e22d..51cda1013 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -8,6 +8,7 @@ use futures::future::try_join_all; use once_cell::sync::{Lazy, OnceCell}; use parking_lot::RwLock; use parking_lot::{Mutex, RawMutex, lock_api::MutexGuard}; +use pgdog_config::Role; use tokio::sync::Notify; use tokio::time::Instant; use tracing::{debug, error}; @@ -488,6 +489,11 @@ impl Pool { *self.inner().lsn_stats.read() } + /// Set pool role returning true if the role changed. + pub(crate) fn set_role(&self, role: Role) -> bool { + self.lock().set_role(role) + } + /// Update pool configuration used in internals. #[cfg(test)] pub(crate) fn update_config(&self, config: Config) { diff --git a/pgdog/src/backend/pool/role.rs b/pgdog/src/backend/pool/role.rs new file mode 100644 index 000000000..949650111 --- /dev/null +++ b/pgdog/src/backend/pool/role.rs @@ -0,0 +1,37 @@ +//! Connection pool role, +//! powered by an atomic that allows +//! it to be changed without re-creating the struct. + +use std::sync::{ + Arc, + atomic::{AtomicU8, Ordering}, +}; + +use pgdog_config::Role; + +/// A pool/LB target role, +/// powered by an atomic so it's easy +/// to change it. +#[derive(Debug, Clone, Default)] +pub(crate) struct PoolRole { + role: Arc, +} + +impl PoolRole { + pub(crate) fn new(role: Role) -> Self { + Self { + role: Arc::new(AtomicU8::new(role.try_into().expect("valid role"))), + } + } + + pub(crate) fn role(&self) -> Role { + let role = self.role.load(Ordering::Relaxed); + role.try_into().expect("valid role") + } + + pub(crate) fn set_role(&self, role: Role) -> bool { + self.role + .swap(role.try_into().expect("valid role"), Ordering::Relaxed) + != u8::try_from(role).expect("valid role") + } +} From 4486d73e9e49e153e80849fc207c0b0bbb6d11cc Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Mon, 3 Aug 2026 21:41:29 -0700 Subject: [PATCH 2/5] clippy --- pgdog-config/src/role_config.rs | 10 ++-------- pgdog/src/backend/pool/inner.rs | 5 ++++- pgdog/src/backend/pool/role.rs | 6 ++---- 3 files changed, 8 insertions(+), 13 deletions(-) diff --git a/pgdog-config/src/role_config.rs b/pgdog-config/src/role_config.rs index f34d013b5..e1d603422 100644 --- a/pgdog-config/src/role_config.rs +++ b/pgdog-config/src/role_config.rs @@ -39,14 +39,8 @@ idle_timeout_replica = 1_000 let users: Users = toml::from_str(&config).unwrap(); assert_eq!(users.users[0].role_config.pool_size_primary, Some(10)); - assert_eq!( - users.users[0].role_config.min_pool_size_primary, - Some(2) - ); - assert_eq!( - users.users[0].role_config.idle_timeout_primary, - Some(500) - ); + assert_eq!(users.users[0].role_config.min_pool_size_primary, Some(2)); + assert_eq!(users.users[0].role_config.idle_timeout_primary, Some(500)); assert_eq!(users.users[0].role_config.idle_timeout_replica, Some(1_000)); } } diff --git a/pgdog/src/backend/pool/inner.rs b/pgdog/src/backend/pool/inner.rs index 1d3fbbfdd..ed30ded65 100644 --- a/pgdog/src/backend/pool/inner.rs +++ b/pgdog/src/backend/pool/inner.rs @@ -698,7 +698,10 @@ mod test { #[test] fn test_role_specific_min_and_idle_timeout() { - let mut inner = Inner::default(); + let mut inner = Inner { + role: Role::Auto, + ..Default::default() + }; inner.config.min = 1; inner.config.min_primary = Some(2); inner.config.min_replica = Some(3); diff --git a/pgdog/src/backend/pool/role.rs b/pgdog/src/backend/pool/role.rs index 949650111..54124341f 100644 --- a/pgdog/src/backend/pool/role.rs +++ b/pgdog/src/backend/pool/role.rs @@ -20,7 +20,7 @@ pub(crate) struct PoolRole { impl PoolRole { pub(crate) fn new(role: Role) -> Self { Self { - role: Arc::new(AtomicU8::new(role.try_into().expect("valid role"))), + role: Arc::new(AtomicU8::new(role.into())), } } @@ -30,8 +30,6 @@ impl PoolRole { } pub(crate) fn set_role(&self, role: Role) -> bool { - self.role - .swap(role.try_into().expect("valid role"), Ordering::Relaxed) - != u8::try_from(role).expect("valid role") + self.role.swap(role.into(), Ordering::Relaxed) != u8::from(role) } } From 6a34cfa6bcdcca129aa9300db8cbce7d922ade5e Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Mon, 3 Aug 2026 21:44:20 -0700 Subject: [PATCH 3/5] comment --- pgdog-config/src/role_config.rs | 27 ++++++++++++++++++++++++++- pgdog/src/backend/pool/role.rs | 5 +++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/pgdog-config/src/role_config.rs b/pgdog-config/src/role_config.rs index e1d603422..af4661071 100644 --- a/pgdog-config/src/role_config.rs +++ b/pgdog-config/src/role_config.rs @@ -23,7 +23,32 @@ pub struct RoleConfig { #[cfg(test)] mod test { - use super::super::Users; + use super::super::{Config, Users}; + + #[test] + fn test_role_config_with_databases() { + let config = r#" +[[databases]] +name = "prod" +host = "127.0.0.1" +pool_size_primary = 10 +pool_size_replica = 20 +min_pool_size_primary = 2 +min_pool_size_replica = 4 +idle_timeout_primary = 500 +idle_timeout_replica = 1_000 + "#; + + let config: Config = toml::from_str(config).unwrap(); + let role_config = &config.databases[0].role_config; + + assert_eq!(role_config.pool_size_primary, Some(10)); + assert_eq!(role_config.pool_size_replica, Some(20)); + assert_eq!(role_config.min_pool_size_primary, Some(2)); + assert_eq!(role_config.min_pool_size_replica, Some(4)); + assert_eq!(role_config.idle_timeout_primary, Some(500)); + assert_eq!(role_config.idle_timeout_replica, Some(1_000)); + } #[test] fn test_role_config_with_users() { diff --git a/pgdog/src/backend/pool/role.rs b/pgdog/src/backend/pool/role.rs index 54124341f..ca958e687 100644 --- a/pgdog/src/backend/pool/role.rs +++ b/pgdog/src/backend/pool/role.rs @@ -18,17 +18,22 @@ pub(crate) struct PoolRole { } impl PoolRole { + /// Create new pool role. + /// + /// The role can be changed atomically. pub(crate) fn new(role: Role) -> Self { Self { role: Arc::new(AtomicU8::new(role.into())), } } + /// Get current role value. pub(crate) fn role(&self) -> Role { let role = self.role.load(Ordering::Relaxed); role.try_into().expect("valid role") } + /// Change role atomically, returning true if the role value changed. pub(crate) fn set_role(&self, role: Role) -> bool { self.role.swap(role.into(), Ordering::Relaxed) != u8::from(role) } From f0960981b4c31f26b863eebb02c1beaaa0b64be3 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Mon, 3 Aug 2026 21:46:06 -0700 Subject: [PATCH 4/5] clippy --- pgdog-config/src/role_config.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgdog-config/src/role_config.rs b/pgdog-config/src/role_config.rs index af4661071..95ef79867 100644 --- a/pgdog-config/src/role_config.rs +++ b/pgdog-config/src/role_config.rs @@ -62,7 +62,7 @@ idle_timeout_primary = 500 idle_timeout_replica = 1_000 "#; - let users: Users = toml::from_str(&config).unwrap(); + let users: Users = toml::from_str(config).unwrap(); assert_eq!(users.users[0].role_config.pool_size_primary, Some(10)); assert_eq!(users.users[0].role_config.min_pool_size_primary, Some(2)); assert_eq!(users.users[0].role_config.idle_timeout_primary, Some(500)); From ff676a0e295d15e77159a0620325b4dacad44d28 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Mon, 3 Aug 2026 21:51:44 -0700 Subject: [PATCH 5/5] schema --- .schema/pgdog.schema.json | 54 +++++++++++++++++++++++++++++++++++++++ .schema/users.schema.json | 54 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+) diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index 745cb39d2..d2cc776d8 100644 --- a/.schema/pgdog.schema.json +++ b/.schema/pgdog.schema.json @@ -416,6 +416,24 @@ "format": "uint64", "minimum": 0 }, + "idle_timeout_primary": { + "description": "Idle connection timeout for the primary database in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 + }, + "idle_timeout_replica": { + "description": "Idle connection timeout for the replica databases in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 + }, "lb_weight": { "description": "Used for weighted load balancing.", "type": "integer", @@ -442,6 +460,24 @@ "format": "uint", "minimum": 0 }, + "min_pool_size_primary": { + "description": "Minimum pool size for the primary database in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 0 + }, + "min_pool_size_replica": { + "description": "Minimum pool size for the replica databases in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 0 + }, "name": { "description": "Name of your database. Clients that connect to PgDog will need to use this name to refer to the database. For multiple entries that are part of the same cluster, use the same value.\n\n", "type": "string" @@ -462,6 +498,24 @@ "format": "uint", "minimum": 0 }, + "pool_size_primary": { + "description": "Maximum pool size for the primary database in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 0 + }, + "pool_size_replica": { + "description": "Maximum pool size for the replica databases in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 0 + }, "pooler_mode": { "description": "Overrides the `pooler_mode` setting. Connections to this database will use this connection pool mode.\n\n", "anyOf": [ diff --git a/.schema/users.schema.json b/.schema/users.schema.json index d9cc9bf4c..3afbf8796 100644 --- a/.schema/users.schema.json +++ b/.schema/users.schema.json @@ -143,6 +143,24 @@ "format": "uint64", "minimum": 0 }, + "idle_timeout_primary": { + "description": "Idle connection timeout for the primary database in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 + }, + "idle_timeout_replica": { + "description": "Idle connection timeout for the replica databases in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 + }, "lock_timeout": { "description": "Lock timeout.\n\nSets the `lock_timeout` on all server connections at connection creation.\nAborts any statement that waits longer than the specified duration to acquire a lock.\nUnlike `statement_timeout`, this only counts time spent waiting for locks, not execution time.\nRecommended for replication destination connections to prevent cross-shard deadlocks\nfrom hanging indefinitely.\n\n**Note:** Nothing is preventing the user from manually changing this setting at runtime,\ne.g., by running `SET lock_timeout TO 0`;\n\n", "type": [ @@ -161,6 +179,24 @@ "format": "uint", "minimum": 0 }, + "min_pool_size_primary": { + "description": "Minimum pool size for the primary database in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 0 + }, + "min_pool_size_replica": { + "description": "Minimum pool size for the replica databases in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 0 + }, "name": { "description": "Name of the user. Clients that connect to PgDog will need to use this username.\n\n", "type": "string" @@ -196,6 +232,24 @@ "format": "uint", "minimum": 0 }, + "pool_size_primary": { + "description": "Maximum pool size for the primary database in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 0 + }, + "pool_size_replica": { + "description": "Maximum pool size for the replica databases in the cluster.", + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 0 + }, "pooler_mode": { "description": "Overrides [`pooler_mode`](https://docs.pgdog.dev/configuration/pgdog.toml/general/) for this user. This allows users in [session mode](https://docs.pgdog.dev/features/session-mode/) to connect to the same PgDog instance as users in [transaction mode](https://docs.pgdog.dev/features/transaction-mode/).\n\n", "anyOf": [