Skip to content

Commit

Permalink
Add support for setting DSCP on client and server sockets
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewhavck committed Jul 12, 2024
1 parent 34fb39e commit 6b60701
Show file tree
Hide file tree
Showing 5 changed files with 50 additions and 3 deletions.
2 changes: 1 addition & 1 deletion .bleep
Original file line number Diff line number Diff line change
@@ -1 +1 @@
bd707d83c6b344fa22ca0e4b61d751acea02f4bc
837db6c7ec2d37abf83f9588be99fda00e2012c3
6 changes: 5 additions & 1 deletion pingora-core/src/connectors/l4.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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() {
Expand Down
12 changes: 11 additions & 1 deletion pingora-core/src/listeners/l4.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<TcpKeepalive>,
/// 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<u8>,
// TODO: allow configuring reuseaddr, backlog, etc. from here?
}

Expand Down Expand Up @@ -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(())
}

Expand Down Expand Up @@ -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(())
}

Expand Down
25 changes: 25 additions & 0 deletions pingora-core/src/protocols/l4/ext.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<u64> {
get_opt_sized::<c_ulonglong>(fd, libc::SOL_SOCKET, libc::SO_COOKIE)
Expand Down
8 changes: 8 additions & 0 deletions pingora-core/src/upstreams/peer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<u8> {
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()
Expand Down Expand Up @@ -301,6 +307,7 @@ pub struct PeerOptions {
pub ca: Option<Arc<Box<[X509]>>>,
pub tcp_keepalive: Option<TcpKeepalive>,
pub tcp_recv_buf: Option<usize>,
pub dscp: Option<u8>,
pub no_header_eos: bool,
pub h2_ping_interval: Option<Duration>,
// how many concurrent h2 stream are allowed in the same connection
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 6b60701

Please sign in to comment.