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::{canonical_socket_addr, types::Version},
19};
20
21/// A client service wrapper that keeps track of its load.
22///
23/// It also keeps track of the peer's reported protocol version.
24#[derive(Debug)]
25pub struct LoadTrackedClient {
26    /// A service representing a connected peer, wrapped in a load tracker.
27    service: PeakEwma<Client>,
28
29    /// The metadata for the connected peer `service`.
30    connection_info: Arc<ConnectionInfo>,
31}
32
33/// Create a new [`LoadTrackedClient`] wrapping the provided `client` service.
34impl From<Client> for LoadTrackedClient {
35    fn from(client: Client) -> Self {
36        let connection_info = client.connection_info.clone();
37
38        let service = PeakEwma::new(
39            client,
40            EWMA_DEFAULT_RTT,
41            EWMA_DECAY_TIME_NANOS,
42            tower::load::CompleteOnResponse::default(),
43        );
44
45        LoadTrackedClient {
46            service,
47            connection_info,
48        }
49    }
50}
51
52impl LoadTrackedClient {
53    /// Retrieve the peer's reported protocol version.
54    pub fn remote_version(&self) -> Version {
55        self.connection_info.remote.version
56    }
57
58    /// Returns true if this peer connected directly to us from `ip`.
59    pub fn is_inbound_direct_from_ip(&self, ip: &IpAddr) -> bool {
60        let expected_ip = canonical_socket_addr(SocketAddr::new(*ip, 0)).ip();
61
62        matches!(
63            self.connection_info.connected_addr,
64            ConnectedAddr::InboundDirect { addr }
65                if canonical_socket_addr(addr.remove_socket_addr_privacy()).ip() == expected_ip
66        )
67    }
68}
69
70impl<Request> Service<Request> for LoadTrackedClient
71where
72    Client: Service<Request>,
73{
74    type Response = <Client as Service<Request>>::Response;
75    type Error = <Client as Service<Request>>::Error;
76    type Future = <PeakEwma<Client> as Service<Request>>::Future;
77
78    fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
79        self.service.poll_ready(context)
80    }
81
82    fn call(&mut self, request: Request) -> Self::Future {
83        self.service.call(request)
84    }
85}
86
87impl Load for LoadTrackedClient {
88    type Metric = <PeakEwma<Client> as Load>::Metric;
89
90    fn load(&self) -> Self::Metric {
91        self.service.load()
92    }
93}
zebra_network/peer/load_tracked_client.rs

zebra_network/peer/
load_tracked_client.rs
