Make transaction pool metrics per-shard (#9185)

As there is actually 1 transaction pool per shard and their metrics mix together.

See #9174 (comment) for some details.

Author: aborg-dev

chain/chunks/src/client.rs

@@ -44,7 +44,6 @@ pub struct ShardedTransactionPool {
 
 impl ShardedTransactionPool {
     pub fn new(rng_seed: RngSeed, pool_size_limit: Option<u64>) -> Self {
-        TransactionPool::init_metrics();
         Self { tx_pools: HashMap::new(), rng_seed, pool_size_limit }
     }
 
@@ -80,7 +79,11 @@ impl ShardedTransactionPool {
 
     fn pool_for_shard(&mut self, shard_id: ShardId) -> &mut TransactionPool {
         self.tx_pools.entry(shard_id).or_insert_with(|| {
-            TransactionPool::new(Self::random_seed(&self.rng_seed, shard_id), self.pool_size_limit)
+            TransactionPool::new(
+                Self::random_seed(&self.rng_seed, shard_id),
+                self.pool_size_limit,
+                &shard_id.to_string(),
+            )
         })
     }
 

chain/pool/src/lib.rs

@@ -4,6 +4,7 @@ use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
 use crate::types::{PoolIterator, PoolKey, TransactionGroup};
 use borsh::BorshSerialize;
 use near_crypto::PublicKey;
+use near_o11y::metrics::prometheus::core::{AtomicI64, GenericGauge};
 use near_primitives::epoch_manager::RngSeed;
 use near_primitives::hash::{hash, CryptoHash};
 use near_primitives::transaction::SignedTransaction;
@@ -39,26 +40,37 @@ pub struct TransactionPool {
     total_transaction_size_limit: Option<u64>,
     /// Total size of transactions in the pool measured in bytes.
     total_transaction_size: u64,
+    /// Metrics tracked for transaction pool.
+    transaction_pool_count_metric: GenericGauge<AtomicI64>,
+    transaction_pool_size_metric: GenericGauge<AtomicI64>,
 }
 
 impl TransactionPool {
-    pub fn new(key_seed: RngSeed, total_transaction_size_limit: Option<u64>) -> Self {
+    pub fn new(
+        key_seed: RngSeed,
+        total_transaction_size_limit: Option<u64>,
+        metrics_label: &str,
+    ) -> Self {
+        let transaction_pool_count_metric =
+            metrics::TRANSACTION_POOL_COUNT.with_label_values(&[metrics_label]);
+        let transaction_pool_size_metric =
+            metrics::TRANSACTION_POOL_SIZE.with_label_values(&[metrics_label]);
+        // A `get()` call initializes a metric even if its value is zero.
+        transaction_pool_count_metric.get();
+        transaction_pool_size_metric.get();
+
         Self {
             key_seed,
             transactions: BTreeMap::new(),
             unique_transactions: HashSet::new(),
             last_used_key: CryptoHash::default(),
             total_transaction_size_limit,
             total_transaction_size: 0,
+            transaction_pool_count_metric,
+            transaction_pool_size_metric,
         }
     }
 
-    pub fn init_metrics() {
-        // A `get()` call initializes a metric even if its value is zero.
-        metrics::TRANSACTION_POOL_COUNT.get();
-        metrics::TRANSACTION_POOL_SIZE.get();
-    }
-
     fn key(&self, account_id: &AccountId, public_key: &PublicKey) -> PoolKey {
         let mut v = public_key.try_to_vec().unwrap();
         v.extend_from_slice(&self.key_seed);
@@ -98,8 +110,8 @@ impl TransactionPool {
             .or_insert_with(Vec::new)
             .push(signed_transaction);
 
-        metrics::TRANSACTION_POOL_COUNT.inc();
-        metrics::TRANSACTION_POOL_SIZE.set(self.total_transaction_size as i64);
+        self.transaction_pool_count_metric.inc();
+        self.transaction_pool_size_metric.set(self.total_transaction_size as i64);
         InsertTransactionResult::Success
     }
 
@@ -150,8 +162,8 @@ impl TransactionPool {
         }
 
         // We can update metrics only once for the whole batch of transactions.
-        metrics::TRANSACTION_POOL_COUNT.set(self.unique_transactions.len() as i64);
-        metrics::TRANSACTION_POOL_SIZE.set(self.total_transaction_size as i64);
+        self.transaction_pool_count_metric.set(self.unique_transactions.len() as i64);
+        self.transaction_pool_size_metric.set(self.total_transaction_size as i64);
     }
 
     /// Returns the number of unique transactions in the pool.
@@ -231,7 +243,7 @@ impl<'a> PoolIterator for PoolIteratorWrapper<'a> {
                 if sorted_group.transactions.is_empty() {
                     for hash in sorted_group.removed_transaction_hashes {
                         if self.pool.unique_transactions.remove(&hash) {
-                            metrics::TRANSACTION_POOL_COUNT.dec();
+                            self.pool.transaction_pool_count_metric.dec();
                         }
                     }
                 } else {
@@ -266,8 +278,8 @@ impl<'a> Drop for PoolIteratorWrapper<'a> {
             }
         }
         // We can update metrics only once for the whole batch of transactions.
-        metrics::TRANSACTION_POOL_COUNT.set(self.pool.unique_transactions.len() as i64);
-        metrics::TRANSACTION_POOL_SIZE.set(self.pool.transaction_size() as i64);
+        self.pool.transaction_pool_count_metric.set(self.pool.unique_transactions.len() as i64);
+        self.pool.transaction_pool_size_metric.set(self.pool.transaction_size() as i64);
     }
 }
 
@@ -313,7 +325,7 @@ mod tests {
         mut transactions: Vec<SignedTransaction>,
         expected_weight: u32,
     ) -> (Vec<u64>, TransactionPool) {
-        let mut pool = TransactionPool::new(TEST_SEED, None);
+        let mut pool = TransactionPool::new(TEST_SEED, None, "");
         let mut rng = thread_rng();
         transactions.shuffle(&mut rng);
         for tx in transactions {
@@ -424,7 +436,7 @@ mod tests {
             })
             .collect::<Vec<_>>();
 
-        let mut pool = TransactionPool::new(TEST_SEED, None);
+        let mut pool = TransactionPool::new(TEST_SEED, None, "");
         let mut rng = thread_rng();
         transactions.shuffle(&mut rng);
         for tx in transactions.clone() {
@@ -536,7 +548,7 @@ mod tests {
 
     #[test]
     fn test_transaction_pool_size() {
-        let mut pool = TransactionPool::new(TEST_SEED, None);
+        let mut pool = TransactionPool::new(TEST_SEED, None, "");
         let transactions = generate_transactions("alice.near", "alice.near", 1, 100);
         let mut total_transaction_size = 0;
         // Adding transactions increases the size.
@@ -560,7 +572,7 @@ mod tests {
         // Each transaction is at least 1 byte in size, so the last transaction will not fit.
         let pool_size_limit =
             transactions.iter().map(|tx| tx.get_size()).sum::<u64>().checked_sub(1).unwrap();
-        let mut pool = TransactionPool::new(TEST_SEED, Some(pool_size_limit));
+        let mut pool = TransactionPool::new(TEST_SEED, Some(pool_size_limit), "");
         for (i, tx) in transactions.iter().cloned().enumerate() {
             if i + 1 < transactions.len() {
                 assert_eq!(pool.insert_transaction(tx), InsertTransactionResult::Success);

chain/pool/src/metrics.rs

@@ -1,18 +1,20 @@
-use near_o11y::metrics::IntGauge;
+use near_o11y::metrics::IntGaugeVec;
 use once_cell::sync::Lazy;
 
-pub static TRANSACTION_POOL_COUNT: Lazy<IntGauge> = Lazy::new(|| {
-    near_o11y::metrics::try_create_int_gauge(
+pub static TRANSACTION_POOL_COUNT: Lazy<IntGaugeVec> = Lazy::new(|| {
+    near_o11y::metrics::try_create_int_gauge_vec(
         "near_transaction_pool_entries",
-        "Total number of transactions currently in the pools tracked by the node",
+        "Total number of transactions currently tracked by the node in a given shard pool",
+        &["shard_id"],
     )
     .unwrap()
 });
 
-pub static TRANSACTION_POOL_SIZE: Lazy<IntGauge> = Lazy::new(|| {
-    near_o11y::metrics::try_create_int_gauge(
+pub static TRANSACTION_POOL_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
+    near_o11y::metrics::try_create_int_gauge_vec(
         "near_transaction_pool_size",
-        "Total size in bytes of transactions currently in the pools tracked by the node",
+        "Total size in bytes of transactions currently tracked by the node in a given shard pool",
+        &["shard_id"],
     )
     .unwrap()
 });