diff --git a/integration/rust/tests/integration/admin/mod.rs b/integration/rust/tests/integration/admin/mod.rs index 2080be470..732c0589a 100644 --- a/integration/rust/tests/integration/admin/mod.rs +++ b/integration/rust/tests/integration/admin/mod.rs @@ -1,6 +1,7 @@ //! Integration tests asserting admin command output over the wire. //! //! Each submodule connects to the live PgDog admin database (`rust::setup::admin_sqlx`). +pub mod reset_stats; pub mod resharding; pub mod show_config; pub mod show_version; diff --git a/integration/rust/tests/integration/admin/reset_stats.rs b/integration/rust/tests/integration/admin/reset_stats.rs new file mode 100644 index 000000000..22e774719 --- /dev/null +++ b/integration/rust/tests/integration/admin/reset_stats.rs @@ -0,0 +1,69 @@ +//! Integration tests for the RESET STATS admin command. +//! +//! One sequential test: RESET and RELOAD both act on global pool statistics, +//! so they must not run concurrently against the same pooler. + +use crate::setup::*; +use rust_decimal::prelude::ToPrimitive; +use sqlx::{Executor, Pool, Postgres, Row}; + +/// Sum of `total_query_count` across every pool reported by SHOW STATS. +/// +/// The column is sent as NUMERIC, so it's decoded via Decimal. +async fn total_query_count(admin: &Pool) -> i64 { + admin + .fetch_all("SHOW STATS") + .await + .unwrap() + .iter() + .map(|row| { + row.get::("total_query_count") + .to_i64() + .unwrap() + }) + .sum() +} + +#[tokio::test] +async fn test_reload_preserves_and_reset_clears_pool_stats() { + let admin = admin_sqlx().await; + + // Bump pool statistics by running queries through the pooler. + let client = connections_sqlx().await.pop().unwrap(); + for _ in 0..20 { + let _: (i64,) = sqlx::query_as("SELECT 1::BIGINT") + .fetch_one(&client) + .await + .unwrap(); + } + + let before = total_query_count(&admin).await; + assert!(before >= 20, "expected client queries to be counted"); + + // RELOAD must not lose the accumulated counters (issue #1281). + admin.execute("RELOAD").await.unwrap(); + let after_reload = total_query_count(&admin).await; + assert!( + after_reload >= before, + "RELOAD must preserve pool statistics (before {before}, after {after_reload})" + ); + + // RESET STATS must zero them. + admin.execute("RESET STATS").await.unwrap(); + let after_reset = total_query_count(&admin).await; + assert_eq!( + after_reset, 0, + "RESET STATS must zero pool query counters (got {after_reset})" + ); + + // Counters keep counting from zero afterwards. + let _: (i64,) = sqlx::query_as("SELECT 1::BIGINT") + .fetch_one(&client) + .await + .unwrap(); + let resumed = total_query_count(&admin).await; + assert!( + resumed > 0, + "statistics must resume counting after RESET STATS" + ); +} diff --git a/pgdog/src/admin/mod.rs b/pgdog/src/admin/mod.rs index 6565801da..a2f29904a 100644 --- a/pgdog/src/admin/mod.rs +++ b/pgdog/src/admin/mod.rs @@ -20,6 +20,7 @@ pub(crate) mod reload; pub(crate) mod replicate; pub(crate) mod reset_prepared; pub(crate) mod reset_query_cache; +pub(crate) mod reset_stats; pub(crate) mod reshard; pub(crate) mod schema_sync; pub(crate) mod server; @@ -64,6 +65,7 @@ pub(crate) use reload::*; pub(crate) use replicate::*; pub(crate) use reset_prepared::*; pub(crate) use reset_query_cache::*; +pub(crate) use reset_stats::*; pub(crate) use reshard::*; pub(crate) use schema_sync::*; pub(crate) use set::*; diff --git a/pgdog/src/admin/parser.rs b/pgdog/src/admin/parser.rs index 771afa22c..53769c36e 100644 --- a/pgdog/src/admin/parser.rs +++ b/pgdog/src/admin/parser.rs @@ -20,6 +20,7 @@ pub(crate) enum ParseResult { ShowQueryCache(ShowQueryCache), ResetPrepared(ResetPrepared), ResetQueryCache(ResetQueryCache), + ResetStats(ResetStats), ShowStats(ShowStats), ShowTransactions(ShowTransactions), ShowMirrors(ShowMirrors), @@ -69,6 +70,7 @@ impl ParseResult { ShowQueryCache(show_query_cache) => show_query_cache.execute().await, ResetPrepared(cmd) => cmd.execute().await, ResetQueryCache(reset_query_cache) => reset_query_cache.execute().await, + ResetStats(cmd) => cmd.execute().await, ShowStats(show_stats) => show_stats.execute().await, ShowTransactions(show_transactions) => show_transactions.execute().await, ShowMirrors(show_mirrors) => show_mirrors.execute().await, @@ -118,6 +120,7 @@ impl ParseResult { ShowQueryCache(show_query_cache) => show_query_cache.name(), ResetPrepared(cmd) => cmd.name(), ResetQueryCache(reset_query_cache) => reset_query_cache.name(), + ResetStats(cmd) => cmd.name(), ShowStats(show_stats) => show_stats.name(), ShowTransactions(show_transactions) => show_transactions.name(), ShowMirrors(show_mirrors) => show_mirrors.name(), @@ -238,6 +241,7 @@ impl Parser { "reset" => match iter.next().ok_or(Error::Syntax)?.trim() { "prepared" => ParseResult::ResetPrepared(ResetPrepared::parse(&sql)?), "query_cache" => ParseResult::ResetQueryCache(ResetQueryCache::parse(&sql)?), + "stats" => ParseResult::ResetStats(ResetStats::parse(&sql)?), command => { debug!("unknown admin show command: '{}'", command); return Err(Error::Syntax); @@ -378,6 +382,12 @@ mod tests { assert!(matches!(result, Ok(ParseResult::ResetPrepared(_)))); } + #[test] + fn parses_reset_stats_command() { + let result = Parser::parse("RESET STATS"); + assert!(matches!(result, Ok(ParseResult::ResetStats(_)))); + } + #[test] fn rejects_unknown_admin_command() { let result = Parser::parse("FOO BAR"); diff --git a/pgdog/src/admin/reset_stats.rs b/pgdog/src/admin/reset_stats.rs new file mode 100644 index 000000000..87739ea42 --- /dev/null +++ b/pgdog/src/admin/reset_stats.rs @@ -0,0 +1,22 @@ +//! RESET STATS. +use crate::backend::databases::databases; + +use super::prelude::*; + +pub struct ResetStats; + +#[async_trait] +impl Command for ResetStats { + fn name(&self) -> String { + "RESET STATS".into() + } + + fn parse(_: &str) -> Result { + Ok(Self) + } + + async fn execute(&self) -> Result, Error> { + databases().reset_stats(); + Ok(vec![]) + } +} diff --git a/pgdog/src/backend/databases.rs b/pgdog/src/backend/databases.rs index 08f67a164..969245961 100644 --- a/pgdog/src/backend/databases.rs +++ b/pgdog/src/backend/databases.rs @@ -418,6 +418,18 @@ impl Databases { Ok(moved) } + /// Reset statistics on all clusters (including mirrors) and their pools, + /// used by the RESET STATS command. + pub(crate) fn reset_stats(&self) { + for cluster in self.databases.values() { + cluster.reset_stats(); + } + + for cluster in self.mirrors.values().flatten() { + cluster.reset_stats(); + } + } + /// Shutdown all pools. fn shutdown(&self) { for cluster in self.all().values() { @@ -949,6 +961,79 @@ mod tests { assert_eq!(bob_mirrors[0].name(), "db1_mirror"); } + #[test] + fn test_reset_stats_resets_primary_and_mirror_clusters() { + // RESET STATS must clear metrics on both the primary clusters and + // their mirrors, not just `self.databases`. + let config = Config { + databases: vec![ + Database { + name: "db1".to_string(), + host: "localhost".to_string(), + port: 5432, + role: Role::Primary, + ..Default::default() + }, + Database { + name: "db1_mirror".to_string(), + host: "localhost".to_string(), + port: 5433, + role: Role::Primary, + ..Default::default() + }, + ], + mirroring: vec![Mirroring { + source_db: "db1".to_string(), + destination_db: "db1_mirror".to_string(), + ..Default::default() + }], + ..Default::default() + }; + + let users = crate::config::Users { + users: vec![ + crate::config::User { + name: "alice".to_string(), + database: "db1".to_string(), + password: Some("pass".to_string()), + ..Default::default() + }, + crate::config::User { + name: "alice".to_string(), + database: "db1_mirror".to_string(), + password: Some("pass".to_string()), + ..Default::default() + }, + ], + ..Default::default() + }; + + let databases = from_config(&ConfigAndUsers { + config, + users, + config_path: std::path::PathBuf::new(), + users_path: std::path::PathBuf::new(), + ..Default::default() + }); + + let primary = databases.cluster(("alice", "db1")).unwrap(); + let mirror = databases + .mirrors(("alice", "db1")) + .unwrap() + .unwrap() + .first() + .cloned() + .unwrap(); + + primary.stats().lock().mirror.total_count = 11; + mirror.stats().lock().mirror.total_count = 7; + + databases.reset_stats(); + + assert_eq!(primary.stats().lock().mirror.total_count, 0); + assert_eq!(mirror.stats().lock().mirror.total_count, 0); + } + #[test] fn test_mirror_user_mismatch_handling() { // Test that mirroring is disabled gracefully when users don't match diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 53fd035f0..157c5c708 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -452,6 +452,23 @@ impl Cluster { /// Move connections from cluster to another, saving them. pub(crate) fn move_conns_to(&self, other: &Cluster) -> Result<(), Error> { + // Ensure no deadlock: locking the same mutex twice on this thread + // (if self and other were the same cluster) would hang forever. + assert!(!Arc::ptr_eq(&self.stats, &other.stats)); + + // Carry cluster-level statistics over so a reload doesn't reset + // mirror and lookup counters. The lookup counters are accumulated + // into the new cluster's own stats Arc (unless it's the same one), + // so lookups served by the new cluster keep counting. + { + let from = self.stats.lock(); + let mut to = other.stats.lock(); + to.mirror = from.mirror; + if !Arc::ptr_eq(&to.lookup, &from.lookup) { + to.lookup.accumulate(&from.lookup); + } + } + for (from, to) in self.shards.iter().zip(other.shards.iter()) { from.move_conns_to(to)?; } @@ -459,6 +476,15 @@ impl Cluster { Ok(()) } + /// Reset statistics collected by this cluster and its pools, + /// used by the RESET STATS command. + pub fn reset_stats(&self) { + self.stats.lock().reset(); + for shard in &self.shards { + shard.reset_stats(); + } + } + /// Cancel a query executed by one of the shards. pub(crate) async fn cancel(&self, id: FrontendPid) -> Result<(), super::super::Error> { for shard in &self.shards { @@ -1224,4 +1250,78 @@ mod test { cluster.query_parser = QueryParserLevel::Off; assert!(!cluster.use_query_parser(&req)); } + + #[test] + fn test_move_conns_to_carries_cluster_metrics() { + let config = ConfigAndUsers::default(); + let old = Cluster::new_test(&config); + let new = Cluster::new_test(&config); + + // Populate the old cluster's metrics. + { + let mut stats = old.stats.lock(); + stats.mirror.total_count = 11; + stats.mirror.mirrored_count = 7; + stats.mirror.dropped_count = 2; + stats.mirror.error_count = 1; + stats + .lookup + .hits + .fetch_add(5, std::sync::atomic::Ordering::Relaxed); + stats + .lookup + .misses + .fetch_add(3, std::sync::atomic::Ordering::Relaxed); + } + + assert!(old.can_move_conns_to(&new)); + old.move_conns_to(&new).unwrap(); + + let stats = new.stats.lock(); + assert_eq!(stats.mirror.total_count, 11); + assert_eq!(stats.mirror.mirrored_count, 7); + assert_eq!(stats.mirror.dropped_count, 2); + assert_eq!(stats.mirror.error_count, 1); + assert_eq!( + stats.lookup.hits.load(std::sync::atomic::Ordering::Relaxed), + 5 + ); + assert_eq!( + stats + .lookup + .misses + .load(std::sync::atomic::Ordering::Relaxed), + 3 + ); + } + + #[test] + fn test_reset_stats_clears_cluster_and_pool_counters() { + let config = ConfigAndUsers::default(); + let cluster = Cluster::new_test(&config); + + // Populate cluster metrics and pool statistics. + { + let mut stats = cluster.stats.lock(); + stats.mirror.total_count = 11; + stats + .lookup + .hits + .fetch_add(5, std::sync::atomic::Ordering::Relaxed); + } + let pool = cluster.shards()[0].pools()[0].clone(); + pool.lock().stats.counts.query_count = 42; + + cluster.reset_stats(); + + { + let stats = cluster.stats.lock(); + assert_eq!(stats.mirror.total_count, 0); + assert_eq!( + stats.lookup.hits.load(std::sync::atomic::Ordering::Relaxed), + 0 + ); + } + assert_eq!(pool.lock().stats.counts.query_count, 0); + } } diff --git a/pgdog/src/backend/pool/cluster_metrics.rs b/pgdog/src/backend/pool/cluster_metrics.rs index 09d34e31e..94f9fab7e 100644 --- a/pgdog/src/backend/pool/cluster_metrics.rs +++ b/pgdog/src/backend/pool/cluster_metrics.rs @@ -80,6 +80,14 @@ pub(crate) struct ClusterMetrics { pub(crate) lookup: Arc, } +impl ClusterMetrics { + /// Reset all counters to zero, used by the RESET STATS command. + pub(super) fn reset(&mut self) { + self.mirror = Counts::default(); + self.lookup.reset(); + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index 1c351dd74..424d00e9c 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -317,6 +317,13 @@ impl Pool { to_guard.paused = true; } + // Carry pool statistics over to the new pool instance. The two + // pools represent the same backend (they were matched by address), + // so counters and averages should survive config reloads; only + // pools that are actually added or removed between configs start + // with fresh statistics. + to_guard.stats = from_guard.stats; + from_guard.online = false; let (idle, taken) = from_guard.move_conns_to(destination); for server in idle { diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index 341481612..030b61edc 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -20,7 +20,7 @@ use crate::config::{LoadBalancingStrategy, ReadWriteSplit, Role}; use crate::net::Parameters; use crate::net::messages::FrontendPid; -use super::{Error, Guard, LoadBalancer, Pool, PoolConfig, Request}; +use super::{Error, Guard, LoadBalancer, Pool, PoolConfig, Request, Stats}; pub(crate) mod failover_signal; pub(crate) mod monitor; @@ -113,6 +113,14 @@ impl Shard { self.lb.can_move_conns_to(&other.lb) } + /// Reset statistics collected by this shard's pools, + /// used by the RESET STATS command. + pub(crate) fn reset_stats(&self) { + for pool in self.pool_iter() { + pool.lock().stats = Stats::default(); + } + } + /// Listen for notifications on channel. pub(crate) async fn listen(&self, channel: &str) -> Result { match self.pub_sub.load_full().deref() { diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index 4201a3c71..b0e3bb6e9 100644 --- a/pgdog/src/backend/pool/test/mod.rs +++ b/pgdog/src/backend/pool/test/mod.rs @@ -1260,3 +1260,34 @@ async fn test_move_conns_to_does_not_pause_destination_when_source_is_not_paused destination.shutdown(); } + +#[tokio::test] +async fn test_move_conns_to_carries_stats_over() { + // Pool statistics (counts and averages) must survive a config reload: + // the destination pool takes over the source pool's stats so SHOW POOLS + // and OpenMetrics gauges don't reset to zero on every reload. + let source = Pool::new_test(); + let destination = Pool::new_test(); + + source.launch(); + destination.launch(); + + { + let mut guard = source.lock(); + guard.stats.counts.query_count = 42; + guard.stats.counts.xact_count = 7; + guard.stats.counts.query_time = Duration::from_millis(500); + guard.stats.calc_averages(Duration::from_secs(1)); + } + + assert!(source.has_compatible_address_with(&destination)); + source.move_conns_to(&destination).unwrap(); + + let stats = destination.lock().stats; + assert_eq!(stats.counts.query_count, 42); + assert_eq!(stats.counts.xact_count, 7); + assert_eq!(stats.counts.query_time, Duration::from_millis(500)); + assert_eq!(stats.averages.query_time, Duration::from_millis(500) / 42); + + destination.shutdown(); +} diff --git a/pgdog/src/frontend/router/sharding/lookup.rs b/pgdog/src/frontend/router/sharding/lookup.rs index aaea671eb..5aebcf36b 100644 --- a/pgdog/src/frontend/router/sharding/lookup.rs +++ b/pgdog/src/frontend/router/sharding/lookup.rs @@ -340,6 +340,33 @@ impl LookupStats { self.lookup_time_us .fetch_add(elapsed.as_micros() as u64, Ordering::Relaxed); } + + /// Zero all counters, used by the RESET STATS command. + pub(crate) fn reset(&self) { + for counter in [ + &self.hits, + &self.misses, + &self.evictions, + &self.lookups, + &self.lookup_time_us, + ] { + counter.store(0, Ordering::Relaxed); + } + } + + /// Add another snapshot of counters into this one, used when a config + /// reload carries lookup statistics over to the new cluster's cache. + pub(crate) fn accumulate(&self, other: &LookupStats) { + for (to, from) in [ + (&self.hits, &other.hits), + (&self.misses, &other.misses), + (&self.evictions, &other.evictions), + (&self.lookups, &other.lookups), + (&self.lookup_time_us, &other.lookup_time_us), + ] { + to.fetch_add(from.load(Ordering::Relaxed), Ordering::Relaxed); + } + } } /// Sharding key lookup cache. Bounded by approximate memory use;