From b7fddfa70de2adb6d4b43acde49f6db22fc2350a Mon Sep 17 00:00:00 2001
From: Anastasia Lubennikova <anastasia@neon.tech>
Date: Tue, 7 Mar 2023 18:47:06 +0200
Subject: [PATCH] Add branch_id field to proxy_io_bytes_per_client metric.
 Since we allow switching endpoints between different branches, it is
 important to use composite key. Otherwise, we may try to calculate delta
 between metric values for two different branches.

---
 proxy/src/metrics.rs | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/proxy/src/metrics.rs b/proxy/src/metrics.rs
index 3b28346872..be22c45836 100644
--- a/proxy/src/metrics.rs
+++ b/proxy/src/metrics.rs
@@ -21,6 +21,7 @@ const PROXY_IO_BYTES_PER_CLIENT: &str = "proxy_io_bytes_per_client";
 #[derive(Eq, Hash, PartialEq, Serialize, Debug)]
 pub struct Ids {
     pub endpoint_id: String,
+    pub branch_id: String,
 }
 
 pub async fn task_main(config: &MetricCollectionConfig) -> anyhow::Result<()> {
@@ -74,12 +75,23 @@ fn gather_proxy_io_bytes_per_client() -> Vec<(Ids, (u64, DateTime<Utc>))> {
                         .find(|l| l.get_name() == "endpoint_id")
                         .unwrap()
                         .get_value();
+                    let branch_id = ms
+                        .get_label()
+                        .iter()
+                        .find(|l| l.get_name() == "branch_id")
+                        .unwrap()
+                        .get_value();
+
                     let value = ms.get_counter().get_value() as u64;
 
-                    debug!("endpoint_id:val - {}: {}", endpoint_id, value);
+                    debug!(
+                        "branch_id {} endpoint_id {} val: {}",
+                        branch_id, endpoint_id, value
+                    );
                     current_metrics.push((
                         Ids {
                             endpoint_id: endpoint_id.to_string(),
+                            branch_id: "".to_string(),
                         },
                         (value, Utc::now()),
                     ));
@@ -131,6 +143,7 @@ async fn collect_metrics_iteration(
                 value,
                 extra: Ids {
                     endpoint_id: curr_key.endpoint_id.clone(),
+                    branch_id: curr_key.branch_id.clone(),
                 },
             })
         })
@@ -172,6 +185,7 @@ async fn collect_metrics_iteration(
                 cached_metrics
                     .entry(Ids {
                         endpoint_id: send_metric.extra.endpoint_id.clone(),
+                        branch_id: send_metric.extra.branch_id.clone(),
                     })
                     // update cached value (add delta) and time
                     .and_modify(|e| {