From b2b159ad9204bf7744b78b8aa062636d1646f4af Mon Sep 17 00:00:00 2001
From: Brooks <brooks@anza.xyz>
Date: Tue, 26 Mar 2024 10:44:23 -0400
Subject: [PATCH] Adds bench for writing accounts to append vecs and hot
 storage (#421)

---
 accounts-db/Cargo.toml                     |  4 +
 accounts-db/benches/bench_accounts_file.rs | 93 ++++++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 accounts-db/benches/bench_accounts_file.rs

diff --git a/accounts-db/Cargo.toml b/accounts-db/Cargo.toml
index d43acdf9e819f3..ff38118806c3c9 100644
--- a/accounts-db/Cargo.toml
+++ b/accounts-db/Cargo.toml
@@ -80,6 +80,10 @@ rustc_version = { workspace = true }
 [features]
 dev-context-only-utils = ["dep:qualifier_attr", "dep:solana-stake-program", "dep:solana-vote-program"]
 
+[[bench]]
+name = "bench_accounts_file"
+harness = false
+
 [[bench]]
 name = "bench_hashing"
 harness = false
diff --git a/accounts-db/benches/bench_accounts_file.rs b/accounts-db/benches/bench_accounts_file.rs
new file mode 100644
index 00000000000000..808f1a630ee7df
--- /dev/null
+++ b/accounts-db/benches/bench_accounts_file.rs
@@ -0,0 +1,93 @@
+#![allow(clippy::arithmetic_side_effects)]
+use {
+    criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion, Throughput},
+    solana_accounts_db::{
+        account_storage::meta::StorableAccountsWithHashesAndWriteVersions,
+        accounts_hash::AccountHash,
+        append_vec::{self, AppendVec},
+        tiered_storage::hot::HotStorageWriter,
+    },
+    solana_sdk::{
+        account::Account, clock::Slot, hash::Hash, pubkey::Pubkey,
+        rent_collector::RENT_EXEMPT_RENT_EPOCH,
+    },
+};
+
+const ACCOUNTS_COUNTS: [usize; 4] = [
+    1,      // the smallest count; will bench overhead
+    100,    // number of accounts written per slot on mnb (with *no* rent rewrites)
+    1_000,  // number of accounts written slot on mnb (with rent rewrites)
+    10_000, // reasonable largest number of accounts written per slot
+];
+
+fn bench_write_accounts_file(c: &mut Criterion) {
+    let mut group = c.benchmark_group("write_accounts_file");
+
+    // most accounts on mnb are 165-200 bytes, so use that here too
+    let space = 200;
+    let lamports = 2_282_880; // the rent-exempt amount for 200 bytes of data
+    let temp_dir = tempfile::tempdir().unwrap();
+
+    for accounts_count in ACCOUNTS_COUNTS {
+        group.throughput(Throughput::Elements(accounts_count as u64));
+
+        let accounts: Vec<_> = std::iter::repeat_with(|| {
+            (
+                Pubkey::new_unique(),
+                Account::new_rent_epoch(
+                    lamports,
+                    space,
+                    &Pubkey::new_unique(),
+                    RENT_EXEMPT_RENT_EPOCH,
+                ),
+            )
+        })
+        .take(accounts_count)
+        .collect();
+        let accounts_refs: Vec<_> = accounts.iter().collect();
+        let accounts_data = (Slot::MAX, accounts_refs.as_slice());
+        let storable_accounts =
+            StorableAccountsWithHashesAndWriteVersions::new_with_hashes_and_write_versions(
+                &accounts_data,
+                vec![AccountHash(Hash::default()); accounts_count],
+                vec![0; accounts_count],
+            );
+
+        group.bench_function(BenchmarkId::new("append_vec", accounts_count), |b| {
+            b.iter_batched_ref(
+                || {
+                    let path = temp_dir.path().join(format!("append_vec_{accounts_count}"));
+                    let file_size = accounts.len() * (space + append_vec::STORE_META_OVERHEAD);
+                    AppendVec::new(&path, true, file_size)
+                },
+                |append_vec| {
+                    let res = append_vec.append_accounts(&storable_accounts, 0).unwrap();
+                    let accounts_written_count = res.len();
+                    assert_eq!(accounts_written_count, accounts_count);
+                },
+                BatchSize::SmallInput,
+            );
+        });
+
+        group.bench_function(BenchmarkId::new("hot_storage", accounts_count), |b| {
+            b.iter_batched_ref(
+                || {
+                    let path = temp_dir
+                        .path()
+                        .join(format!("hot_storage_{accounts_count}"));
+                    _ = std::fs::remove_file(&path);
+                    HotStorageWriter::new(path).unwrap()
+                },
+                |hot_storage| {
+                    let res = hot_storage.write_accounts(&storable_accounts, 0).unwrap();
+                    let accounts_written_count = res.len();
+                    assert_eq!(accounts_written_count, accounts_count);
+                },
+                BatchSize::SmallInput,
+            );
+        });
+    }
+}
+
+criterion_group!(benches, bench_write_accounts_file);
+criterion_main!(benches);