Fix connection stats handling and add RESET STATS command

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;

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<Postgres>) -> i64 {
+    admin
+        .fetch_all("SHOW STATS")
+        .await
+        .unwrap()
+        .iter()
+        .map(|row| {
+            row.get::<rust_decimal::Decimal, _>("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"
+    );
+}

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::*;

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");

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<Self, Error> {
+        Ok(Self)
+    }
+
+    async fn execute(&self) -> Result<Vec<Message>, Error> {
+        databases().reset_stats();
+        Ok(vec![])
+    }
+}

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

pgdog/src/backend/pool/cluster.rs

@@ -452,13 +452,39 @@ 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)?;
         }
 
         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);
+    }
 }

pgdog/src/backend/pool/cluster_metrics.rs

@@ -80,6 +80,14 @@ pub(crate) struct ClusterMetrics {
     pub(crate) lookup: Arc<LookupStats>,
 }
 
+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::*;