From b7028a1dc59e7f41b4d0cb2eed9dccbd89e44ad4 Mon Sep 17 00:00:00 2001
From: Brennan <brennan.watt@anza.xyz>
Date: Thu, 31 Oct 2024 08:03:25 -0700
Subject: [PATCH] more status cache benches (#3368)

---
 runtime/benches/status_cache.rs | 77 ++++++++++++++++++++++++++++++++-
 1 file changed, 76 insertions(+), 1 deletion(-)

diff --git a/runtime/benches/status_cache.rs b/runtime/benches/status_cache.rs
index 17d509b310d87a..f2d9f8997428b8 100644
--- a/runtime/benches/status_cache.rs
+++ b/runtime/benches/status_cache.rs
@@ -3,10 +3,12 @@ extern crate test;
 
 use {
     bincode::serialize,
+    rand::{rngs::SmallRng, Rng, SeedableRng},
+    solana_accounts_db::ancestors::Ancestors,
     solana_runtime::{bank::BankStatusCache, status_cache::*},
     solana_sdk::{
         hash::{hash, Hash},
-        signature::Signature,
+        signature::{Signature, SIGNATURE_BYTES},
     },
     test::Bencher,
 };
@@ -49,3 +51,76 @@ fn bench_status_cache_root_slot_deltas(bencher: &mut Bencher) {
 
     bencher.iter(|| test::black_box(status_cache.root_slot_deltas()));
 }
+
+fn fill_status_cache(status_cache: &mut BankStatusCache, max_cache_entries: u64, num_txs: usize) {
+    for slot in 0..max_cache_entries {
+        let blockhash = Hash::new_unique();
+        fill_status_cache_slot(status_cache, &blockhash, slot, num_txs);
+    }
+}
+
+fn fill_status_cache_slot(
+    status_cache: &mut BankStatusCache,
+    blockhash: &Hash,
+    slot: u64,
+    num_txs: usize,
+) {
+    for _ in 0..num_txs {
+        let tx_hash = Hash::new_unique();
+        status_cache.insert(blockhash, tx_hash, slot, Ok(()));
+    }
+}
+
+#[bench]
+fn bench_status_cache_check_and_insert(bencher: &mut Bencher) {
+    // Fill up the status cache to better match what intense runtime usage would
+    // look like.
+    let max_cache_entries = MAX_CACHE_ENTRIES as u64;
+    let mut status_cache = BankStatusCache::default();
+    fill_status_cache(&mut status_cache, max_cache_entries - 1, 100_000);
+
+    // Manually fill the last slot so we can save off the blockhash to use for
+    // querying and inserting into.
+    let blockhash = Hash::new_unique();
+    fill_status_cache_slot(&mut status_cache, &blockhash, max_cache_entries, 100_000);
+
+    let slot = max_cache_entries + 1;
+    let ancestors = Ancestors::from((slot - 32..slot).collect::<Vec<u64>>());
+
+    // Pre-generate unique tx_hashes so we don't spend benchmark time generating
+    // them.
+    let batch_size = 1_000;
+    let mut tx_hashes = Vec::with_capacity(batch_size);
+    let mut rng = SmallRng::seed_from_u64(0);
+    for _ in 0..batch_size {
+        let mut sigbytes = [0u8; SIGNATURE_BYTES];
+        rng.fill(&mut sigbytes);
+        tx_hashes.push(Signature::from(sigbytes));
+    }
+
+    bencher.iter(|| {
+        for tx_hash in &tx_hashes {
+            if status_cache
+                .get_status(*tx_hash, &blockhash, &ancestors)
+                .is_none()
+            {
+                status_cache.insert(&blockhash, *tx_hash, slot, Ok(()));
+            }
+        }
+    });
+}
+
+#[bench]
+fn bench_status_cache_add_roots(bencher: &mut Bencher) {
+    // Fill up the status cache to better match what intense runtime usage would
+    // look like.
+    let max_cache_entries = MAX_CACHE_ENTRIES as u64;
+    let mut status_cache = BankStatusCache::default();
+    fill_status_cache(&mut status_cache, max_cache_entries, 100_000);
+    let start_slot = max_cache_entries + 1;
+    bencher.iter(|| {
+        for root in start_slot..start_slot + max_cache_entries {
+            status_cache.add_root(root);
+        }
+    });
+}