From 25a0d254c124fdae834f7ba03a96e807fb94e549 Mon Sep 17 00:00:00 2001
From: Anastasios Kichidis <akihidis@gmail.com>
Date: Fri, 29 Jul 2022 17:38:15 +0100
Subject: [PATCH] refactor metrics

---
 network/src/metrics.rs | 16 ++++++++++------
 network/src/primary.rs |  4 ++--
 network/src/worker.rs  |  5 ++++-
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/network/src/metrics.rs b/network/src/metrics.rs
index 80bd8d478..5881e9341 100644
--- a/network/src/metrics.rs
+++ b/network/src/metrics.rs
@@ -1,6 +1,5 @@
 // Copyright (c) 2022, Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
-
 use prometheus::{default_registry, register_int_gauge_vec_with_registry, IntGaugeVec, Registry};
 use std::sync::Arc;
 
@@ -20,7 +19,7 @@ impl PrimaryNetworkMetrics {
             network_available_tasks: register_int_gauge_vec_with_registry!(
                 "primary_network_available_tasks",
                 "The number of available tasks to run in the network connector",
-                &["module", "network"],
+                &["module", "network", "address"],
                 registry
             )
             .unwrap(),
@@ -52,7 +51,7 @@ impl WorkerNetworkMetrics {
             network_available_tasks: register_int_gauge_vec_with_registry!(
                 "worker_network_concurrent_tasks",
                 "The number of available tasks to run in the network connector",
-                &["module", "network"],
+                &["module", "network", "address"],
                 registry
             )
             .unwrap(),
@@ -107,10 +106,14 @@ impl<N: NetworkMetrics> Metrics<N> {
         self.network_type.clone()
     }
 
-    pub fn set_network_available_tasks(&self, value: i64) {
+    pub fn set_network_available_tasks(&self, value: i64, addr: Option<String>) {
         self.metrics_handler
             .network_available_tasks()
-            .with_label_values(&[self.module_tag.as_str(), self.network_type.as_str()])
+            .with_label_values(&[
+                self.module_tag.as_str(),
+                self.network_type.as_str(),
+                addr.map_or("".to_string(), |a| a).as_str(),
+            ])
             .set(value);
     }
 }
@@ -132,12 +135,13 @@ mod test {
         };
 
         // WHEN update metrics
-        metrics.set_network_available_tasks(14);
+        metrics.set_network_available_tasks(14, Some("127.0.0.1".to_string()));
 
         // THEN registry should be updated with expected tag
         let mut m = HashMap::new();
         m.insert("module", "demo_handler");
         m.insert("network", "primary");
+        m.insert("address", "127.0.0.1");
         assert_eq!(
             metrics
                 .metrics_handler
diff --git a/network/src/primary.rs b/network/src/primary.rs
index c222c3b9e..81d7678b8 100644
--- a/network/src/primary.rs
+++ b/network/src/primary.rs
@@ -232,7 +232,7 @@ impl PrimaryNetwork {
             for (addr, executor) in &self.executors {
                 let available = executor.available_capacity();
 
-                m.set_network_available_tasks(available as i64);
+                m.set_network_available_tasks(available as i64, Some(addr.to_string()));
 
                 if available == 0 {
                     warn!("Executor in network:{} and module:{} available tasks is 0 for client address: {}", m.network_type(), m.module_tag(), addr);
@@ -336,7 +336,7 @@ impl PrimaryToWorkerNetwork {
 
     fn update_metrics(&self) {
         if let Some(m) = &self.metrics {
-            m.set_network_available_tasks(self.executor.available_capacity() as i64);
+            m.set_network_available_tasks(self.executor.available_capacity() as i64, None);
         }
     }
 }
diff --git a/network/src/worker.rs b/network/src/worker.rs
index 8ec1afe22..e5a3055f9 100644
--- a/network/src/worker.rs
+++ b/network/src/worker.rs
@@ -200,7 +200,10 @@ impl WorkerNetwork {
             for (addr, executor) in &self.executors {
                 let available = executor.available_capacity();
 
-                m.set_network_available_tasks(executor.available_capacity() as i64);
+                m.set_network_available_tasks(
+                    executor.available_capacity() as i64,
+                    Some(addr.to_string()),
+                );
 
                 if available == 0 {
                     warn!("Executor in network:{} and module:{} available tasks is 0 for client address: {}", m.network_type(), m.module_tag(), addr);