[proxy] Enable keepalives for all tcp connections (#1448)

proxy/Cargo.toml

@@ -22,8 +22,9 @@ rustls = "0.19.1"
 scopeguard = "1.1.0"
 serde = "1"
 serde_json = "1"
+socket2 = "0.4.4"
 thiserror = "1.0"
-tokio = { version = "1.11", features = ["macros"] }
+tokio = { version = "1.17", features = ["macros"] }
 tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }
 tokio-rustls = "0.22.0"
 

proxy/src/compute.rs

@@ -41,6 +41,7 @@ impl DatabaseInfo {
         let host_port = format!("{}:{}", self.host, self.port);
         let socket = TcpStream::connect(host_port).await?;
         let socket_addr = socket.peer_addr()?;
+        socket2::SockRef::from(&socket).set_keepalive(true)?;
 
         Ok((socket_addr, socket))
     }

proxy/src/proxy.rs

@@ -50,6 +50,10 @@ pub async fn thread_main(
         println!("proxy has shut down");
     }
 
+    // When set for the server socket, the keepalive setting
+    // will be inherited by all accepted client sockets.
+    socket2::SockRef::from(&listener).set_keepalive(true)?;
+
     let cancel_map = Arc::new(CancelMap::default());
     loop {
         let (socket, peer_addr) = listener.accept().await?;
@@ -367,4 +371,24 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn keepalive_is_inherited() -> anyhow::Result<()> {
+        use tokio::net::{TcpListener, TcpStream};
+
+        let listener = TcpListener::bind("127.0.0.1:0").await?;
+        let port = listener.local_addr()?.port();
+        socket2::SockRef::from(&listener).set_keepalive(true)?;
+
+        let t = tokio::spawn(async move {
+            let (client, _) = listener.accept().await?;
+            let keepalive = socket2::SockRef::from(&client).keepalive()?;
+            anyhow::Ok(keepalive)
+        });
+
+        let _ = TcpStream::connect(("127.0.0.1", port)).await?;
+        assert!(t.await??, "keepalive should be inherited");
+
+        Ok(())
+    }
 }