Skip to main content

msrt_udp/
client.rs

1//! Connected UDP client adapter.
2
3use std::io::ErrorKind;
4use std::net::SocketAddr;
5use std::time::Instant;
6
7use msrt::endpoint::{ClientEndpoint, EndpointPoll, EngineConfig, PeerState};
8use tokio::net::{ToSocketAddrs, UdpSocket};
9
10use crate::error::Result;
11use crate::event::UdpClientEvent;
12
13const RX_BYTES: usize = 2048;
14const TX_BYTES: usize = 256;
15
16/// Connected UDP client adapter.
17#[derive(Debug)]
18pub struct UdpClient {
19    socket: UdpSocket,
20    endpoint: ClientEndpoint,
21    start: Instant,
22    rx_buf: [u8; RX_BYTES],
23    tx_buf: [u8; TX_BYTES],
24}
25
26impl UdpClient {
27    /// Binds a local UDP socket and connects it to `remote`.
28    pub async fn bind<A, R>(local: A, remote: R) -> Result<Self>
29    where
30        A: ToSocketAddrs,
31        R: ToSocketAddrs,
32    {
33        Self::bind_with_config(local, remote, EngineConfig::default()).await
34    }
35
36    /// Binds a local UDP socket, connects it to `remote`, and uses `config`.
37    pub async fn bind_with_config<A, R>(local: A, remote: R, config: EngineConfig) -> Result<Self>
38    where
39        A: ToSocketAddrs,
40        R: ToSocketAddrs,
41    {
42        let socket = UdpSocket::bind(local).await?;
43        socket.connect(remote).await?;
44        Ok(Self::from_socket(socket, config))
45    }
46
47    /// Creates a client from an already connected UDP socket.
48    pub fn from_socket(socket: UdpSocket, config: EngineConfig) -> Self {
49        Self {
50            socket,
51            endpoint: ClientEndpoint::new(config),
52            start: Instant::now(),
53            rx_buf: [0; RX_BYTES],
54            tx_buf: [0; TX_BYTES],
55        }
56    }
57
58    /// Returns the local socket address.
59    pub fn local_addr(&self) -> Result<SocketAddr> {
60        Ok(self.socket.local_addr()?)
61    }
62
63    /// Returns the remote socket address.
64    pub fn peer_addr(&self) -> Result<SocketAddr> {
65        Ok(self.socket.peer_addr()?)
66    }
67
68    /// Returns the endpoint peer state.
69    pub fn peer_state(&self) -> PeerState {
70        self.endpoint.peer().state()
71    }
72
73    /// Returns a shared reference to the UDP socket.
74    pub const fn socket(&self) -> &UdpSocket {
75        &self.socket
76    }
77
78    /// Returns a mutable reference to the UDP socket.
79    pub fn socket_mut(&mut self) -> &mut UdpSocket {
80        &mut self.socket
81    }
82
83    /// Consumes the adapter and returns the UDP socket.
84    pub fn into_socket(self) -> UdpSocket {
85        self.socket
86    }
87
88    /// Starts a fresh client session.
89    pub fn connect(&mut self) -> Result<()> {
90        self.endpoint.connect(self.now_ms())?;
91        Ok(())
92    }
93
94    /// Drops the active client session.
95    pub fn disconnect(&mut self) {
96        self.endpoint.disconnect();
97    }
98
99    /// Queues an application message.
100    pub fn send(&mut self, message: &[u8]) -> Result<bool> {
101        Ok(self.endpoint.send(message)?.is_some())
102    }
103
104    /// Receives currently available UDP datagrams and feeds them into MSRT.
105    pub fn receive_available(&mut self) -> Result<usize> {
106        let mut datagrams = 0;
107        loop {
108            match self.socket.try_recv(&mut self.rx_buf) {
109                Ok(n) => {
110                    datagrams += 1;
111                    let now_ms = self.now_ms();
112                    let _ = self.endpoint.receive(now_ms, &self.rx_buf[..n]);
113                }
114                Err(error) if error.kind() == ErrorKind::WouldBlock => return Ok(datagrams),
115                Err(error) if error.kind() == ErrorKind::Interrupted => continue,
116                Err(error) => return Err(error.into()),
117            }
118        }
119    }
120
121    /// Polls one adapter event and sends pending UDP datagrams.
122    pub async fn poll(&mut self) -> Result<UdpClientEvent> {
123        let now_ms = self.now_ms();
124        loop {
125            match self.endpoint.poll(now_ms, &mut self.tx_buf)? {
126                EndpointPoll::Transmit { bytes, .. } => {
127                    let _ = self.socket.send(bytes).await?;
128                }
129                EndpointPoll::Message(message) => return Ok(UdpClientEvent::Message(message)),
130                EndpointPoll::SendFailed(failed) => return Ok(UdpClientEvent::SendFailed(failed)),
131                EndpointPoll::Idle => return Ok(UdpClientEvent::Idle),
132            }
133        }
134    }
135
136    /// Runs `receive_available` followed by `poll`.
137    pub async fn tick(&mut self) -> Result<UdpClientEvent> {
138        if let Err(error) = self.receive_available() {
139            if let Some(kind) = recoverable_transport_error(&error) {
140                return Ok(UdpClientEvent::TransportUnavailable { kind });
141            }
142            return Err(error);
143        }
144
145        match self.poll().await {
146            Ok(event) => Ok(event),
147            Err(error) => {
148                if let Some(kind) = recoverable_transport_error(&error) {
149                    return Ok(UdpClientEvent::TransportUnavailable { kind });
150                }
151                Err(error)
152            }
153        }
154    }
155
156    fn now_ms(&self) -> u64 {
157        self.start
158            .elapsed()
159            .as_millis()
160            .try_into()
161            .unwrap_or(u64::MAX)
162    }
163}
164
165fn recoverable_transport_error(error: &crate::Error) -> Option<ErrorKind> {
166    let crate::Error::Io(error) = error else {
167        return None;
168    };
169
170    match error.kind() {
171        ErrorKind::ConnectionRefused
172        | ErrorKind::ConnectionReset
173        | ErrorKind::ConnectionAborted
174        | ErrorKind::NotConnected
175        | ErrorKind::TimedOut => Some(error.kind()),
176        _ => None,
177    }
178}