Add the ability to set TCP recv buf size

4 files changed, 55 insertions, 2 deletions

diff --git a/.bleep b/.bleep
index 6426a33..9cae03e 100644
--- a/.bleep
+++ b/.bleep
@@ -1 +1 @@
-0460a7c567a3dcb3cc2f9940a8c8bc3307036fb2
\ No newline at end of file
+54aea7e241fff10bb31cf067348afb8c139fa39c
\ No newline at end of file
diff --git a/pingora-core/src/connectors/l4.rs b/pingora-core/src/connectors/l4.rs
index 3db2771..95c32ea 100644
--- a/pingora-core/src/connectors/l4.rs
+++ b/pingora-core/src/connectors/l4.rs
@@ -18,7 +18,9 @@ use rand::seq::SliceRandom;
 use std::net::SocketAddr as InetSocketAddr;
 use std::os::unix::io::AsRawFd;
 
-use crate::protocols::l4::ext::{connect as tcp_connect, connect_uds, set_tcp_keepalive};
+use crate::protocols::l4::ext::{
+    connect as tcp_connect, connect_uds, set_recv_buf, set_tcp_keepalive,
+};
 use crate::protocols::l4::socket::SocketAddr;
 use crate::protocols::l4::stream::Stream;
 use crate::protocols::{GetSocketDigest, SocketDigest};
@@ -53,6 +55,10 @@ where
                         debug!("Setting tcp keepalive");
                         set_tcp_keepalive(&socket, ka)?;
                     }
+                    if let Some(recv_buf) = peer.tcp_recv_buf() {
+                        debug!("Setting recv buf size");
+                        set_recv_buf(socket.as_raw_fd(), recv_buf)?;
+                    }
                     Ok(socket.into())
                 }
                 Err(e) => {
diff --git a/pingora-core/src/protocols/l4/ext.rs b/pingora-core/src/protocols/l4/ext.rs
index a30e26a..525f94a 100644
--- a/pingora-core/src/protocols/l4/ext.rs
+++ b/pingora-core/src/protocols/l4/ext.rs
@@ -211,6 +211,18 @@ pub fn get_tcp_info(_fd: RawFd) -> io::Result<TCP_INFO> {
     Ok(unsafe { TCP_INFO::new() })
 }
 
+/// Set the TCP receive buffer size. See SO_RCVBUF.
+#[cfg(target_os = "linux")]
+pub fn set_recv_buf(fd: RawFd, val: usize) -> Result<()> {
+    set_opt(fd, libc::SOL_SOCKET, libc::SO_RCVBUF, val as c_int)
+        .or_err(ConnectError, "failed to set SO_RCVBUF")
+}
+
+#[cfg(not(target_os = "linux"))]
+pub fn set_recv_buf(_fd: RawFd) -> io::Result<()> {
+    Ok(())
+}
+
 /*
  * this extension is needed until the following are addressed
  * https://github.com/tokio-rs/tokio/issues/1543
@@ -295,3 +307,31 @@ pub fn set_tcp_keepalive(stream: &TcpStream, ka: &TcpKeepalive) -> Result<()> {
     // TODO: check localhost or if keepalive is already set
     set_keepalive(fd, ka).or_err(ConnectError, "failed to set keepalive")
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_set_recv_buf() {
+        use tokio::net::TcpSocket;
+        let socket = TcpSocket::new_v4().unwrap();
+        set_recv_buf(socket.as_raw_fd(), 102400).unwrap();
+
+        let mut recv_size: c_int = 0;
+        let mut size = std::mem::size_of::<c_int>() as u32;
+        get_opt(
+            socket.as_raw_fd(),
+            libc::SOL_SOCKET,
+            libc::SO_RCVBUF,
+            &mut recv_size,
+            &mut size,
+        )
+        .unwrap();
+
+        if cfg!(target_os = "linux") {
+            // kernel doubles whatever is set
+            assert_eq!(recv_size, 102400 * 2);
+        }
+    }
+}
diff --git a/pingora-core/src/upstreams/peer.rs b/pingora-core/src/upstreams/peer.rs
index 51f8aa1..f39ea5b 100644
--- a/pingora-core/src/upstreams/peer.rs
+++ b/pingora-core/src/upstreams/peer.rs
@@ -163,6 +163,11 @@ pub trait Peer: Display + Clone {
         self.get_peer_options().and_then(|o| o.h2_ping_interval)
     }
 
+    /// The size of the TCP receive buffer should be limited to. See SO_RCVBUF for more details.
+    fn tcp_recv_buf(&self) -> Option<usize> {
+        self.get_peer_options().and_then(|o| o.tcp_recv_buf)
+    }
+
     fn matches_fd<V: AsRawFd>(&self, fd: V) -> bool {
         self.address().check_fd_match(fd)
     }
@@ -271,6 +276,7 @@ pub struct PeerOptions {
     pub alpn: ALPN,
     pub ca: Option<Arc<Box<[X509]>>>,
     pub tcp_keepalive: Option<TcpKeepalive>,
+    pub tcp_recv_buf: Option<usize>,
     pub no_header_eos: bool,
     pub h2_ping_interval: Option<Duration>,
     // how many concurrent h2 stream are allowed in the same connection
@@ -301,6 +307,7 @@ impl PeerOptions {
             alpn: ALPN::H1,
             ca: None,
             tcp_keepalive: None,
+            tcp_recv_buf: None,
             no_header_eos: false,
             h2_ping_interval: None,
             max_h2_streams: 1,