solana_udp_client/nonblocking/
udp_client.rs1use {
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}