Skip to content

Commit ca4431b

Browse files
authored
Add idle client in transaction configuration (#380)
* Add idle client in transaction configuration * fmt * Update docs * trigger build * Add tests * Make the config dynamic from reloads * fmt * comments * trigger build * fix config.md * remove error
1 parent d66b377 commit ca4431b

File tree

5 files changed

+111
-4
lines changed

5 files changed

+111
-4
lines changed

CONFIG.md

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,14 @@ default: 30000 # milliseconds
4949

5050
How long an idle connection with a server is left open (ms).
5151

52+
### idle_client_in_transaction_timeout
53+
```
54+
path: general.idle_client_in_transaction_timeout
55+
default: 0 # milliseconds
56+
```
57+
58+
How long a client is allowed to be idle while in a transaction (ms).
59+
5260
### healthcheck_timeout
5361
```
5462
path: general.healthcheck_timeout

pgcat.toml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,9 @@ connect_timeout = 5000 # milliseconds
2323
# How long an idle connection with a server is left open (ms).
2424
idle_timeout = 30000 # milliseconds
2525

26+
# How long a client is allowed to be idle while in a transaction (ms).
27+
idle_client_in_transaction_timeout = 0 # milliseconds
28+
2629
# How much time to give the health check query to return with a result (ms).
2730
healthcheck_timeout = 1000 # milliseconds
2831

src/client.rs

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ use tokio::sync::broadcast::Receiver;
1212
use tokio::sync::mpsc::Sender;
1313

1414
use crate::admin::{generate_server_info_for_admin, handle_admin};
15-
use crate::config::{get_config, Address, PoolMode};
15+
use crate::config::{get_config, get_idle_client_in_transaction_timeout, Address, PoolMode};
1616
use crate::constants::*;
1717

1818
use crate::messages::*;
@@ -859,6 +859,11 @@ where
859859

860860
let mut initial_message = Some(message);
861861

862+
let idle_client_timeout_duration = match get_idle_client_in_transaction_timeout() {
863+
0 => tokio::time::Duration::MAX,
864+
timeout => tokio::time::Duration::from_millis(timeout),
865+
};
866+
862867
// Transaction loop. Multiple queries can be issued by the client here.
863868
// The connection belongs to the client until the transaction is over,
864869
// or until the client disconnects if we are in session mode.
@@ -870,15 +875,26 @@ where
870875
None => {
871876
trace!("Waiting for message inside transaction or in session mode");
872877

873-
match read_message(&mut self.read).await {
874-
Ok(message) => message,
875-
Err(err) => {
878+
match tokio::time::timeout(
879+
idle_client_timeout_duration,
880+
read_message(&mut self.read),
881+
)
882+
.await
883+
{
884+
Ok(Ok(message)) => message,
885+
Ok(Err(err)) => {
876886
// Client disconnected inside a transaction.
877887
// Clean up the server and re-use it.
878888
server.checkin_cleanup().await?;
879889

880890
return Err(err);
881891
}
892+
Err(_) => {
893+
// Client idle in transaction timeout
894+
error_response(&mut self.write, "idle transaction timeout").await?;
895+
error!("Client idle in transaction timeout: {{ pool_name: {:?}, username: {:?}, shard: {:?}, role: \"{:?}\"}}", self.pool_name.clone(), self.username.clone(), query_router.shard(), query_router.role());
896+
break;
897+
}
882898
}
883899
}
884900
Some(message) => {

src/config.rs

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,9 @@ pub struct General {
197197
#[serde(default = "General::default_ban_time")]
198198
pub ban_time: i64,
199199

200+
#[serde(default = "General::default_idle_client_in_transaction_timeout")]
201+
pub idle_client_in_transaction_timeout: u64,
202+
200203
#[serde(default = "General::default_worker_threads")]
201204
pub worker_threads: usize,
202205

@@ -260,6 +263,10 @@ impl General {
260263
pub fn default_worker_threads() -> usize {
261264
4
262265
}
266+
267+
pub fn default_idle_client_in_transaction_timeout() -> u64 {
268+
0
269+
}
263270
}
264271

265272
impl Default for General {
@@ -276,6 +283,7 @@ impl Default for General {
276283
healthcheck_delay: Self::default_healthcheck_delay(),
277284
ban_time: Self::default_ban_time(),
278285
worker_threads: Self::default_worker_threads(),
286+
idle_client_in_transaction_timeout: Self::default_idle_client_in_transaction_timeout(),
279287
tcp_keepalives_idle: Self::default_tcp_keepalives_idle(),
280288
tcp_keepalives_count: Self::default_tcp_keepalives_count(),
281289
tcp_keepalives_interval: Self::default_tcp_keepalives_interval(),
@@ -655,6 +663,13 @@ impl From<&Config> for std::collections::HashMap<String, String> {
655663
config.general.healthcheck_delay.to_string(),
656664
),
657665
("ban_time".to_string(), config.general.ban_time.to_string()),
666+
(
667+
"idle_client_in_transaction_timeout".to_string(),
668+
config
669+
.general
670+
.idle_client_in_transaction_timeout
671+
.to_string(),
672+
),
658673
];
659674

660675
r.append(&mut static_settings);
@@ -666,6 +681,10 @@ impl Config {
666681
/// Print current configuration.
667682
pub fn show(&self) {
668683
info!("Ban time: {}s", self.general.ban_time);
684+
info!(
685+
"Idle client in transaction timeout: {}ms",
686+
self.general.idle_client_in_transaction_timeout
687+
);
669688
info!("Worker threads: {}", self.general.worker_threads);
670689
info!(
671690
"Healthcheck timeout: {}ms",
@@ -819,6 +838,12 @@ pub fn get_config() -> Config {
819838
(*(*CONFIG.load())).clone()
820839
}
821840

841+
pub fn get_idle_client_in_transaction_timeout() -> u64 {
842+
(*(*CONFIG.load()))
843+
.general
844+
.idle_client_in_transaction_timeout
845+
}
846+
822847
/// Parse the configuration file located at the path.
823848
pub async fn parse(path: &str) -> Result<(), Error> {
824849
let mut contents = String::new();
@@ -889,6 +914,7 @@ mod test {
889914
assert_eq!(get_config().path, "pgcat.toml".to_string());
890915

891916
assert_eq!(get_config().general.ban_time, 60);
917+
assert_eq!(get_config().general.idle_client_in_transaction_timeout, 0);
892918
assert_eq!(get_config().general.idle_timeout, 30000);
893919
assert_eq!(get_config().pools.len(), 2);
894920
assert_eq!(get_config().pools["sharded_db"].shards.len(), 3);

tests/ruby/misc_spec.rb

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -309,4 +309,58 @@
309309
end
310310
end
311311
end
312+
313+
describe "Idle client timeout" do
314+
context "idle transaction timeout set to 0" do
315+
before do
316+
current_configs = processes.pgcat.current_config
317+
correct_idle_client_transaction_timeout = current_configs["general"]["idle_client_in_transaction_timeout"]
318+
puts(current_configs["general"]["idle_client_in_transaction_timeout"])
319+
320+
current_configs["general"]["idle_client_in_transaction_timeout"] = 0
321+
322+
processes.pgcat.update_config(current_configs) # with timeout 0
323+
processes.pgcat.reload_config
324+
end
325+
326+
it "Allow client to be idle in transaction" do
327+
conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
328+
conn.async_exec("BEGIN")
329+
conn.async_exec("SELECT 1")
330+
sleep(2)
331+
conn.async_exec("COMMIT")
332+
conn.close
333+
end
334+
end
335+
336+
context "idle transaction timeout set to 500ms" do
337+
before do
338+
current_configs = processes.pgcat.current_config
339+
correct_idle_client_transaction_timeout = current_configs["general"]["idle_client_in_transaction_timeout"]
340+
current_configs["general"]["idle_client_in_transaction_timeout"] = 500
341+
342+
processes.pgcat.update_config(current_configs) # with timeout 500
343+
processes.pgcat.reload_config
344+
end
345+
346+
it "Allow client to be idle in transaction below timeout" do
347+
conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
348+
conn.async_exec("BEGIN")
349+
conn.async_exec("SELECT 1")
350+
sleep(0.4) # below 500ms
351+
conn.async_exec("COMMIT")
352+
conn.close
353+
end
354+
355+
it "Error when client idle in transaction time exceeds timeout" do
356+
conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
357+
conn.async_exec("BEGIN")
358+
conn.async_exec("SELECT 1")
359+
sleep(1) # above 500ms
360+
expect{ conn.async_exec("COMMIT") }.to raise_error(PG::SystemError, /idle transaction timeout/)
361+
conn.async_exec("SELECT 1") # should be able to send another query
362+
conn.close
363+
end
364+
end
365+
end
312366
end

0 commit comments

Comments
 (0)