Optimize TpuConnection and its implementations and refactor connection-cache to not use dyn in order to enable those changes (#23877)

This commit is contained in:
ryleung-solana
2022-03-24 11:40:26 -04:00
committed by GitHub
parent 5b916961b5
commit 82945ba973
11 changed files with 266 additions and 48 deletions

View File

@@ -41,6 +41,7 @@ solana-faucet = { path = "../faucet", version = "=1.11.0" }
solana-measure = { path = "../measure", version = "=1.11.0" }
solana-net-utils = { path = "../net-utils", version = "=1.11.0" }
solana-sdk = { path = "../sdk", version = "=1.11.0" }
solana-streamer = { path = "../streamer", version = "=1.11.0" }
solana-transaction-status = { path = "../transaction-status", version = "=1.11.0" }
solana-version = { path = "../version", version = "=1.11.0" }
solana-vote-program = { path = "../programs/vote", version = "=1.11.0" }

View File

@@ -3,6 +3,7 @@ use {
quic_client::QuicTpuConnection, tpu_connection::TpuConnection, udp_client::UdpTpuConnection,
},
lazy_static::lazy_static,
solana_sdk::{transaction::VersionedTransaction, transport::TransportError},
std::{
collections::{hash_map::Entry, BTreeMap, HashMap},
net::{SocketAddr, UdpSocket},
@@ -13,9 +14,15 @@ use {
// Should be non-zero
static MAX_CONNECTIONS: usize = 64;
#[derive(Clone)]
enum Connection {
Udp(Arc<UdpTpuConnection>),
Quic(Arc<QuicTpuConnection>),
}
struct ConnMap {
// Keeps track of the connection associated with an addr and the last time it was used
map: HashMap<SocketAddr, (Arc<dyn TpuConnection + 'static + Sync + Send>, u64)>,
map: HashMap<SocketAddr, (Connection, u64)>,
// Helps to find the least recently used connection. The search and inserts are O(log(n))
// but since we're bounding the size of the collections, this should be constant
// (and hopefully negligible) time. In theory, we can do this in constant time
@@ -55,7 +62,7 @@ pub fn set_use_quic(use_quic: bool) {
#[allow(dead_code)]
// TODO: see https://github.com/solana-labs/solana/issues/23661
// remove lazy_static and optimize and refactor this
pub fn get_connection(addr: &SocketAddr) -> Arc<dyn TpuConnection + 'static + Sync + Send> {
fn get_connection(addr: &SocketAddr) -> Connection {
let mut map = (*CONNECTION_MAP).lock().unwrap();
let ticks = map.ticks;
let use_quic = map.use_quic;
@@ -71,10 +78,10 @@ pub fn get_connection(addr: &SocketAddr) -> Arc<dyn TpuConnection + 'static + Sy
// TODO: see https://github.com/solana-labs/solana/issues/23659
// make it configurable (e.g. via the command line) whether to use UDP or Quic
let conn: Arc<dyn TpuConnection + 'static + Sync + Send> = if use_quic {
Arc::new(QuicTpuConnection::new(send_socket, *addr))
let conn = if use_quic {
Connection::Quic(Arc::new(QuicTpuConnection::new(send_socket, *addr)))
} else {
Arc::new(UdpTpuConnection::new(send_socket, *addr))
Connection::Udp(Arc::new(UdpTpuConnection::new(send_socket, *addr)))
};
entry.insert((conn.clone(), ticks));
@@ -101,13 +108,69 @@ pub fn get_connection(addr: &SocketAddr) -> Arc<dyn TpuConnection + 'static + Sy
conn
}
// TODO: see https://github.com/solana-labs/solana/issues/23851
// use enum_dispatch and get rid of this tedious code.
// The main blocker to using enum_dispatch right now is that
// the it doesn't work with static methods like TpuConnection::new
// which is used by thin_client. This will be eliminated soon
// once thin_client is moved to using this connection cache.
// Once that is done, we will migrate to using enum_dispatch
// This will be done in a followup to
// https://github.com/solana-labs/solana/pull/23817
pub fn send_wire_transaction_batch(
packets: &[&[u8]],
addr: &SocketAddr,
) -> Result<(), TransportError> {
let conn = get_connection(addr);
match conn {
Connection::Udp(conn) => conn.send_wire_transaction_batch(packets),
Connection::Quic(conn) => conn.send_wire_transaction_batch(packets),
}
}
pub fn send_wire_transaction(
wire_transaction: &[u8],
addr: &SocketAddr,
) -> Result<(), TransportError> {
let conn = get_connection(addr);
match conn {
Connection::Udp(conn) => conn.send_wire_transaction(wire_transaction),
Connection::Quic(conn) => conn.send_wire_transaction(wire_transaction),
}
}
pub fn serialize_and_send_transaction(
transaction: &VersionedTransaction,
addr: &SocketAddr,
) -> Result<(), TransportError> {
let conn = get_connection(addr);
match conn {
Connection::Udp(conn) => conn.serialize_and_send_transaction(transaction),
Connection::Quic(conn) => conn.serialize_and_send_transaction(transaction),
}
}
pub fn par_serialize_and_send_transaction_batch(
transactions: &[VersionedTransaction],
addr: &SocketAddr,
) -> Result<(), TransportError> {
let conn = get_connection(addr);
match conn {
Connection::Udp(conn) => conn.par_serialize_and_send_transaction_batch(transactions),
Connection::Quic(conn) => conn.par_serialize_and_send_transaction_batch(transactions),
}
}
#[cfg(test)]
mod tests {
use {
crate::connection_cache::{get_connection, CONNECTION_MAP, MAX_CONNECTIONS},
crate::{
connection_cache::{get_connection, Connection, CONNECTION_MAP, MAX_CONNECTIONS},
tpu_connection::TpuConnection,
},
rand::{Rng, SeedableRng},
rand_chacha::ChaChaRng,
std::net::SocketAddr,
std::net::{IpAddr, SocketAddr},
};
fn get_addr(rng: &mut ChaChaRng) -> SocketAddr {
@@ -121,6 +184,13 @@ mod tests {
addr_str.parse().expect("Invalid address")
}
fn ip(conn: Connection) -> IpAddr {
match conn {
Connection::Udp(conn) => conn.tpu_addr().ip(),
Connection::Quic(conn) => conn.tpu_addr().ip(),
}
}
#[test]
fn test_connection_cache() {
// Allow the test to run deterministically
@@ -136,7 +206,7 @@ mod tests {
// be lazy and not connect until first use or handle connection errors somehow
// (without crashing, as would be required in a real practical validator)
let first_addr = get_addr(&mut rng);
assert!(get_connection(&first_addr).tpu_addr().ip() == first_addr.ip());
assert!(ip(get_connection(&first_addr)) == first_addr.ip());
let addrs = (0..MAX_CONNECTIONS)
.into_iter()
.map(|_| {
@@ -149,7 +219,7 @@ mod tests {
let map = (*CONNECTION_MAP).lock().unwrap();
addrs.iter().for_each(|a| {
let conn = map.map.get(a).expect("Address not found");
assert!(a.ip() == conn.0.tpu_addr().ip());
assert!(a.ip() == ip(conn.0.clone()));
});
assert!(map.map.get(&first_addr).is_none());

View File

@@ -63,19 +63,22 @@ impl TpuConnection for QuicTpuConnection {
&self.client.addr
}
fn send_wire_transaction(&self, wire_transaction: &[u8]) -> TransportResult<()> {
fn send_wire_transaction<T>(&self, wire_transaction: T) -> TransportResult<()>
where
T: AsRef<[u8]>,
{
let _guard = self.client.runtime.enter();
let send_buffer = self.client.send_buffer(wire_transaction);
self.client.runtime.block_on(send_buffer)?;
Ok(())
}
fn send_wire_transaction_batch(
&self,
wire_transaction_batch: &[Vec<u8>],
) -> TransportResult<()> {
fn send_wire_transaction_batch<T>(&self, buffers: &[T]) -> TransportResult<()>
where
T: AsRef<[u8]>,
{
let _guard = self.client.runtime.enter();
let send_batch = self.client.send_batch(wire_transaction_batch);
let send_batch = self.client.send_batch(buffers);
self.client.runtime.block_on(send_batch)?;
Ok(())
}
@@ -158,12 +161,18 @@ impl QuicClient {
}
}
pub async fn send_buffer(&self, data: &[u8]) -> Result<(), ClientErrorKind> {
self._send_buffer(data).await?;
pub async fn send_buffer<T>(&self, data: T) -> Result<(), ClientErrorKind>
where
T: AsRef<[u8]>,
{
self._send_buffer(data.as_ref()).await?;
Ok(())
}
pub async fn send_batch(&self, buffers: &[Vec<u8>]) -> Result<(), ClientErrorKind> {
pub async fn send_batch<T>(&self, buffers: &[T]) -> Result<(), ClientErrorKind>
where
T: AsRef<[u8]>,
{
// Start off by "testing" the connection by sending the first transaction
// This will also connect to the server if not already connected
// and reconnect and retry if the first send attempt failed
@@ -178,7 +187,7 @@ impl QuicClient {
if buffers.is_empty() {
return Ok(());
}
let connection = self._send_buffer(&buffers[0]).await?;
let connection = self._send_buffer(buffers[0].as_ref()).await?;
// Used to avoid dereferencing the Arc multiple times below
// by just getting a reference to the NewConnection once
@@ -192,7 +201,7 @@ impl QuicClient {
join_all(
buffs
.into_iter()
.map(|buf| Self::_send_buffer_using_conn(buf, connection_ref)),
.map(|buf| Self::_send_buffer_using_conn(buf.as_ref(), connection_ref)),
)
});

View File

@@ -613,7 +613,7 @@ impl<C: 'static + TpuConnection> AsyncClient for ThinClient<C> {
fn async_send_batch(&self, transactions: Vec<Transaction>) -> TransportResult<()> {
let batch: Vec<VersionedTransaction> = transactions.into_iter().map(Into::into).collect();
self.tpu_connection()
.par_serialize_and_send_transaction_batch(&batch)?;
.par_serialize_and_send_transaction_batch(&batch[..])?;
Ok(())
}

View File

@@ -1,13 +1,11 @@
use {
rayon::iter::{IntoParallelRefIterator, ParallelIterator},
rayon::iter::{IntoParallelIterator, ParallelIterator},
solana_sdk::{transaction::VersionedTransaction, transport::Result as TransportResult},
std::net::{SocketAddr, UdpSocket},
};
pub trait TpuConnection {
fn new(client_socket: UdpSocket, tpu_addr: SocketAddr) -> Self
where
Self: Sized;
fn new(client_socket: UdpSocket, tpu_addr: SocketAddr) -> Self;
fn tpu_addr(&self) -> &SocketAddr;
@@ -20,26 +18,23 @@ pub trait TpuConnection {
self.send_wire_transaction(&wire_transaction)
}
fn send_wire_transaction(&self, wire_transaction: &[u8]) -> TransportResult<()>;
fn send_wire_transaction<T>(&self, wire_transaction: T) -> TransportResult<()>
where
T: AsRef<[u8]>;
fn par_serialize_and_send_transaction_batch(
&self,
transaction_batch: &[VersionedTransaction],
transactions: &[VersionedTransaction],
) -> TransportResult<()> {
let wire_transaction_batch: Vec<_> = transaction_batch
.par_iter()
let buffers = transactions
.into_par_iter()
.map(|tx| bincode::serialize(&tx).expect("serialize Transaction in send_batch"))
.collect();
self.send_wire_transaction_batch(&wire_transaction_batch)
.collect::<Vec<_>>();
self.send_wire_transaction_batch(&buffers)
}
fn send_wire_transaction_batch(
&self,
wire_transaction_batch: &[Vec<u8>],
) -> TransportResult<()> {
for wire_transaction in wire_transaction_batch {
self.send_wire_transaction(wire_transaction)?;
}
Ok(())
}
fn send_wire_transaction_batch<T>(&self, buffers: &[T]) -> TransportResult<()>
where
T: AsRef<[u8]>;
}

View File

@@ -3,7 +3,9 @@
use {
crate::tpu_connection::TpuConnection,
core::iter::repeat,
solana_sdk::transport::Result as TransportResult,
solana_streamer::sendmmsg::batch_send,
std::net::{SocketAddr, UdpSocket},
};
@@ -24,8 +26,20 @@ impl TpuConnection for UdpTpuConnection {
&self.addr
}
fn send_wire_transaction(&self, wire_transaction: &[u8]) -> TransportResult<()> {
self.socket.send_to(wire_transaction, self.addr)?;
fn send_wire_transaction<T>(&self, wire_transaction: T) -> TransportResult<()>
where
T: AsRef<[u8]>,
{
self.socket.send_to(wire_transaction.as_ref(), self.addr)?;
Ok(())
}
fn send_wire_transaction_batch<T>(&self, buffers: &[T]) -> TransportResult<()>
where
T: AsRef<[u8]>,
{
let pkts: Vec<_> = buffers.iter().zip(repeat(self.tpu_addr())).collect();
batch_send(&self.socket, &pkts)?;
Ok(())
}
}