From dae240d30c16be7b59e9fae3b1600a34826790b3 Mon Sep 17 00:00:00 2001
From: Lev Kokotov <levkk@users.noreply.github.com>
Date: Mon, 6 Nov 2023 12:18:52 -0800
Subject: [PATCH] Add connet_timeout and idle_timeout to the user (#634)

* Add connect_timeout to the user

* Allow user to override connect timeout

* version

* lock

* Add both timeouts to the user
---
 Cargo.lock              |  2 +-
 Cargo.toml              |  2 +-
 pgcat.toml              |  2 ++
 src/auth_passthrough.rs |  2 ++
 src/config.rs           | 22 ++++++++++++++++++++++
 src/pool.rs             | 14 ++++++++++----
 6 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index ff7fddf9..e5043975 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1020,7 +1020,7 @@ checksum = "9b2a4787296e9989611394c33f193f676704af1686e70b8f8033ab5ba9a35a94"
 
 [[package]]
 name = "pgcat"
-version = "1.1.2-dev1"
+version = "1.1.2-dev2"
 dependencies = [
  "arc-swap",
  "async-trait",
diff --git a/Cargo.toml b/Cargo.toml
index b79270ea..6485622a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pgcat"
-version = "1.1.2-dev1"
+version = "1.1.2-dev2"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/pgcat.toml b/pgcat.toml
index 841649ee..9e19c13b 100644
--- a/pgcat.toml
+++ b/pgcat.toml
@@ -301,6 +301,8 @@ username = "other_user"
 password = "other_user"
 pool_size = 21
 statement_timeout = 15000
+connect_timeout = 1000
+idle_timeout = 1000
 
 # Shard configs are structured as pool.<pool_name>.shards.<shard_id>
 # Each shard config contains a list of servers that make up the shard
diff --git a/src/auth_passthrough.rs b/src/auth_passthrough.rs
index fc0f6dc6..159847ed 100644
--- a/src/auth_passthrough.rs
+++ b/src/auth_passthrough.rs
@@ -79,6 +79,8 @@ impl AuthPassthrough {
             pool_mode: None,
             server_lifetime: None,
             min_pool_size: None,
+            connect_timeout: None,
+            idle_timeout: None,
         };
 
         let user = &address.username;
diff --git a/src/config.rs b/src/config.rs
index 3d140b9b..ef7952f2 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -216,6 +216,8 @@ pub struct User {
     pub server_lifetime: Option<u64>,
     #[serde(default)] // 0
     pub statement_timeout: u64,
+    pub connect_timeout: Option<u64>,
+    pub idle_timeout: Option<u64>,
 }
 
 impl Default for User {
@@ -230,6 +232,8 @@ impl Default for User {
             statement_timeout: 0,
             pool_mode: None,
             server_lifetime: None,
+            connect_timeout: None,
+            idle_timeout: None,
         }
     }
 }
@@ -1307,6 +1311,24 @@ impl Config {
                         None => "default".to_string(),
                     }
                 );
+                info!(
+                    "[pool: {}][user: {}] Connection timeout: {}",
+                    pool_name,
+                    user.1.username,
+                    match user.1.connect_timeout {
+                        Some(connect_timeout) => format!("{}ms", connect_timeout),
+                        None => "not set".to_string(),
+                    }
+                );
+                info!(
+                    "[pool: {}][user: {}] Idle timeout: {}",
+                    pool_name,
+                    user.1.username,
+                    match user.1.idle_timeout {
+                        Some(idle_timeout) => format!("{}ms", idle_timeout),
+                        None => "not set".to_string(),
+                    }
+                );
             }
         }
     }
diff --git a/src/pool.rs b/src/pool.rs
index 0de16c2a..751f1875 100644
--- a/src/pool.rs
+++ b/src/pool.rs
@@ -436,14 +436,20 @@ impl ConnectionPool {
                             pool_config.prepared_statements_cache_size,
                         );
 
-                        let connect_timeout = match pool_config.connect_timeout {
+                        let connect_timeout = match user.connect_timeout {
                             Some(connect_timeout) => connect_timeout,
-                            None => config.general.connect_timeout,
+                            None => match pool_config.connect_timeout {
+                                Some(connect_timeout) => connect_timeout,
+                                None => config.general.connect_timeout,
+                            },
                         };
 
-                        let idle_timeout = match pool_config.idle_timeout {
+                        let idle_timeout = match user.idle_timeout {
                             Some(idle_timeout) => idle_timeout,
-                            None => config.general.idle_timeout,
+                            None => match pool_config.idle_timeout {
+                                Some(idle_timeout) => idle_timeout,
+                                None => config.general.idle_timeout,
+                            },
                         };
 
                         let server_lifetime = match user.server_lifetime {