zebra_network/peer/
load_tracked_client.rs

1//! A peer connection service wrapper type to handle load tracking and provide access to the
2//! reported protocol version.
3
4use std::{
5    net::{IpAddr, SocketAddr},
6    sync::Arc,
7    task::{Context, Poll},
8};
9
10use tower::{
11    load::{Load, PeakEwma},
12    Service,
13};
14
15use crate::{
16    constants::{EWMA_DECAY_TIME_NANOS, EWMA_DEFAULT_RTT},
17    peer::{Client, ConnectedAddr, ConnectionInfo},
18    protocol::external::{
19        canonical_socket_addr,
20        types::{PeerServices, Version},
21    },
22};
23
24/// A client service wrapper that keeps track of its load.
25///
26/// It also keeps track of the peer's reported protocol version.
27#[derive(Debug)]
28pub struct LoadTrackedClient {
29    /// A service representing a connected peer, wrapped in a load tracker.
30    service: PeakEwma<Client>,
31
32    /// The metadata for the connected peer `service`.
33    connection_info: Arc<ConnectionInfo>,
34}
35
36/// Create a new [`LoadTrackedClient`] wrapping the provided `client` service.
37impl From<Client> for LoadTrackedClient {
38    fn from(client: Client) -> Self {
39        let connection_info = client.connection_info.clone();
40
41        let service = PeakEwma::new(
42            client,
43            EWMA_DEFAULT_RTT,
44            EWMA_DECAY_TIME_NANOS,
45            tower::load::CompleteOnResponse::default(),
46        );
47
48        LoadTrackedClient {
49            service,
50            connection_info,
51        }
52    }
53}
54
55impl LoadTrackedClient {
56    /// Retrieve the peer's reported protocol version.
57    pub fn remote_version(&self) -> Version {
58        self.connection_info.remote.version
59    }
60
61    /// Retrieve the services the peer advertised in its `version` message.
62    pub fn remote_services(&self) -> PeerServices {
63        self.connection_info.remote.services
64    }
65
66    /// Returns the shared metadata allocation identifying this connection.
67    pub(crate) fn connection_info(&self) -> &Arc<ConnectionInfo> {
68        &self.connection_info
69    }
70
71    /// Returns true if this peer connected directly to us from `ip`.
72    pub fn is_inbound_direct_from_ip(&self, ip: &IpAddr) -> bool {
73        let expected_ip = canonical_socket_addr(SocketAddr::new(*ip, 0)).ip();
74
75        matches!(
76            self.connection_info.connected_addr,
77            ConnectedAddr::InboundDirect { addr }
78                if canonical_socket_addr(addr.remove_socket_addr_privacy()).ip() == expected_ip
79        )
80    }
81}
82
83impl<Request> Service<Request> for LoadTrackedClient
84where
85    Client: Service<Request>,
86{
87    type Response = <Client as Service<Request>>::Response;
88    type Error = <Client as Service<Request>>::Error;
89    type Future = <PeakEwma<Client> as Service<Request>>::Future;
90
91    fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
92        self.service.poll_ready(context)
93    }
94
95    fn call(&mut self, request: Request) -> Self::Future {
96        self.service.call(request)
97    }
98}
99
100impl Load for LoadTrackedClient {
101    type Metric = <PeakEwma<Client> as Load>::Metric;
102
103    fn load(&self) -> Self::Metric {
104        self.service.load()
105    }
106}
zebra_network/peer/load_tracked_client.rs

zebra_network/peer/
load_tracked_client.rs
