perf fixes && fix disconnect server (#63)

Author: vadv

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "pg_doorman"
-version = "2.4.3"
+version = "2.5.0"
 edition = "2021"
 rust-version = "1.87.0"
 license = "MIT"

Cargo.toml

@@ -4,6 +4,16 @@ title: Changelog
 
 # Changelog
 
+### 2.5.0 <small>Nov 18, 2025</small> { id="2.5.0" }
+
+**Improvements:**
+- Reworked the statistics collection system, yielding up to 20% performance gain on fast queries.
+- Improved detection of `SAVEPOINT` usage, allowing the auto-rollback feature to be applied in more situations.
+
+**Bug Fixes / Behavior:**
+- Less aggressive behavior on write errors when sending a response to the client: the server connection is no longer immediately marked as "bad" and evicted from the pool. We now read the remaining server response and clean up its state, returning the connection to the pool in a clean state. This improves performance during client reconnections.
+
+
 ### 2.4.3 <small>Nov 15, 2025</small> { id="2.4.3" }
 
 **Bug Fixes:**

documentation/docs/changelog.md

@@ -175,6 +175,20 @@ syslog_prog_name = "pg_doorman"
 # For example, it can be the address mask 10.0.0.0/8. By default, access is allowed from any address.
 hba = []
 
+[prometheus]
+# Prometheus metrics exporter settings. See documentation/docs/reference/prometheus.md
+# Enable or disable the Prometheus exporter.
+# Default: false
+enabled = false
+
+# Host to listen on for the metrics HTTP endpoint.
+# Default: "0.0.0.0"
+host = "0.0.0.0"
+
+# Port for the metrics HTTP endpoint.
+# Default: 9127
+port = 9127
+
 [pools]
 # Each record in the pool is the name of the virtual database that the pg-doorman client can connect to.
 [pools.exampledb]

pg_doorman.toml

@@ -138,29 +138,15 @@ where
         "free_clients".to_string(),
         client_stats
             .keys()
-            .filter(|client_id| {
-                client_stats
-                    .get(client_id)
-                    .unwrap()
-                    .state
-                    .load(Ordering::Relaxed)
-                    == CLIENT_STATE_IDLE
-            })
+            .filter(|client_id| client_stats.get(client_id).unwrap().state() == CLIENT_STATE_IDLE)
             .count()
             .to_string(),
     ]));
     res.put(data_row(&vec![
         "used_clients".to_string(),
         client_stats
             .keys()
-            .filter(|client_id| {
-                client_stats
-                    .get(client_id)
-                    .unwrap()
-                    .state
-                    .load(Ordering::Relaxed)
-                    == CLIENT_STATE_ACTIVE
-            })
+            .filter(|client_id| client_stats.get(client_id).unwrap().state() == CLIENT_STATE_ACTIVE)
             .count()
             .to_string(),
     ]));
@@ -172,29 +158,15 @@ where
         "free_servers".to_string(),
         server_stats
             .keys()
-            .filter(|server_id| {
-                server_stats
-                    .get(server_id)
-                    .unwrap()
-                    .state
-                    .load(Ordering::Relaxed)
-                    == SERVER_STATE_IDLE
-            })
+            .filter(|server_id| server_stats.get(server_id).unwrap().state() == SERVER_STATE_IDLE)
             .count()
             .to_string(),
     ]));
     res.put(data_row(&vec![
         "used_servers".to_string(),
         server_stats
             .keys()
-            .filter(|server_id| {
-                server_stats
-                    .get(server_id)
-                    .unwrap()
-                    .state
-                    .load(Ordering::Relaxed)
-                    == SERVER_STATE_ACTIVE
-            })
+            .filter(|server_id| server_stats.get(server_id).unwrap().state() == SERVER_STATE_ACTIVE)
             .count()
             .to_string(),
     ]));

src/admin.rs

@@ -201,12 +201,18 @@ fn local_trust_is_ignored_for_tcp() {
     // For SCRAM password
     let ci_scram = ci_from_hba(hba, false);
     let pw_scram = format!("{}secret", SCRAM_SHA_256);
-    assert_eq!(eval_hba_for_pool_password(&pw_scram, &ci_scram), CheckResult::Deny);
+    assert_eq!(
+        eval_hba_for_pool_password(&pw_scram, &ci_scram),
+        CheckResult::Deny
+    );
 
     // For MD5 password
     let ci_md5 = ci_from_hba(hba, false);
     let pw_md5 = format!("{}hash", MD5_PASSWORD_PREFIX);
-    assert_eq!(eval_hba_for_pool_password(&pw_md5, &ci_md5), CheckResult::Deny);
+    assert_eq!(
+        eval_hba_for_pool_password(&pw_md5, &ci_md5),
+        CheckResult::Deny
+    );
 }
 
 #[test]
