solana_streamer/
quic_socket.rs

1//! This module defines [`QuicSocket`], which allows selecting between kernel UDP and AF_XDP-backed
2//! QUIC socket configurations.
3use {
4    agave_xdp::{
5        ecn_codepoint::EcnCodepoint as XdpEcnCodepoint,
6        transmitter::{BytesTxPacket, XdpSender},
7    },
8    bytes::Bytes,
9    crossbeam_channel::TrySendError,
10    quinn::{
11        AsyncUdpSocket, Runtime, TokioRuntime, UdpPoller,
12        udp::{EcnCodepoint as QuinnEcnCodepoint, RecvMeta, Transmit},
13    },
14    std::{
15        fmt::{self, Debug},
16        io::{self, IoSliceMut},
17        net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
18        pin::Pin,
19        sync::Arc,
20        task::{Context, Poll},
21    },
22};
23
24/// [`QuicSocket`] is an enum for selecting between a kernel UDP socket and an AF_XDP-backed
25/// socket for QUIC communication.
26#[derive(Debug)]
27pub enum QuicSocket {
28    /// A QUIC socket that uses AF_XDP for sending and a kernel UDP socket for receiving.
29    Xdp(QuicXdpSocketParts),
30    /// A QUIC socket that uses kernel UDP socket for both sending and receiving.
31    Kernel(std::net::UdpSocket),
32}
33
34impl From<std::net::UdpSocket> for QuicSocket {
35    fn from(socket: std::net::UdpSocket) -> Self {
36        QuicSocket::Kernel(socket)
37    }
38}
39
40impl QuicSocket {
41    pub fn with_xdp(
42        socket: std::net::UdpSocket,
43        fallback_src_ip: Ipv4Addr,
44        xdp_sender: XdpSender,
45    ) -> Self {
46        Self::Xdp(QuicXdpSocketParts {
47            socket,
48            fallback_src_ip,
49            xdp_sender,
50        })
51    }
52
53    #[cfg(feature = "dev-context-only-utils")]
54    pub fn local_addr(&self) -> io::Result<SocketAddr> {
55        match self {
56            QuicSocket::Xdp(parts) => parts.socket.local_addr(),
57            QuicSocket::Kernel(socket) => socket.local_addr(),
58        }
59    }
60}
61
62/// [`QuicXdpSocketParts`] wraps the resources required to construct an AF_XDP-backed QUIC socket.
63///
64/// It carries both an [`XdpSender`] and a [`std::net::UdpSocket`], rather than constructing an
65/// [`QuicXdpTxSocket`] directly, because the underlying sockets can only be created when a Tokio
66/// runtime is present. `fallback_src_ip` is used when the local address of `socket` is a
67/// wildcard address.
68pub struct QuicXdpSocketParts {
69    pub socket: std::net::UdpSocket,
70    pub fallback_src_ip: Ipv4Addr,
71    pub xdp_sender: XdpSender,
72}
73
74impl Debug for QuicXdpSocketParts {
75    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
76        f.debug_struct("QuicXdpSocketParts")
77            .field("socket", &self.socket)
78            .finish()
79    }
80}
81
82/// [`QuicXdpTxSocket`] uses AF_XDP for egress traffic and `UdpSocket` for ingress traffic.
83///
84/// For egress traffic, it employs an underlying `QuicXdpSender` for non-local destinations. For
85/// destinations owned by the local host (routed via `lo`, including loopback and local interface
86/// IPs), it falls back to a kernel `UdpSocket`.
87pub(crate) struct QuicXdpTxSocket {
88    udp_socket: Arc<dyn AsyncUdpSocket>,
89    xdp_sender: QuicXdpSender,
90    local_ips: Vec<Ipv4Addr>,
91}
92
93impl QuicXdpTxSocket {
94    pub(crate) fn new(
95        socket: std::net::UdpSocket,
96        fallback_src_ip: Ipv4Addr,
97        xdp_sender: XdpSender,
98    ) -> io::Result<Self> {
99        let src_addr = socket.local_addr()?;
100        let SocketAddr::V4(src_addr) = src_addr else {
101            return Err(io::Error::new(
102                io::ErrorKind::InvalidInput,
103                "Only IPv4 addresses are supported",
104            ));
105        };
106        // if local address is wildcard, override it with fallback_src_ip.
107        let src_addr = if src_addr.ip().is_unspecified() {
108            SocketAddrV4::new(fallback_src_ip, src_addr.port())
109        } else {
110            src_addr
111        };
112
113        // Collect local interface IPs once at construction time. We do not refresh them if
114        // interface addresses change later. This is a low-risk tradeoff because local-destination
115        // egress is expected to be rare: only RPC sendTransaction traffic or local testing.
116        let local_ips = collect_local_ipv4_ips()?;
117
118        Ok(Self {
119            udp_socket: TokioRuntime.wrap_udp_socket(socket)?,
120            xdp_sender: QuicXdpSender::new(xdp_sender, src_addr),
121            local_ips,
122        })
123    }
124
125    fn should_use_kernel_udp(&self, dst: SocketAddr) -> bool {
126        dst.ip().is_loopback() || matches!(dst.ip(), IpAddr::V4(ip) if self.local_ips.contains(&ip))
127    }
128}
129
130impl fmt::Debug for QuicXdpTxSocket {
131    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
132        f.debug_struct("QuicXdpTxSocket")
133            .field("local_addr", &self.udp_socket.local_addr())
134            .finish_non_exhaustive()
135    }
136}
137
138impl AsyncUdpSocket for QuicXdpTxSocket {
139    fn create_io_poller(self: Arc<Self>) -> Pin<Box<dyn UdpPoller>> {
140        // The kernel UDP socket poller is always returned here, ignoring the XDP sender. This
141        // implementation is correct under the following assumptions:
142        // 1. When egress AF_XDP is enabled, the kernel UDP socket is used rarely and only for local
143        //    destinations, so it should almost always be writable.
144        // 2. `QuicXdpSender` can almost always enqueue.
145        //
146        // A rare mismatch is still possible: if the UDP socket is not writable while
147        // `QuicXdpSender` could enqueue, throughput may be temporarily suboptimal until the UDP
148        // socket becomes writable. The reverse mismatch is also possible: the UDP poller is ready
149        // but the selected `QuicXdpSender` channel is full. In this case `try_send` fails with
150        // `WouldBlock`, and the caller invokes `poll_writable` again before retrying.
151        self.udp_socket.clone().create_io_poller()
152    }
153
154    /// Attempts to send the given [`Transmit`].
155    ///
156    /// For non-local destinations uses AF_XDP, otherwise kernel UDP.
157    ///
158    /// If enqueueing fails after some datagrams were already enqueued, this method returns
159    /// `Err(WouldBlock)`. The caller may retry the whole transmit, which can cause duplicate
160    /// datagrams to be sent for the already enqueued chunks. QUIC packet numbers make this
161    /// protocol-safe, but duplicates can still degrade throughput and congestion behavior. This
162    /// implementation therefore assumes the AF_XDP channel is rarely (ideally never) full.
163    fn try_send(&self, t: &Transmit<'_>) -> io::Result<()> {
164        if self.should_use_kernel_udp(t.destination) {
165            return self.udp_socket.try_send(t);
166        }
167        if t.destination.is_ipv6() {
168            return Err(io::Error::new(
169                io::ErrorKind::InvalidInput,
170                "IPv6 destination addresses are not supported for AF_XDP sends",
171            ));
172        }
173        let src_ip = match t.src_ip {
174            Some(IpAddr::V4(ip)) => Some(ip),
175            Some(IpAddr::V6(_)) => {
176                return Err(io::Error::new(
177                    io::ErrorKind::InvalidInput,
178                    "IPv6 source addresses are not supported",
179                ));
180            }
181            None => None,
182        };
183
184        debug_assert!(
185            t.segment_size.is_none(),
186            "GSO segmentation is disabled for AF_XDP sends, but segment_size is {:?}",
187            t.segment_size
188        );
189
190        let payload = Bytes::copy_from_slice(t.contents);
191        match self
192            .xdp_sender
193            .try_send(src_ip, t.destination, t.ecn, payload)
194        {
195            Ok(()) => Ok(()),
196            Err(TrySendError::Full(_)) => Err(io::ErrorKind::WouldBlock.into()),
197            Err(TrySendError::Disconnected(_)) => Err(io::ErrorKind::BrokenPipe.into()),
198        }
199    }
200
201    fn poll_recv(
202        &self,
203        cx: &mut Context,
204        bufs: &mut [IoSliceMut<'_>],
205        meta: &mut [RecvMeta],
206    ) -> Poll<io::Result<usize>> {
207        self.udp_socket.poll_recv(cx, bufs, meta)
208    }
209
210    fn local_addr(&self) -> io::Result<SocketAddr> {
211        self.udp_socket.local_addr()
212    }
213
214    fn max_transmit_segments(&self) -> usize {
215        // no GSO batches, so each transmit describes exactly one datagram
216        1
217    }
218
219    fn max_receive_segments(&self) -> usize {
220        self.udp_socket.max_receive_segments()
221    }
222
223    fn may_fragment(&self) -> bool {
224        false
225    }
226}
227
228/// [`QuicXdpSender`] wraps [`XdpSender`] and provides destination-based sender selection.
229///
230/// This wrapper maps each remote IP to a stable XDP sender index. Keeping packets for the same
231/// remote host on one TX queue avoids queue-induced reordering within packet bursts.
232struct QuicXdpSender {
233    xdp_sender: XdpSender,
234    src_addr: SocketAddrV4,
235}
236
237impl QuicXdpSender {
238    fn new(xdp_sender: XdpSender, src_addr: SocketAddrV4) -> Self {
239        Self {
240            xdp_sender,
241            src_addr,
242        }
243    }
244
245    fn try_send(
246        &self,
247        src_ip: Option<Ipv4Addr>,
248        destination: SocketAddr,
249        ecn: Option<QuinnEcnCodepoint>,
250        payload: Bytes,
251    ) -> Result<(), TrySendError<BytesTxPacket>> {
252        // Keep packets for the same remote IP on the same XDP TX queue when there is more than
253        // one queue. Avoid hashing entirely for the common single-sender case.
254        let sender_key = if self.xdp_sender.len() == 1 {
255            0
256        } else {
257            fold_xor(destination_ip_key(&destination)) as usize
258        };
259
260        let src_ip = src_ip.unwrap_or(*self.src_addr.ip());
261        // Respect Quinn's per-packet source IP, used for wildcard-bound sockets, while keeping the
262        // port from `self.src_addr`.
263        let src_addr = SocketAddrV4::new(src_ip, self.src_addr.port());
264        let ecn = ecn.map(quinn_ecn_to_xdp);
265
266        let mut packet = BytesTxPacket::new(src_addr, destination, ecn, payload);
267        packet.set_allow_mtu_overflow(true);
268        self.xdp_sender.try_send(sender_key, packet)
269    }
270}
271
272/// Collects IPv4 addresses assigned to local network interfaces.
273#[cfg(target_os = "linux")]
274fn collect_local_ipv4_ips() -> io::Result<Vec<Ipv4Addr>> {
275    use nix::ifaddrs::getifaddrs;
276
277    let mut ips = Vec::new();
278    for ifa in getifaddrs().map_err(io::Error::other)? {
279        let Some(addr) = ifa.address else { continue };
280        if let Some(v4) = addr.as_sockaddr_in() {
281            let ip = v4.ip();
282            if !ips.contains(&ip) {
283                ips.push(ip);
284            }
285        }
286    }
287    Ok(ips)
288}
289
290#[cfg(not(target_os = "linux"))]
291fn collect_local_ipv4_ips() -> io::Result<Vec<Ipv4Addr>> {
292    Ok(Vec::new())
293}
294
295#[inline]
296const fn quinn_ecn_to_xdp(ecn: QuinnEcnCodepoint) -> XdpEcnCodepoint {
297    match ecn {
298        QuinnEcnCodepoint::Ect0 => XdpEcnCodepoint::Ect0,
299        QuinnEcnCodepoint::Ect1 => XdpEcnCodepoint::Ect1,
300        QuinnEcnCodepoint::Ce => XdpEcnCodepoint::Ce,
301    }
302}
303
304#[inline]
305fn fold_xor(mut x: u32) -> u32 {
306    x ^= x >> 16;
307    x ^= x >> 8;
308    x
309}
310
311#[inline]
312fn destination_ip_key(destination: &SocketAddr) -> u32 {
313    match destination {
314        SocketAddr::V4(destination) => u32::from(*destination.ip()),
315        SocketAddr::V6(_) => unreachable!("IPv6 destinations are rejected before AF_XDP send"),
316    }
317}
solana_streamer/quic_socket.rs

solana_streamer/
quic_socket.rs
