1use std::io;
32use std::net::{SocketAddr, UdpSocket};
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub enum DgramBackend {
37 Udp,
39 IoUring,
41 Wire,
45 Demux,
49}
50
51#[allow(clippy::large_enum_variant)]
56enum Inner {
57 Udp(UdpSocket),
58 #[cfg(target_os = "linux")]
59 IoUring(linux_iou::IoUringDgram),
60 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
61 Wire(wire_backend::WireDgram),
62 Demux(DemuxDgram),
65}
66
67pub type DemuxQueue =
72 std::sync::Arc<std::sync::Mutex<std::collections::VecDeque<(Vec<u8>, SocketAddr, Option<i128>)>>>;
73
74pub fn new_demux_queue() -> DemuxQueue {
76 std::sync::Arc::new(std::sync::Mutex::new(std::collections::VecDeque::new()))
77}
78
79struct DemuxDgram {
85 real: std::sync::Arc<UdpSocket>,
86 queue: DemuxQueue,
87 peer: std::sync::Mutex<Option<SocketAddr>>,
90 sent: Option<std::sync::Arc<std::sync::atomic::AtomicU64>>,
93 pop_attempts: std::sync::atomic::AtomicU64,
95 pop_yields: std::sync::atomic::AtomicU64,
96}
97
98impl DemuxDgram {
99 fn recv_with_kts(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Option<i128>)> {
100 self.pop_attempts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
101 match self.queue.lock().unwrap().pop_front() {
102 Some((data, from, kts)) => {
103 self.pop_yields.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
104 let n = data.len().min(buf.len());
105 buf[..n].copy_from_slice(&data[..n]);
106 Ok((n, from, kts))
107 }
108 None => Err(io::Error::new(io::ErrorKind::WouldBlock, "demux queue empty")),
111 }
112 }
113
114 fn connect(&self, addr: SocketAddr) {
115 *self.peer.lock().unwrap() = Some(addr);
116 }
117
118 fn count_fwd(&self, buf: &[u8]) {
119 if let Some(c) = &self.sent
123 && matches!(buf.first(), Some(&(1 | 10 | 11)))
124 {
125 c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
126 }
127 }
128
129 fn send_to(&self, buf: &[u8], addr: SocketAddr) -> io::Result<usize> {
130 self.count_fwd(buf);
131 self.real.send_to(buf, addr)
132 }
133
134 fn send(&self, buf: &[u8]) -> io::Result<usize> {
135 self.count_fwd(buf);
136 match *self.peer.lock().unwrap() {
137 Some(p) => self.real.send_to(buf, p),
138 None => Err(io::Error::new(io::ErrorKind::NotConnected, "demux send before connect")),
139 }
140 }
141}
142
143pub struct DgramSock {
146 inner: Inner,
147}
148
149impl DgramSock {
150 pub fn wrap(sock: UdpSocket) -> Self {
154 enable_rx_timestamp(&sock);
155 let forced = std::env::var("SUBETHA_DGRAM").ok();
156 if forced.as_deref() == Some("udp") {
157 return Self { inner: Inner::Udp(sock) };
158 }
159
160 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
170 {
171 let wire_if = std::env::var("SUBETHA_WIRE_IFNAME").ok();
172 let force_wire = forced.as_deref() == Some("wire");
173 let auto_wire = forced.is_none()
174 && wire_if.as_deref().map(wire_link_fast_enough).unwrap_or(false);
175 if force_wire || auto_wire {
176 let port = sock.local_addr().map(|a| a.port()).unwrap_or(0);
177 match wire_backend::WireDgram::from_env(port) {
178 Ok(w) => {
179 if force_wire
180 && !wire_if.as_deref().map(wire_link_fast_enough).unwrap_or(false)
181 {
182 eprintln!(
183 "SUBETHA_DGRAM=wire forced on a link below the {:.0} Gbit/s \
184 gate; the kernel-bypass path is net-negative below line rate",
185 wire_min_link_bps() as f64 / 1e9
186 );
187 }
188 return Self { inner: Inner::Wire(w) };
189 }
190 Err(e) => eprintln!(
191 "Wire backend requested but unavailable ({e}); falling back"
192 ),
193 }
194 }
195 }
196
197 #[cfg(target_os = "linux")]
198 {
199 let force_ring = forced.as_deref() == Some("iouring");
200 match linux_iou::IoUringDgram::new(sock) {
201 Ok(d) => Self { inner: Inner::IoUring(d) },
202 Err((sock, e)) => {
203 if force_ring {
204 eprintln!(
205 "SUBETHA_DGRAM=iouring requested but the ring is unavailable \
206 ({e}); using plain UDP"
207 );
208 }
209 Self { inner: Inner::Udp(sock) }
210 }
211 }
212 }
213
214 #[cfg(not(target_os = "linux"))]
215 {
216 drop(forced);
217 Self { inner: Inner::Udp(sock) }
218 }
219 }
220
221 pub fn demux(real: std::sync::Arc<UdpSocket>, queue: DemuxQueue) -> Self {
227 Self::demux_inner(real, queue, None)
228 }
229
230 pub fn demux_counted(
233 real: std::sync::Arc<UdpSocket>,
234 queue: DemuxQueue,
235 sent: std::sync::Arc<std::sync::atomic::AtomicU64>,
236 ) -> Self {
237 Self::demux_inner(real, queue, Some(sent))
238 }
239
240 fn demux_inner(
241 real: std::sync::Arc<UdpSocket>,
242 queue: DemuxQueue,
243 sent: Option<std::sync::Arc<std::sync::atomic::AtomicU64>>,
244 ) -> Self {
245 Self {
246 inner: Inner::Demux(DemuxDgram {
247 real,
248 queue,
249 peer: std::sync::Mutex::new(None),
250 sent,
251 pop_attempts: std::sync::atomic::AtomicU64::new(0),
252 pop_yields: std::sync::atomic::AtomicU64::new(0),
253 }),
254 }
255 }
256
257 pub fn demux_probe(&self) -> Option<(u64, u64, u64, u64)> {
262 match &self.inner {
263 Inner::Demux(d) => Some((
264 d.pop_attempts.load(std::sync::atomic::Ordering::Relaxed),
265 d.pop_yields.load(std::sync::atomic::Ordering::Relaxed),
266 std::sync::Arc::as_ptr(&d.queue) as usize as u64,
267 d.queue.lock().unwrap().len() as u64,
268 )),
269 _ => None,
270 }
271 }
272
273 pub fn from_udp(sock: UdpSocket) -> Self {
281 Self { inner: Inner::Udp(sock) }
282 }
283
284 pub fn as_udp(&self) -> Option<&UdpSocket> {
289 match &self.inner {
290 Inner::Udp(s) => Some(s),
291 _ => None,
292 }
293 }
294
295 pub fn connect(&self, addr: SocketAddr) -> io::Result<()> {
299 match &self.inner {
300 Inner::Udp(s) => s.connect(addr),
301 Inner::Demux(d) => {
302 d.connect(addr);
303 Ok(())
304 }
305 #[allow(unreachable_patterns)]
306 _ => Err(io::Error::new(
307 io::ErrorKind::Unsupported,
308 "connect on a non-Udp/Demux backend",
309 )),
310 }
311 }
312
313 pub fn send(&self, buf: &[u8]) -> io::Result<usize> {
315 match &self.inner {
316 Inner::Udp(s) => s.send(buf),
317 Inner::Demux(d) => d.send(buf),
318 #[allow(unreachable_patterns)]
319 _ => Err(io::Error::new(
320 io::ErrorKind::Unsupported,
321 "send on a non-Udp/Demux backend",
322 )),
323 }
324 }
325
326 pub fn recv(&self, buf: &mut [u8]) -> io::Result<usize> {
329 match &self.inner {
330 Inner::Udp(s) => s.recv(buf),
331 Inner::Demux(d) => d.recv_with_kts(buf).map(|(n, _, _)| n),
332 #[allow(unreachable_patterns)]
333 _ => Err(io::Error::new(
334 io::ErrorKind::Unsupported,
335 "recv on a non-Udp/Demux backend",
336 )),
337 }
338 }
339
340 pub fn backend(&self) -> DgramBackend {
342 match &self.inner {
343 Inner::Udp(_) => DgramBackend::Udp,
344 #[cfg(target_os = "linux")]
345 Inner::IoUring(_) => DgramBackend::IoUring,
346 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
347 Inner::Wire(_) => DgramBackend::Wire,
348 Inner::Demux(_) => DgramBackend::Demux,
349 }
350 }
351
352 pub fn send_to(&self, buf: &[u8], addr: SocketAddr) -> io::Result<usize> {
353 match &self.inner {
354 Inner::Udp(s) => s.send_to(buf, addr),
355 #[cfg(target_os = "linux")]
356 Inner::IoUring(d) => d.send_to(buf, addr),
357 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
358 Inner::Wire(d) => d.send_to(buf, addr),
359 Inner::Demux(d) => d.send_to(buf, addr),
360 }
361 }
362
363 #[cfg(target_os = "linux")]
367 fn raw_fd(&self) -> Option<i32> {
368 use std::os::unix::io::AsRawFd;
369 match &self.inner {
370 Inner::Udp(s) => Some(s.as_raw_fd()),
371 Inner::IoUring(d) => Some(d.raw_fd()),
372 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
373 Inner::Wire(_) => None,
374 Inner::Demux(_) => None,
375 }
376 }
377
378 pub fn send_gso(&self, batch: &[u8], seg_size: u16, addr: SocketAddr) -> io::Result<()> {
387 #[cfg(target_os = "linux")]
388 if let Some(fd) = self.raw_fd() {
389 return linux_gso_send(fd, batch, seg_size, addr);
390 }
391 let n = (seg_size as usize).max(1);
392 let mut off = 0;
393 while off < batch.len() {
394 let end = (off + n).min(batch.len());
395 self.send_to(&batch[off..end], addr)?;
396 off = end;
397 }
398 Ok(())
399 }
400
401 pub fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
402 match &self.inner {
403 Inner::Udp(s) => s.recv_from(buf),
404 #[cfg(target_os = "linux")]
405 Inner::IoUring(d) => d.recv_with_kts(buf).map(|(n, a, _)| (n, a)),
406 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
407 Inner::Wire(d) => d.recv_with_kts(buf).map(|(n, a, _)| (n, a)),
408 Inner::Demux(d) => d.recv_with_kts(buf).map(|(n, a, _)| (n, a)),
409 }
410 }
411
412 pub fn recv_with_kts(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Option<i128>)> {
414 match &self.inner {
415 Inner::Udp(s) => udp_recv_with_kts(s, buf),
416 #[cfg(target_os = "linux")]
417 Inner::IoUring(d) => d.recv_with_kts(buf),
418 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
419 Inner::Wire(d) => d.recv_with_kts(buf),
420 Inner::Demux(d) => d.recv_with_kts(buf),
421 }
422 }
423
424 pub fn local_addr(&self) -> io::Result<SocketAddr> {
425 match &self.inner {
426 Inner::Udp(s) => s.local_addr(),
427 #[cfg(target_os = "linux")]
428 Inner::IoUring(d) => d.local_addr(),
429 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
430 Inner::Wire(d) => d.local_addr(),
431 Inner::Demux(d) => d.real.local_addr(),
432 }
433 }
434
435 pub fn set_nonblocking(&self, nb: bool) -> io::Result<()> {
436 match &self.inner {
437 Inner::Udp(s) => s.set_nonblocking(nb),
438 #[cfg(target_os = "linux")]
439 Inner::IoUring(d) => d.set_nonblocking(nb),
440 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
441 Inner::Wire(d) => d.set_nonblocking(nb),
442 Inner::Demux(_) => Ok(()),
445 }
446 }
447}
448
449#[cfg(target_os = "linux")]
455fn enable_rx_timestamp(sock: &UdpSocket) {
456 use std::os::fd::AsRawFd;
457 let on: libc::c_int = 1;
458 unsafe {
459 libc::setsockopt(
460 sock.as_raw_fd(),
461 libc::SOL_SOCKET,
462 libc::SO_TIMESTAMPNS,
463 &on as *const libc::c_int as *const libc::c_void,
464 std::mem::size_of::<libc::c_int>() as libc::socklen_t,
465 );
466 }
467}
468
469#[cfg(not(target_os = "linux"))]
470fn enable_rx_timestamp(_sock: &UdpSocket) {}
471
472#[cfg(target_os = "linux")]
474pub(crate) fn udp_recv_with_kts(sock: &UdpSocket, buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Option<i128>)> {
475 use std::os::fd::AsRawFd;
476 let mut iov = libc::iovec {
477 iov_base: buf.as_mut_ptr() as *mut libc::c_void,
478 iov_len: buf.len(),
479 };
480 let mut name: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
481 let mut control = [0u8; 64];
482 let mut msg: libc::msghdr = unsafe { std::mem::zeroed() };
483 msg.msg_name = &mut name as *mut _ as *mut libc::c_void;
484 msg.msg_namelen = std::mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;
485 msg.msg_iov = &mut iov;
486 msg.msg_iovlen = 1;
487 msg.msg_control = control.as_mut_ptr() as *mut libc::c_void;
488 msg.msg_controllen = control.len();
489 let n = unsafe { libc::recvmsg(sock.as_raw_fd(), &mut msg, 0) };
490 if n < 0 {
491 return Err(io::Error::last_os_error());
492 }
493 let kts = parse_timestamp(&msg);
494 let from = unsafe { sockaddr_to_socketaddr(&name) }
495 .ok_or_else(|| io::Error::other("non-IP source"))?;
496 Ok((n as usize, from, kts))
497}
498
499#[cfg(not(target_os = "linux"))]
500pub(crate) fn udp_recv_with_kts(sock: &UdpSocket, buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Option<i128>)> {
501 let (n, from) = sock.recv_from(buf)?;
502 Ok((n, from, None))
503}
504
505#[cfg(target_os = "linux")]
506fn parse_timestamp(msg: &libc::msghdr) -> Option<i128> {
507 unsafe {
508 let mut cmsg = libc::CMSG_FIRSTHDR(msg);
509 while !cmsg.is_null() {
510 if (*cmsg).cmsg_level == libc::SOL_SOCKET && (*cmsg).cmsg_type == libc::SCM_TIMESTAMPNS {
511 let ts = (libc::CMSG_DATA(cmsg) as *const libc::timespec).read_unaligned();
512 return Some(ts.tv_sec as i128 * 1_000_000_000 + ts.tv_nsec as i128);
513 }
514 cmsg = libc::CMSG_NXTHDR(msg, cmsg);
515 }
516 }
517 None
518}
519
520#[cfg(target_os = "linux")]
521unsafe fn sockaddr_to_socketaddr(name: &libc::sockaddr_storage) -> Option<SocketAddr> {
522 use std::net::{Ipv4Addr, Ipv6Addr};
523 match name.ss_family as libc::c_int {
524 libc::AF_INET => {
525 let sin = unsafe { &*(name as *const libc::sockaddr_storage as *const libc::sockaddr_in) };
526 let ip = Ipv4Addr::from(sin.sin_addr.s_addr.to_ne_bytes());
527 Some(SocketAddr::new(ip.into(), u16::from_be(sin.sin_port)))
528 }
529 libc::AF_INET6 => {
530 let sin6 = unsafe { &*(name as *const libc::sockaddr_storage as *const libc::sockaddr_in6) };
531 let ip = Ipv6Addr::from(sin6.sin6_addr.s6_addr);
532 Some(SocketAddr::new(ip.into(), u16::from_be(sin6.sin6_port)))
533 }
534 _ => None,
535 }
536}
537
538#[cfg(target_os = "linux")]
541fn socketaddr_to_sockaddr(addr: SocketAddr) -> (libc::sockaddr_storage, libc::socklen_t) {
542 let mut storage: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
543 match addr {
544 SocketAddr::V4(a) => {
545 let sin = unsafe { &mut *(&mut storage as *mut _ as *mut libc::sockaddr_in) };
546 sin.sin_family = libc::AF_INET as libc::sa_family_t;
547 sin.sin_port = a.port().to_be();
548 sin.sin_addr.s_addr = u32::from_ne_bytes(a.ip().octets());
549 (storage, std::mem::size_of::<libc::sockaddr_in>() as libc::socklen_t)
550 }
551 SocketAddr::V6(a) => {
552 let sin6 = unsafe { &mut *(&mut storage as *mut _ as *mut libc::sockaddr_in6) };
553 sin6.sin6_family = libc::AF_INET6 as libc::sa_family_t;
554 sin6.sin6_port = a.port().to_be();
555 sin6.sin6_addr.s6_addr = a.ip().octets();
556 (storage, std::mem::size_of::<libc::sockaddr_in6>() as libc::socklen_t)
557 }
558 }
559}
560
561#[cfg(target_os = "linux")]
565fn linux_gso_send(fd: i32, batch: &[u8], seg_size: u16, addr: SocketAddr) -> io::Result<()> {
566 const UDP_SEGMENT: libc::c_int = 103;
568 const SOL_UDP: libc::c_int = 17;
569 let (storage, addrlen) = socketaddr_to_sockaddr(addr);
570 let mut iov = libc::iovec {
571 iov_base: batch.as_ptr() as *mut libc::c_void,
572 iov_len: batch.len(),
573 };
574 let cmsg_space = unsafe { libc::CMSG_SPACE(2) } as usize;
575 let mut cbuf = vec![0u8; cmsg_space];
576 let mut msg: libc::msghdr = unsafe { std::mem::zeroed() };
577 msg.msg_name = &storage as *const _ as *mut libc::c_void;
578 msg.msg_namelen = addrlen;
579 msg.msg_iov = &mut iov;
580 msg.msg_iovlen = 1;
581 msg.msg_control = cbuf.as_mut_ptr() as *mut libc::c_void;
582 msg.msg_controllen = cmsg_space as _;
583 unsafe {
586 let c = libc::CMSG_FIRSTHDR(&msg);
587 (*c).cmsg_level = SOL_UDP;
588 (*c).cmsg_type = UDP_SEGMENT;
589 (*c).cmsg_len = libc::CMSG_LEN(2) as _;
590 std::ptr::copy_nonoverlapping(&seg_size as *const u16 as *const u8, libc::CMSG_DATA(c), 2);
591 loop {
592 let r = libc::sendmsg(fd, &msg, 0);
593 if r >= 0 {
594 return Ok(());
595 }
596 let e = io::Error::last_os_error();
597 if e.kind() == io::ErrorKind::Interrupted {
598 continue;
599 }
600 return Err(e);
601 }
602 }
603}
604
605#[cfg(target_os = "linux")]
610mod linux_iou {
611 use super::{parse_timestamp, sockaddr_to_socketaddr, socketaddr_to_sockaddr};
612 use std::collections::VecDeque;
613 use std::io;
614 use std::net::{SocketAddr, UdpSocket};
615 use std::os::fd::AsRawFd;
616
617 use parking_lot::Mutex;
618
619 use io_uring::{opcode, types, IoUring};
620
621 const RECV_DEPTH: usize = 32;
622 const SEND_DEPTH: usize = 32;
623 const FRAME_CAP: usize = 2048;
624 const CONTROL_CAP: usize = 64;
625 const SEND_TAG: u64 = 1 << 32; struct RecvCtx {
632 addr: libc::sockaddr_storage,
633 iov: libc::iovec,
634 control: [u8; CONTROL_CAP],
635 msghdr: libc::msghdr,
636 buf: [u8; FRAME_CAP],
637 }
638
639 impl RecvCtx {
640 fn boxed() -> Box<Self> {
641 let mut b: Box<Self> = Box::new(unsafe { std::mem::zeroed() });
642 b.refresh();
643 b
644 }
645 fn refresh(&mut self) {
646 self.iov.iov_base = self.buf.as_mut_ptr() as *mut libc::c_void;
647 self.iov.iov_len = FRAME_CAP;
648 self.msghdr.msg_name = std::ptr::addr_of_mut!(self.addr) as *mut libc::c_void;
649 self.msghdr.msg_namelen = std::mem::size_of::<libc::sockaddr_storage>() as u32;
650 self.msghdr.msg_iov = std::ptr::addr_of_mut!(self.iov);
651 self.msghdr.msg_iovlen = 1;
652 self.msghdr.msg_control = self.control.as_mut_ptr() as *mut libc::c_void;
653 self.msghdr.msg_controllen = CONTROL_CAP;
654 }
655 }
656
657 struct SendCtx {
658 addr: libc::sockaddr_storage,
659 iov: libc::iovec,
660 msghdr: libc::msghdr,
661 buf: [u8; FRAME_CAP],
662 }
663
664 struct State {
665 ring: IoUring,
666 #[allow(clippy::vec_box)]
671 recv: Vec<Box<RecvCtx>>,
672 #[allow(clippy::vec_box)]
673 send: Vec<Box<SendCtx>>,
674 ready: VecDeque<usize>,
675 ready_len: Vec<usize>,
676 free_send: Vec<usize>,
677 }
678
679 pub struct IoUringDgram {
680 sock: UdpSocket,
681 fd: i32,
682 st: Mutex<State>,
683 }
684
685 unsafe impl Send for IoUringDgram {}
689 unsafe impl Sync for IoUringDgram {}
690
691 impl IoUringDgram {
692 pub fn new(sock: UdpSocket) -> Result<Self, (UdpSocket, io::Error)> {
693 let entries = ((RECV_DEPTH + SEND_DEPTH) * 2).next_power_of_two() as u32;
694 let ring = match IoUring::new(entries) {
695 Ok(r) => r,
696 Err(e) => return Err((sock, e)),
697 };
698 let fd = sock.as_raw_fd();
699 let recv: Vec<Box<RecvCtx>> = (0..RECV_DEPTH).map(|_| RecvCtx::boxed()).collect();
700 let send: Vec<Box<SendCtx>> =
701 (0..SEND_DEPTH).map(|_| Box::new(unsafe { std::mem::zeroed() })).collect();
702 let st = State {
703 ring,
704 recv,
705 send,
706 ready: VecDeque::new(),
707 ready_len: vec![0usize; RECV_DEPTH],
708 free_send: (0..SEND_DEPTH).collect(),
709 };
710 let me = Self { sock, fd, st: Mutex::new(st) };
711 if let Err(e) = me.submit_all_recv() {
712 return Err((me.sock, e));
713 }
714 Ok(me)
715 }
716
717 pub fn local_addr(&self) -> io::Result<SocketAddr> {
718 self.sock.local_addr()
719 }
720
721 pub fn raw_fd(&self) -> i32 {
725 self.fd
726 }
727
728 pub fn set_nonblocking(&self, nb: bool) -> io::Result<()> {
729 self.sock.set_nonblocking(nb)
730 }
731
732 fn submit_all_recv(&self) -> io::Result<()> {
733 let mut st = self.st.lock();
734 for i in 0..st.recv.len() {
735 self.push_recv(&mut st, i)?;
736 }
737 st.ring.submit()?;
738 Ok(())
739 }
740
741 fn push_recv(&self, st: &mut State, idx: usize) -> io::Result<()> {
742 st.recv[idx].refresh();
743 let msg: *mut libc::msghdr = std::ptr::addr_of_mut!(st.recv[idx].msghdr);
744 let e = opcode::RecvMsg::new(types::Fd(self.fd), msg)
745 .build()
746 .user_data(idx as u64);
747 unsafe {
750 st.ring
751 .submission()
752 .push(&e)
753 .map_err(|_| io::Error::other("io_uring SQ full (recv)"))?;
754 }
755 Ok(())
756 }
757
758 fn reap(&self, st: &mut State) {
759 let mut completed: Vec<(u64, i32)> = Vec::new();
760 for cqe in st.ring.completion() {
761 completed.push((cqe.user_data(), cqe.result()));
762 }
763 for (ud, res) in completed {
764 if ud >= SEND_TAG {
765 st.free_send.push((ud - SEND_TAG) as usize);
766 } else {
767 let idx = ud as usize;
768 if res >= 0 {
769 st.ready_len[idx] = res as usize;
770 st.ready.push_back(idx);
771 } else {
772 self.push_recv(st, idx).ok();
773 }
774 }
775 }
776 }
777
778 pub fn recv_with_kts(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Option<i128>)> {
779 let mut st = self.st.lock();
780 self.reap(&mut st);
781 if st.ready.is_empty() {
782 st.ring.submit()?;
783 self.reap(&mut st);
784 if st.ready.is_empty() {
785 return Err(io::Error::from(io::ErrorKind::WouldBlock));
786 }
787 }
788 let idx = st.ready.pop_front().unwrap();
789 let n = st.ready_len[idx];
790 let copy = n.min(buf.len());
791 let from;
792 let kts;
793 {
794 let ctx = &st.recv[idx];
795 buf[..copy].copy_from_slice(&ctx.buf[..copy]);
796 from = unsafe { sockaddr_to_socketaddr(&ctx.addr) }
797 .ok_or_else(|| io::Error::other("non-IP source"))?;
798 kts = parse_timestamp(&ctx.msghdr);
799 }
800 self.push_recv(&mut st, idx)?;
801 st.ring.submit()?;
802 Ok((copy, from, kts))
803 }
804
805 pub fn send_to(&self, data: &[u8], addr: SocketAddr) -> io::Result<usize> {
806 if data.len() > FRAME_CAP {
807 return Err(io::Error::other("datagram exceeds frame cap"));
808 }
809 let mut st = self.st.lock();
810 self.reap(&mut st);
811 if st.free_send.is_empty() {
812 st.ring.submit()?;
813 self.reap(&mut st);
814 if st.free_send.is_empty() {
815 return self.sock.send_to(data, addr);
816 }
817 }
818 let idx = st.free_send.pop().unwrap();
819 let (sa, sa_len) = socketaddr_to_sockaddr(addr);
820 {
821 let ctx = &mut st.send[idx];
822 ctx.buf[..data.len()].copy_from_slice(data);
823 ctx.addr = sa;
824 ctx.iov.iov_base = ctx.buf.as_mut_ptr() as *mut libc::c_void;
825 ctx.iov.iov_len = data.len();
826 ctx.msghdr.msg_name = std::ptr::addr_of_mut!(ctx.addr) as *mut libc::c_void;
827 ctx.msghdr.msg_namelen = sa_len;
828 ctx.msghdr.msg_iov = std::ptr::addr_of_mut!(ctx.iov);
829 ctx.msghdr.msg_iovlen = 1;
830 }
831 let msg: *const libc::msghdr = std::ptr::addr_of!(st.send[idx].msghdr);
832 let e = opcode::SendMsg::new(types::Fd(self.fd), msg)
833 .build()
834 .user_data(SEND_TAG | idx as u64);
835 unsafe {
838 st.ring
839 .submission()
840 .push(&e)
841 .map_err(|_| io::Error::other("io_uring SQ full (send)"))?;
842 }
843 st.ring.submit()?;
844 Ok(data.len())
845 }
846 }
847}
848
849#[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
872pub fn link_speed_bps(ifname: &str) -> Option<u64> {
873 let phys = ifname.strip_prefix("netmap:").unwrap_or(ifname);
876 if phys.starts_with("vale") || phys.contains(':') || phys.contains('{') {
877 return None;
878 }
879 #[cfg(target_os = "linux")]
880 {
881 let mbps: i64 = std::fs::read_to_string(format!("/sys/class/net/{phys}/speed"))
883 .ok()?
884 .trim()
885 .parse()
886 .ok()?;
887 (mbps > 0).then_some(mbps as u64 * 1_000_000)
888 }
889 #[cfg(any(target_os = "freebsd", target_os = "macos"))]
890 {
891 link_speed_baudrate(phys)
892 }
893}
894
895#[cfg(all(feature = "wire-locale", any(target_os = "freebsd", target_os = "macos")))]
898fn link_speed_baudrate(ifname: &str) -> Option<u64> {
899 use std::ffi::CStr;
900 let mut ifap: *mut libc::ifaddrs = std::ptr::null_mut();
901 if unsafe { libc::getifaddrs(&mut ifap) } != 0 {
902 return None;
903 }
904 let mut speed = None;
905 let mut cur = ifap;
906 while !cur.is_null() {
907 let ifa = unsafe { &*cur };
911 if !ifa.ifa_name.is_null() && !ifa.ifa_addr.is_null() && !ifa.ifa_data.is_null() {
912 let name = unsafe { CStr::from_ptr(ifa.ifa_name) }.to_string_lossy();
913 let family = unsafe { (*ifa.ifa_addr).sa_family } as i32;
914 if name == ifname && family == libc::AF_LINK {
915 let baud =
916 unsafe { (*(ifa.ifa_data as *const libc::if_data)).ifi_baudrate };
917 if baud > 0 {
918 #[cfg(target_os = "freebsd")]
920 {
921 speed = Some(baud);
922 }
923 #[cfg(target_os = "macos")]
924 {
925 speed = Some(u64::from(baud));
926 }
927 }
928 }
929 }
930 cur = ifa.ifa_next;
931 }
932 unsafe { libc::freeifaddrs(ifap) };
933 speed
934}
935
936#[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
942fn wire_min_link_bps() -> u64 {
943 std::env::var("SUBETHA_WIRE_MIN_GBPS")
944 .ok()
945 .and_then(|s| s.trim().parse::<f64>().ok())
946 .map(|g| (g * 1e9) as u64)
947 .unwrap_or(10_000_000_000)
948}
949
950#[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
955pub fn wire_gate_admits(ifname: &str) -> bool {
956 wire_link_fast_enough(ifname)
957}
958
959#[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
960fn wire_link_fast_enough(ifname: &str) -> bool {
961 let threshold = wire_min_link_bps();
962 if threshold == 0 {
963 return true;
964 }
965 link_speed_bps(ifname).map(|bps| bps >= threshold).unwrap_or(false)
966}
967
968#[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
969mod wire_backend {
970 use super::*;
971 use std::net::{IpAddr, Ipv4Addr};
972
973 use parking_lot::Mutex;
974
975 use crate::locale_wire::WireSocket;
976
977 const ETH_HDR: usize = 14;
978 const IP_HDR: usize = 20;
979 const UDP_HDR: usize = 8;
980 const HDRS: usize = ETH_HDR + IP_HDR + UDP_HDR; fn ipv4_csum(h: &[u8]) -> u16 {
983 let mut sum = 0u32;
984 let mut i = 0;
985 while i + 1 < h.len() {
986 sum += u16::from_be_bytes([h[i], h[i + 1]]) as u32;
987 i += 2;
988 }
989 while (sum >> 16) != 0 {
990 sum = (sum & 0xffff) + (sum >> 16);
991 }
992 !(sum as u16)
993 }
994
995 #[allow(clippy::too_many_arguments)]
998 fn build_frame(
999 dst_mac: [u8; 6],
1000 src_mac: [u8; 6],
1001 src_ip: [u8; 4],
1002 dst_ip: [u8; 4],
1003 src_port: u16,
1004 dst_port: u16,
1005 payload: &[u8],
1006 out: &mut Vec<u8>,
1007 ) {
1008 out.clear();
1009 out.resize(HDRS + payload.len(), 0);
1010 out[0..6].copy_from_slice(&dst_mac);
1011 out[6..12].copy_from_slice(&src_mac);
1012 out[12..14].copy_from_slice(&0x0800u16.to_be_bytes());
1013 out[14] = 0x45;
1014 out[16..18].copy_from_slice(&((IP_HDR + UDP_HDR + payload.len()) as u16).to_be_bytes());
1015 out[22] = 64;
1016 out[23] = 17;
1017 out[26..30].copy_from_slice(&src_ip);
1018 out[30..34].copy_from_slice(&dst_ip);
1019 let c = ipv4_csum(&out[14..34]);
1020 out[24..26].copy_from_slice(&c.to_be_bytes());
1021 out[34..36].copy_from_slice(&src_port.to_be_bytes());
1022 out[36..38].copy_from_slice(&dst_port.to_be_bytes());
1023 out[38..40].copy_from_slice(&((UDP_HDR + payload.len()) as u16).to_be_bytes());
1024 out[42..].copy_from_slice(payload);
1025 }
1026
1027 fn parse_frame(f: &[u8]) -> Option<([u8; 4], u16, u16, usize, usize)> {
1030 if f.len() < HDRS || f[12..14] != 0x0800u16.to_be_bytes() {
1031 return None;
1032 }
1033 let ihl = (f[14] & 0x0f) as usize * 4;
1034 if ihl < IP_HDR || f[23] != 17 {
1035 return None;
1036 }
1037 let l4 = ETH_HDR + ihl;
1038 if f.len() < l4 + UDP_HDR {
1039 return None;
1040 }
1041 let src_ip = [f[26], f[27], f[28], f[29]];
1042 let src_port = u16::from_be_bytes([f[l4], f[l4 + 1]]);
1043 let dst_port = u16::from_be_bytes([f[l4 + 2], f[l4 + 3]]);
1044 let udp_len = u16::from_be_bytes([f[l4 + 4], f[l4 + 5]]) as usize;
1045 let pstart = l4 + UDP_HDR;
1046 let plen = udp_len.saturating_sub(UDP_HDR).min(f.len() - pstart);
1047 Some((src_ip, src_port, dst_port, pstart, plen))
1048 }
1049
1050 fn parse_mac(s: &str) -> Option<[u8; 6]> {
1051 let mut mac = [0u8; 6];
1052 let parts: Vec<&str> = s.split(':').collect();
1053 if parts.len() != 6 {
1054 return None;
1055 }
1056 for (i, p) in parts.iter().enumerate() {
1057 mac[i] = u8::from_str_radix(p, 16).ok()?;
1058 }
1059 Some(mac)
1060 }
1061
1062 fn parse_ipv4(s: &str) -> Option<[u8; 4]> {
1063 s.parse::<Ipv4Addr>().ok().map(|a| a.octets())
1064 }
1065
1066 pub struct WireDgram {
1072 wire: Mutex<WireSocket>,
1073 scratch: Mutex<Vec<u8>>,
1074 local_ip: [u8; 4],
1075 local_mac: [u8; 6],
1076 peer_mac: [u8; 6],
1077 local_port: u16,
1078 }
1079
1080 impl WireDgram {
1081 pub fn from_env(local_port: u16) -> io::Result<Self> {
1082 let getv = |k: &str| -> io::Result<String> {
1083 std::env::var(k).map_err(|_| io::Error::other(format!("{k} unset")))
1084 };
1085 let ifname = getv("SUBETHA_WIRE_IFNAME")?;
1086 let local_ip = parse_ipv4(&getv("SUBETHA_WIRE_LOCAL_IP")?)
1087 .ok_or_else(|| io::Error::other("bad SUBETHA_WIRE_LOCAL_IP"))?;
1088 let local_mac = parse_mac(&getv("SUBETHA_WIRE_LOCAL_MAC")?)
1089 .ok_or_else(|| io::Error::other("bad SUBETHA_WIRE_LOCAL_MAC"))?;
1090 let peer_mac = parse_mac(&getv("SUBETHA_WIRE_PEER_MAC")?)
1091 .ok_or_else(|| io::Error::other("bad SUBETHA_WIRE_PEER_MAC"))?;
1092 let wire = WireSocket::bind(&ifname, 0)?;
1093 Ok(Self {
1094 wire: Mutex::new(wire),
1095 scratch: Mutex::new(Vec::with_capacity(HDRS + 2048)),
1096 local_ip,
1097 local_mac,
1098 peer_mac,
1099 local_port,
1100 })
1101 }
1102
1103 pub fn local_addr(&self) -> io::Result<SocketAddr> {
1104 Ok(SocketAddr::new(IpAddr::V4(Ipv4Addr::from(self.local_ip)), self.local_port))
1105 }
1106
1107 pub fn set_nonblocking(&self, _nb: bool) -> io::Result<()> {
1108 Ok(())
1110 }
1111
1112 pub fn send_to(&self, buf: &[u8], addr: SocketAddr) -> io::Result<usize> {
1113 let dst_ip = match addr.ip() {
1114 IpAddr::V4(v) => v.octets(),
1115 IpAddr::V6(_) => return Err(io::Error::other("wire backend is IPv4-only")),
1116 };
1117 let mut scratch = self.scratch.lock();
1118 build_frame(
1119 self.peer_mac,
1120 self.local_mac,
1121 self.local_ip,
1122 dst_ip,
1123 self.local_port,
1124 addr.port(),
1125 buf,
1126 &mut scratch,
1127 );
1128 self.wire.lock().send_frame(&scratch)?;
1129 Ok(buf.len())
1130 }
1131
1132 pub fn recv_with_kts(
1133 &self,
1134 buf: &mut [u8],
1135 ) -> io::Result<(usize, SocketAddr, Option<i128>)> {
1136 let mut fb = [0u8; 2048];
1137 let mut wire = self.wire.lock();
1138 loop {
1139 let n = wire.recv_frame(&mut fb, 0)?;
1141 if n == 0 {
1142 return Err(io::Error::from(io::ErrorKind::WouldBlock));
1143 }
1144 if let Some((src_ip, src_port, dst_port, pstart, plen)) = parse_frame(&fb[..n])
1147 && dst_port == self.local_port
1148 {
1149 let copy = plen.min(buf.len());
1150 buf[..copy].copy_from_slice(&fb[pstart..pstart + copy]);
1151 let from = SocketAddr::new(IpAddr::V4(Ipv4Addr::from(src_ip)), src_port);
1152 return Ok((copy, from, None));
1153 }
1154 }
1156 }
1157 }
1158}