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_streamer::nonblocking::sendmmsg::batch_send, solana_transaction_error::TransportResult,
8    std::net::SocketAddr, tokio::net::UdpSocket,
9};
10
11pub struct UdpClientConnection {
12    pub socket: UdpSocket,
13    pub addr: SocketAddr,
14}
15
16impl UdpClientConnection {
17    pub fn new_from_addr(socket: std::net::UdpSocket, server_addr: SocketAddr) -> Self {
18        socket.set_nonblocking(true).unwrap();
19        let socket = UdpSocket::from_std(socket).unwrap();
20        Self {
21            socket,
22            addr: server_addr,
23        }
24    }
25}
26
27#[async_trait]
28impl ClientConnection for UdpClientConnection {
29    fn server_addr(&self) -> &SocketAddr {
30        &self.addr
31    }
32
33    async fn send_data(&self, buffer: &[u8]) -> TransportResult<()> {
34        self.socket.send_to(buffer, self.addr).await?;
35        Ok(())
36    }
37
38    async fn send_data_batch(&self, buffers: &[Vec<u8>]) -> TransportResult<()> {
39        let pkts: Vec<_> = buffers.iter().zip(repeat(self.server_addr())).collect();
40        batch_send(&self.socket, &pkts).await?;
41        Ok(())
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use {
48        super::*,
49        solana_net_utils::sockets::{
50            bind_to_async, bind_to_with_config, unique_port_range_for_tests,
51            SocketConfiguration as SocketConfig,
52        },
53        solana_packet::{Packet, PACKET_DATA_SIZE},
54        solana_streamer::nonblocking::recvmmsg::recv_mmsg,
55        std::net::{IpAddr, Ipv4Addr},
56        tokio::net::UdpSocket,
57    };
58
59    async fn check_send_one(connection: &UdpClientConnection, reader: &UdpSocket) {
60        let packet = vec![111u8; PACKET_DATA_SIZE];
61        connection.send_data(&packet).await.unwrap();
62        let mut packets = vec![Packet::default(); 32];
63        let recv = recv_mmsg(reader, &mut packets[..]).await.unwrap();
64        assert_eq!(1, recv);
65    }
66
67    async fn check_send_batch(connection: &UdpClientConnection, reader: &UdpSocket) {
68        let packets: Vec<_> = (0..32).map(|_| vec![0u8; PACKET_DATA_SIZE]).collect();
69        connection.send_data_batch(&packets).await.unwrap();
70        let mut packets = vec![Packet::default(); 32];
71        let recv = recv_mmsg(reader, &mut packets[..]).await.unwrap();
72        assert_eq!(32, recv);
73    }
74
75    #[tokio::test]
76    async fn test_send_from_addr() {
77        let mut port_range = unique_port_range_for_tests(4);
78        let socket = bind_to_with_config(
79            IpAddr::V4(Ipv4Addr::UNSPECIFIED),
80            port_range.next().unwrap(),
81            SocketConfig::default(),
82        )
83        .unwrap();
84
85        let reader_ip = IpAddr::V4(Ipv4Addr::LOCALHOST);
86        let reader_port = port_range.next().unwrap();
87        let connection =
88            UdpClientConnection::new_from_addr(socket, SocketAddr::new(reader_ip, reader_port));
89
90        let reader = bind_to_async(reader_ip, reader_port).await.expect("bind");
91        check_send_one(&connection, &reader).await;
92        check_send_batch(&connection, &reader).await;
93    }
94}