Repository: mobilecoinfoundation/mobilecoin
Pull request #4034: add balance metrics to fog-test-client
Author: login: eranrund
Created: 2025-05-19T18:51:07Z | Merged: 2025-05-19T20:12:03Z

https://sentz.atlassian.net/browse/SENTZ-3824 requested that we add metrics for wallet balances.
I tested this locally and the metrics look correct:
```
test_client_balance{client_index="0",token_id="0"} 8320000000911
test_client_balance{client_index="1",token_id="0"} 6918000000062
test_client_balance_block_count{client_index="0"} 3837232
test_client_balance_block_count{client_index="1"} 3837232
```


Diff:
diff --git a/Cargo.lock b/Cargo.lock
index 0330dc5283..6a2572573f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -620,9 +620,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
 
 [[package]]
 name = "cc"
-version = "1.2.11"
+version = "1.2.23"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e4730490333d58093109dc02c23174c3f4d490998c3fed3cc8e82d57afedb9cf"
+checksum = "5f4ac86a9e5bc1e2b3449ab9d7d3a6a405e3d1bb28d7b9be8614f55846ae3766"
 dependencies = [
  "shlex",
 ]
@@ -792,9 +792,9 @@ dependencies = [
 
 [[package]]
 name = "cmake"
-version = "0.1.49"
+version = "0.1.54"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "db34956e100b30725f2eb215f90d4871051239535632f84fea3bc92722c66b7c"
+checksum = "e7caa3f9de89ddbe2c607f4101924c5abec803763ae9534e4f4d7d8f84aa81f0"
 dependencies = [
  "cc",
 ]
@@ -4648,6 +4648,7 @@ dependencies = [
  "mc-util-uri",
  "more-asserts",
  "once_cell",
+ "prometheus",
  "rand",
  "retry",
  "serde",
diff --git a/fog/test-client/Cargo.toml b/fog/test-client/Cargo.toml
index cda0d934b4..f1b66b5ed8 100644
--- a/fog/test-client/Cargo.toml
+++ b/fog/test-client/Cargo.toml
@@ -41,6 +41,7 @@ lazy_static = "1.4"
 maplit = "1.0"
 more-asserts = "0.3"
 once_cell = "1.19"
+prometheus = "0.13"
 rand = "0.8"
 retry = "2.0"
 serde = "1"
diff --git a/fog/test-client/src/counters.rs b/fog/test-client/src/counters.rs
index 657a01ace7..8de10f223b 100644
--- a/fog/test-client/src/counters.rs
+++ b/fog/test-client/src/counters.rs
@@ -2,7 +2,12 @@
 
 //! Prometheus metrics, interesting when the test runs continuously
 
-use mc_util_metrics::{register_histogram, Histogram, IntCounter, IntGauge, OpMetrics};
+use mc_transaction_core::TokenId;
+use mc_util_metrics::{
+    register_histogram, Collector, Desc, Histogram, IntCounter, IntGauge, IntGaugeVec,
+    MetricFamily, OpMetrics, Opts,
+};
+use std::collections::HashMap;
 
 // Histogram buckets used for reporting the TX_CONFIRMED_TIME and
 // TX_RECEIVED_TIME to prometheus
@@ -29,6 +34,9 @@ lazy_static::lazy_static! {
     /// Counter group
     pub static ref OP_COUNTERS: OpMetrics = OpMetrics::new_and_registered("fog_test_client");
 
+    /// Metrics for each individual client.
+    pub static ref CLIENT_METRICS: ClientMetrics = ClientMetrics::new_and_registered();
+
     /// Time in seconds that it takes for the source client to build a transaction (including fog interactions)
     pub static ref TX_BUILD_TIME: Histogram =
         register_histogram!("fog_test_client_tx_build_time", "Time for source client to build a transaction, including fog interactions", TX_TIME_BUCKETS_SHORT.to_vec()).unwrap();
@@ -108,3 +116,79 @@ lazy_static::lazy_static! {
     /// This is updated after every transfer attempt.
     pub static ref LAST_POLLING_SUCCESSFUL: IntGauge = OP_COUNTERS.gauge("last_polling_successful");
 }
+
+/// Metrics for each individual client.
+/// Clients are identified by their index in the config.
+#[derive(Clone)]
+pub struct ClientMetrics {
+    balance_gauges: IntGaugeVec,
+    balance_block_count_gauges: IntGaugeVec,
+}
+
+impl Default for ClientMetrics {
+    fn default() -> Self {
+        ClientMetrics {
+            balance_gauges: IntGaugeVec::new(
+                Opts::new(
+                    "test_client_balance",
+                    "Balance of each individual test client",
+                ),
+                &["client_index", "token_id"],
+            )
+            .unwrap(),
+            balance_block_count_gauges: IntGaugeVec::new(
+                Opts::new(
+                    "test_client_balance_block_count",
+                    "Block count at which the balance was last updated",
+                ),
+                &["client_index"],
+            )
+            .unwrap(),
+        }
+    }
+}
+
+impl ClientMetrics {
+    /// Create a new ClientMetrics instance and register it with the Prometheus
+    /// crate.
+    pub fn new_and_registered() -> Self {
+        let client_metrics = Self::default();
+        prometheus::register(Box::new(client_metrics.clone()))
+            .expect("ClientMetrics registration on Prometheus failed.");
+        client_metrics
+    }
+
+    /// Update the balance and block count for a given client.
+    pub fn update_balance(
+        &self,
+        client_index: usize,
+        balances: &HashMap<TokenId, u64>,
+        block_count: impl Into<u64>,
+    ) {
+        let block_count = block_count.into();
+        for (token_id, balance) in balances {
+            self.balance_gauges
+                .with_label_values(&[&client_index.to_string(), &token_id.to_string()])
+                .set(*balance as i64);
+        }
+        self.balance_block_count_gauges
+            .with_label_values(&[&client_index.to_string()])
+            .set(block_count as i64);
+    }
+}
+
+impl Collector for ClientMetrics {
+    fn desc(&self) -> Vec<&Desc> {
+        let mut metrics = Vec::with_capacity(2);
+        metrics.extend(self.balance_gauges.desc());
+        metrics.extend(self.balance_block_count_gauges.desc());
+        metrics
+    }
+
+    fn collect(&self) -> Vec<MetricFamily> {
+        let mut metrics = Vec::with_capacity(2);
+        metrics.extend(self.balance_gauges.collect());
+        metrics.extend(self.balance_block_count_gauges.collect());
+        metrics
+    }
+}
diff --git a/fog/test-client/src/test_client.rs b/fog/test-client/src/test_client.rs
index 8bf7dcfbf7..1dafe42ff3 100644
--- a/fog/test-client/src/test_client.rs
+++ b/fog/test-client/src/test_client.rs
@@ -5,7 +5,10 @@
 //! - Fog canary (sends transactions in prod to alert if it fails, and collect
 //!   timings)
 
-use crate::{counters, error::TestClientError};
+use crate::{
+    counters::{self, CLIENT_METRICS},
+    error::TestClientError,
+};
 
 use hex_fmt::HexList;
 use maplit::hashmap;
@@ -296,7 +299,9 @@ impl TestClient {
     fn transfer(
         &self,
         source_client: &mut Client,
+        source_client_index: usize,
         target_client: &Client,
+        _target_client_index: usize,
         token_id: TokenId,
     ) -> Result<TransferData, TestClientError> {
         self.tx_info.clear();
@@ -310,11 +315,12 @@ impl TestClient {
 
         // First do a balance check to flush out any spent txos
         let tracer = tracer!();
-        tracer.in_span("pre_transfer_balance_check", |_cx| {
+        let (balances, block_count) = tracer.in_span("pre_transfer_balance_check", |_cx| {
             source_client
                 .check_balance()
                 .map_err(TestClientError::CheckBalance)
         })?;
+        CLIENT_METRICS.update_balance(source_client_index, &balances, block_count);
 
         let mut rng = McRng;
         assert!(target_address.fog_report_url().is_some());
@@ -429,6 +435,8 @@ impl TestClient {
     /// balance.
     ///
     /// Arguments:
+    /// * client: The client to use for this check
+    /// * client_index: The index of the client in the list of clients
     /// * block_index: The block_index containing new transactions that must be
     ///   in the balance
     /// * expected_balance: The expected balance to compute after this
@@ -436,6 +444,7 @@ impl TestClient {
     fn ensure_expected_balance_after_block(
         &self,
         client: &mut Client,
+        client_index: usize,
         block_index: BlockIndex,
         expected_balances: HashMap<TokenId, u64>,
     ) -> Result<(), TestClientError> {
@@ -446,6 +455,7 @@ impl TestClient {
             let (new_balances, new_block_count) = client
                 .check_balance()
                 .map_err(TestClientError::CheckBalance)?;
+            CLIENT_METRICS.update_balance(client_index, &new_balances, new_block_count);
 
             // Wait for client cursor to include the index where the transaction landed.
             if u64::from(new_block_count) > block_index {
@@ -567,6 +577,7 @@ impl TestClient {
                 let (src_balances, src_cursor) = source_client_lk
                     .check_balance()
                     .map_err(TestClientError::CheckBalance)?;
+                CLIENT_METRICS.update_balance(source_client_index, &src_balances, src_cursor);
                 let src_balance = src_balances.get(&token_id).cloned().unwrap_or_default();
 
                 log::info!(
@@ -580,8 +591,8 @@ impl TestClient {
                 let (tgt_balances, tgt_cursor) = target_client_lk
                     .check_balance()
                     .map_err(TestClientError::CheckBalance)?;
+                CLIENT_METRICS.update_balance(target_client_index, &tgt_balances, tgt_cursor);
                 let tgt_balance = tgt_balances.get(&token_id).cloned().unwrap_or_default();
-
                 log::info!(
                     self.logger,
                     "client {} has a TokenId({}) balance of {} after {} blocks",
@@ -599,7 +610,13 @@ impl TestClient {
         )?;
 
         let transfer_start = std::time::SystemTime::now();
-        let transfer_data = self.transfer(&mut source_client_lk, &target_client_lk, token_id)?;
+        let transfer_data = self.transfer(
+            &mut source_client_lk,
+            source_client_index,
+            &target_client_lk,
+            target_client_index,
+            token_id,
+        )?;
 
         let mut span = block_span_builder(&tracer, "test_iteration", transfer_data.block_count)
             .with_start_time(transfer_start)
@@ -626,6 +643,7 @@ impl TestClient {
 
         let mut receive_tx_worker = ReceiveTxWorker::new(
             target_client,
+            target_client_index,
             hashmap! { token_id => tgt_balance },
             hashmap! { token_id => tgt_balance + self.policy.transfer_amount },
             self.policy.clone(),
@@ -652,6 +670,7 @@ impl TestClient {
         tracer.in_span("ensure_expected_balance_after_block", |_cx| {
             self.ensure_expected_balance_after_block(
                 &mut source_client_lk,
+                source_client_index,
                 transaction_appeared,
                 hashmap! { token_id => src_balance - self.policy.transfer_amount - transfer_data.fee.value },
             )
@@ -749,7 +768,9 @@ impl TestClient {
     fn atomic_swap(
         &self,
         source_client: &mut Client,
+        source_client_index: usize,
         target_client: &mut Client,
+        target_client_index: usize,
         token_id1: TokenId,
         token_id2: TokenId,
         is_partial_fill: bool,
@@ -780,16 +801,19 @@ impl TestClient {
 
         // First do a balance check to flush out any spent txos
         let tracer = tracer!();
-        tracer.in_span("pre_swap_src_balance_check", |_cx| {
+        let (src_balances, src_cursor) = tracer.in_span("pre_swap_src_balance_check", |_cx| {
             source_client
                 .check_balance()
                 .map_err(TestClientError::CheckBalance)
         })?;
-        tracer.in_span("pre_swap_dst_balance_check", |_cx| {
+        CLIENT_METRICS.update_balance(source_client_index, &src_balances, src_cursor);
+
+        let (tgt_balances, tgt_cursor) = tracer.in_span("pre_swap_dst_balance_check", |_cx| {
             target_client
                 .check_balance()
                 .map_err(TestClientError::CheckBalance)
         })?;
+        CLIENT_METRICS.update_balance(target_client_index, &tgt_balances, tgt_cursor);
 
         assert!(target_address.fog_report_url().is_some());
 
@@ -874,9 +898,9 @@ impl TestClient {
         token_id2: TokenId,
         is_partial_fill: bool,
         source_client: Arc<Mutex<Client>>,
-        _source_client_index: usize,
+        source_client_index: usize,
         target_client: Arc<Mutex<Client>>,
-        _target_client_index: usize,
+        target_client_index: usize,
     ) -> Result<(), TestClientError> {
         self.tx_info.clear();
         let tracer = tracer!();
@@ -887,13 +911,15 @@ impl TestClient {
         let (src_balances, tgt_balances) = tracer.in_span(
             "test_atomic_swap_pre_checks",
             |_cx| -> Result<(_, _), TestClientError> {
-                let (src_balances, _src_cursor) = source_client_lk
+                let (src_balances, src_cursor) = source_client_lk
                     .check_balance()
                     .map_err(TestClientError::CheckBalance)?;
+                CLIENT_METRICS.update_balance(source_client_index, &src_balances, src_cursor);
 
-                let (tgt_balances, _tgt_cursor) = target_client_lk
+                let (tgt_balances, tgt_cursor) = target_client_lk
                     .check_balance()
                     .map_err(TestClientError::CheckBalance)?;
+                CLIENT_METRICS.update_balance(target_client_index, &tgt_balances, tgt_cursor);
 
                 Ok((src_balances, tgt_balances))
             },
@@ -902,7 +928,9 @@ impl TestClient {
         let transfer_start = std::time::SystemTime::now();
         let transfer_data = self.atomic_swap(
             &mut source_client_lk,
+            source_client_index,
             &mut target_client_lk,
+            target_client_index,
             token_id1,
             token_id2,
             is_partial_fill,
@@ -926,6 +954,7 @@ impl TestClient {
         drop(target_client_lk);
         let mut receive_tx_worker = ReceiveTxWorker::new(
             target_client,
+            target_client_index,
             tgt_balances,
             expected_tgt_balances,
             self.policy.clone(),
@@ -959,6 +988,7 @@ impl TestClient {
         tracer.in_span("ensure_expected_balance_after_block", |_cx| {
             self.ensure_expected_balance_after_block(
                 &mut source_client_lk,
+                source_client_index,
                 transaction_appeared,
                 expected_src_balance,
             )
@@ -1221,6 +1251,7 @@ impl ReceiveTxWorker {
     /// * logger
     pub fn new(
         client: Arc<Mutex<Client>>,
+        client_index: usize,
         current_balances: HashMap<TokenId, u64>,
         expected_balances: HashMap<TokenId, u64>,
         policy: TestClientPolicy,
@@ -1260,6 +1291,7 @@ impl ReceiveTxWorker {
                     let (new_balances, new_block_count) = client
                         .check_balance()
                         .map_err(TestClientError::CheckBalance)?;
+                    CLIENT_METRICS.update_balance(client_index, &new_balances, new_block_count);
 
                     if balance_match(&expected_balances, &new_balances) {
                         counters::TX_RECEIVED_TIME.observe(start.elapsed().as_secs_f64());


Reviews:
user: login: nick-mobilecoin
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-19T19:41:29Z
user: login: eranrund
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T19:42:48Z

Review comments:
diff_hunk: @@ -108,3 +116,76 @@ lazy_static::lazy_static! {
     /// This is updated after every transfer attempt.
     pub static ref LAST_POLLING_SUCCESSFUL: IntGauge = OP_COUNTERS.gauge("last_polling_successful");
 }
+
+/// Metrics for each individual client.
+/// Clients are identified by their index in the config.
+#[derive(Clone)]
+pub struct ClientMetrics {
+    balance_gauges: IntGaugeVec,
+    balance_block_count_gauges: IntGaugeVec,
+}
+
+impl Default for ClientMetrics {
+    fn default() -> Self {
+        ClientMetrics {
+            balance_gauges: IntGaugeVec::new(
+                Opts::new(
+                    "test_client_balance",
+                    "Balance of each individual test client",
+                ),
+                &["client_index", "token_id"],
+            )
+            .unwrap(),
+            balance_block_count_gauges: IntGaugeVec::new(
+                Opts::new(
+                    "test_client_balance_block_count",
+                    "Block count at which the balance was last updated",
+                ),
+                &["client_index"],
+            )
+            .unwrap(),
+        }
+    }
+}
+
+impl ClientMetrics {
+    /// Create a new ClientMetrics instance and register it with the Prometheus
+    /// crate.
+    pub fn new_and_registered() -> Self {
+        let client_metrics = Self::default();
+        prometheus::register(Box::new(client_metrics.clone()))
+            .expect("ClientMetrics registration on Prometheus failed.");
+        client_metrics
+    }
+
+    /// Update the balance and block count for a given client.
+    pub fn update_balance(
+        &self,
+        client_index: usize,
+        balances: &HashMap<TokenId, u64>,
+        block_count: impl Into<u64>,
+    ) {
+        let block_count = block_count.into();
+        for (token_id, balance) in balances {
+            self.balance_gauges
+                .with_label_values(&[&client_index.to_string(), &token_id.to_string()])
+                .set(*balance as i64);
+        }
+        self.balance_block_count_gauges
+            .with_label_values(&[&client_index.to_string()])
+            .set(block_count as i64);
+    }
+}
+
+impl Collector for ClientMetrics {
+    fn desc(&self) -> Vec<&Desc> {
+        self.balance_gauges.desc()
+    }
path: fog/test-client/src/counters.rs
user: login: nick-mobilecoin
body: Missing the bock count
```suggestion
    fn desc(&self) -> Vec<&Desc> {
        let mut metrics = Vec::with_capacity(2);
        metrics.extend(self.balance_gauges.desc());
        metrics.extend(self.balance_block_count_gauges.desc());
        metrics
    }
```
created_at: 2025-05-19T19:36:17Z
author_association: COLLABORATOR
diff_hunk: @@ -108,3 +116,76 @@ lazy_static::lazy_static! {
     /// This is updated after every transfer attempt.
     pub static ref LAST_POLLING_SUCCESSFUL: IntGauge = OP_COUNTERS.gauge("last_polling_successful");
 }
+
+/// Metrics for each individual client.
+/// Clients are identified by their index in the config.
+#[derive(Clone)]
+pub struct ClientMetrics {
+    balance_gauges: IntGaugeVec,
+    balance_block_count_gauges: IntGaugeVec,
+}
+
+impl Default for ClientMetrics {
+    fn default() -> Self {
+        ClientMetrics {
+            balance_gauges: IntGaugeVec::new(
+                Opts::new(
+                    "test_client_balance",
+                    "Balance of each individual test client",
+                ),
+                &["client_index", "token_id"],
+            )
+            .unwrap(),
+            balance_block_count_gauges: IntGaugeVec::new(
+                Opts::new(
+                    "test_client_balance_block_count",
+                    "Block count at which the balance was last updated",
+                ),
+                &["client_index"],
+            )
+            .unwrap(),
+        }
+    }
+}
+
+impl ClientMetrics {
+    /// Create a new ClientMetrics instance and register it with the Prometheus
+    /// crate.
+    pub fn new_and_registered() -> Self {
+        let client_metrics = Self::default();
+        prometheus::register(Box::new(client_metrics.clone()))
+            .expect("ClientMetrics registration on Prometheus failed.");
+        client_metrics
+    }
+
+    /// Update the balance and block count for a given client.
+    pub fn update_balance(
+        &self,
+        client_index: usize,
+        balances: &HashMap<TokenId, u64>,
+        block_count: impl Into<u64>,
+    ) {
+        let block_count = block_count.into();
+        for (token_id, balance) in balances {
+            self.balance_gauges
+                .with_label_values(&[&client_index.to_string(), &token_id.to_string()])
+                .set(*balance as i64);
+        }
+        self.balance_block_count_gauges
+            .with_label_values(&[&client_index.to_string()])
+            .set(block_count as i64);
+    }
+}
+
+impl Collector for ClientMetrics {
+    fn desc(&self) -> Vec<&Desc> {
+        self.balance_gauges.desc()
+    }
path: fog/test-client/src/counters.rs
user: login: eranrund
body: Good catch!
created_at: 2025-05-19T19:42:48Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-05-19T18:55:06Z
body: ⚠️ Downstream repo mobilecoinofficial/full-service failed to build. Check actions status for details.
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-05-19T19:48:57Z
body: ⚠️ Downstream repo mobilecoinofficial/full-service failed to build. Check actions status for details.
author_association: NONE