use crate::{
constants,
counters::{self},
logging::*,
peer::{Peer, PeerNotification, PeerRequest},
transport::{
Connection, ConnectionId, ConnectionMetadata, TSocket as TransportTSocket,
TRANSPORT_TIMEOUT,
},
ProtocolId,
};
use aptos_config::network_id::NetworkContext;
use aptos_logger::prelude::*;
use aptos_rate_limiter::rate_limit::TokenBucketRateLimiter;
use aptos_time_service::{TimeService, TimeServiceTrait};
use aptos_types::{network_address::NetworkAddress, PeerId};
use channel::{self, aptos_channel, message_queues::QueueStyle};
use futures::{
channel::oneshot,
io::{AsyncRead, AsyncWrite, AsyncWriteExt},
sink::SinkExt,
stream::StreamExt,
};
use netcore::transport::{ConnectionOrigin, Transport};
use short_hex_str::AsShortHexStr;
use std::{
collections::{hash_map::Entry, HashMap},
marker::PhantomData,
net::{IpAddr, Ipv4Addr},
sync::Arc,
time::Duration,
};
use tokio::runtime::Handle;
pub mod builder;
pub mod conn_notifs_channel;
mod error;
mod senders;
#[cfg(test)]
mod tests;
mod transport;
mod types;
pub use self::error::PeerManagerError;
use crate::{
application::storage::PeerMetadataStorage,
peer_manager::transport::{TransportHandler, TransportRequest},
protocols::network::SerializedRequest,
};
use aptos_config::config::{PeerRole, PeerSet};
use aptos_infallible::RwLock;
pub use senders::*;
pub use types::*;
pub type IpAddrTokenBucketLimiter = TokenBucketRateLimiter<IpAddr>;
pub struct PeerManager<TTransport, TSocket>
where
TTransport: Transport,
TSocket: AsyncRead + AsyncWrite,
{
network_context: NetworkContext,
executor: Handle,
time_service: TimeService,
listen_addr: NetworkAddress,
transport_handler: Option<TransportHandler<TTransport, TSocket>>,
active_peers: HashMap<
PeerId,
(
ConnectionMetadata,
aptos_channel::Sender<ProtocolId, PeerRequest>,
),
>,
peer_metadata_storage: Arc<PeerMetadataStorage>,
trusted_peers: Arc<RwLock<PeerSet>>,
requests_rx: aptos_channel::Receiver<(PeerId, ProtocolId), PeerManagerRequest>,
upstream_handlers:
HashMap<ProtocolId, aptos_channel::Sender<(PeerId, ProtocolId), PeerManagerNotification>>,
connection_event_handlers: Vec<conn_notifs_channel::Sender>,
transport_reqs_tx: channel::Sender<TransportRequest>,
transport_notifs_tx: channel::Sender<TransportNotification<TSocket>>,
connection_reqs_rx: aptos_channel::Receiver<PeerId, ConnectionRequest>,
transport_notifs_rx: channel::Receiver<TransportNotification<TSocket>>,
outstanding_disconnect_requests:
HashMap<ConnectionId, oneshot::Sender<Result<(), PeerManagerError>>>,
phantom_transport: PhantomData<TTransport>,
max_concurrent_network_reqs: usize,
channel_size: usize,
max_frame_size: usize,
inbound_connection_limit: usize,
inbound_rate_limiters: IpAddrTokenBucketLimiter,
outbound_rate_limiters: IpAddrTokenBucketLimiter,
}
impl<TTransport, TSocket> PeerManager<TTransport, TSocket>
where
TTransport: Transport<Output = Connection<TSocket>> + Send + 'static,
TSocket: TransportTSocket,
{
#[allow(clippy::too_many_arguments)]
pub fn new(
executor: Handle,
time_service: TimeService,
transport: TTransport,
network_context: NetworkContext,
listen_addr: NetworkAddress,
peer_metadata_storage: Arc<PeerMetadataStorage>,
trusted_peers: Arc<RwLock<PeerSet>>,
requests_rx: aptos_channel::Receiver<(PeerId, ProtocolId), PeerManagerRequest>,
connection_reqs_rx: aptos_channel::Receiver<PeerId, ConnectionRequest>,
upstream_handlers: HashMap<
ProtocolId,
aptos_channel::Sender<(PeerId, ProtocolId), PeerManagerNotification>,
>,
connection_event_handlers: Vec<conn_notifs_channel::Sender>,
channel_size: usize,
max_concurrent_network_reqs: usize,
max_frame_size: usize,
inbound_connection_limit: usize,
inbound_rate_limiters: IpAddrTokenBucketLimiter,
outbound_rate_limiters: IpAddrTokenBucketLimiter,
) -> Self {
let (transport_notifs_tx, transport_notifs_rx) = channel::new(
channel_size,
&counters::PENDING_CONNECTION_HANDLER_NOTIFICATIONS,
);
let (transport_reqs_tx, transport_reqs_rx) =
channel::new(channel_size, &counters::PENDING_PEER_MANAGER_DIAL_REQUESTS);
let transport_notifs_tx_clone = transport_notifs_tx.clone();
let _guard = executor.enter();
let (transport_handler, listen_addr) = TransportHandler::new(
network_context,
time_service.clone(),
transport,
listen_addr,
transport_reqs_rx,
transport_notifs_tx_clone,
);
Self {
network_context,
executor,
time_service,
listen_addr,
transport_handler: Some(transport_handler),
active_peers: HashMap::new(),
peer_metadata_storage,
trusted_peers,
requests_rx,
connection_reqs_rx,
transport_reqs_tx,
transport_notifs_tx,
transport_notifs_rx,
outstanding_disconnect_requests: HashMap::new(),
phantom_transport: PhantomData,
upstream_handlers,
connection_event_handlers,
max_concurrent_network_reqs,
channel_size,
max_frame_size,
inbound_connection_limit,
inbound_rate_limiters,
outbound_rate_limiters,
}
}
pub fn update_connected_peers_metrics(&self) {
let total = self.active_peers.len();
let inbound = self
.active_peers
.iter()
.filter(|(_, (metadata, _))| metadata.origin == ConnectionOrigin::Inbound)
.count();
let outbound = total.saturating_sub(inbound);
counters::connections(&self.network_context, ConnectionOrigin::Inbound).set(inbound as i64);
counters::connections(&self.network_context, ConnectionOrigin::Outbound)
.set(outbound as i64);
}
fn sample_connected_peers(&self) {
sample!(SampleRate::Duration(Duration::from_secs(60)), {
let peers: Vec<_> = self
.active_peers
.values()
.map(|(connection, _)| {
(
connection.remote_peer_id,
connection.addr.clone(),
connection.origin,
)
})
.collect();
info!(
NetworkSchema::new(&self.network_context),
peers = ?peers,
"Current connected peers"
)
});
}
pub fn listen_addr(&self) -> &NetworkAddress {
&self.listen_addr
}
pub async fn start(mut self) {
info!(
NetworkSchema::new(&self.network_context),
"Start listening for incoming connections on {}", self.listen_addr
);
self.start_connection_listener();
loop {
::futures::select! {
connection_event = self.transport_notifs_rx.select_next_some() => {
self.handle_connection_event(connection_event);
}
connection_request = self.connection_reqs_rx.select_next_some() => {
self.handle_outbound_connection_request(connection_request).await;
}
request = self.requests_rx.select_next_some() => {
self.handle_outbound_request(request).await;
}
complete => {
break;
}
}
}
warn!(
NetworkSchema::new(&self.network_context),
"PeerManager actor terminated"
);
}
fn handle_connection_event(&mut self, event: TransportNotification<TSocket>) {
trace!(
NetworkSchema::new(&self.network_context),
transport_notification = format!("{:?}", event),
"{} TransportNotification::{:?}",
self.network_context,
event
);
self.sample_connected_peers();
match event {
TransportNotification::NewConnection(mut conn) => {
match conn.metadata.origin {
ConnectionOrigin::Outbound => {
conn.metadata.role = self
.trusted_peers
.read()
.get(&conn.metadata.remote_peer_id)
.map_or(PeerRole::Unknown, |auth_context| auth_context.role);
if conn.metadata.role == PeerRole::Unknown {
warn!(
NetworkSchema::new(&self.network_context)
.connection_metadata_with_address(&conn.metadata),
"{} Outbound connection made with unknown peer role: {}",
self.network_context,
conn.metadata
)
}
}
ConnectionOrigin::Inbound => {
if conn.metadata.role == PeerRole::Unknown {
let unknown_inbound_conns = self
.active_peers
.iter()
.filter(|(peer_id, (metadata, _))| {
metadata.origin == ConnectionOrigin::Inbound
&& self
.trusted_peers
.read()
.get(peer_id)
.map_or(true, |peer| peer.role == PeerRole::Unknown)
})
.count();
if !self
.active_peers
.contains_key(&conn.metadata.remote_peer_id)
&& unknown_inbound_conns + 1 > self.inbound_connection_limit
{
info!(
NetworkSchema::new(&self.network_context)
.connection_metadata_with_address(&conn.metadata),
"{} Connection rejected due to connection limit: {}",
self.network_context,
conn.metadata
);
counters::connections_rejected(
&self.network_context,
conn.metadata.origin,
)
.inc();
self.disconnect(conn);
return;
}
}
}
}
info!(
NetworkSchema::new(&self.network_context)
.connection_metadata_with_address(&conn.metadata),
"{} New connection established: {}", self.network_context, conn.metadata
);
self.add_peer(conn);
self.update_connected_peers_metrics();
}
TransportNotification::Disconnected(lost_conn_metadata, reason) => {
info!(
NetworkSchema::new(&self.network_context)
.connection_metadata_with_address(&lost_conn_metadata),
disconnection_reason = reason,
"{} Connection {} closed due to {}",
self.network_context,
lost_conn_metadata,
reason
);
let peer_id = lost_conn_metadata.remote_peer_id;
if let Entry::Occupied(entry) = self.active_peers.entry(peer_id) {
let (conn_metadata, _) = entry.get();
if conn_metadata.connection_id == lost_conn_metadata.connection_id {
entry.remove();
self.peer_metadata_storage.remove_connection(
self.network_context.network_id(),
&lost_conn_metadata,
)
}
}
self.update_connected_peers_metrics();
if let Some(oneshot_tx) = self
.outstanding_disconnect_requests
.remove(&lost_conn_metadata.connection_id)
{
if let Err(send_err) = oneshot_tx.send(Ok(())) {
info!(
NetworkSchema::new(&self.network_context),
error = ?send_err,
"{} Failed to notify upstream client of closed connection for peer {}: {:?}",
self.network_context,
peer_id,
send_err
);
}
}
let ip_addr = lost_conn_metadata
.addr
.find_ip_addr()
.unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED));
if !self.active_peers.contains_key(&peer_id) {
let notif = ConnectionNotification::LostPeer(
lost_conn_metadata,
self.network_context,
reason,
);
self.send_conn_notification(peer_id, notif);
}
self.inbound_rate_limiters.try_garbage_collect_key(&ip_addr);
self.outbound_rate_limiters
.try_garbage_collect_key(&ip_addr);
}
}
}
async fn handle_outbound_connection_request(&mut self, request: ConnectionRequest) {
trace!(
NetworkSchema::new(&self.network_context),
peer_manager_request = request,
"{} PeerManagerRequest::{:?}",
self.network_context,
request
);
self.sample_connected_peers();
match request {
ConnectionRequest::DialPeer(requested_peer_id, addr, response_tx) => {
if let Some((curr_connection, _)) = self.active_peers.get(&requested_peer_id) {
let error = PeerManagerError::AlreadyConnected(curr_connection.addr.clone());
debug!(
NetworkSchema::new(&self.network_context)
.connection_metadata_with_address(curr_connection),
"{} Already connected to Peer {} with connection {:?}. Not dialing address {}",
self.network_context,
requested_peer_id.short_str(),
curr_connection,
addr
);
if let Err(send_err) = response_tx.send(Err(error)) {
info!(
NetworkSchema::new(&self.network_context)
.remote_peer(&requested_peer_id),
"{} Failed to notify that peer is already connected for Peer {}: {:?}",
self.network_context,
requested_peer_id.short_str(),
send_err
);
}
} else {
let request = TransportRequest::DialPeer(requested_peer_id, addr, response_tx);
self.transport_reqs_tx.send(request).await.unwrap();
};
}
ConnectionRequest::DisconnectPeer(peer_id, resp_tx) => {
if let Some((conn_metadata, sender)) = self.active_peers.remove(&peer_id) {
let connection_id = conn_metadata.connection_id;
self.peer_metadata_storage
.remove_connection(self.network_context.network_id(), &conn_metadata);
drop(sender);
self.outstanding_disconnect_requests
.insert(connection_id, resp_tx);
} else {
info!(
NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
"{} Connection with peer: {} was already closed",
self.network_context,
peer_id.short_str(),
);
if let Err(err) = resp_tx.send(Err(PeerManagerError::NotConnected(peer_id))) {
info!(
NetworkSchema::new(&self.network_context),
error = ?err,
"{} Failed to notify that connection was already closed for Peer {}: {:?}",
self.network_context,
peer_id,
err
);
}
}
}
}
}
async fn handle_outbound_request(&mut self, request: PeerManagerRequest) {
trace!(
NetworkSchema::new(&self.network_context),
peer_manager_request = request,
"{} PeerManagerRequest::{:?}",
self.network_context,
request
);
self.sample_connected_peers();
let (peer_id, protocol_id, peer_request) = match request {
PeerManagerRequest::SendDirectSend(peer_id, msg) => {
(peer_id, msg.protocol_id(), PeerRequest::SendDirectSend(msg))
}
PeerManagerRequest::SendRpc(peer_id, req) => {
(peer_id, req.protocol_id(), PeerRequest::SendRpc(req))
}
};
if let Some((conn_metadata, sender)) = self.active_peers.get_mut(&peer_id) {
if let Err(err) = sender.push(protocol_id, peer_request) {
info!(
NetworkSchema::new(&self.network_context).connection_metadata(conn_metadata),
protocol_id = %protocol_id,
error = ?err,
"{} Failed to forward outbound message to downstream actor. Error: {:?}",
self.network_context, err
);
}
} else {
warn!(
NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
protocol_id = %protocol_id,
"{} Can't send message to peer. Peer {} is currently not connected",
self.network_context,
peer_id.short_str()
);
}
}
fn start_connection_listener(&mut self) {
let transport_handler = self
.transport_handler
.take()
.expect("Transport handler already taken");
self.executor.spawn(transport_handler.listen());
}
fn simultaneous_dial_tie_breaking(
own_peer_id: PeerId,
remote_peer_id: PeerId,
existing_origin: ConnectionOrigin,
new_origin: ConnectionOrigin,
) -> bool {
match (existing_origin, new_origin) {
(ConnectionOrigin::Inbound, ConnectionOrigin::Inbound) => true,
(ConnectionOrigin::Outbound, ConnectionOrigin::Outbound) => true,
(ConnectionOrigin::Inbound, ConnectionOrigin::Outbound) => remote_peer_id < own_peer_id,
(ConnectionOrigin::Outbound, ConnectionOrigin::Inbound) => own_peer_id < remote_peer_id,
}
}
fn disconnect(&mut self, connection: Connection<TSocket>) {
let network_context = self.network_context;
let time_service = self.time_service.clone();
let drop_fut = async move {
let mut connection = connection;
let peer_id = connection.metadata.remote_peer_id;
if let Err(e) = time_service
.timeout(TRANSPORT_TIMEOUT, connection.socket.close())
.await
{
error!(
NetworkSchema::new(&network_context)
.remote_peer(&peer_id),
error = %e,
"{} Closing connection with Peer {} failed with error: {}",
network_context,
peer_id.short_str(),
e
);
};
};
self.executor.spawn(drop_fut);
}
fn add_peer(&mut self, connection: Connection<TSocket>) {
let conn_meta = connection.metadata.clone();
let peer_id = conn_meta.remote_peer_id;
if self.network_context.peer_id() == peer_id {
debug_assert!(false, "Self dials shouldn't happen");
warn!(
NetworkSchema::new(&self.network_context)
.connection_metadata_with_address(&conn_meta),
"Received self-dial, disconnecting it"
);
self.disconnect(connection);
return;
}
let mut send_new_peer_notification = true;
if let Entry::Occupied(active_entry) = self.active_peers.entry(peer_id) {
let (curr_conn_metadata, _) = active_entry.get();
if Self::simultaneous_dial_tie_breaking(
self.network_context.peer_id(),
peer_id,
curr_conn_metadata.origin,
conn_meta.origin,
) {
let (_, peer_handle) = active_entry.remove();
drop(peer_handle);
info!(
NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
"{} Closing existing connection with Peer {} to mitigate simultaneous dial",
self.network_context,
peer_id.short_str()
);
send_new_peer_notification = false;
} else {
info!(
NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
"{} Closing incoming connection with Peer {} to mitigate simultaneous dial",
self.network_context,
peer_id.short_str()
);
self.disconnect(connection);
return;
}
}
let ip_addr = connection
.metadata
.addr
.find_ip_addr()
.unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED));
let inbound_rate_limiter = self.inbound_rate_limiters.bucket(ip_addr);
let outbound_rate_limiter = self.outbound_rate_limiters.bucket(ip_addr);
let (peer_reqs_tx, peer_reqs_rx) = aptos_channel::new(
QueueStyle::FIFO,
self.channel_size,
Some(&counters::PENDING_NETWORK_REQUESTS),
);
let (peer_notifs_tx, peer_notifs_rx) = aptos_channel::new(
QueueStyle::FIFO,
self.channel_size,
Some(&counters::PENDING_NETWORK_NOTIFICATIONS),
);
let peer = Peer::new(
self.network_context,
self.executor.clone(),
self.time_service.clone(),
connection,
self.transport_notifs_tx.clone(),
peer_reqs_rx,
peer_notifs_tx,
Duration::from_millis(constants::INBOUND_RPC_TIMEOUT_MS),
constants::MAX_CONCURRENT_INBOUND_RPCS,
constants::MAX_CONCURRENT_OUTBOUND_RPCS,
self.max_frame_size,
Some(inbound_rate_limiter),
Some(outbound_rate_limiter),
);
self.executor.spawn(peer.start());
self.spawn_peer_network_events_handler(peer_id, peer_notifs_rx);
self.active_peers
.insert(peer_id, (conn_meta.clone(), peer_reqs_tx));
self.peer_metadata_storage
.insert_connection(self.network_context.network_id(), conn_meta.clone());
if send_new_peer_notification {
let notif = ConnectionNotification::NewPeer(conn_meta, self.network_context);
self.send_conn_notification(peer_id, notif);
}
}
fn send_conn_notification(&mut self, peer_id: PeerId, notification: ConnectionNotification) {
for handler in self.connection_event_handlers.iter_mut() {
if let Err(e) = handler.push(peer_id, notification.clone()) {
warn!(
NetworkSchema::new(&self.network_context)
.remote_peer(&peer_id),
error = ?e,
connection_notification = notification,
"{} Failed to send notification {} to handler for peer: {}. Error: {:?}",
self.network_context,
notification,
peer_id.short_str(),
e
);
}
}
}
fn spawn_peer_network_events_handler(
&self,
peer_id: PeerId,
network_events: aptos_channel::Receiver<ProtocolId, PeerNotification>,
) {
let mut upstream_handlers = self.upstream_handlers.clone();
let network_context = self.network_context;
self.executor.spawn(network_events.for_each_concurrent(
self.max_concurrent_network_reqs,
move |inbound_event| {
handle_inbound_request(
network_context,
inbound_event,
peer_id,
&mut upstream_handlers,
);
futures::future::ready(())
},
));
}
}
fn handle_inbound_request(
network_context: NetworkContext,
inbound_event: PeerNotification,
peer_id: PeerId,
upstream_handlers: &mut HashMap<
ProtocolId,
aptos_channel::Sender<(PeerId, ProtocolId), PeerManagerNotification>,
>,
) {
let (protocol_id, notification) = match inbound_event {
PeerNotification::RecvMessage(msg) => (
msg.protocol_id(),
PeerManagerNotification::RecvMessage(peer_id, msg),
),
PeerNotification::RecvRpc(req) => (
req.protocol_id(),
PeerManagerNotification::RecvRpc(peer_id, req),
),
};
if let Some(handler) = upstream_handlers.get_mut(&protocol_id) {
if let Err(err) = handler.push((peer_id, protocol_id), notification) {
warn!(
NetworkSchema::new(&network_context),
error = ?err,
protocol_id = protocol_id,
"{} Upstream handler unable to handle message for protocol: {}. Error: {:?}",
network_context, protocol_id, err
);
}
} else {
debug!(
NetworkSchema::new(&network_context),
protocol_id = protocol_id,
message = format!("{:?}", notification),
"{} Received network message for unregistered protocol: {:?}",
network_context,
notification,
);
}
}