@@ -217,10 +223,16 @@ fn local_then_host_rule_behaves_like_host_only() {
     // MD5 password should be allowed by the host md5 rule
     let ci_md5 = ci_from_hba(hba, false);
     let pw_md5 = format!("{}hash", MD5_PASSWORD_PREFIX);
-    assert_eq!(eval_hba_for_pool_password(&pw_md5, &ci_md5), CheckResult::Allow);
+    assert_eq!(
+        eval_hba_for_pool_password(&pw_md5, &ci_md5),
+        CheckResult::Allow
+    );
 
     // SCRAM password should be allowed to proceed (scram NotMatched, md5 Allow → overall Allow)
     let ci_scram = ci_from_hba(hba, false);
     let pw_scram = format!("{}secret", SCRAM_SHA_256);
-    assert_eq!(eval_hba_for_pool_password(&pw_scram, &ci_scram), CheckResult::Allow);
+    assert_eq!(
+        eval_hba_for_pool_password(&pw_scram, &ci_scram),
+        CheckResult::Allow
+    );
 }

src/auth/hba_eval_tests.rs

@@ -11,7 +11,7 @@ use std::marker::Unpin;
 
 // External crate imports
 use crate::auth::hba::CheckResult;
-use log::{error, warn, info};
+use log::{error, info, warn};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 // Internal crate imports
 use crate::auth::jwt::get_user_name_from_jwt;
@@ -53,7 +53,10 @@ where
         if client_identifier.hba_md5 == CheckResult::Trust
             || client_identifier.hba_scram == CheckResult::Trust
         {
-            info!("HBA trust for admin user: {username_from_parameters} from: {:?}.", client_identifier.addr);
+            info!(
+                "HBA trust for admin user: {username_from_parameters} from: {:?}.",
+                client_identifier.addr
+            );
             return Ok((false, ServerParameters::admin(), false));
         }
         if client_identifier.hba_md5 == CheckResult::Deny

src/auth/mod.rs

@@ -902,7 +902,7 @@ where
                     Some((process_id, secret_key, address, port)) => {
                         {
                             let mut cancel_guard = CANCELED_PIDS.lock();
-                            cancel_guard.push(*process_id);
+                            cancel_guard.insert(*process_id);
                         }
                         (*process_id, *secret_key, address.clone(), *port)
                     }
@@ -945,20 +945,15 @@ where
                 self.stats.disconnect();
                 return Ok(());
             }
