From 669bc43bcd4ccdff79ba4d9738554a221201a994 Mon Sep 17 00:00:00 2001
From: HaoranYi <haoran.yi@gmail.com>
Date: Mon, 23 Oct 2023 09:48:17 -0500
Subject: [PATCH] Buffer account's fields for hash (#33788)

* buffer accounts field for hash

* use smallvec to allocate hash buffer on stack

* sort deps

* more opt

* clippy

---------

Co-authored-by: HaoranYi <haoran.yi@solana.com>
---
 Cargo.lock                     |  5 +++--
 Cargo.toml                     |  1 +
 accounts-db/Cargo.toml         |  1 +
 accounts-db/src/accounts_db.rs | 39 +++++++++++++++++++++++++---------
 programs/sbf/Cargo.lock        |  5 +++--
 5 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index ff5576d984..cdd3278af3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5048,9 +5048,9 @@ checksum = "9def91fd1e018fe007022791f865d0ccc9b3a0d5001e01aabb8b40e46000afb5"
 
 [[package]]
 name = "smallvec"
-version = "1.7.0"
+version = "1.11.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1ecab6c735a6bb4139c0caafd0cc3635748bbb3acf4550e8138122099251f309"
+checksum = "942b4a808e05215192e39f4ab80813e599068285906cc91aa64f923db842bd5a"
 
 [[package]]
 name = "smpl_jwt"
@@ -5214,6 +5214,7 @@ dependencies = [
  "seqlock",
  "serde",
  "serde_derive",
+ "smallvec",
  "solana-accounts-db",
  "solana-bucket-map",
  "solana-config-program",
diff --git a/Cargo.toml b/Cargo.toml
index 21fc3fcfb2..53889bf2f4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -298,6 +298,7 @@ sha2 = "0.10.8"
 sha3 = "0.10.4"
 signal-hook = "0.3.17"
 siphasher = "0.3.11"
+smallvec = "1.11.1"
 smpl_jwt = "0.7.1"
 socket2 = "0.5.5"
 soketto = "0.7"
diff --git a/accounts-db/Cargo.toml b/accounts-db/Cargo.toml
index 0cafcfe3e4..b4fcceea00 100644
--- a/accounts-db/Cargo.toml
+++ b/accounts-db/Cargo.toml
@@ -43,6 +43,7 @@ regex = { workspace = true }
 seqlock = { workspace = true }
 serde = { workspace = true, features = ["rc"] }
 serde_derive = { workspace = true }
+smallvec = { workspace = true }
 solana-bucket-map = { workspace = true }
 solana-config-program = { workspace = true }
 solana-frozen-abi = { workspace = true }
diff --git a/accounts-db/src/accounts_db.rs b/accounts-db/src/accounts_db.rs
index 3bc7a651d4..711c893dfe 100644
--- a/accounts-db/src/accounts_db.rs
+++ b/accounts-db/src/accounts_db.rs
@@ -79,6 +79,7 @@ use {
     rand::{thread_rng, Rng},
     rayon::{prelude::*, ThreadPool},
     serde::{Deserialize, Serialize},
+    smallvec::SmallVec,
     solana_measure::{measure::Measure, measure_us},
     solana_nohash_hasher::IntSet,
     solana_rayon_threadlimit::get_thread_count,
@@ -6274,31 +6275,49 @@ impl AccountsDb {
         }
         let mut hasher = blake3::Hasher::new();
 
-        hasher.update(&lamports.to_le_bytes());
+        // allocate 128 bytes buffer on the stack
+        const BUF_SIZE: usize = 128;
+        const TOTAL_FIELD_SIZE: usize = 8 /* lamports */ + 8 /* slot */ + 8 /* rent_epoch */ + 1 /* exec_flag */ + 32 /* owner_key */ + 32 /* pubkey */;
+        const DATA_SIZE_CAN_FIT: usize = BUF_SIZE - TOTAL_FIELD_SIZE;
+
+        let mut buffer = SmallVec::<[u8; BUF_SIZE]>::new();
+
+        // collect lamports, slot, rent_epoch into buffer to hash
+        buffer.extend_from_slice(&lamports.to_le_bytes());
 
         match include_slot {
             IncludeSlotInHash::IncludeSlot => {
                 // upon feature activation, stop including slot# in the account hash
-                hasher.update(&slot.to_le_bytes());
+                buffer.extend_from_slice(&slot.to_le_bytes());
             }
             IncludeSlotInHash::RemoveSlot => {}
             IncludeSlotInHash::IrrelevantAssertOnUse => {
                 panic!("IncludeSlotInHash is irrelevant, but we are calculating hash");
             }
         }
+        buffer.extend_from_slice(&rent_epoch.to_le_bytes());
 
-        hasher.update(&rent_epoch.to_le_bytes());
+        if data.len() > DATA_SIZE_CAN_FIT {
+            // For larger accounts whose data can't fit into the buffer, update the hash now.
+            hasher.update(&buffer);
+            buffer.clear();
 
-        hasher.update(data);
-
-        if executable {
-            hasher.update(&[1u8; 1]);
+            // hash account's data
+            hasher.update(data);
         } else {
-            hasher.update(&[0u8; 1]);
+            // For small accounts whose data can fit into the buffer, append it to the buffer.
+            buffer.extend_from_slice(data);
         }
 
-        hasher.update(owner.as_ref());
-        hasher.update(pubkey.as_ref());
+        // collect exec_flag, owner, pubkey into buffer to hash
+        if executable {
+            buffer.push(1_u8);
+        } else {
+            buffer.push(0_u8);
+        }
+        buffer.extend_from_slice(owner.as_ref());
+        buffer.extend_from_slice(pubkey.as_ref());
+        hasher.update(&buffer);
 
         AccountHash(Hash::new_from_array(hasher.finalize().into()))
     }
diff --git a/programs/sbf/Cargo.lock b/programs/sbf/Cargo.lock
index d1f3d72cc3..5b9ac7556e 100644
--- a/programs/sbf/Cargo.lock
+++ b/programs/sbf/Cargo.lock
@@ -4402,9 +4402,9 @@ checksum = "c111b5bd5695e56cffe5129854aa230b39c93a305372fdbb2668ca2394eea9f8"
 
 [[package]]
 name = "smallvec"
-version = "1.6.1"
+version = "1.11.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fe0f37c9e8f3c5a4a66ad655a93c74daac4ad00c441533bf5c6e7990bb42604e"
+checksum = "942b4a808e05215192e39f4ab80813e599068285906cc91aa64f923db842bd5a"
 
 [[package]]
 name = "smpl_jwt"
@@ -4517,6 +4517,7 @@ dependencies = [
  "seqlock",
  "serde",
  "serde_derive",
+ "smallvec",
  "solana-bucket-map",
  "solana-config-program",
  "solana-frozen-abi",