sova_udp/
lib.rs

1//! UDP listeners as [`BackgroundService`](sova_core::BackgroundService).
2
3use sova_core::extend::{wait_shutdown, BoxFuture, StateMap};
4use sova_core::{BackgroundService, Shutdown};
5use std::net::SocketAddr;
6use std::sync::Arc;
7use tokio::net::UdpSocket;
8
9pub type UdpHandler =
10    Arc<dyn Fn(SocketAddr, Vec<u8>, Arc<UdpSocket>) -> BoxFuture<()> + Send + Sync>;
11
12/// Bind a UDP socket and invoke `handler` for each datagram until shutdown.
13pub struct UdpService {
14    addr: SocketAddr,
15    handler: UdpHandler,
16    name: String,
17}
18
19impl UdpService {
20    pub fn new<F, Fut>(addr: SocketAddr, handler: F) -> Self
21    where
22        F: Fn(SocketAddr, Vec<u8>, Arc<UdpSocket>) -> Fut + Send + Sync + 'static,
23        Fut: std::future::Future<Output = ()> + Send + 'static,
24    {
25        Self {
26            addr,
27            handler: Arc::new(move |a, b, s| Box::pin(handler(a, b, s))),
28            name: format!("udp:{addr}"),
29        }
30    }
31
32    pub fn echo(addr: SocketAddr) -> Self {
33        Self::new(addr, |_peer, data, sock| async move {
34            let _ = sock.send_to(&data, _peer).await;
35        })
36    }
37}
38
39impl BackgroundService for UdpService {
40    fn name(&self) -> &str {
41        &self.name
42    }
43
44    fn run(
45        self: Box<Self>,
46        _state: Arc<StateMap>,
47        shutdown: Shutdown,
48    ) -> BoxFuture<()> {
49        Box::pin(async move {
50            let sock = match UdpSocket::bind(self.addr).await {
51                Ok(s) => Arc::new(s),
52                Err(e) => {
53                    tracing::error!(error = %e, addr = %self.addr, "udp bind failed");
54                    return;
55                }
56            };
57            tracing::info!(addr = %self.addr, "udp listening");
58            let mut buf = vec![0u8; 65535];
59            loop {
60                tokio::select! {
61                    _ = wait_shutdown(shutdown.clone()) => break,
62                    res = sock.recv_from(&mut buf) => {
63                        match res {
64                            Ok((n, peer)) => {
65                                let data = buf[..n].to_vec();
66                                (self.handler)(peer, data, Arc::clone(&sock)).await;
67                            }
68                            Err(e) => {
69                                tracing::debug!(error = %e, "udp recv");
70                                break;
71                            }
72                        }
73                    }
74                }
75            }
76        })
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use std::time::Duration;
84
85    #[tokio::test]
86    async fn echo_roundtrip() {
87        let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
88        let sock = UdpSocket::bind(addr).await.unwrap();
89        let addr = sock.local_addr().unwrap();
90        drop(sock);
91
92        let (tx, shutdown) = sova_core::shutdown_channel();
93        let svc = Box::new(UdpService::echo(addr));
94        let handle = tokio::spawn(svc.run(Arc::new(StateMap::new()), shutdown));
95
96        tokio::time::sleep(Duration::from_millis(50)).await;
97        let client = UdpSocket::bind("127.0.0.1:0").await.unwrap();
98        client.send_to(b"ping", addr).await.unwrap();
99        let mut buf = [0u8; 16];
100        let (n, _) = tokio::time::timeout(Duration::from_secs(1), client.recv_from(&mut buf))
101            .await
102            .unwrap()
103            .unwrap();
104        assert_eq!(&buf[..n], b"ping");
105        let _ = tx.send(true);
106        let _ = handle.await;
107    }
108
109    #[tokio::test]
110    async fn shutdown_stops_recv_loop() {
111        let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
112        let sock = UdpSocket::bind(addr).await.unwrap();
113        let addr = sock.local_addr().unwrap();
114        drop(sock);
115
116        let (tx, shutdown) = sova_core::shutdown_channel();
117        let svc = Box::new(UdpService::echo(addr));
118        let handle = tokio::spawn(svc.run(Arc::new(StateMap::new()), shutdown));
119
120        tokio::time::sleep(Duration::from_millis(50)).await;
121        let _ = tx.send(true);
122        tokio::time::timeout(Duration::from_secs(1), handle)
123            .await
124            .expect("service did not stop after shutdown")
125            .expect("join error");
126    }
127
128    #[tokio::test]
129    async fn max_datagram_size_accepted() {
130        let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
131        let sock = UdpSocket::bind(addr).await.unwrap();
132        let addr = sock.local_addr().unwrap();
133        drop(sock);
134
135        let (tx, shutdown) = sova_core::shutdown_channel();
136        let svc = Box::new(UdpService::echo(addr));
137        let handle = tokio::spawn(svc.run(Arc::new(StateMap::new()), shutdown));
138
139        tokio::time::sleep(Duration::from_millis(50)).await;
140        let client = UdpSocket::bind("127.0.0.1:0").await.unwrap();
141        let payload = vec![0xABu8; 8192];
142        client.send_to(&payload, addr).await.unwrap();
143        let mut buf = vec![0u8; 8192];
144        let (n, _) = tokio::time::timeout(Duration::from_secs(1), client.recv_from(&mut buf))
145            .await
146            .unwrap()
147            .unwrap();
148        assert_eq!(&buf[..n], &payload[..]);
149        let _ = tx.send(true);
150        let _ = handle.await;
151    }
152}
sova_udp/lib.rs

sova_udp/
lib.rs
