solana_client/
connection_cache.rs

1pub use solana_connection_cache::connection_cache::Protocol;
2use {
3    solana_connection_cache::{
4        client_connection::ClientConnection,
5        connection_cache::{
6            BaseClientConnection, ConnectionCache as BackendConnectionCache, ConnectionPool,
7            NewConnectionConfig,
8        },
9    },
10    solana_keypair::Keypair,
11    solana_pubkey::Pubkey,
12    solana_quic_client::{QuicConfig, QuicConnectionManager, QuicPool},
13    solana_quic_definitions::NotifyKeyUpdate,
14    solana_streamer::streamer::StakedNodes,
15    solana_transaction_error::TransportResult,
16    solana_udp_client::{UdpConfig, UdpConnectionManager, UdpPool},
17    std::{
18        net::{IpAddr, Ipv4Addr, SocketAddr, UdpSocket},
19        sync::{Arc, RwLock},
20    },
21};
22
23const DEFAULT_CONNECTION_POOL_SIZE: usize = 4;
24const DEFAULT_CONNECTION_CACHE_USE_QUIC: bool = true;
25
26/// A thin wrapper over connection-cache/ConnectionCache to ease
27/// construction of the ConnectionCache for code dealing both with udp and quic.
28/// For the scenario only using udp or quic, use connection-cache/ConnectionCache directly.
29pub enum ConnectionCache {
30    Quic(Arc<BackendConnectionCache<QuicPool, QuicConnectionManager, QuicConfig>>),
31    Udp(Arc<BackendConnectionCache<UdpPool, UdpConnectionManager, UdpConfig>>),
32}
33
34type QuicBaseClientConnection = <QuicPool as ConnectionPool>::BaseClientConnection;
35type UdpBaseClientConnection = <UdpPool as ConnectionPool>::BaseClientConnection;
36
37pub enum BlockingClientConnection {
38    Quic(Arc<<QuicBaseClientConnection as BaseClientConnection>::BlockingClientConnection>),
39    Udp(Arc<<UdpBaseClientConnection as BaseClientConnection>::BlockingClientConnection>),
40}
41
42pub enum NonblockingClientConnection {
43    Quic(Arc<<QuicBaseClientConnection as BaseClientConnection>::NonblockingClientConnection>),
44    Udp(Arc<<UdpBaseClientConnection as BaseClientConnection>::NonblockingClientConnection>),
45}
46
47impl NotifyKeyUpdate for ConnectionCache {
48    fn update_key(&self, key: &Keypair) -> Result<(), Box<dyn std::error::Error>> {
49        match self {
50            Self::Udp(_) => Ok(()),
51            Self::Quic(backend) => backend.update_key(key),
52        }
53    }
54}
55
56impl ConnectionCache {
57    pub fn new(name: &'static str) -> Self {
58        if DEFAULT_CONNECTION_CACHE_USE_QUIC {
59            let cert_info = (&Keypair::new(), IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)));
60            ConnectionCache::new_with_client_options(
61                name,
62                DEFAULT_CONNECTION_POOL_SIZE,
63                None, // client_endpoint
64                Some(cert_info),
65                None, // stake_info
66            )
67        } else {
68            ConnectionCache::with_udp(name, DEFAULT_CONNECTION_POOL_SIZE)
69        }
70    }
71
72    /// Create a quic connection_cache
73    pub fn new_quic(name: &'static str, connection_pool_size: usize) -> Self {
74        Self::new_with_client_options(name, connection_pool_size, None, None, None)
75    }
76
77    /// Create a quic connection_cache with more client options
78    pub fn new_with_client_options(
79        name: &'static str,
80        connection_pool_size: usize,
81        client_socket: Option<UdpSocket>,
82        cert_info: Option<(&Keypair, IpAddr)>,
83        stake_info: Option<(&Arc<RwLock<StakedNodes>>, &Pubkey)>,
84    ) -> Self {
85        // The minimum pool size is 1.
86        let connection_pool_size = 1.max(connection_pool_size);
87        let mut config = QuicConfig::new().unwrap();
88        if let Some(cert_info) = cert_info {
89            config.update_client_certificate(cert_info.0, cert_info.1);
90        }
91        if let Some(client_socket) = client_socket {
92            config.update_client_endpoint(client_socket);
93        }
94        if let Some(stake_info) = stake_info {
95            config.set_staked_nodes(stake_info.0, stake_info.1);
96        }
97        let connection_manager = QuicConnectionManager::new_with_connection_config(config);
98        let cache =
99            BackendConnectionCache::new(name, connection_manager, connection_pool_size).unwrap();
100        Self::Quic(Arc::new(cache))
101    }
102
103    #[inline]
104    pub fn protocol(&self) -> Protocol {
105        match self {
106            Self::Quic(_) => Protocol::QUIC,
107            Self::Udp(_) => Protocol::UDP,
108        }
109    }
110
111    pub fn with_udp(name: &'static str, connection_pool_size: usize) -> Self {
112        // The minimum pool size is 1.
113        let connection_pool_size = 1.max(connection_pool_size);
114        let connection_manager = UdpConnectionManager::default();
115        let cache =
116            BackendConnectionCache::new(name, connection_manager, connection_pool_size).unwrap();
117        Self::Udp(Arc::new(cache))
118    }
119
120    pub fn use_quic(&self) -> bool {
121        matches!(self, Self::Quic(_))
122    }
123
124    pub fn get_connection(&self, addr: &SocketAddr) -> BlockingClientConnection {
125        match self {
126            Self::Quic(cache) => BlockingClientConnection::Quic(cache.get_connection(addr)),
127            Self::Udp(cache) => BlockingClientConnection::Udp(cache.get_connection(addr)),
128        }
129    }
130
131    pub fn get_nonblocking_connection(&self, addr: &SocketAddr) -> NonblockingClientConnection {
132        match self {
133            Self::Quic(cache) => {
134                NonblockingClientConnection::Quic(cache.get_nonblocking_connection(addr))
135            }
136            Self::Udp(cache) => {
137                NonblockingClientConnection::Udp(cache.get_nonblocking_connection(addr))
138            }
139        }
140    }
141}
142
143macro_rules! dispatch {
144    ($(#[$meta:meta])* $vis:vis fn $name:ident$(<$($t:ident: $cons:ident + ?Sized),*>)?(&self $(, $arg:ident: $ty:ty)*) $(-> $out:ty)?) => {
145        #[inline]
146        $(#[$meta])*
147        $vis fn $name$(<$($t: $cons + ?Sized),*>)?(&self $(, $arg:$ty)*) $(-> $out)? {
148            match self {
149                Self::Quic(this) => this.$name($($arg, )*),
150                Self::Udp(this) => this.$name($($arg, )*),
151            }
152        }
153    };
154    ($(#[$meta:meta])* $vis:vis fn $name:ident$(<$($t:ident: $cons:ident + ?Sized),*>)?(&mut self $(, $arg:ident: $ty:ty)*) $(-> $out:ty)?) => {
155        #[inline]
156        $(#[$meta])*
157        $vis fn $name$(<$($t: $cons + ?Sized),*>)?(&mut self $(, $arg:$ty)*) $(-> $out)? {
158            match self {
159                Self::Quic(this) => this.$name($($arg, )*),
160                Self::Udp(this) => this.$name($($arg, )*),
161            }
162        }
163    };
164}
165
166pub(crate) use dispatch;
167
168impl ClientConnection for BlockingClientConnection {
169    dispatch!(fn server_addr(&self) -> &SocketAddr);
170    dispatch!(fn send_data(&self, buffer: &[u8]) -> TransportResult<()>);
171    dispatch!(fn send_data_async(&self, buffer: Vec<u8>) -> TransportResult<()>);
172    dispatch!(fn send_data_batch(&self, buffers: &[Vec<u8>]) -> TransportResult<()>);
173    dispatch!(fn send_data_batch_async(&self, buffers: Vec<Vec<u8>>) -> TransportResult<()>);
174}
175
176#[async_trait::async_trait]
177impl solana_connection_cache::nonblocking::client_connection::ClientConnection
178    for NonblockingClientConnection
179{
180    dispatch!(fn server_addr(&self) -> &SocketAddr);
181
182    async fn send_data(&self, buffer: &[u8]) -> TransportResult<()> {
183        match self {
184            Self::Quic(cache) => Ok(cache.send_data(buffer).await?),
185            Self::Udp(cache) => Ok(cache.send_data(buffer).await?),
186        }
187    }
188
189    async fn send_data_batch(&self, buffers: &[Vec<u8>]) -> TransportResult<()> {
190        match self {
191            Self::Quic(cache) => Ok(cache.send_data_batch(buffers).await?),
192            Self::Udp(cache) => Ok(cache.send_data_batch(buffers).await?),
193        }
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use {
200        super::*,
201        crate::connection_cache::ConnectionCache,
202        solana_net_utils::sockets::{bind_to, localhost_port_range_for_tests},
203        std::net::{IpAddr, Ipv4Addr, SocketAddr},
204    };
205
206    #[test]
207    fn test_connection_with_specified_client_endpoint() {
208        let port_range = localhost_port_range_for_tests();
209        let mut port_range = port_range.0..port_range.1;
210        let client_socket =
211            bind_to(IpAddr::V4(Ipv4Addr::LOCALHOST), port_range.next().unwrap()).unwrap();
212        let connection_cache = ConnectionCache::new_with_client_options(
213            "connection_cache_test",
214            1,                   // connection_pool_size
215            Some(client_socket), // client_endpoint
216            None,                // cert_info
217            None,                // stake_info
218        );
219
220        // server port 1:
221        let port1 = port_range.next().unwrap();
222        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port1);
223        let conn = connection_cache.get_connection(&addr);
224        assert_eq!(conn.server_addr().port(), port1);
225
226        // server port 2:
227        let port2 = port_range.next().unwrap();
228        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port2);
229        let conn = connection_cache.get_connection(&addr);
230        assert_eq!(conn.server_addr().port(), port2);
231    }
232}
solana_client/connection_cache.rs

solana_client/
connection_cache.rs
