fix: put Sync into its own ClientRequest when spliced, feat: add client_connection_recovery setting (#764)

- fix: put `Sync` into its own `ClientRequest`. This affects sharded
spliced requests only (multiple `Execute`/`Bind` in one request).
Previously, Sync was only attached to the last request only. Now, it's
sent individually to all shards, making sure all shards execute
requests.
- feat: add `client_connection_recovery` setting. Allows to disconnect
clients on hard to recover connection errors, like pool timeouts.

```toml
[general]
client_connection_recovery = "drop"
```

Author: levkk

integration/rust/tests/integration/connection_recovery.rs

@@ -0,0 +1,196 @@
+use rust::setup::{admin_sqlx, admin_tokio};
+use serial_test::serial;
+use sqlx::{Executor, Row};
+use std::time::Duration;
+use tokio::time::sleep;
+use tokio_postgres::NoTls;
+
+async fn connection_tokio(db: &str) -> tokio_postgres::Client {
+    let (client, connection) = tokio_postgres::connect(
+        &format!(
+            "host=127.0.0.1 user=pgdog dbname={} password=pgdog port=6432",
+            db
+        ),
+        NoTls,
+    )
+    .await
+    .unwrap();
+
+    tokio::spawn(async move {
+        if let Err(e) = connection.await {
+            eprintln!("connection error: {}", e);
+        }
+    });
+
+    client
+}
+
+async fn get_pool_ids(database: &str) -> (i64, i64) {
+    let admin = admin_sqlx().await;
+    let pools = admin.fetch_all("SHOW POOLS").await.unwrap();
+
+    let primary_id = pools
+        .iter()
+        .find(|p| {
+            p.get::<String, &str>("database") == database
+                && p.get::<String, &str>("user") == "pgdog"
+                && p.get::<String, &str>("role") == "primary"
+        })
+        .map(|p| p.get::<i64, &str>("id"))
+        .unwrap();
+
+    let replica_id = pools
+        .iter()
+        .find(|p| {
+            p.get::<String, &str>("database") == database
+                && p.get::<String, &str>("user") == "pgdog"
+                && p.get::<String, &str>("role") == "replica"
+        })
+        .map(|p| p.get::<i64, &str>("id"))
+        .unwrap();
+
+    (primary_id, replica_id)
+}
+
+async fn ban_pools(database: &str) {
+    let admin = admin_sqlx().await;
+    let (primary_id, replica_id) = get_pool_ids(database).await;
+
+    admin
+        .execute(format!("BAN {}", primary_id).as_str())
+        .await
+        .unwrap();
+    admin
+        .execute(format!("BAN {}", replica_id).as_str())
+        .await
+        .unwrap();
+}
+
+async fn unban_pools(database: &str) {
+    let admin = admin_sqlx().await;
+    let (primary_id, replica_id) = get_pool_ids(database).await;
+
+    admin
+        .execute(format!("UNBAN {}", primary_id).as_str())
+        .await
+        .unwrap();
+    admin
+        .execute(format!("UNBAN {}", replica_id).as_str())
+        .await
+        .unwrap();
+}
+
+#[tokio::test]
+#[serial]
+async fn test_client_connection_recovery_default() {
+    let admin = admin_tokio().await;
+
+    // Ensure recover mode (default)
+    admin
+        .simple_query("SET client_connection_recovery TO 'recover'")
+        .await
+        .unwrap();
+
+    // Give pools time to reinitialize
+    sleep(Duration::from_millis(200)).await;
+
+    // Connection that we'll test recovery on
+    let conn = connection_tokio("pgdog").await;
+
+    // Ban both primary and replica pools to force Banned error
+    ban_pools("pgdog").await;
+
+    // Give ban time to take effect
+    sleep(Duration::from_millis(50)).await;
+
+    // This query should fail with Banned error because both pools are banned
+    conn.simple_query("BEGIN").await.unwrap();
+    let result = conn.simple_query("SELECT 1").await;
+    assert!(result.is_err(), "Expected error when pools are banned");
+    let err = result.unwrap_err();
+    let err_str = err.to_string().to_lowercase();
+    assert!(
+        err_str.contains("all replicas down"),
+        "Expected 'all replicas down' error, got: {}",
+        err
+    );
+
+    // Unban pools
+    unban_pools("pgdog").await;
+
+    // Give unban time to take effect
+    sleep(Duration::from_millis(100)).await;
+
+    // In recovery mode, the connection should still be usable
+    let result = conn.simple_query("BEGIN").await;
+    assert!(
+        result.is_ok(),
+        "Expected query to succeed after recovery, got: {:?}",
+        result.err()
+    );
+    let result = conn.simple_query("SELECT 1").await;
+    assert!(
+        result.is_ok(),
+        "Expected SELECT to succeed after recovery, got: {:?}",
+        result.err()
+    );
+    conn.simple_query("COMMIT").await.unwrap();
+
+    // Reset settings
+    admin.simple_query("RELOAD").await.unwrap();
+}
+
+#[tokio::test]
+#[serial]
+async fn test_client_connection_recovery_drop() {
+    let admin = admin_tokio().await;
+
+    // Set drop mode - client should be disconnected on recoverable errors
+    admin
+        .simple_query("SET client_connection_recovery TO 'drop'")
+        .await
+        .unwrap();
+
+    // Give pools time to reinitialize
+    sleep(Duration::from_millis(200)).await;
+
+    // Connection that we'll test drop behavior on
+    let conn = connection_tokio("pgdog").await;
+
+    // Ban both pools to force Banned error
+    ban_pools("pgdog").await;
+
+    // Give ban time to take effect
+    sleep(Duration::from_millis(50)).await;
+
+    // This query should fail because pools are banned
+    conn.simple_query("BEGIN").await.unwrap();
+    let result = conn.simple_query("SELECT 1").await;
+    assert!(result.is_err(), "Expected error when pools are banned");
+
+    // Unban pools
+    unban_pools("pgdog").await;
+
+    // Give unban time to take effect
+    sleep(Duration::from_millis(100)).await;
+
+    // In drop mode, the connection should be disconnected
+    // Subsequent queries should fail because the connection is closed
+    let result = conn.simple_query("BEGIN").await;
+    assert!(
+        result.is_err(),
+        "Expected query to fail because connection was dropped"
+    );
+
+    // Verify that the error indicates connection was closed
+    let err = result.unwrap_err();
+    let err_str = err.to_string().to_lowercase();
+    assert!(
+        err_str.contains("connection closed"),
+        "Expected 'connection closed' error, got: {}",
+        err
+    );
+
+    // Reset settings
+    admin.simple_query("RELOAD").await.unwrap();
+}

integration/rust/tests/integration/mod.rs

@@ -4,6 +4,7 @@ pub mod auto_id;
 pub mod avg;
 pub mod ban;
 pub mod client_ids;
+pub mod connection_recovery;
 pub mod cross_shard_disabled;
 pub mod distinct;
 pub mod explain;

pgdog-config/src/general.rs

@@ -182,6 +182,9 @@ pub struct General {
     /// Connection cleanup algorithm.
     #[serde(default = "General::connection_recovery")]
     pub connection_recovery: ConnectionRecovery,
+    /// Client connection recovery
+    #[serde(default = "General::client_connection_recovery")]
+    pub client_connection_recovery: ConnectionRecovery,
     /// LSN check interval.
     #[serde(default = "General::lsn_check_interval")]
     pub lsn_check_interval: u64,
@@ -270,6 +273,7 @@ impl Default for General {
             server_lifetime: Self::server_lifetime(),
             stats_period: Self::stats_period(),
             connection_recovery: Self::connection_recovery(),
+            client_connection_recovery: Self::client_connection_recovery(),
             lsn_check_interval: Self::lsn_check_interval(),
             lsn_check_timeout: Self::lsn_check_timeout(),
             lsn_check_delay: Self::lsn_check_delay(),
@@ -601,6 +605,10 @@ impl General {
         Self::env_enum_or_default("PGDOG_CONNECTION_RECOVERY")
     }
 
+    pub fn client_connection_recovery() -> ConnectionRecovery {
+        Self::env_enum_or_default("PGDOG_CLIENT_CONNECTION_RECOVERY")
+    }
+
     fn stats_period() -> u64 {
         Self::env_or_default("PGDOG_STATS_PERIOD", 15_000)
     }

pgdog/src/admin/set.rs

@@ -164,6 +164,22 @@ impl Command for Set {
                 config.config.general.reload_schema_on_ddl = Self::from_json(&self.value)?;
             }
 
+            "connection_recovery" => {
+                config.config.general.connection_recovery = Self::from_json(&self.value)?;
+            }
+
+            "client_connection_recovery" => {
+                config.config.general.client_connection_recovery = Self::from_json(&self.value)?;
+            }
+
+            "default_pool_size" => {
+                config.config.general.default_pool_size = self.value.parse()?;
+            }
+
+            "connect_timeout" => {
+                config.config.general.connect_timeout = self.value.parse()?;
+            }
+
             _ => return Err(Error::Syntax),
         }
 

pgdog/src/backend/pool/cluster.rs

@@ -71,6 +71,7 @@ pub struct Cluster {
     pub_sub_channel_size: usize,
     query_parser: QueryParserLevel,
     connection_recovery: ConnectionRecovery,
+    client_connection_recovery: ConnectionRecovery,
     query_parser_engine: QueryParserEngine,
     reload_schema_on_ddl: bool,
 }
@@ -144,6 +145,7 @@ pub struct ClusterConfig<'a> {
     pub query_parser: QueryParserLevel,
     pub query_parser_engine: QueryParserEngine,
     pub connection_recovery: ConnectionRecovery,
+    pub client_connection_recovery: ConnectionRecovery,
     pub lsn_check_interval: Duration,
     pub reload_schema_on_ddl: bool,
 }
@@ -192,6 +194,7 @@ impl<'a> ClusterConfig<'a> {
             query_parser: general.query_parser,
             query_parser_engine: general.query_parser_engine,
             connection_recovery: general.connection_recovery,
+            client_connection_recovery: general.client_connection_recovery,
             lsn_check_interval: Duration::from_millis(general.lsn_check_interval),
             reload_schema_on_ddl: general.reload_schema_on_ddl,
         }
@@ -225,6 +228,7 @@ impl Cluster {
             pub_sub_channel_size,
             query_parser,
             connection_recovery,
+            client_connection_recovery,
             lsn_check_interval,
             query_parser_engine,
             reload_schema_on_ddl,
@@ -272,6 +276,7 @@ impl Cluster {
             pub_sub_channel_size,
             query_parser,
             connection_recovery,
+            client_connection_recovery,
             query_parser_engine,
             reload_schema_on_ddl,
         }
@@ -379,6 +384,10 @@ impl Cluster {
         &self.connection_recovery
     }
 
+    pub fn client_connection_recovery(&self) -> &ConnectionRecovery {
+        &self.client_connection_recovery
+    }
+
     pub fn dry_run(&self) -> bool {
         self.dry_run
     }

pgdog/src/backend/pool/connection/binding.rs

@@ -160,7 +160,14 @@ impl Binding {
                     result?;
                 }
 
-                state.update(shards_sent, client_request.route());
+                // For Sync-only requests, update shards count but don't reset counters.
+                // Sync needs correct shards for ReadyForQuery counting, but we must
+                // preserve buffered CommandComplete from previous queries.
+                if client_request.is_sync_only() {
+                    state.update_shards(shards_sent);
+                } else {
+                    state.update(shards_sent, client_request.route());
+                }
 
                 Ok(())
             }

pgdog/src/backend/pool/connection/multi_shard/mod.rs

@@ -81,6 +81,13 @@ impl MultiShard {
         self.route = route.clone();
     }
 
+    /// Update only the shards count without resetting counters.
+    /// Used for Sync-only requests where we need correct ReadyForQuery counting
+    /// but must preserve buffered CommandComplete from previous queries.
+    pub(super) fn update_shards(&mut self, shards: usize) {
+        self.shards = shards;
+    }
+
     pub(super) fn reset(&mut self) {
         self.counters = Counters::default();
         self.buffer.reset();