diff --git a/zebra-network/src/peer/handshake.rs b/zebra-network/src/peer/handshake.rs index 3761ddc6c..7dc467983 100644 --- a/zebra-network/src/peer/handshake.rs +++ b/zebra-network/src/peer/handshake.rs @@ -109,8 +109,13 @@ where let fut = async move { debug!("connecting to remote peer"); - let mut stream = - Framed::new(tcp_stream, Codec::builder().for_network(network).finish()); + let mut stream = Framed::new( + tcp_stream, + Codec::builder() + .for_network(network) + .with_metrics_label(addr.ip().to_string()) + .finish(), + ); let local_nonce = Nonce::default(); nonces diff --git a/zebra-network/src/protocol/external/codec.rs b/zebra-network/src/protocol/external/codec.rs index 65f10f876..7cee5c19f 100644 --- a/zebra-network/src/protocol/external/codec.rs +++ b/zebra-network/src/protocol/external/codec.rs @@ -45,6 +45,8 @@ pub struct Builder { version: Version, /// The maximum allowable message length. max_len: usize, + /// An optional label to use for reporting metrics. + metrics_label: Option, } impl Codec { @@ -54,6 +56,7 @@ impl Codec { network: Network::Mainnet, version: constants::CURRENT_VERSION, max_len: MAX_PROTOCOL_MESSAGE_LEN, + metrics_label: None, } } @@ -91,6 +94,12 @@ impl Builder { self.max_len = len; self } + + /// Configure the codec for the given peer address. + pub fn with_metrics_label(mut self, metrics_label: String) -> Self { + self.metrics_label = Some(metrics_label); + self + } } // ======== Encoding ========= @@ -113,6 +122,10 @@ impl Encoder for Codec { return Err(Parse("body length exceeded maximum size")); } + if let Some(label) = self.builder.metrics_label.clone() { + metrics::counter!("bytes.written", (body.len() + HEADER_LEN) as u64, "addr" => label); + } + use Message::*; // Note: because all match arms must have // the same type, and the array length is @@ -325,6 +338,10 @@ impl Decoder for Codec { return Err(Parse("body length exceeded maximum size")); } + if let Some(label) = self.builder.metrics_label.clone() { + metrics::counter!("bytes.read", (body_len + HEADER_LEN) as u64, "addr" => label); + } + // Reserve buffer space for the expected body and the following header. src.reserve(body_len + HEADER_LEN);