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_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}