udp_echo/
udp_echo.rs

1//! A connectionless sibling of `echo.rs`: a thread-per-core **UDP** echo server.
2//!
3//! `SO_REUSEPORT` works for datagram sockets too, so each pinned worker binds the
4//! same address and the kernel load-balances incoming datagrams across the group
5//! by flow hash — the same trick the TCP examples use, one layer down. There is
6//! no accept loop and no connection; each worker just loops `recv_from` ->
7//! `send_to` on its own socket, borrowing a datagram buffer from this thread's
8//! [`Pool`] for each packet.
9//!
10//! The crate's `bind_reuseport` is TCP-only, so the few lines of `socket2` for a
11//! `SO_REUSEPORT` *datagram* socket live here in the example.
12//!
13//! ```text
14//! cargo run --release --example udp_echo -- [ADDR] [CAPACITY]
15//!   ADDR       bind address               default 0.0.0.0:7000
16//!   CAPACITY   datagram buffers per worker default 1024
17//!
18//! # try it (sends one datagram, prints the echo):
19//! echo -n hello | nc -u -w1 127.0.0.1 7000
20//! ```
21
22use std::{io, net::SocketAddr, thread};
23
24use compio::{BufResult, net::UdpSocket, runtime::Runtime};
25use compio_pool::{ManageConnection, Pool, cpu};
26use socket2::{Domain, Protocol, Socket, Type};
27
28/// Hands out reusable 64 KiB datagram buffers — one large enough for any single
29/// UDP payload — one per packet in flight.
30struct Buffers;
31
32impl ManageConnection for Buffers {
33    type Connection = Vec<u8>;
34    type Error = io::Error;
35
36    async fn connect(&self) -> io::Result<Vec<u8>> {
37        Ok(Vec::with_capacity(64 * 1024))
38    }
39
40    async fn is_valid(&self, _buf: &mut Vec<u8>) -> io::Result<()> {
41        Ok(())
42    }
43
44    fn has_broken(&self, _buf: &mut Vec<u8>) -> bool {
45        false
46    }
47}
48
49/// A `SO_REUSEPORT` UDP socket, the datagram counterpart of
50/// [`compio_pool::bind_reuseport`]. Both reuse flags go on before `bind`, which
51/// is what lets every worker bind the same address and have the kernel split
52/// datagrams across them.
53fn bind_reuseport_udp(addr: SocketAddr) -> io::Result<std::net::UdpSocket> {
54    let socket = Socket::new(Domain::for_address(addr), Type::DGRAM, Some(Protocol::UDP))?;
55    socket.set_reuse_address(true)?;
56    socket.set_reuse_port(true)?;
57    socket.bind(&addr.into())?;
58    Ok(socket.into())
59}
60
61fn main() -> io::Result<()> {
62    let addr: SocketAddr = std::env::args()
63        .nth(1)
64        .unwrap_or_else(|| "0.0.0.0:7000".into())
65        .parse()
66        .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("bad address: {e}")))?;
67    let capacity: u32 = std::env::args()
68        .nth(2)
69        .and_then(|s| s.parse().ok())
70        .unwrap_or(1024);
71
72    let pool = Pool::builder().max_size(capacity).build(Buffers);
73
74    let cores = cpu::cores();
75    if cores.is_empty() {
76        return Err(io::Error::new(io::ErrorKind::Unsupported, "no cores to run on"));
77    }
78    println!(
79        "udp_echo on {addr}: {} workers, capacity {capacity}/worker",
80        cores.len()
81    );
82
83    let mut handles = Vec::new();
84    for (index, core) in cores.into_iter().enumerate() {
85        let pool = pool.clone();
86        let handle = thread::Builder::new()
87            .name(format!("worker/{index}"))
88            .spawn(move || worker(core, addr, pool))?;
89        handles.push(handle);
90    }
91    for handle in handles {
92        let _ = handle.join();
93    }
94    Ok(())
95}
96
97/// One pinned thread: its own runtime, its own `SO_REUSEPORT` datagram socket,
98/// and its own [`LocalPool`] of datagram buffers.
99fn worker(core: cpu::CoreId, addr: SocketAddr, pool: Pool<Buffers>) {
100    cpu::pin_current(core);
101
102    let runtime = Runtime::builder().build().expect("build compio runtime");
103    runtime.block_on(async move {
104        let local = pool.local();
105
106        let std_socket = bind_reuseport_udp(addr).expect("bind SO_REUSEPORT UDP");
107        let socket = UdpSocket::from_std(std_socket).expect("wrap socket in ring");
108
109        // No accept loop for a datagram socket: one task owns the socket and
110        // echoes each packet. Borrow a buffer per datagram so the working set is
111        // bounded and reused rather than reallocated.
112        loop {
113            let mut lease = match local.get().await {
114                Ok(lease) => lease,
115                Err(_) => continue,
116            };
117
118            let mut buf = std::mem::take(&mut *lease);
119            buf.clear();
120            let BufResult(res, buf) = socket.recv_from(buf).await;
121            let (buf, peer) = match res {
122                // After `recv_from` the buffer's length is the datagram size, so
123                // sending it straight back echoes exactly those bytes.
124                Ok((_n, peer)) => (buf, peer),
125                Err(_) => {
126                    *lease = buf;
127                    continue;
128                }
129            };
130
131            let BufResult(_sent, buf) = socket.send_to(buf, peer).await;
132            *lease = buf; // return the buffer to the pool on drop
133        }
134    });
135}
udp_echo/udp_echo.rs

udp_echo/
udp_echo.rs
