From d61d2dcda089f83627726c57181e34e7f7f00864 Mon Sep 17 00:00:00 2001 From: meskill <8974488+meskill@users.noreply.github.com> Date: Thu, 6 Aug 2026 19:29:35 +0000 Subject: [PATCH 1/3] feat(prepared_statements): add new ttl config --- .schema/pgdog.schema.json | 16 + example.pgdog.toml | 16 + pgdog-config/src/general.rs | 151 +++++++ pgdog-stats/src/pool.rs | 33 +- pgdog/src/backend/pool/config.rs | 34 +- pgdog/src/backend/pool/guard.rs | 36 +- pgdog/src/backend/pool/pool_impl.rs | 5 +- pgdog/src/backend/pool/test/mod.rs | 5 +- pgdog/src/backend/prepared_statements.rs | 421 +++++++++++++++--- pgdog/src/backend/server.rs | 120 +++-- pgdog/src/frontend/prepared_statements/mod.rs | 19 +- pgdog/src/util.rs | 2 + pgdog/src/util/time.rs | 103 +++++ 13 files changed, 861 insertions(+), 100 deletions(-) create mode 100644 pgdog/src/util/time.rs diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index a334988f3..4f0cf6bd3 100644 --- a/.schema/pgdog.schema.json +++ b/.schema/pgdog.schema.json @@ -83,6 +83,8 @@ "port": 6432, "prepared_statements": "extended", "prepared_statements_limit": 9223372036854775807, + "prepared_statements_ttl": 300000, + "prepared_statements_ttl_jitter": 30000, "pub_sub_channel_size": 0, "query_cache_limit": 1000, "query_log": null, @@ -1028,6 +1030,20 @@ "default": 9223372036854775807, "minimum": 0 }, + "prepared_statements_ttl": { + "description": "How long a prepared statement is allowed to stay prepared on a server connection, in milliseconds.\n\n**Note:** Expired statements are closed and prepared again the next time they are used. This stops stale execution plans from staying in Postgres. Set to `0` to let statements stay prepared forever.\n\n_Default:_ `300000` (5 minutes)\n\n", + "type": "integer", + "format": "uint64", + "default": 300000, + "minimum": 0 + }, + "prepared_statements_ttl_jitter": { + "description": "Maximum random adjustment applied to `prepared_statements_ttl` per prepared\nstatement, in milliseconds. Each statement expires at a point sampled uniformly\nfrom `[ttl - jitter, ttl + jitter]`, chosen once when the statement is prepared.\n\n**Note:** Clients run their whole statement set together, so a small jitter\nleaves them expiring in lockstep for hours. Keep this at a meaningful fraction\nof the TTL.\n\n_Default:_ `30000` (30 seconds)\n\n", + "type": "integer", + "format": "uint64", + "default": 30000, + "minimum": 0 + }, "pub_sub_channel_size": { "description": "Enables support for pub/sub and configures the size of the background task queue.\n\n", "type": "integer", diff --git a/example.pgdog.toml b/example.pgdog.toml index d110dd659..c812d2c2b 100644 --- a/example.pgdog.toml +++ b/example.pgdog.toml @@ -197,6 +197,22 @@ query_parser = "on" # Default: unlimited # prepared_statements_limit = 1_000 +# How long a prepared statement can stay prepared on a Postgres server +# connection. Expired statements are closed and prepared again the next +# time they are used, so execution plans don't get stale. +# +# Set to 0 to let statements stay prepared forever. +# +# Default: 300_000 (5 minutes) +# +prepared_statements_ttl = 300_000 +# Random spread applied to prepared_statements_ttl, per statement. Clients +# usually prepare their whole statement set at once, so without this they +# would all expire at the same moment, over and over. +# +# Default: 30_000 (30 seconds) +# +prepared_statements_ttl_jitter = 30_000 # Limit on the number of queries cached in the Abstract Syntax Tree # cache used for query routing and sharding. # diff --git a/pgdog-config/src/general.rs b/pgdog-config/src/general.rs index 5630ef054..d2d3768a5 100644 --- a/pgdog-config/src/general.rs +++ b/pgdog-config/src/general.rs @@ -382,6 +382,30 @@ pub struct General { #[serde(default = "General::prepared_statements_limit")] pub prepared_statements_limit: usize, + /// How long a prepared statement is allowed to stay prepared on a server connection, in milliseconds. + /// + /// **Note:** Expired statements are closed and prepared again the next time they are used. This stops stale execution plans from staying in Postgres. Set to `0` to let statements stay prepared forever. + /// + /// _Default:_ `300000` (5 minutes) + /// + /// + #[serde(default = "General::default_prepared_statements_ttl")] + pub prepared_statements_ttl: u64, + + /// Maximum random adjustment applied to `prepared_statements_ttl` per prepared + /// statement, in milliseconds. Each statement expires at a point sampled uniformly + /// from `[ttl - jitter, ttl + jitter]`, chosen once when the statement is prepared. + /// + /// **Note:** Clients run their whole statement set together, so a small jitter + /// leaves them expiring in lockstep for hours. Keep this at a meaningful fraction + /// of the TTL. + /// + /// _Default:_ `30000` (30 seconds) + /// + /// + #[serde(default = "General::default_prepared_statements_ttl_jitter")] + pub prepared_statements_ttl_jitter: u64, + /// Limit on the number of statements saved in the statement cache used to accelerate query parsing. /// /// _Default:_ `50000` @@ -887,6 +911,8 @@ impl Default for General { regex_parser_limit: Self::regex_parser_limit(), query_parser_engine: QueryParserEngine::default(), prepared_statements_limit: Self::prepared_statements_limit(), + prepared_statements_ttl: Self::default_prepared_statements_ttl(), + prepared_statements_ttl_jitter: Self::default_prepared_statements_ttl_jitter(), query_cache_limit: Self::query_cache_limit(), passthrough_auth: Self::default_passthrough_auth(), connect_timeout: Self::default_connect_timeout(), @@ -1142,6 +1168,30 @@ impl General { self.dns_ttl.map(Duration::from_millis) } + /// How long a statement can stay prepared on a server connection. + /// + /// `None` means statements never expire. + pub fn prepared_statements_ttl(&self) -> Option { + let ttl = self.prepared_statements_ttl; + if ttl == 0 || ttl >= crate::MAX_DURATION.as_millis() as u64 { + None + } else { + Some(Duration::from_millis(ttl)) + } + } + + /// Random spread applied to [`Self::prepared_statements_ttl`]. + /// + /// Clamped to just below the TTL. + pub fn prepared_statements_ttl_jitter(&self) -> Duration { + let ttl = self + .prepared_statements_ttl + .min(crate::MAX_DURATION.as_millis() as u64) + .saturating_sub(1); + + Duration::from_millis(self.prepared_statements_ttl_jitter.min(ttl)) + } + pub fn client_idle_timeout(&self) -> Duration { Duration::from_millis(self.client_idle_timeout) } @@ -1381,6 +1431,20 @@ impl General { Self::env_or_default("PGDOG_PREPARED_STATEMENTS_LIMIT", i64::MAX as usize) } + pub fn default_prepared_statements_ttl() -> u64 { + Self::env_or_default( + "PGDOG_PREPARED_STATEMENTS_TTL", + Duration::from_secs(300).as_millis() as u64, + ) + } + + pub fn default_prepared_statements_ttl_jitter() -> u64 { + Self::env_or_default( + "PGDOG_PREPARED_STATEMENTS_TTL_JITTER", + Duration::from_secs(30).as_millis() as u64, + ) + } + pub fn query_cache_limit() -> usize { Self::env_or_default("PGDOG_QUERY_CACHE_LIMIT", 1_000) } @@ -1494,6 +1558,85 @@ mod tests { use super::*; use crate::test_utils::*; + #[test] + fn test_prepared_statements_ttl_defaults() { + let general = General::default(); + + assert_eq!( + general.prepared_statements_ttl(), + Some(Duration::from_millis(300_000)) + ); + assert_eq!( + general.prepared_statements_ttl_jitter(), + Duration::from_millis(30_000) + ); + } + + #[test] + fn test_prepared_statements_ttl_is_read_in_millis() { + let general = General { + prepared_statements_ttl: 3_600_000, + prepared_statements_ttl_jitter: 5_000, + ..Default::default() + }; + + assert_eq!( + general.prepared_statements_ttl(), + Some(Duration::from_millis(3_600_000)) + ); + assert_eq!( + general.prepared_statements_ttl_jitter(), + Duration::from_millis(5_000) + ); + } + + #[test] + fn test_prepared_statements_ttl_disabled() { + for ttl in [0, crate::MAX_DURATION.as_millis() as u64] { + let general = General { + prepared_statements_ttl: ttl, + ..Default::default() + }; + + assert_eq!( + general.prepared_statements_ttl(), + None, + "ttl {ttl} should disable expiration" + ); + } + } + + #[test] + fn test_prepared_statements_ttl_jitter_is_clamped_below_ttl() { + let general = General { + prepared_statements_ttl: 10_000, + prepared_statements_ttl_jitter: u64::MAX, + ..Default::default() + }; + + assert_eq!( + general.prepared_statements_ttl_jitter(), + Duration::from_millis(9_999) + ); + } + + #[test] + fn test_prepared_statements_ttl_jitter_never_reaches_the_ttl() { + for ttl in [0, 1, 2, 300_000] { + let general = General { + prepared_statements_ttl: ttl, + prepared_statements_ttl_jitter: u64::MAX, + ..Default::default() + }; + + let jitter = general.prepared_statements_ttl_jitter().as_millis() as u64; + assert!( + jitter < ttl.max(1), + "jitter {jitter} must stay below ttl {ttl}" + ); + } + } + #[test] fn test_sharding_lookup_cache_size() { let general = General::default(); @@ -1820,6 +1963,8 @@ mod tests { let _guard = set_env_var("PGDOG_BROADCAST_PORT", "7432"); let _guard = set_env_var("PGDOG_OPENMETRICS_PORT", "9090"); let _guard = set_env_var("PGDOG_PREPARED_STATEMENTS_LIMIT", "1000"); + let _guard = set_env_var("PGDOG_PREPARED_STATEMENTS_TTL", "3600000"); + let _guard = set_env_var("PGDOG_PREPARED_STATEMENTS_TTL_JITTER", "5000"); let _guard = set_env_var("PGDOG_QUERY_CACHE_LIMIT", "500"); let _guard = set_env_var("PGDOG_CONNECT_ATTEMPTS", "3"); let _guard = set_env_var("PGDOG_MIRROR_QUEUE", "256"); @@ -1832,6 +1977,8 @@ mod tests { assert_eq!(General::broadcast_port(), 7432); assert_eq!(General::openmetrics_port(), Some(9090)); assert_eq!(General::prepared_statements_limit(), 1000); + assert_eq!(General::default_prepared_statements_ttl(), 3600000); + assert_eq!(General::default_prepared_statements_ttl_jitter(), 5000); assert_eq!(General::query_cache_limit(), 500); assert_eq!(General::connect_attempts(), 3); assert_eq!(General::mirror_queue(), 256); @@ -1844,6 +1991,8 @@ mod tests { let _guard = remove_env_var("PGDOG_BROADCAST_PORT"); let _guard = remove_env_var("PGDOG_OPENMETRICS_PORT"); let _guard = remove_env_var("PGDOG_PREPARED_STATEMENTS_LIMIT"); + let _guard = remove_env_var("PGDOG_PREPARED_STATEMENTS_TTL"); + let _guard = remove_env_var("PGDOG_PREPARED_STATEMENTS_TTL_JITTER"); let _guard = remove_env_var("PGDOG_QUERY_CACHE_LIMIT"); let _guard = remove_env_var("PGDOG_CONNECT_ATTEMPTS"); let _guard = remove_env_var("PGDOG_MIRROR_QUEUE"); @@ -1856,6 +2005,8 @@ mod tests { assert_eq!(General::broadcast_port(), General::port() + 1); assert_eq!(General::openmetrics_port(), None); assert_eq!(General::prepared_statements_limit(), i64::MAX as usize); + assert_eq!(General::default_prepared_statements_ttl(), 300_000); + assert_eq!(General::default_prepared_statements_ttl_jitter(), 30_000); assert_eq!(General::query_cache_limit(), 1_000); assert_eq!(General::connect_attempts(), 1); assert_eq!(General::mirror_queue(), 128); diff --git a/pgdog-stats/src/pool.rs b/pgdog-stats/src/pool.rs index a993c8ec2..2cb2224d7 100644 --- a/pgdog-stats/src/pool.rs +++ b/pgdog-stats/src/pool.rs @@ -285,6 +285,30 @@ pub struct State { pub lsn_stats: LsnStats, } +/// How a server connection handles prepared statements. +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq)] +pub struct PreparedStatementsConfig { + /// Which statements PgDog keeps prepared on the connection. + pub level: PreparedStatements, + /// Maximum prepared statements per connection. + pub limit: usize, + /// How long a statement can keep a cached plan. `None` never expires. + pub ttl: Option, + /// Random spread applied to `ttl`, per statement. + pub ttl_jitter: Duration, +} + +impl Default for PreparedStatementsConfig { + fn default() -> Self { + Self { + level: PreparedStatements::default(), + limit: usize::MAX, + ttl: None, + ttl_jitter: Duration::ZERO, + } + } +} + /// Pool configuration. #[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq)] pub struct Config { @@ -353,8 +377,8 @@ pub struct Config { pub pooler_mode: PoolerMode, /// Read only mode. pub read_only: bool, - /// Maximum prepared statements per connection. - pub prepared_statements_limit: usize, + /// Prepared statements config. + pub prepared_statements: PreparedStatementsConfig, /// Stats averaging period. pub stats_period: Duration, /// Recovery algo. @@ -371,8 +395,6 @@ pub struct Config { pub resharding_only: bool, /// LB weight. pub lb_weight: u8, - /// Prepared statements level. - pub prepared_statements_level: PreparedStatements, } pub struct RoleSpecificConfig { @@ -424,7 +446,7 @@ impl Default for Config { replication_mode: false, pooler_mode: PoolerMode::default(), read_only: false, - prepared_statements_limit: usize::MAX, + prepared_statements: PreparedStatementsConfig::default(), stats_period: Duration::from_millis(15_000), dns_ttl: Duration::from_millis(60_000), connection_recovery: ConnectionRecovery::Recover, @@ -434,7 +456,6 @@ impl Default for Config { role_detection: false, resharding_only: false, lb_weight: 255, - prepared_statements_level: PreparedStatements::default(), } } } diff --git a/pgdog/src/backend/pool/config.rs b/pgdog/src/backend/pool/config.rs index c1fcfb0bd..db9a592e0 100644 --- a/pgdog/src/backend/pool/config.rs +++ b/pgdog/src/backend/pool/config.rs @@ -122,7 +122,12 @@ impl Config { read_only: user .read_only .unwrap_or(database.read_only.unwrap_or_default()), - prepared_statements_limit: general.prepared_statements_limit, + prepared_statements: pgdog_stats::PreparedStatementsConfig { + level: general.prepared_statements, + limit: general.prepared_statements_limit, + ttl: general.prepared_statements_ttl(), + ttl_jitter: general.prepared_statements_ttl_jitter(), + }, stats_period: Duration::from_millis(general.stats_period), bannable: !is_only_replica, connection_recovery: general.connection_recovery, @@ -132,7 +137,6 @@ impl Config { role_detection: database.role == Role::Auto, resharding_only: database.resharding_only, lb_weight: database.lb_weight, - prepared_statements_level: general.prepared_statements, ..Default::default() }, } @@ -165,6 +169,32 @@ mod test { assert!(config.role_detection); } + #[test] + fn test_prepared_statements_config_from_general() { + let general = General { + prepared_statements_ttl: 60_000, + prepared_statements_limit: 10, + ..Default::default() + }; + + let config = Config::new( + &general, + &create_database(Role::Primary), + &User::default(), + false, + ); + + assert_eq!( + config.prepared_statements.ttl, + Some(Duration::from_millis(60_000)) + ); + assert_eq!(config.prepared_statements.limit, 10); + assert_eq!( + config.prepared_statements.level, + general.prepared_statements + ); + } + #[test] fn test_user_takes_precedence_over_database() { let general = General::default(); diff --git a/pgdog/src/backend/pool/guard.rs b/pgdog/src/backend/pool/guard.rs index 03d27cf76..96357dd23 100644 --- a/pgdog/src/backend/pool/guard.rs +++ b/pgdog/src/backend/pool/guard.rs @@ -240,6 +240,7 @@ mod test { use std::time::Duration; use pgdog_config::pooling::ConnectionRecovery; + use pgdog_stats::PreparedStatementsConfig; use tokio::time::Instant; use crate::util::{safe_sleep, safe_timeout}; @@ -391,7 +392,12 @@ mod test { Box::new(test_server().await), Instant::now(), ); - server.prepared_statements_mut().set_capacity(1); + server + .prepared_statements_mut() + .configure(PreparedStatementsConfig { + limit: 1, + ..Default::default() + }); for i in 0..5 { server @@ -455,7 +461,12 @@ mod test { Box::new(test_server().await), Instant::now(), ); - server.prepared_statements_mut().set_capacity(1); + server + .prepared_statements_mut() + .configure(PreparedStatementsConfig { + limit: 1, + ..Default::default() + }); server .send( @@ -503,7 +514,12 @@ mod test { Box::new(test_server().await), Instant::now(), ); - server.prepared_statements_mut().set_capacity(1); + server + .prepared_statements_mut() + .configure(PreparedStatementsConfig { + limit: 1, + ..Default::default() + }); server .send( @@ -549,7 +565,12 @@ mod test { Box::new(test_server().await), Instant::now(), ); - server.prepared_statements_mut().set_capacity(1); + server + .prepared_statements_mut() + .configure(PreparedStatementsConfig { + limit: 1, + ..Default::default() + }); server .send( @@ -708,7 +729,12 @@ mod test { Box::new(test_server().await), Instant::now(), ); - server.prepared_statements_mut().set_capacity(1); + server + .prepared_statements_mut() + .configure(PreparedStatementsConfig { + limit: 1, + ..Default::default() + }); server .send( diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index 5dc56235e..29ca7e0bd 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -175,10 +175,7 @@ impl Pool { server .prepared_statements_mut() - .set_capacity(self.inner.config.prepared_statements_limit); - server - .prepared_statements_mut() - .set_prepared_statements_level(self.inner.config.prepared_statements_level); + .configure(self.inner.config.prepared_statements); server.set_pooler_mode(self.inner.config.pooler_mode); match self diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index 0d213ded9..afa9ec861 100644 --- a/pgdog/src/backend/pool/test/mod.rs +++ b/pgdog/src/backend/pool/test/mod.rs @@ -48,7 +48,10 @@ pub fn pool_with_prepared_capacity(capacity: usize) -> Pool { inner: pgdog_stats::Config { max: 1, min: 1, - prepared_statements_limit: capacity, + prepared_statements: pgdog_stats::PreparedStatementsConfig { + limit: capacity, + ..Default::default() + }, ..Config::default().inner }, }; diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index dde70b9d4..3bfe44605 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -1,6 +1,11 @@ use lru::LruCache; -use std::{collections::VecDeque, sync::Arc}; +use std::{ + collections::VecDeque, + sync::Arc, + time::{Duration, Instant}, +}; +use crate::util::time::deadline; use crate::{ frontend::{self, prepared_statements::GlobalCache}, net::{ @@ -10,7 +15,7 @@ use crate::{ }, }; use parking_lot::RwLock; -use pgdog_config::PreparedStatements as PreparedStatementsLevel; +use pgdog_stats::PreparedStatementsConfig; use super::{Error, Oids}; use super::{ @@ -18,20 +23,71 @@ use super::{ state::ExecutionCode, }; -/// Approximate memory used by a String. +/// Rough size of one local cache entry. Ignores the LRU node itself, +/// so it undercounts a little. #[inline] -fn str_mem(s: &str) -> usize { - s.len() + std::mem::size_of::() +fn entry_mem(s: &str) -> usize { + s.len() + std::mem::size_of::() + std::mem::size_of::() +} + +/// A statement info prepared on this connection. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct LocalStatement { + /// When this statement should be replanned + deadline: Option, +} + +impl LocalStatement { + fn new(ttl: Option, jitter: Duration) -> Self { + Self { + deadline: ttl.map(|ttl| deadline(ttl, jitter)), + } + } + + pub fn deadline(&self) -> Option { + self.deadline + } + + /// Check for expired + /// + /// If the check is called and deadline is not set then it's marked as expired + /// to cover the case when the TTL was set after the statement creation + pub fn expired(&self, now: Instant) -> bool { + self.deadline.is_none_or(|deadline| deadline <= now) + } +} + +/// A statement that has to be run before client messages +#[derive(Debug, Clone, PartialEq)] +pub struct Prepare { + /// Some if statement was prepared previously, but has expired since + close: Option, + parse: ProtocolMessage, +} + +impl Prepare { + /// The stale statement to close first, if the name is taken. + pub fn close(&self) -> Option<&ProtocolMessage> { + self.close.as_ref() + } + + pub fn parse(&self) -> &ProtocolMessage { + &self.parse + } + + fn anonymize(&mut self) { + self.parse.anonymize(); + } } #[derive(Debug, Clone, PartialEq)] pub enum HandleResult { - Forward, Drop, - Prepend(ProtocolMessage), + Forward, Rewrite(ProtocolMessage), + Prepend(Prepare), PrependRewrite { - prepend: ProtocolMessage, + prepend: Prepare, rewrite: ProtocolMessage, }, } @@ -44,15 +100,14 @@ pub enum HandleResult { #[derive(Debug)] pub struct PreparedStatements { global_cache: Arc>, - local_cache: LruCache, + local_cache: LruCache, state: ProtocolState, // Prepared statements being prepared now on the connection. parses: VecDeque, // Describes being executed now on the connection. describes: VecDeque, - capacity: usize, + config: PreparedStatementsConfig, memory_used: usize, - level: PreparedStatementsLevel, oids: Arc, } @@ -72,27 +127,26 @@ impl PreparedStatements { state: ProtocolState::default(), parses: VecDeque::new(), describes: VecDeque::new(), - capacity: usize::MAX, + config: PreparedStatementsConfig::default(), memory_used: 0, - level: PreparedStatementsLevel::default(), oids, } } - /// Set maximum prepared statements capacity. + /// Apply the pool's prepared statement settings. #[inline] - pub fn set_capacity(&mut self, capacity: usize) { - self.capacity = capacity; + pub fn configure(&mut self, config: PreparedStatementsConfig) { + self.config = config; } - #[inline] - pub fn set_prepared_statements_level(&mut self, level: PreparedStatementsLevel) { - self.level = level; + /// Current prepared statement settings. + pub fn config(&self) -> PreparedStatementsConfig { + self.config } /// Get prepared statements capacity. pub fn capacity(&self) -> usize { - self.capacity + self.config.limit } /// Force the server to ignore the response to this message. @@ -118,10 +172,13 @@ impl PreparedStatements { let message = self.check_prepared(bind.statement())?; match message { Some(mut message) => { + if message.close.is_some() { + self.state.add_ignore('3'); + } self.state.add_ignore('1'); self.parses.push_back(bind.statement().to_string()); self.state.add('2'); - if self.level.rewrite_anonymous() { + if self.config.level.rewrite_anonymous() { message.anonymize(); let mut bind = bind.clone(); bind.anonymize(); @@ -136,7 +193,7 @@ impl PreparedStatements { None => { self.state.add('2'); - if self.level.rewrite_anonymous() { + if self.config.level.rewrite_anonymous() { let mut bind = bind.clone(); bind.anonymize(); return Ok(HandleResult::Rewrite(ProtocolMessage::Bind(bind))); @@ -153,12 +210,15 @@ impl PreparedStatements { match message { Some(mut message) => { + if message.close.is_some() { + self.state.add_ignore('3'); + } self.state.add_ignore('1'); self.parses.push_back(describe.statement().to_string()); self.state.add(ExecutionCode::DescriptionOrNothing); // t self.state.add(ExecutionCode::DescriptionOrNothing); // T - if self.level.rewrite_anonymous() { + if self.config.level.rewrite_anonymous() { // Save the RowDescription because // we don't actually save prepared statements in the server // anymore so they can be different every time. @@ -181,7 +241,7 @@ impl PreparedStatements { self.state.add(ExecutionCode::DescriptionOrNothing); // T self.describes.push_back(describe.statement().to_string()); - if self.level.rewrite_anonymous() { + if self.config.level.rewrite_anonymous() { let mut describe = describe.clone(); describe.anonymize(); return Ok(HandleResult::Rewrite(ProtocolMessage::Describe( @@ -224,7 +284,7 @@ impl PreparedStatements { // The client is sending named prepared statements, // but we're in ExtendedAnonymous mode so we rewrite // them to anonymous to avoid storing them in Postgres. - if self.level.rewrite_anonymous() { + if self.config.level.rewrite_anonymous() { parse.anonymize(); rewritten = true; } @@ -314,6 +374,18 @@ impl PreparedStatements { } } + // The close statement that is ignored and we have the parse for + // means we're repreparing the statement right now, so + // drop from cache first and let it be readded later on ParseComplete + '3' if matches!(action, Action::Ignore) => { + // ok, pop_front -> push_front just to avoid borrowing issues + // and not to copy the name just to remove by name + if let Some(name) = self.parses.pop_front() { + self.remove(&name); + self.parses.push_front(name); + } + } + 'G' => { self.state.prepend('G'); // Next thing we'll see is a CopyFail or CopyDone. } @@ -332,7 +404,7 @@ impl PreparedStatements { // Reset cache, forcing all Bind/Execute, Describe, solo requests // to always re-prepare the statement next time it's sent. - if !self.has_more_messages() && self.level.rewrite_anonymous() { + if !self.has_more_messages() && self.config.level.rewrite_anonymous() { self.clear(); } @@ -363,17 +435,38 @@ impl PreparedStatements { self.state.out_of_sync() } - fn check_prepared(&mut self, name: &str) -> Result, Error> { - if !self.contains(name) && !self.parses.iter().any(|s| s == name) { - let parse = self.parse(name); - if let Some(parse) = parse { - Ok(Some(ProtocolMessage::Parse(parse))) - } else { - Ok(None) - } - } else { - Ok(None) + /// Check the prepared state to identify if we need + /// to run something before actual client's requests + fn check_prepared(&mut self, name: &str) -> Result, Error> { + // Ignore if we already have a Parse in progress. + if self.parses.iter().any(|s| s == name) { + return Ok(None); + } + + let entry = self.local_cache.get(name); + let expired = + self.config.ttl.is_some() && entry.is_some_and(|entry| entry.expired(Instant::now())); + + if entry.is_some() && !expired { + return Ok(None); } + + // Nothing to prepare it from, so leave whatever is there alone. + let Some(parse) = self.parse(name) else { + return Ok(None); + }; + + Ok(Some(Prepare { + // Postgres still has the expired statement under this name, so + // we need to close it first to reprepare. + // + // The entry stays in the cache until CloseComplete confirms the + // drop: an error earlier in the batch makes Postgres skip our + // Close, and dropping it here would leave us re-preparing a name + // it still holds. + close: expired.then(|| ProtocolMessage::Close(Close::named(name))), + parse: ProtocolMessage::Parse(parse), + })) } /// The server has prepared this statement already. @@ -381,10 +474,19 @@ impl PreparedStatements { self.local_cache.promote(name) } - /// Indicate this statement is prepared on the connection. + #[cfg(test)] + fn statement(&self, name: &str) -> Option<&LocalStatement> { + self.local_cache.peek(name) + } + pub fn prepared(&mut self, name: &str) { - self.memory_used += str_mem(name); - self.local_cache.push(name.to_owned(), ()); + let statement = LocalStatement::new(self.config.ttl, self.config.ttl_jitter); + + // Cache is unbounded, so anything handed back is the old entry + // for this same name, never an eviction. Only new names cost us. + if self.local_cache.push(name.to_owned(), statement).is_none() { + self.memory_used += entry_mem(name); + } } /// How much memory is used by this structure, approx. @@ -418,7 +520,7 @@ impl PreparedStatements { /// or failed to parse. pub(crate) fn remove(&mut self, name: &str) -> bool { if self.local_cache.pop(name).is_some() { - self.memory_used = self.memory_used.saturating_sub(str_mem(name)); + self.memory_used = self.memory_used.saturating_sub(entry_mem(name)); true } else { false @@ -461,12 +563,12 @@ impl PreparedStatements { #[must_use] pub fn ensure_capacity(&mut self) -> Vec { let mut close = vec![]; - while self.local_cache.len() > self.capacity { + while self.local_cache.len() > self.config.limit { let candidate = self.local_cache.pop_lru(); if let Some((name, _)) = candidate { close.push(Close::named(&name)); - self.memory_used = self.memory_used.saturating_sub(str_mem(&name)); + self.memory_used = self.memory_used.saturating_sub(entry_mem(&name)); } } @@ -523,29 +625,240 @@ impl PreparedStatements { } #[cfg(test)] -mod test { +pub(crate) mod test { use super::*; use crate::frontend::PreparedStatements as FrontendPreparedStatements; use crate::net::{ - Bind, Describe, Execute, Message, Parse, ProtocolMessage, Query, Sync, bind::Parameter, - messages::ReadyForQuery, + Bind, Describe, ErrorResponse, Execute, Message, Parse, ProtocolMessage, Query, Sync, + bind::Parameter, messages::ReadyForQuery, }; use pgdog_config::PreparedStatements as PreparedStatementsLevel; /// Build a PreparedStatements instance configured for ExtendedAnonymous mode. fn new_extended_anonymous() -> PreparedStatements { - let mut ps = PreparedStatements::default(); - ps.set_prepared_statements_level(PreparedStatementsLevel::ExtendedAnonymous); - ps + new_with_level(PreparedStatementsLevel::ExtendedAnonymous) } /// Build a PreparedStatements instance configured for Extended (default) mode. fn new_extended() -> PreparedStatements { + new_with_level(PreparedStatementsLevel::Extended) + } + + fn new_with_level(level: PreparedStatementsLevel) -> PreparedStatements { let mut ps = PreparedStatements::default(); - ps.set_prepared_statements_level(PreparedStatementsLevel::Extended); + ps.configure(PreparedStatementsConfig { + level, + ..ps.config() + }); + ps + } + + const TTL: Duration = Duration::from_secs(300); + + fn new_with_ttl() -> PreparedStatements { + let mut ps = new_extended(); + ps.configure(PreparedStatementsConfig { + ttl: Some(TTL), + ..ps.config() + }); ps } + pub(crate) fn prepare_expired(ps: &mut PreparedStatements, name: &str) { + let config = ps.config(); + ps.configure(PreparedStatementsConfig { + ttl: Some(Duration::ZERO), + ttl_jitter: Duration::ZERO, + ..config + }); + ps.prepared(name); + ps.configure(config); + } + + macro_rules! assert_close_and_parse { + ($result:expr, $name:expr) => { + match $result { + HandleResult::Prepend(prepare) => { + assert_eq!( + prepare.close(), + Some(&ProtocolMessage::Close(Close::named($name))) + ); + assert!(matches!(prepare.parse(), ProtocolMessage::Parse(_))); + } + other => panic!("expected Prepend carrying a Close, got {other:?}"), + } + }; + } + + macro_rules! assert_parse_without_close { + ($result:expr) => { + match $result { + HandleResult::Prepend(prepare) => { + assert_eq!(prepare.close(), None); + assert!(matches!(prepare.parse(), ProtocolMessage::Parse(_))); + } + other => panic!("expected Prepend without a Close, got {other:?}"), + } + }; + } + + fn bind(name: &str) -> ProtocolMessage { + ProtocolMessage::Bind(Bind::new_statement(name)) + } + + #[test] + fn bind_prepares_an_unknown_statement_without_a_close() { + let name = insert_global("ttl_unknown", "SELECT $1::bigint"); + let mut ps = new_with_ttl(); + + assert_parse_without_close!(ps.handle(&bind(&name)).unwrap()); + } + + #[test] + fn bind_leaves_a_statement_within_its_ttl_alone() { + let name = insert_global("ttl_fresh", "SELECT $1::bigint"); + let mut ps = new_with_ttl(); + ps.prepared(&name); + + assert_eq!(ps.handle(&bind(&name)).unwrap(), HandleResult::Forward); + assert!(ps.contains(&name)); + } + + #[test] + fn bind_closes_a_statement_past_its_ttl() { + let name = insert_global("ttl_expired", "SELECT $1::bigint"); + let mut ps = new_with_ttl(); + prepare_expired(&mut ps, &name); + + assert_close_and_parse!(ps.handle(&bind(&name)).unwrap(), &name); + } + + #[test] + fn bind_closes_a_statement_prepared_before_the_ttl_was_set() { + let name = insert_global("ttl_enabled_later", "SELECT $1::bigint"); + let mut ps = new_extended(); + ps.prepared(&name); + + assert_eq!(ps.config().ttl, None); + assert_eq!(ps.statement(&name).unwrap().deadline(), None); + + ps.configure(PreparedStatementsConfig { + ttl: Some(TTL), + ..ps.config() + }); + + assert_close_and_parse!(ps.handle(&bind(&name)).unwrap(), &name); + } + + #[test] + fn bind_leaves_an_expired_statement_alone_when_the_ttl_is_disabled() { + let name = insert_global("ttl_disabled", "SELECT $1::bigint"); + let mut ps = new_extended(); + prepare_expired(&mut ps, &name); + + assert_eq!(ps.config().ttl, None); + assert!(ps.statement(&name).unwrap().expired(Instant::now())); + + assert_eq!(ps.handle(&bind(&name)).unwrap(), HandleResult::Forward); + assert!(ps.contains(&name)); + } + + #[test] + fn bind_leaves_an_expired_statement_alone_when_it_cannot_be_re_prepared() { + let mut ps = new_with_ttl(); + prepare_expired(&mut ps, "not_in_global_cache"); + + assert_eq!( + ps.handle(&bind("not_in_global_cache")).unwrap(), + HandleResult::Forward + ); + assert!(ps.contains("not_in_global_cache")); + } + + #[test] + fn bind_closes_a_statement_past_its_ttl_only_once() { + let name = insert_global("ttl_in_flight", "SELECT $1::bigint"); + let mut ps = new_with_ttl(); + prepare_expired(&mut ps, &name); + + assert_close_and_parse!(ps.handle(&bind(&name)).unwrap(), &name); + assert_eq!(ps.handle(&bind(&name)).unwrap(), HandleResult::Forward); + } + + #[test] + fn describe_closes_a_statement_past_its_ttl() { + let name = insert_global("ttl_describe", "SELECT $1::bigint"); + let mut ps = new_with_ttl(); + prepare_expired(&mut ps, &name); + + let describe = ProtocolMessage::Describe(Describe::new_statement(&name)); + assert_close_and_parse!(ps.handle(&describe).unwrap(), &name); + } + + #[test] + fn close_confirmed_then_failed_parse_leaves_no_stale_entry() { + let name = insert_global("ttl_close_then_error", "SELECT $1::bigint"); + let mut ps = new_with_ttl(); + prepare_expired(&mut ps, &name); + + // Bind prepends Close + Parse for the expired statement. + assert_close_and_parse!(ps.handle(&bind(&name)).unwrap(), &name); + + // Postgres closes the statement, then rejects the Parse. The client + // never asked for the Close, so its reply stays with us. + let mut close_complete = Message::new(CloseComplete.to_bytes()); + assert!(!ps.forward(&mut close_complete).unwrap()); + + let mut error = Message::new(ErrorResponse::syntax("boom").to_bytes()); + assert!(ps.forward(&mut error).unwrap()); + + // The name is gone from the server, so the cache must not claim it. + assert!(!ps.contains(&name)); + } + + #[test] + fn error_before_the_close_keeps_the_entry() { + let name = insert_global("ttl_error_then_close", "SELECT $1::bigint"); + let mut ps = new_with_ttl(); + prepare_expired(&mut ps, &name); + + assert_close_and_parse!(ps.handle(&bind(&name)).unwrap(), &name); + + // An earlier message failed, so Postgres skipped our Close and still + // holds the statement. + let mut error = Message::new(ErrorResponse::syntax("boom").to_bytes()); + assert!(ps.forward(&mut error).unwrap()); + + assert!(ps.contains(&name)); + + // The next Bind must close it again before re-preparing. + assert_close_and_parse!(ps.handle(&bind(&name)).unwrap(), &name); + } + + #[test] + fn a_portal_close_does_not_drop_a_pending_statement() { + let name = insert_global("ttl_portal_close", "SELECT $1::bigint"); + let mut ps = new_with_ttl(); + prepare_expired(&mut ps, &name); + + // The client closes a portal, then binds the expired statement. + let portal = ProtocolMessage::Close(Close::portal("p")); + assert_eq!(ps.handle(&portal).unwrap(), HandleResult::Forward); + assert_close_and_parse!(ps.handle(&bind(&name)).unwrap(), &name); + + // Postgres answers the portal Close, then a message between it and our + // Close fails, so our Close never runs. The client asked for this + // Close, so its reply goes back. + let mut close_complete = Message::new(CloseComplete.to_bytes()); + assert!(ps.forward(&mut close_complete).unwrap()); + + let mut error = Message::new(ErrorResponse::syntax("boom").to_bytes()); + assert!(ps.forward(&mut error).unwrap()); + + // Postgres still holds the statement, so the cache must too. + assert!(ps.contains(&name)); + } + /// Insert a prepared statement into the global cache so check_prepared can find it. fn insert_global(name: &str, query: &str) -> String { let parse = Parse::named(name, query); @@ -649,7 +962,11 @@ mod test { let bind = Bind::new_statement(&name); let result = ps.handle(&ProtocolMessage::Bind(bind)).unwrap(); match result { - HandleResult::Prepend(ProtocolMessage::Parse(p)) => { + HandleResult::Prepend(prepare) => { + assert_eq!(prepare.close(), None); + let ProtocolMessage::Parse(p) = prepare.parse() else { + panic!("expected prepend to be Parse"); + }; assert_eq!(p.query(), "SELECT $1"); } other => panic!("expected Prepend(Parse), got {:?}", other), @@ -665,7 +982,7 @@ mod test { match result { HandleResult::PrependRewrite { prepend, rewrite } => { // The prepended Parse should be anonymized. - if let ProtocolMessage::Parse(p) = &prepend { + if let ProtocolMessage::Parse(p) = prepend.parse() { assert!(p.anonymous(), "prepended Parse should be anonymous"); assert_eq!(p.query(), "SELECT $1"); } else { @@ -707,7 +1024,7 @@ mod test { let result = ps.handle(&ProtocolMessage::Describe(describe)).unwrap(); match result { HandleResult::PrependRewrite { prepend, rewrite } => { - if let ProtocolMessage::Parse(p) = &prepend { + if let ProtocolMessage::Parse(p) = prepend.parse() { assert!(p.anonymous(), "prepended Parse should be anonymous"); } else { panic!("expected prepend to be Parse"); diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index f574ec87b..043dba2fc 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -14,7 +14,8 @@ use tracing::{debug, error, info, trace, warn}; use super::{ ConnectReason, DisconnectReason, Error, Oids, PreparedStatements, ServerOptions, Stats, - pool::Address, prepared_statements::HandleResult, + pool::Address, + prepared_statements::{HandleResult, Prepare}, }; use crate::{ auth::{md5, scram::Client}, @@ -492,26 +493,30 @@ impl Server { let result = self.prepared_statements.handle(message)?; - let queue = match result { - HandleResult::Drop => [None, None], - HandleResult::Prepend(ref prepare) => [Some(prepare), Some(message)], - HandleResult::Forward => [Some(message), None], - HandleResult::Rewrite(ref message) => [Some(message), None], - HandleResult::PrependRewrite { - ref prepend, - ref rewrite, - } => [Some(prepend), Some(rewrite)], - }; - - for message in queue.iter().flatten() { - trace!("{:#?} >>> [{}]", message, self.addr()); + match &result { + HandleResult::Drop => {} + HandleResult::Forward => self.send_stream(message).await?, + HandleResult::Rewrite(rewrite) => self.send_stream(rewrite).await?, + HandleResult::Prepend(prepare) => { + self.send_prepare(prepare).await?; + self.send_stream(message).await?; + } + HandleResult::PrependRewrite { prepend, rewrite } => { + self.send_prepare(prepend).await?; + self.send_stream(rewrite).await?; + } } - for message in queue.into_iter().flatten() { - self.send_stream(message).await?; + Ok(()) + } + + /// Close a stale prepared statement, if any, then prepare the new one. + async fn send_prepare(&mut self, prepare: &Prepare) -> Result<(), Error> { + if let Some(close) = prepare.close() { + self.send_stream(close).await?; } - Ok(()) + self.send_stream(prepare.parse()).await } /// Send a message to Postgres and force us to ignore its respose in [`Self::read`]. @@ -529,6 +534,8 @@ impl Server { /// Send message to Postgres, checking for any errors /// and setting the server state accordingly. async fn send_stream(&mut self, message: &ProtocolMessage) -> Result<(), Error> { + trace!("{:#?} >>> [{}]", message, self.addr()); + match self.stream().send(message).await { Ok(sent) => self.stats.send(sent, message.code() as u8), Err(err) => { @@ -1320,6 +1327,7 @@ pub mod test { use std::time::SystemTime; use bytes::{BufMut, BytesMut}; + use pgdog_stats::PreparedStatementsConfig; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::TcpListener, @@ -2765,7 +2773,12 @@ pub mod test { let in_sync = server.fetch_all::("SELECT 1::bigint").await.unwrap(); assert_eq!(in_sync[0], 1); - server.prepared_statements.set_capacity(3); + server + .prepared_statements + .configure(PreparedStatementsConfig { + limit: 3, + ..Default::default() + }); assert_eq!(server.prepared_statements.capacity(), 3); server @@ -2809,10 +2822,27 @@ pub mod test { } } - /// Names Postgres itself thinks are prepared on this connection. - pub(crate) async fn prepared_in_postgres(server: &mut Server) -> Vec { + /// What Postgres itself thinks is prepared on this connection. + #[derive(Debug, Clone, PartialEq, Eq)] + pub(crate) struct PreparedInPostgres { + pub name: String, + pub prepared_at: String, + } + + impl From for PreparedInPostgres { + fn from(value: DataRow) -> Self { + Self { + name: value.get_text(0).unwrap_or_default(), + prepared_at: value.get_text(1).unwrap_or_default(), + } + } + } + + pub(crate) async fn prepared_in_postgres(server: &mut Server) -> Vec { server - .fetch_all::("SELECT name FROM pg_prepared_statements") + .fetch_all::( + "SELECT name, prepare_time::text FROM pg_prepared_statements ORDER BY name", + ) .await .unwrap() } @@ -2838,6 +2868,41 @@ pub mod test { server.fetch_all::(request).await.unwrap() } + #[tokio::test] + async fn test_prepared_statement_is_re_prepared_once_its_ttl_runs_out() { + use crate::backend::prepared_statements::test::prepare_expired; + + let parse = Parse::named("ttl_client", "SELECT $1::bigint"); + let (_, name) = crate::frontend::PreparedStatements::global() + .write() + .insert(&parse); + + let mut server = test_server().await; + let statements = server.prepared_statements_mut(); + statements.configure(PreparedStatementsConfig { + ttl: Some(Duration::from_secs(300)), + ..statements.config() + }); + + assert_eq!(execute_prepared(&mut server, &name, b"1").await, [1]); + let prepared = prepared_in_postgres(&mut server).await; + assert_eq!(prepared.len(), 1); + assert_eq!(prepared[0].name, name); + + assert_eq!(execute_prepared(&mut server, &name, b"1").await, [1]); + assert_eq!(prepared_in_postgres(&mut server).await, prepared); + + prepare_expired(server.prepared_statements_mut(), &name); + + assert_eq!(execute_prepared(&mut server, &name, b"1").await, [1]); + let re_prepared = prepared_in_postgres(&mut server).await; + assert_eq!(re_prepared.len(), 1); + assert_eq!(re_prepared[0].name, name); + assert_ne!(re_prepared, prepared); + + assert!(server.done()); + } + #[tokio::test] async fn test_deallocate_all_clears_cache() { let mut server = test_server().await; @@ -3975,9 +4040,12 @@ pub mod test { /// Set a server's prepared_statements level to ExtendedAnonymous. fn set_extended_anonymous(server: &mut Server) { use pgdog_config::PreparedStatements as PSLevel; - server - .prepared_statements_mut() - .set_prepared_statements_level(PSLevel::ExtendedAnonymous); + let statements = server.prepared_statements_mut(); + let config = statements.config(); + statements.configure(PreparedStatementsConfig { + level: PSLevel::ExtendedAnonymous, + ..config + }); } #[tokio::test] @@ -4245,7 +4313,9 @@ pub mod test { use crate::net::bind::Parameter; let mut server = test_server().await; set_extended_anonymous(&mut server); - server.prepared_statements_mut().set_capacity(3); + let statements = server.prepared_statements_mut(); + let config = statements.config(); + statements.configure(PreparedStatementsConfig { limit: 3, ..config }); // Send many different "named" statements. // Because they're all anonymized, no named statements are stored in Postgres, diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index f41dacc83..988dba535 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -194,11 +194,20 @@ pub fn run_maintenance() { #[cfg(test)] mod test { + use crate::backend::Server; use crate::backend::server::test::{execute_prepared, prepared_in_postgres, test_server}; use crate::net::messages::Bind; use super::*; + async fn prepared_names(server: &mut Server) -> Vec { + prepared_in_postgres(server) + .await + .into_iter() + .map(|statement| statement.name) + .collect() + } + #[tokio::test] async fn test_close_unused_does_not_reuse_names() { let mut client = PreparedStatements::new(); @@ -209,13 +218,13 @@ mod test { let mut server = test_server().await; assert_eq!(execute_prepared(&mut server, first_name, b"1").await, [1]); - assert_eq!(prepared_in_postgres(&mut server).await, [first_name]); + assert_eq!(prepared_names(&mut server).await, [first_name]); client.close("client_a"); PreparedStatements::global().write().close_unused(0); assert!(PreparedStatements::global().read().is_empty()); - assert_eq!(prepared_in_postgres(&mut server).await, [first_name]); + assert_eq!(prepared_names(&mut server).await, [first_name]); let mut second = Parse::named("client_b", "SELECT $1::bigint + 100"); client.insert(&mut second); @@ -238,16 +247,16 @@ mod test { let mut warm = test_server().await; assert_eq!(execute_prepared(&mut warm, name, b"1").await, [1]); - assert_eq!(prepared_in_postgres(&mut warm).await, [name]); + assert_eq!(prepared_names(&mut warm).await, [name]); PreparedStatements::global().write().close_unused(0); assert_eq!(execute_prepared(&mut warm, name, b"1").await, [1]); let mut cold = test_server().await; - assert!(prepared_in_postgres(&mut cold).await.is_empty()); + assert!(prepared_names(&mut cold).await.is_empty()); assert_eq!(execute_prepared(&mut cold, name, b"1").await, [1]); - assert_eq!(prepared_in_postgres(&mut cold).await, [name]); + assert_eq!(prepared_names(&mut cold).await, [name]); } #[test] diff --git a/pgdog/src/util.rs b/pgdog/src/util.rs index 2e3a53f55..a81b902ee 100644 --- a/pgdog/src/util.rs +++ b/pgdog/src/util.rs @@ -1,5 +1,7 @@ //! What's a project without a util module. +pub mod time; + use chrono::{DateTime, Local, Utc}; use once_cell::sync::Lazy; use rand::{Rng, distr::Alphanumeric}; diff --git a/pgdog/src/util/time.rs b/pgdog/src/util/time.rs new file mode 100644 index 000000000..b517b8cd6 --- /dev/null +++ b/pgdog/src/util/time.rs @@ -0,0 +1,103 @@ +//! Deadlines and jitter. + +use rand::Rng; +use std::time::{Duration, Instant}; + +/// Random point in `[base - jitter, base + jitter]`, so things created +/// together don't all expire on the same tick. +/// +/// Jitter has millisecond granularity: anything smaller is treated as zero. +pub fn jitter_duration(base: Duration, jitter: Duration) -> Duration { + if jitter.is_zero() { + return base; + } + + // make sure to clamp the jitter + let jitter = jitter.as_millis().min(i64::MAX as u128) as i64; + if jitter == 0 { + return base; + } + + let offset = rand::rng().random_range(-jitter..=jitter); + let magnitude = Duration::from_millis(offset.unsigned_abs()); + + if offset >= 0 { + base.saturating_add(magnitude) + } else { + base.saturating_sub(magnitude) + } +} + +/// Calculate the deadline from now +/// +/// # Panics +/// +/// Panics on a `ttl` big enough to overflow the clock. Config caps it +/// long before that. +pub fn deadline(ttl: Duration, jitter: Duration) -> Instant { + Instant::now() + jitter_duration(ttl, jitter) +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_jitter_duration_zero_returns_base() { + let base = Duration::from_secs(60); + assert_eq!(jitter_duration(base, Duration::ZERO), base); + } + + #[test] + fn test_jitter_duration_below_a_millisecond_returns_base() { + let base = Duration::from_secs(60); + assert_eq!(jitter_duration(base, Duration::from_micros(500)), base); + } + + #[test] + fn test_jitter_duration_survives_an_absurd_jitter() { + let base = Duration::from_secs(60); + let jitter = Duration::from_millis(u64::MAX); + assert!(jitter_duration(base, jitter) <= base.saturating_add(jitter)); + } + + #[test] + fn test_jitter_duration_stays_within_bounds() { + let base = Duration::from_secs(60); + let jitter = Duration::from_secs(10); + let mut below = false; + let mut above = false; + + for _ in 0..1_000 { + let sampled = jitter_duration(base, jitter); + assert!( + sampled >= base - jitter && sampled <= base + jitter, + "sampled {sampled:?} outside [{:?}, {:?}]", + base - jitter, + base + jitter + ); + below |= sampled < base; + above |= sampled > base; + } + + assert!(below, "never sampled below base"); + assert!(above, "never sampled above base"); + } + + #[test] + fn test_jitter_duration_saturates_at_zero() { + let base = Duration::from_millis(10); + let jitter = Duration::from_millis(100); + for _ in 0..1_000 { + let sampled = jitter_duration(base, jitter); + assert!(sampled <= base + jitter, "sampled {sampled:?} too large"); + } + } + + #[test] + fn test_largest_configurable_ttl_fits_in_an_instant() { + // Largest ttl the config will actually let through. + let ttl = Duration::from_millis(i64::MAX as u64 - 1); + assert!(Instant::now().checked_add(ttl).is_some()); + } +} From dec0966e3253a27a1795d8cb40597dab82bc38f1 Mon Sep 17 00:00:00 2001 From: meskill <8974488+meskill@users.noreply.github.com> Date: Wed, 12 Aug 2026 18:40:49 +0000 Subject: [PATCH 2/3] fix visibility --- .schema/pgdog.schema.json | 11 +++--- example.pgdog.toml | 4 +-- pgdog-config/src/general.rs | 44 ++++++++++-------------- pgdog/src/backend/pool/config.rs | 2 +- pgdog/src/backend/prepared_statements.rs | 20 +++++------ pgdog/src/backend/server.rs | 4 +-- pgdog/src/util.rs | 2 +- pgdog/src/util/time.rs | 4 +-- 8 files changed, 41 insertions(+), 50 deletions(-) diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index 4f0cf6bd3..16f9b57b1 100644 --- a/.schema/pgdog.schema.json +++ b/.schema/pgdog.schema.json @@ -83,7 +83,7 @@ "port": 6432, "prepared_statements": "extended", "prepared_statements_limit": 9223372036854775807, - "prepared_statements_ttl": 300000, + "prepared_statements_ttl": null, "prepared_statements_ttl_jitter": 30000, "pub_sub_channel_size": 0, "query_cache_limit": 1000, @@ -1031,10 +1031,13 @@ "minimum": 0 }, "prepared_statements_ttl": { - "description": "How long a prepared statement is allowed to stay prepared on a server connection, in milliseconds.\n\n**Note:** Expired statements are closed and prepared again the next time they are used. This stops stale execution plans from staying in Postgres. Set to `0` to let statements stay prepared forever.\n\n_Default:_ `300000` (5 minutes)\n\n", - "type": "integer", + "description": "How long a prepared statement is allowed to stay prepared on a server connection, in milliseconds.\n\n**Note:** Expired statements are closed and prepared again the next time they are used. This stops stale execution plans from staying in Postgres. Omit or set to `0` to let statements stay prepared forever.\n\n_Default:_ `None` (disabled)\n\n", + "type": [ + "integer", + "null" + ], "format": "uint64", - "default": 300000, + "default": null, "minimum": 0 }, "prepared_statements_ttl_jitter": { diff --git a/example.pgdog.toml b/example.pgdog.toml index c812d2c2b..923007b39 100644 --- a/example.pgdog.toml +++ b/example.pgdog.toml @@ -201,9 +201,9 @@ prepared_statements_limit = 1_000 # connection. Expired statements are closed and prepared again the next # time they are used, so execution plans don't get stale. # -# Set to 0 to let statements stay prepared forever. +# Omit or set to 0 to let statements stay prepared forever. # -# Default: 300_000 (5 minutes) +# Default: none (disabled) # prepared_statements_ttl = 300_000 # Random spread applied to prepared_statements_ttl, per statement. Clients diff --git a/pgdog-config/src/general.rs b/pgdog-config/src/general.rs index d2d3768a5..bdfdc0287 100644 --- a/pgdog-config/src/general.rs +++ b/pgdog-config/src/general.rs @@ -384,13 +384,13 @@ pub struct General { /// How long a prepared statement is allowed to stay prepared on a server connection, in milliseconds. /// - /// **Note:** Expired statements are closed and prepared again the next time they are used. This stops stale execution plans from staying in Postgres. Set to `0` to let statements stay prepared forever. + /// **Note:** Expired statements are closed and prepared again the next time they are used. This stops stale execution plans from staying in Postgres. Omit or set to `0` to let statements stay prepared forever. /// - /// _Default:_ `300000` (5 minutes) + /// _Default:_ `None` (disabled) /// /// #[serde(default = "General::default_prepared_statements_ttl")] - pub prepared_statements_ttl: u64, + pub prepared_statements_ttl: Option, /// Maximum random adjustment applied to `prepared_statements_ttl` per prepared /// statement, in milliseconds. Each statement expires at a point sampled uniformly @@ -1172,12 +1172,9 @@ impl General { /// /// `None` means statements never expire. pub fn prepared_statements_ttl(&self) -> Option { - let ttl = self.prepared_statements_ttl; - if ttl == 0 || ttl >= crate::MAX_DURATION.as_millis() as u64 { - None - } else { - Some(Duration::from_millis(ttl)) - } + self.prepared_statements_ttl + .filter(|ttl| *ttl > 0 && *ttl < crate::MAX_DURATION.as_millis() as u64) + .map(Duration::from_millis) } /// Random spread applied to [`Self::prepared_statements_ttl`]. @@ -1186,6 +1183,7 @@ impl General { pub fn prepared_statements_ttl_jitter(&self) -> Duration { let ttl = self .prepared_statements_ttl + .unwrap_or(0) .min(crate::MAX_DURATION.as_millis() as u64) .saturating_sub(1); @@ -1431,11 +1429,8 @@ impl General { Self::env_or_default("PGDOG_PREPARED_STATEMENTS_LIMIT", i64::MAX as usize) } - pub fn default_prepared_statements_ttl() -> u64 { - Self::env_or_default( - "PGDOG_PREPARED_STATEMENTS_TTL", - Duration::from_secs(300).as_millis() as u64, - ) + fn default_prepared_statements_ttl() -> Option { + Self::env_option("PGDOG_PREPARED_STATEMENTS_TTL") } pub fn default_prepared_statements_ttl_jitter() -> u64 { @@ -1562,20 +1557,17 @@ mod tests { fn test_prepared_statements_ttl_defaults() { let general = General::default(); - assert_eq!( - general.prepared_statements_ttl(), - Some(Duration::from_millis(300_000)) - ); + assert_eq!(general.prepared_statements_ttl(), None); assert_eq!( general.prepared_statements_ttl_jitter(), - Duration::from_millis(30_000) + Duration::from_millis(0) ); } #[test] fn test_prepared_statements_ttl_is_read_in_millis() { let general = General { - prepared_statements_ttl: 3_600_000, + prepared_statements_ttl: Some(3_600_000), prepared_statements_ttl_jitter: 5_000, ..Default::default() }; @@ -1592,7 +1584,7 @@ mod tests { #[test] fn test_prepared_statements_ttl_disabled() { - for ttl in [0, crate::MAX_DURATION.as_millis() as u64] { + for ttl in [None, Some(0), Some(crate::MAX_DURATION.as_millis() as u64)] { let general = General { prepared_statements_ttl: ttl, ..Default::default() @@ -1601,7 +1593,7 @@ mod tests { assert_eq!( general.prepared_statements_ttl(), None, - "ttl {ttl} should disable expiration" + "ttl {ttl:?} should disable expiration" ); } } @@ -1609,7 +1601,7 @@ mod tests { #[test] fn test_prepared_statements_ttl_jitter_is_clamped_below_ttl() { let general = General { - prepared_statements_ttl: 10_000, + prepared_statements_ttl: Some(10_000), prepared_statements_ttl_jitter: u64::MAX, ..Default::default() }; @@ -1624,7 +1616,7 @@ mod tests { fn test_prepared_statements_ttl_jitter_never_reaches_the_ttl() { for ttl in [0, 1, 2, 300_000] { let general = General { - prepared_statements_ttl: ttl, + prepared_statements_ttl: Some(ttl), prepared_statements_ttl_jitter: u64::MAX, ..Default::default() }; @@ -1977,7 +1969,7 @@ mod tests { assert_eq!(General::broadcast_port(), 7432); assert_eq!(General::openmetrics_port(), Some(9090)); assert_eq!(General::prepared_statements_limit(), 1000); - assert_eq!(General::default_prepared_statements_ttl(), 3600000); + assert_eq!(General::default_prepared_statements_ttl(), Some(3600000)); assert_eq!(General::default_prepared_statements_ttl_jitter(), 5000); assert_eq!(General::query_cache_limit(), 500); assert_eq!(General::connect_attempts(), 3); @@ -2005,7 +1997,7 @@ mod tests { assert_eq!(General::broadcast_port(), General::port() + 1); assert_eq!(General::openmetrics_port(), None); assert_eq!(General::prepared_statements_limit(), i64::MAX as usize); - assert_eq!(General::default_prepared_statements_ttl(), 300_000); + assert_eq!(General::default_prepared_statements_ttl(), None); assert_eq!(General::default_prepared_statements_ttl_jitter(), 30_000); assert_eq!(General::query_cache_limit(), 1_000); assert_eq!(General::connect_attempts(), 1); diff --git a/pgdog/src/backend/pool/config.rs b/pgdog/src/backend/pool/config.rs index db9a592e0..507959c3d 100644 --- a/pgdog/src/backend/pool/config.rs +++ b/pgdog/src/backend/pool/config.rs @@ -172,7 +172,7 @@ mod test { #[test] fn test_prepared_statements_config_from_general() { let general = General { - prepared_statements_ttl: 60_000, + prepared_statements_ttl: Some(60_000), prepared_statements_limit: 10, ..Default::default() }; diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 3bfe44605..5c59a352f 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -32,7 +32,7 @@ fn entry_mem(s: &str) -> usize { /// A statement info prepared on this connection. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct LocalStatement { +struct LocalStatement { /// When this statement should be replanned deadline: Option, } @@ -44,22 +44,18 @@ impl LocalStatement { } } - pub fn deadline(&self) -> Option { - self.deadline - } - /// Check for expired /// /// If the check is called and deadline is not set then it's marked as expired /// to cover the case when the TTL was set after the statement creation - pub fn expired(&self, now: Instant) -> bool { + fn expired(&self, now: Instant) -> bool { self.deadline.is_none_or(|deadline| deadline <= now) } } /// A statement that has to be run before client messages #[derive(Debug, Clone, PartialEq)] -pub struct Prepare { +pub(super) struct Prepare { /// Some if statement was prepared previously, but has expired since close: Option, parse: ProtocolMessage, @@ -67,11 +63,11 @@ pub struct Prepare { impl Prepare { /// The stale statement to close first, if the name is taken. - pub fn close(&self) -> Option<&ProtocolMessage> { + pub(super) fn close(&self) -> Option<&ProtocolMessage> { self.close.as_ref() } - pub fn parse(&self) -> &ProtocolMessage { + pub(super) fn parse(&self) -> &ProtocolMessage { &self.parse } @@ -81,7 +77,7 @@ impl Prepare { } #[derive(Debug, Clone, PartialEq)] -pub enum HandleResult { +pub(super) enum HandleResult { Drop, Forward, Rewrite(ProtocolMessage), @@ -165,7 +161,7 @@ impl PreparedStatements { } /// Handle extended protocol message. - pub fn handle(&mut self, request: &ProtocolMessage) -> Result { + pub(super) fn handle(&mut self, request: &ProtocolMessage) -> Result { match request { ProtocolMessage::Bind(bind) => { if !bind.anonymous() { @@ -740,7 +736,7 @@ pub(crate) mod test { ps.prepared(&name); assert_eq!(ps.config().ttl, None); - assert_eq!(ps.statement(&name).unwrap().deadline(), None); + assert!(ps.statement(&name).unwrap().expired(Instant::now())); ps.configure(PreparedStatementsConfig { ttl: Some(TTL), diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 043dba2fc..e47171ee8 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -2825,8 +2825,8 @@ pub mod test { /// What Postgres itself thinks is prepared on this connection. #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct PreparedInPostgres { - pub name: String, - pub prepared_at: String, + pub(crate) name: String, + pub(crate) prepared_at: String, } impl From for PreparedInPostgres { diff --git a/pgdog/src/util.rs b/pgdog/src/util.rs index a81b902ee..14d6778c7 100644 --- a/pgdog/src/util.rs +++ b/pgdog/src/util.rs @@ -1,6 +1,6 @@ //! What's a project without a util module. -pub mod time; +pub(crate) mod time; use chrono::{DateTime, Local, Utc}; use once_cell::sync::Lazy; diff --git a/pgdog/src/util/time.rs b/pgdog/src/util/time.rs index b517b8cd6..64af76234 100644 --- a/pgdog/src/util/time.rs +++ b/pgdog/src/util/time.rs @@ -7,7 +7,7 @@ use std::time::{Duration, Instant}; /// together don't all expire on the same tick. /// /// Jitter has millisecond granularity: anything smaller is treated as zero. -pub fn jitter_duration(base: Duration, jitter: Duration) -> Duration { +pub(crate) fn jitter_duration(base: Duration, jitter: Duration) -> Duration { if jitter.is_zero() { return base; } @@ -34,7 +34,7 @@ pub fn jitter_duration(base: Duration, jitter: Duration) -> Duration { /// /// Panics on a `ttl` big enough to overflow the clock. Config caps it /// long before that. -pub fn deadline(ttl: Duration, jitter: Duration) -> Instant { +pub(crate) fn deadline(ttl: Duration, jitter: Duration) -> Instant { Instant::now() + jitter_duration(ttl, jitter) } From 8b8e72994cec9cc5375b6acc4627c690f9025c2e Mon Sep 17 00:00:00 2001 From: meskill <8974488+meskill@users.noreply.github.com> Date: Thu, 13 Aug 2026 14:49:36 +0000 Subject: [PATCH 3/3] enable ttl for tests --- integration/pgdog.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/integration/pgdog.toml b/integration/pgdog.toml index d5fd64d75..2536ceb75 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -12,6 +12,8 @@ openmetrics_port = 9090 openmetrics_namespace = "pgdog_" prepared_statements_limit = 500 prepared_statements = "extended" +prepared_statements_ttl = 5_000 +prepared_statements_ttl_jitter = 1_000 expanded_explain = true dns_ttl = 1_000 query_cache_limit = 500