solana_udp_client/nonblocking/
udp_client.rs

1//! Simple UDP client that communicates with the given UDP port with UDP and provides
2//! an interface for sending data
3
4use {
5    async_trait::async_trait, core::iter::repeat,
6    solana_connection_cache::nonblocking::client_connection::ClientConnection,
7    solana_sdk::transport::Result as TransportResult,
8    solana_streamer::nonblocking::sendmmsg::batch_send, std::net::SocketAddr,
9    tokio::net::UdpSocket,
10};
11
12pub struct UdpClientConnection {
13    pub socket: UdpSocket,
14    pub addr: SocketAddr,
15}
16
17impl UdpClientConnection {
18    pub fn new_from_addr(socket: std::net::UdpSocket, server_addr: SocketAddr) -> Self {
19        socket.set_nonblocking(true).unwrap();
20        let socket = UdpSocket::from_std(socket).unwrap();
21        Self {
22            socket,
23            addr: server_addr,
24        }
25    }
26}
27
28#[async_trait]
29impl ClientConnection for UdpClientConnection {
30    fn server_addr(&self) -> &SocketAddr {
31        &self.addr
32    }
33
34    async fn send_data(&self, buffer: &[u8]) -> TransportResult<()> {
35        self.socket.send_to(buffer, self.addr).await?;
36        Ok(())
37    }
38
39    async fn send_data_batch(&self, buffers: &[Vec<u8>]) -> TransportResult<()> {
40        let pkts: Vec<_> = buffers.iter().zip(repeat(self.server_addr())).collect();
41        batch_send(&self.socket, &pkts).await?;
42        Ok(())
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use {
49        super::*,
50        solana_sdk::packet::{Packet, PACKET_DATA_SIZE},
51        solana_streamer::nonblocking::recvmmsg::recv_mmsg,
52        std::net::{IpAddr, Ipv4Addr},
53        tokio::net::UdpSocket,
54    };
55
56    async fn check_send_one(connection: &UdpClientConnection, reader: &UdpSocket) {
57        let packet = vec![111u8; PACKET_DATA_SIZE];
58        connection.send_data(&packet).await.unwrap();
59        let mut packets = vec![Packet::default(); 32];
60        let recv = recv_mmsg(reader, &mut packets[..]).await.unwrap();
61        assert_eq!(1, recv);
62    }
63
64    async fn check_send_batch(connection: &UdpClientConnection, reader: &UdpSocket) {
65        let packets: Vec<_> = (0..32).map(|_| vec![0u8; PACKET_DATA_SIZE]).collect();
66        connection.send_data_batch(&packets).await.unwrap();
67        let mut packets = vec![Packet::default(); 32];
68        let recv = recv_mmsg(reader, &mut packets[..]).await.unwrap();
69        assert_eq!(32, recv);
70    }
71
72    #[tokio::test]
73    async fn test_send_from_addr() {
74        let addr_str = "0.0.0.0:50100";
75        let addr = addr_str.parse().unwrap();
76        let socket =
77            solana_net_utils::bind_with_any_port(IpAddr::V4(Ipv4Addr::UNSPECIFIED)).unwrap();
78        let connection = UdpClientConnection::new_from_addr(socket, addr);
79        let reader = UdpSocket::bind(addr_str).await.expect("bind");
80        check_send_one(&connection, &reader).await;
81        check_send_batch(&connection, &reader).await;
82    }
83}