Skip to main content

msrt_udp/
server.rs

1//! Multi-peer UDP server adapter.
2
3use std::io::ErrorKind;
4use std::net::SocketAddr;
5use std::time::Instant;
6
7use msrt::endpoint::{EndpointPoll, EngineConfig, PeerState, ServerEndpoint};
8use tokio::net::{ToSocketAddrs, UdpSocket};
9
10use crate::error::{Error, Result};
11use crate::event::UdpServerEvent;
12
13const RX_BYTES: usize = 2048;
14const TX_BYTES: usize = 256;
15
16/// Multi-peer UDP server adapter.
17#[derive(Debug)]
18pub struct UdpServer<const N: usize> {
19    socket: UdpSocket,
20    endpoint: ServerEndpoint<SocketAddr, N>,
21    start: Instant,
22    rx_buf: [u8; RX_BYTES],
23    tx_buf: [u8; TX_BYTES],
24}
25
26impl<const N: usize> UdpServer<N> {
27    /// Binds a UDP server socket using default MSRT config.
28    pub async fn bind<A>(local: A) -> Result<Self>
29    where
30        A: ToSocketAddrs,
31    {
32        Self::bind_with_config(local, EngineConfig::default()).await
33    }
34
35    /// Binds a UDP server socket using `config`.
36    pub async fn bind_with_config<A>(local: A, config: EngineConfig) -> Result<Self>
37    where
38        A: ToSocketAddrs,
39    {
40        let socket = UdpSocket::bind(local).await?;
41        Ok(Self::from_socket(socket, config))
42    }
43
44    /// Creates a server from an existing UDP socket.
45    pub fn from_socket(socket: UdpSocket, config: EngineConfig) -> Self {
46        Self {
47            socket,
48            endpoint: ServerEndpoint::new(config),
49            start: Instant::now(),
50            rx_buf: [0; RX_BYTES],
51            tx_buf: [0; TX_BYTES],
52        }
53    }
54
55    /// Returns the local socket address.
56    pub fn local_addr(&self) -> Result<SocketAddr> {
57        Ok(self.socket.local_addr()?)
58    }
59
60    /// Returns a shared reference to the UDP socket.
61    pub const fn socket(&self) -> &UdpSocket {
62        &self.socket
63    }
64
65    /// Returns a mutable reference to the UDP socket.
66    pub fn socket_mut(&mut self) -> &mut UdpSocket {
67        &mut self.socket
68    }
69
70    /// Consumes the adapter and returns the UDP socket.
71    pub fn into_socket(self) -> UdpSocket {
72        self.socket
73    }
74
75    /// Returns the peer state for `peer`.
76    pub fn peer_state(&mut self, peer: SocketAddr) -> Option<PeerState> {
77        self.endpoint.peer_mut(peer).map(|slot| slot.state())
78    }
79
80    /// Returns the currently accepted peers.
81    pub fn peers(&self) -> impl Iterator<Item = SocketAddr> + '_ {
82        self.endpoint.peers().map(|peer| *peer.peer_id())
83    }
84
85    /// Queues an application message for `peer`.
86    pub fn send_to(&mut self, peer: SocketAddr, message: &[u8]) -> Result<bool> {
87        match self.endpoint.send(peer, message) {
88            Some(Ok(Some(_))) => Ok(true),
89            Some(Ok(None)) | None => Ok(false),
90            Some(Err(error)) => Err(error.into()),
91        }
92    }
93
94    /// Disconnects a peer and frees its server slot.
95    pub fn disconnect(&mut self, peer: SocketAddr) -> bool {
96        self.endpoint.disconnect(peer)
97    }
98
99    /// Disconnects every peer idle for at least `timeout_ms`.
100    pub fn disconnect_idle(&mut self, timeout_ms: u64) -> usize {
101        self.endpoint.disconnect_idle(self.now_ms(), timeout_ms)
102    }
103
104    /// Receives currently available UDP datagrams and feeds them into MSRT.
105    ///
106    /// Unknown peers are accepted automatically. If the fixed-capacity peer
107    /// table is full, `Error::Accept` is returned.
108    pub fn receive_available(&mut self) -> Result<usize> {
109        let mut datagrams = 0;
110        loop {
111            match self.socket.try_recv_from(&mut self.rx_buf) {
112                Ok((n, peer)) => {
113                    datagrams += 1;
114                    let now_ms = self.now_ms();
115                    if self.endpoint.peer_mut(peer).is_none() {
116                        self.endpoint.accept(peer, now_ms)?;
117                    }
118                    let _ = self.endpoint.receive(peer, now_ms, &self.rx_buf[..n]);
119                }
120                Err(error) if error.kind() == ErrorKind::WouldBlock => return Ok(datagrams),
121                Err(error) if error.kind() == ErrorKind::Interrupted => continue,
122                Err(error) => return Err(error.into()),
123            }
124        }
125    }
126
127    /// Polls one adapter event and sends pending UDP datagrams.
128    pub async fn poll(&mut self) -> Result<UdpServerEvent> {
129        let now_ms = self.now_ms();
130        let peers: Vec<SocketAddr> = self.peers().collect();
131
132        for peer in peers {
133            loop {
134                let poll = self
135                    .endpoint
136                    .poll(peer, now_ms, &mut self.tx_buf)
137                    .ok_or_else(|| Error::Io(ErrorKind::NotFound.into()))??;
138
139                match poll {
140                    EndpointPoll::Transmit { bytes, .. } => {
141                        let _ = self.socket.send_to(bytes, peer).await?;
142                    }
143                    EndpointPoll::Message(message) => {
144                        return Ok(UdpServerEvent::Message { peer, message });
145                    }
146                    EndpointPoll::SendFailed(failed) => {
147                        return Ok(UdpServerEvent::SendFailed { peer, failed });
148                    }
149                    EndpointPoll::Idle => break,
150                }
151            }
152        }
153
154        Ok(UdpServerEvent::Idle)
155    }
156
157    /// Runs `receive_available` followed by `poll`.
158    pub async fn tick(&mut self) -> Result<UdpServerEvent> {
159        let _ = self.receive_available()?;
160        self.poll().await
161    }
162
163    fn now_ms(&self) -> u64 {
164        self.start
165            .elapsed()
166            .as_millis()
167            .try_into()
168            .unwrap_or(u64::MAX)
169    }
170}