From 8a55087d8edd51e4fa44ffc7318521ed7a73ba75 Mon Sep 17 00:00:00 2001
From: Tao Zhu <82401714+tao-stones@users.noreply.github.com>
Date: Wed, 8 May 2024 18:22:11 -0600
Subject: [PATCH] Use ahash in cost tracker (#1252)

* add bench for cost-tracker

* use `ahash` in cost_tracker
---
 Cargo.lock                         |  2 +
 cost-model/Cargo.toml              |  5 ++
 cost-model/benches/cost_tracker.rs | 81 ++++++++++++++++++++++++++++++
 cost-model/src/cost_tracker.rs     |  7 ++-
 programs/sbf/Cargo.lock            |  1 +
 5 files changed, 94 insertions(+), 2 deletions(-)
 create mode 100644 cost-model/benches/cost_tracker.rs

diff --git a/Cargo.lock b/Cargo.lock
index 4b5f146a683f5c..c7621915c1b8c2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -6041,6 +6041,8 @@ dependencies = [
 name = "solana-cost-model"
 version = "2.0.0"
 dependencies = [
+ "ahash 0.8.10",
+ "itertools",
  "lazy_static",
  "log",
  "rustc_version 0.4.0",
diff --git a/cost-model/Cargo.toml b/cost-model/Cargo.toml
index 4a8b159bbf4cb6..cf13a3377376ed 100644
--- a/cost-model/Cargo.toml
+++ b/cost-model/Cargo.toml
@@ -10,6 +10,7 @@ license = { workspace = true }
 edition = { workspace = true }
 
 [dependencies]
+ahash = { workspace = true }
 lazy_static = { workspace = true }
 log = { workspace = true }
 solana-address-lookup-table-program = { workspace = true }
@@ -31,6 +32,7 @@ crate-type = ["lib"]
 name = "solana_cost_model"
 
 [dev-dependencies]
+itertools = { workspace = true }
 solana-logger = { workspace = true }
 solana-sdk = { workspace = true, features = ["dev-context-only-utils"] }
 static_assertions = { workspace = true }
@@ -41,3 +43,6 @@ targets = ["x86_64-unknown-linux-gnu"]
 
 [build-dependencies]
 rustc_version = { workspace = true }
+
+[[bench]]
+name = "cost_tracker"
diff --git a/cost-model/benches/cost_tracker.rs b/cost-model/benches/cost_tracker.rs
new file mode 100644
index 00000000000000..ff1f453643b655
--- /dev/null
+++ b/cost-model/benches/cost_tracker.rs
@@ -0,0 +1,81 @@
+#![feature(test)]
+extern crate test;
+use {
+    itertools::Itertools,
+    solana_cost_model::{
+        cost_tracker::CostTracker,
+        transaction_cost::{TransactionCost, UsageCostDetails},
+    },
+    solana_sdk::pubkey::Pubkey,
+    test::Bencher,
+};
+
+struct BenchSetup {
+    cost_tracker: CostTracker,
+    tx_costs: Vec<TransactionCost>,
+}
+
+fn setup(num_transactions: usize, contentious_transactions: bool) -> BenchSetup {
+    let mut cost_tracker = CostTracker::default();
+    // set cost_tracker with max limits to stretch testing
+    cost_tracker.set_limits(u64::MAX, u64::MAX, u64::MAX);
+
+    let max_accounts_per_tx = 128;
+    let pubkey = Pubkey::new_unique();
+    let tx_costs = (0..num_transactions)
+        .map(|_| {
+            let mut usage_cost_details = UsageCostDetails::default();
+            (0..max_accounts_per_tx).for_each(|_| {
+                let writable_account_key = if contentious_transactions {
+                    pubkey
+                } else {
+                    Pubkey::new_unique()
+                };
+                usage_cost_details
+                    .writable_accounts
+                    .push(writable_account_key)
+            });
+            usage_cost_details.programs_execution_cost = 9999;
+            TransactionCost::Transaction(usage_cost_details)
+        })
+        .collect_vec();
+
+    BenchSetup {
+        cost_tracker,
+        tx_costs,
+    }
+}
+
+#[bench]
+fn bench_cost_tracker_non_contentious_transaction(bencher: &mut Bencher) {
+    let BenchSetup {
+        mut cost_tracker,
+        tx_costs,
+    } = setup(1024, false);
+
+    bencher.iter(|| {
+        for tx_cost in tx_costs.iter() {
+            if cost_tracker.try_add(tx_cost).is_err() {
+                break;
+            } // stop when hit limits
+            cost_tracker.update_execution_cost(tx_cost, 0); // update execution cost down to zero
+        }
+    });
+}
+
+#[bench]
+fn bench_cost_tracker_contentious_transaction(bencher: &mut Bencher) {
+    let BenchSetup {
+        mut cost_tracker,
+        tx_costs,
+    } = setup(1024, true);
+
+    bencher.iter(|| {
+        for tx_cost in tx_costs.iter() {
+            if cost_tracker.try_add(tx_cost).is_err() {
+                break;
+            } // stop when hit limits
+            cost_tracker.update_execution_cost(tx_cost, 0); // update execution cost down to zero
+        }
+    });
+}
diff --git a/cost-model/src/cost_tracker.rs b/cost-model/src/cost_tracker.rs
index af28a2052dd48d..a6341ed8e276ee 100644
--- a/cost-model/src/cost_tracker.rs
+++ b/cost-model/src/cost_tracker.rs
@@ -62,7 +62,7 @@ pub struct CostTracker {
     account_cost_limit: u64,
     block_cost_limit: u64,
     vote_cost_limit: u64,
-    cost_by_writable_accounts: HashMap<Pubkey, u64>,
+    cost_by_writable_accounts: HashMap<Pubkey, u64, ahash::RandomState>,
     block_cost: u64,
     vote_cost: u64,
     transaction_count: u64,
@@ -88,7 +88,10 @@ impl Default for CostTracker {
             account_cost_limit: MAX_WRITABLE_ACCOUNT_UNITS,
             block_cost_limit: MAX_BLOCK_UNITS,
             vote_cost_limit: MAX_VOTE_UNITS,
-            cost_by_writable_accounts: HashMap::with_capacity(WRITABLE_ACCOUNTS_PER_BLOCK),
+            cost_by_writable_accounts: HashMap::with_capacity_and_hasher(
+                WRITABLE_ACCOUNTS_PER_BLOCK,
+                ahash::RandomState::new(),
+            ),
             block_cost: 0,
             vote_cost: 0,
             transaction_count: 0,
diff --git a/programs/sbf/Cargo.lock b/programs/sbf/Cargo.lock
index e68400f38109a7..a81821ebb36b2a 100644
--- a/programs/sbf/Cargo.lock
+++ b/programs/sbf/Cargo.lock
@@ -5013,6 +5013,7 @@ dependencies = [
 name = "solana-cost-model"
 version = "2.0.0"
 dependencies = [
+ "ahash 0.8.10",
  "lazy_static",
  "log",
  "rustc_version",