From b41e33e066dc18ff0a1309a9ccf7357d3a2a7b0a Mon Sep 17 00:00:00 2001
From: Alfredo Garcia <oxarbitrage@gmail.com>
Date: Fri, 14 Aug 2020 19:50:26 -0300
Subject: [PATCH] Bytes read and bytes written metrics (#901)

* add bytes read and written metrics

* Apply suggestions from code review

Co-authored-by: Jane Lusby <jlusby42@gmail.com>

* store address as string

* Apply suggestions from code review

Co-authored-by: Henry de Valence <hdevalence@hdevalence.ca>

* change addr to label

Co-authored-by: Henry de Valence <hdevalence@hdevalence.ca>

* remove newline

Co-authored-by: Jane Lusby <jlusby42@gmail.com>
Co-authored-by: Henry de Valence <hdevalence@hdevalence.ca>
---
 zebra-network/src/peer/handshake.rs          |  9 +++++++--
 zebra-network/src/protocol/external/codec.rs | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

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<String>,
 }
 
 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);