-            tokio::select! {
-                _ = self.shutdown.recv() => {
-                    if !self.admin {
-                        warn!("Dropping client {:?} because connection pooler is shutting down", self.addr);
-                        error_response_terminal(
-                            &mut self.write,
-                            "pooler is shut down now",
-                            "58006"
-                        ).await?;
-                        self.stats.disconnect();
-                        return Ok(());
-                    }
-                },
-                _ = tokio::task::yield_now() => {}
+            if self.shutdown.try_recv().is_ok() && !self.admin {
+                warn!(
+                    "Dropping client {:?} because connection pooler is shutting down",
+                    self.addr
+                );
+                error_response_terminal(&mut self.write, "pooler is shut down now", "58006")
+                    .await?;
+                self.stats.disconnect();
+                return Ok(());
             }
             // Handle admin database queries.
             if self.admin {
@@ -977,16 +972,16 @@ where
 
             match message[0] as char {
                 'Q' => {
-                    if self.pooler_check_query_request_vec.eq(&message.to_vec()) {
+                    if self.pooler_check_query_request_vec.as_slice() == &message[..] {
                         // This is the first message in the transaction, since we are responding with 'IZ',
                         // then we can not expect a server connection and immediately send answer and exit transaction loop.
                         write_all_flush(&mut self.write, &check_query_response()).await?;
                         continue;
                     }
                     if message.len() < 40 && message.len() > QUERY_DEALLOCATE.len() + 5 {
                         // Do not pass simple query with deallocate, as it will run on an unknown server.
-                        let query = message[5..QUERY_DEALLOCATE.len() + 5].to_vec();
-                        if QUERY_DEALLOCATE.eq(&query) {
+                        let query = &message[5..QUERY_DEALLOCATE.len() + 5];
+                        if QUERY_DEALLOCATE == query {
                             write_all_flush(&mut self.write, &deallocate_response()).await?;
                             continue;
                         }
@@ -1046,7 +1041,7 @@ where
                             {
                                 let mut guard = CANCELED_PIDS.lock();
                                 if guard.contains(&conn.get_process_id()) {
-                                    guard.retain(|&id| id != conn.get_process_id());
+                                    guard.remove(&conn.get_process_id());
                                     conn.mark_bad("because was canceled");
                                     continue; // try to find another server.
                                 }
@@ -1855,12 +1850,12 @@ where
             self.stats.active_write();
             if let Err(err_write) = write_all_flush(&mut self.write, &response).await {
                 server.wait_available().await;
-                // Lazy error formatting
-                let mut msg = String::with_capacity(64);
-                use std::fmt::Write;
-                let _ = write!(msg, "flush to client {} {:?}", self.addr, err_write);
-                server.mark_bad(&msg);
-                return Err(err_write);
+                if server.is_async() || server.in_copy_mode() {
+                    server.mark_bad(
+                        format!("flush to client {} {:?}", self.addr, err_write).as_str(),
+                    );
+                    return Err(err_write);
+                }
             }
 
             self.stats.active_idle();

src/client.rs

@@ -13,7 +13,6 @@ use std::hash::{Hash, Hasher};
 use std::mem;
 use std::net::IpAddr;
 use std::path::Path;
-use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 use tokio::fs::File;
 use tokio::io::AsyncReadExt;
@@ -51,8 +50,6 @@ pub struct Address {
     pub pool_name: String,
     /// Address stats
     pub stats: Arc<AddressStats>,
-    /// Number of errors encountered since last successful checkout
-    pub error_count: Arc<AtomicU64>,
 }
 
 impl Default for Address {
@@ -66,7 +63,6 @@ impl Default for Address {
             password: String::from("password"),
             pool_name: String::from("pool_name"),
             stats: Arc::new(AddressStats::default()),
-            error_count: Arc::new(AtomicU64::new(0)),
         }
     }
 }
@@ -111,17 +107,6 @@ impl Address {
     pub fn name(&self) -> String {
         self.pool_name.clone() + "-" + &*self.virtual_pool_id.to_string()
     }
-    pub fn error_count(&self) -> u64 {
-        self.error_count.load(Ordering::Relaxed)
-    }
-
-    pub fn increment_error_count(&self) {
-        self.error_count.fetch_add(1, Ordering::Relaxed);
-    }
-
-    pub fn reset_error_count(&self) {
-        self.error_count.store(0, Ordering::Relaxed);
-    }
 }
 
 /// Pool mode:

src/config.rs

@@ -4,10 +4,10 @@ use log::{error, info, warn};
 use lru::LruCache;
 use once_cell::sync::Lazy;
 use parking_lot::Mutex;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Display, Formatter};
 use std::num::NonZeroUsize;
-use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -31,8 +31,8 @@ pub type PoolMap = HashMap<PoolIdentifierVirtual, ConnectionPool>;
 /// This is atomic and safe and read-optimized.
 /// The pool is recreated dynamically when the config is reloaded.
 pub static POOLS: Lazy<ArcSwap<PoolMap>> = Lazy::new(|| ArcSwap::from_pointee(HashMap::default()));
-pub static CANCELED_PIDS: Lazy<Arc<Mutex<Vec<ProcessId>>>> =
-    Lazy::new(|| Arc::new(Mutex::new(Vec::new())));
+pub static CANCELED_PIDS: Lazy<Arc<Mutex<HashSet<ProcessId>>>> =
+    Lazy::new(|| Arc::new(Mutex::new(HashSet::new())));
 
 pub type PreparedStatementCacheType = Arc<Mutex<PreparedStatementCache>>;
 pub type ServerParametersType = Arc<tokio::sync::Mutex<ServerParameters>>;
@@ -247,7 +247,6 @@ impl ConnectionPool {
                         password: user.password.clone(),
                         pool_name: pool_name.clone(),
                         stats: Arc::new(AddressStats::default()),
-                        error_count: Arc::new(AtomicU64::new(0)),
                     };
 
                     let prepared_statements_cache_size = match config.general.prepared_statements {