From 52bb2b6f39c04da03528983a0475b4f79b21251b Mon Sep 17 00:00:00 2001
From: Satya Vusirikala <satyasatya123456@gmail.com>
Date: Wed, 12 Jun 2024 09:04:48 -0700
Subject: [PATCH] Update counters

---
 mempool/src/core_mempool/mempool.rs           | 52 +++++++++++++------
 mempool/src/core_mempool/transaction_store.rs | 10 ++--
 2 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/mempool/src/core_mempool/mempool.rs b/mempool/src/core_mempool/mempool.rs
index 850bdc054328e..cee33f209db89 100644
--- a/mempool/src/core_mempool/mempool.rs
+++ b/mempool/src/core_mempool/mempool.rs
@@ -315,9 +315,13 @@ impl Mempool {
         counters::MEMPOOL_TOTAL_NUM_TXNS.observe(self.transactions.total_num_transactions() as f64);
         counters::MEMPOOL_REQUESTED_TXNS_IN_GET_BATCH.observe(max_txns as f64);
         counters::MEMPOOL_REQUESTED_BYTES_IN_GET_BATCH.observe(max_bytes as f64);
-        let mempool_total_txns_excluding_progressing = self
-            .transactions
-            .total_num_transactions_excluding(&exclude_transactions);
+        let mempool_total_txns_excluding_progressing =
+            self.transactions.total_num_transactions_excluding(
+                &exclude_transactions
+                    .keys()
+                    .cloned()
+                    .collect::<Vec<TransactionSummary>>(),
+            );
         counters::MEMPOOL_TXNS_EXCLUDING_PROGRESING
             .observe(mempool_total_txns_excluding_progressing as f64);
         let start_time = Instant::now();
@@ -532,22 +536,38 @@ impl Mempool {
             },
         );
 
-        let mut actual_remaining_txns = 0;
-        for txn in self.transactions.iter_queue() {
-            let txn_ptr = TxnPointer::from(txn);
-            if exclude_transactions.contains_key(&txn_ptr) {
-                continue;
-            }
-            if block.iter().any(|t| {
-                t.sender() == txn_ptr.sender && t.sequence_number() == txn_ptr.sequence_number
-            }) {
-                continue;
-            }
-            actual_remaining_txns += 1;
+        // let mut actual_remaining_txns = 0;
+        // for txn in self.transactions.iter_queue() {
+        //     let txn_ptr = TxnPointer::from(txn);
+        //     if exclude_transactions.contains_key(&txn_ptr) {
+        //         continue;
+        //     }
+        //     if block.iter().any(|t| {
+        //         t.sender() == txn_ptr.sender && t.sequence_number() == txn_ptr.sequence_number
+        //     }) {
+        //         continue;
+        //     }
+        //     actual_remaining_txns += 1;
+        // }
+
+        let mut total_exclude_transactions: Vec<TransactionSummary> = exclude_transactions
+            .keys()
+            .cloned()
+            .collect::<Vec<TransactionSummary>>();
+        for txn in block.iter() {
+            let txn_ptr = TransactionSummary {
+                sender: txn.sender(),
+                hash: txn.committed_hash(),
+                sequence_number: txn.sequence_number(),
+            };
+            total_exclude_transactions.push(txn_ptr);
         }
+        let actual_remaining_txns = self
+            .transactions
+            .total_num_transactions_excluding(&total_exclude_transactions);
         counters::MEMPOOL_ACTUAL_REMAINING_TXNS.observe(actual_remaining_txns as f64);
         counters::MEMPOOL_ACTUAL_REMAINING_TXNS_SAME_AS_SKIPPED.observe(
-            if actual_remaining_txns == skipped.len() {
+            if actual_remaining_txns == (skipped.len() as u64) {
                 1.0
             } else {
                 0.0
diff --git a/mempool/src/core_mempool/transaction_store.rs b/mempool/src/core_mempool/transaction_store.rs
index 66d8b9d2567aa..acecf87c6bf3b 100644
--- a/mempool/src/core_mempool/transaction_store.rs
+++ b/mempool/src/core_mempool/transaction_store.rs
@@ -18,7 +18,7 @@ use crate::{
     shared_mempool::types::MultiBucketTimelineIndexIds,
 };
 use aptos_config::config::MempoolConfig;
-use aptos_consensus_types::common::{TransactionInProgress, TransactionSummary};
+use aptos_consensus_types::common::TransactionSummary;
 use aptos_crypto::HashValue;
 use aptos_logger::{prelude::*, Level};
 use aptos_types::{
@@ -28,7 +28,7 @@ use aptos_types::{
 };
 use std::{
     cmp::max,
-    collections::{BTreeMap, HashMap, HashSet},
+    collections::{HashMap, HashSet},
     mem::size_of,
     ops::Bound,
     time::{Duration, Instant, SystemTime},
@@ -201,12 +201,8 @@ impl TransactionStore {
 
     pub(crate) fn total_num_transactions_excluding(
         &self,
-        excluded_transactions: &BTreeMap<TransactionSummary, TransactionInProgress>,
+        excluded_transactions: &[TransactionSummary],
     ) -> u64 {
-        let excluded_transactions = excluded_transactions
-            .iter()
-            .map(|(txn_summary, _)| txn_summary)
-            .collect::<HashSet<_>>();
         self.transactions
             .values()
             .map(|txns| {