From 6b60701a06f7853f3e283f25c1d6e6cea9ad6017 Mon Sep 17 00:00:00 2001 From: Andrew Hauck Date: Wed, 10 Jul 2024 21:47:19 -0700 Subject: [PATCH] Add support for setting DSCP on client and server sockets --- .bleep | 2 +- pingora-core/src/connectors/l4.rs | 6 +++++- pingora-core/src/listeners/l4.rs | 12 +++++++++++- pingora-core/src/protocols/l4/ext.rs | 25 +++++++++++++++++++++++++ pingora-core/src/upstreams/peer.rs | 8 ++++++++ 5 files changed, 50 insertions(+), 3 deletions(-) diff --git a/.bleep b/.bleep index 48a93eee..8439eacf 100644 --- a/.bleep +++ b/.bleep @@ -1 +1 @@ -bd707d83c6b344fa22ca0e4b61d751acea02f4bc \ No newline at end of file +837db6c7ec2d37abf83f9588be99fda00e2012c3 \ No newline at end of file diff --git a/pingora-core/src/connectors/l4.rs b/pingora-core/src/connectors/l4.rs index ffbf084d..449ea4c8 100644 --- a/pingora-core/src/connectors/l4.rs +++ b/pingora-core/src/connectors/l4.rs @@ -19,7 +19,7 @@ use std::net::SocketAddr as InetSocketAddr; use std::os::unix::io::AsRawFd; use crate::protocols::l4::ext::{ - connect_uds, connect_with as tcp_connect, set_recv_buf, set_tcp_fastopen_connect, + connect_uds, connect_with as tcp_connect, set_dscp, set_recv_buf, set_tcp_fastopen_connect, }; use crate::protocols::l4::socket::SocketAddr; use crate::protocols::l4::stream::Stream; @@ -47,6 +47,10 @@ where debug!("Setting recv buf size"); set_recv_buf(socket.as_raw_fd(), recv_buf)?; } + if let Some(dscp) = peer.dscp() { + debug!("Setting dscp"); + set_dscp(socket.as_raw_fd(), dscp)?; + } Ok(()) }); let conn_res = match peer.connection_timeout() { diff --git a/pingora-core/src/listeners/l4.rs b/pingora-core/src/listeners/l4.rs index 42adc87e..748037a0 100644 --- a/pingora-core/src/listeners/l4.rs +++ b/pingora-core/src/listeners/l4.rs @@ -25,7 +25,7 @@ use std::os::unix::net::UnixListener as StdUnixListener; use std::time::Duration; use tokio::net::TcpSocket; -use crate::protocols::l4::ext::set_tcp_fastopen_backlog; +use crate::protocols::l4::ext::{set_dscp, set_tcp_fastopen_backlog}; use crate::protocols::l4::listener::Listener; pub use crate::protocols::l4::stream::Stream; use crate::protocols::TcpKeepalive; @@ -76,6 +76,9 @@ pub struct TcpSocketOptions { /// Enable TCP keepalive on accepted connections. /// See the [man page](https://man7.org/linux/man-pages/man7/tcp.7.html) for more information. pub tcp_keepalive: Option, + /// Specifies the server should set the following DSCP value on outgoing connections. + /// See the [RFC](https://datatracker.ietf.org/doc/html/rfc2474) for more details. + pub dscp: Option, // TODO: allow configuring reuseaddr, backlog, etc. from here? } @@ -150,6 +153,10 @@ fn apply_tcp_socket_options(sock: &TcpSocket, opt: Option<&TcpSocketOptions>) -> if let Some(backlog) = opt.tcp_fastopen { set_tcp_fastopen_backlog(sock.as_raw_fd(), backlog)?; } + + if let Some(dscp) = opt.dscp { + set_dscp(sock.as_raw_fd(), dscp)?; + } Ok(()) } @@ -280,6 +287,9 @@ impl ListenerEndpoint { if let Some(ka) = op.tcp_keepalive.as_ref() { stream.set_keepalive(ka)?; } + if let Some(dscp) = op.dscp { + set_dscp(stream.as_raw_fd(), dscp)?; + } Ok(()) } diff --git a/pingora-core/src/protocols/l4/ext.rs b/pingora-core/src/protocols/l4/ext.rs index f68cdcf1..56af5220 100644 --- a/pingora-core/src/protocols/l4/ext.rs +++ b/pingora-core/src/protocols/l4/ext.rs @@ -275,6 +275,31 @@ pub fn set_tcp_fastopen_backlog(_fd: RawFd, _backlog: usize) -> Result<()> { Ok(()) } +#[cfg(target_os = "linux")] +pub fn set_dscp(fd: RawFd, value: u8) -> Result<()> { + use super::socket::SocketAddr; + use pingora_error::OkOrErr; + + let sock = SocketAddr::from_raw_fd(fd, false); + let addr = sock + .as_ref() + .and_then(|s| s.as_inet()) + .or_err(SocketError, "failed to set dscp, invalid IP socket")?; + + if addr.is_ipv6() { + set_opt(fd, libc::IPPROTO_IPV6, libc::IPV6_TCLASS, value as c_int) + .or_err(SocketError, "failed to set dscp (IPV6_TCLASS)") + } else { + set_opt(fd, libc::IPPROTO_IP, libc::IP_TOS, value as c_int) + .or_err(SocketError, "failed to set dscp (IP_TOS)") + } +} + +#[cfg(not(target_os = "linux"))] +pub fn set_dscp(_fd: RawFd, _value: u8) -> Result<()> { + Ok(()) +} + #[cfg(target_os = "linux")] pub fn get_socket_cookie(fd: RawFd) -> io::Result { get_opt_sized::(fd, libc::SOL_SOCKET, libc::SO_COOKIE) diff --git a/pingora-core/src/upstreams/peer.rs b/pingora-core/src/upstreams/peer.rs index 78e59ede..d0c8125a 100644 --- a/pingora-core/src/upstreams/peer.rs +++ b/pingora-core/src/upstreams/peer.rs @@ -172,6 +172,12 @@ pub trait Peer: Display + Clone { self.get_peer_options().and_then(|o| o.tcp_recv_buf) } + /// The DSCP value that should be applied to the send side of this connection. + /// See the [RFC](https://datatracker.ietf.org/doc/html/rfc2474) for more details. + fn dscp(&self) -> Option { + self.get_peer_options().and_then(|o| o.dscp) + } + /// Whether to enable TCP fast open. fn tcp_fast_open(&self) -> bool { self.get_peer_options() @@ -301,6 +307,7 @@ pub struct PeerOptions { pub ca: Option>>, pub tcp_keepalive: Option, pub tcp_recv_buf: Option, + pub dscp: Option, pub no_header_eos: bool, pub h2_ping_interval: Option, // how many concurrent h2 stream are allowed in the same connection @@ -334,6 +341,7 @@ impl PeerOptions { ca: None, tcp_keepalive: None, tcp_recv_buf: None, + dscp: None, no_header_eos: false, h2_ping_interval: None, max_h2_streams: 1,