Introduce tcp_keepalives to PgCat (#315)

We have encountered a case where PgCat pools were stuck following a database incident. Our best understanding at this point is that the PgCat -> Postgres connections died silently and because Tokio defaults to disabling keepalives, connections in the pool were marked as busy forever. Only when we deployed PgCat did we see recovery.

This PR introduces tcp_keepalives to PgCat. This sets the defaults to be

keepalives_idle: 5        # seconds
keepalives_interval: 5 # seconds
keepalives_count: 5    # a count
These settings can detect the death of an idle connection within 30 seconds of its death. Please note that the connection can remain idle forever (from an application perspective) as long as the keepalive packets are flowing so disconnection will only occur if the other end is not acknowledging keepalive packets (keepalive packet acks are handled by the OS, the application does not need to do anything). I plan to add tcp_user_timeout in a follow-up PR.

src/config.rs

@@ -160,6 +160,13 @@ pub struct General {
     #[serde(default = "General::default_idle_timeout")]
     pub idle_timeout: u64,
 
+    #[serde(default = "General::default_tcp_keepalives_idle")]
+    pub tcp_keepalives_idle: u64,
+    #[serde(default = "General::default_tcp_keepalives_count")]
+    pub tcp_keepalives_count: u32,
+    #[serde(default = "General::default_tcp_keepalives_interval")]
+    pub tcp_keepalives_interval: u64,
+
     #[serde(default)] // False
     pub log_client_connections: bool,
 
@@ -203,6 +210,21 @@ impl General {
         1000
     }
 
+    // These keepalive defaults should detect a dead connection within 30 seconds.
+    // Tokio defaults to disabling keepalives which keeps dead connections around indefinitely.
+    // This can lead to permenant server pool exhaustion
+    pub fn default_tcp_keepalives_idle() -> u64 {
+        5 // 5 seconds
+    }
+
+    pub fn default_tcp_keepalives_count() -> u32 {
+        5 // 5 time
+    }
+
+    pub fn default_tcp_keepalives_interval() -> u64 {
+        5 // 5 seconds
+    }
+
     pub fn default_idle_timeout() -> u64 {
         60000 // 10 minutes
     }
@@ -242,6 +264,9 @@ impl Default for General {
             healthcheck_delay: Self::default_healthcheck_delay(),
             ban_time: Self::default_ban_time(),
             worker_threads: Self::default_worker_threads(),
+            tcp_keepalives_idle: Self::default_tcp_keepalives_idle(),
+            tcp_keepalives_count: Self::default_tcp_keepalives_count(),
+            tcp_keepalives_interval: Self::default_tcp_keepalives_interval(),
             log_client_connections: false,
             log_client_disconnections: false,
             autoreload: false,

src/messages.rs

@@ -1,14 +1,18 @@
 /// Helper functions to send one-off protocol messages
 /// and handle TcpStream (TCP socket).
 use bytes::{Buf, BufMut, BytesMut};
+use log::error;
 use md5::{Digest, Md5};
+use socket2::{SockRef, TcpKeepalive};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpStream;
 
+use crate::config::get_config;
 use crate::errors::Error;
 use std::collections::HashMap;
 use std::io::{BufRead, Cursor};
 use std::mem;
+use std::time::Duration;
 
 /// Postgres data type mappings
 /// used in RowDescription ('T') message.
@@ -550,6 +554,26 @@ pub fn server_parameter_message(key: &str, value: &str) -> BytesMut {
     server_info
 }
 
+pub fn configure_socket(stream: &TcpStream) {
+    let sock_ref = SockRef::from(stream);
+    let conf = get_config();
+
+    match sock_ref.set_keepalive(true) {
+        Ok(_) => {
+            match sock_ref.set_tcp_keepalive(
+                &TcpKeepalive::new()
+                    .with_interval(Duration::from_secs(conf.general.tcp_keepalives_interval))
+                    .with_retries(conf.general.tcp_keepalives_count)
+                    .with_time(Duration::from_secs(conf.general.tcp_keepalives_idle)),
+            ) {
+                Ok(_) => (),
+                Err(err) => error!("Could not configure socket: {}", err),
+            }
+        }
+        Err(err) => error!("Could not configure socket: {}", err),
+    }
+}
+
 pub trait BytesMutReader {
     fn read_string(&mut self) -> Result<String, Error>;
 }

src/server.rs

@@ -92,6 +92,7 @@ impl Server {
                     )));
                 }
             };
+        configure_socket(&stream);
 
         trace!("Sending StartupMessage");
 
@@ -368,6 +369,7 @@ impl Server {
                 return Err(Error::SocketError(format!("Error reading cancel message")));
             }
         };
+        configure_socket(&stream);
 
         debug!("Sending CancelRequest");