feat(l1): code cache rocksdb (#5307)

**Motivation**

Adds a code cache to rocksdb, evicting the least used entries if size
goes over 64mb

---------

Co-authored-by: Javier Chatruc <[email protected]>
Co-authored-by: Edgar <[email protected]>
Co-authored-by: Tomás Grüner <[email protected]>

Author: 4 people

Cargo.lock

@@ -65,6 +65,21 @@ impl Code {
         }
         targets
     }
+
+    /// Estimates the size of the Code struct in bytes
+    /// (including stack size and heap allocation).
+    ///
+    /// Note: This is an estimation and may not be exact.
+    ///
+    /// # Returns
+    ///
+    /// usize - Estimated size in bytes
+    pub fn size(&self) -> usize {
+        let hash_size = size_of::<H256>();
+        let bytes_size = size_of::<Bytes>();
+        let vec_size = size_of::<Vec<u32>>() + self.jump_targets.len() * size_of::<u32>();
+        hash_size + bytes_size + vec_size
+    }
 }
 
 impl AsRef<Bytes> for Code {

crates/common/types/account.rs

@@ -28,6 +28,7 @@ tokio = { workspace = true, optional = true, features = ["rt"] }
 bincode = "1.3.3"
 qfilter = "0.2.5"
 rayon.workspace = true
+lru = "0.16.2"
 
 [features]
 default = []

crates/l2/prover/src/guest_program/src/risc0/Cargo.lock

@@ -14,15 +14,18 @@ use ethrex_common::{
     },
 };
 use ethrex_trie::{Nibbles, Node, Trie};
+use lru::LruCache;
 use rocksdb::{
     BlockBasedOptions, BoundColumnFamily, ColumnFamilyDescriptor, DBWithThreadMode, MultiThreaded,
     Options, WriteBatch, checkpoint::Checkpoint,
 };
+use rustc_hash::FxBuildHasher;
 use std::{
     collections::HashSet,
     path::{Path, PathBuf},
     sync::{
         Arc, Mutex,
+        atomic::AtomicU64,
         mpsc::{SyncSender, sync_channel},
     },
 };
@@ -149,13 +152,40 @@ enum FKVGeneratorControlMessage {
     Continue,
 }
 
-#[derive(Debug, Clone)]
+// 64mb
+const CODE_CACHE_MAX_SIZE: u64 = 64 * 1024 * 1024;
+type CodeCache = LruCache<H256, Code, FxBuildHasher>;
+
+#[derive(Debug)]
 pub struct Store {
     db: Arc<DBWithThreadMode<MultiThreaded>>,
     trie_cache: Arc<Mutex<Arc<TrieLayerCache>>>,
     flatkeyvalue_control_tx: std::sync::mpsc::SyncSender<FKVGeneratorControlMessage>,
     trie_update_worker_tx: TriedUpdateWorkerTx,
     last_computed_flatkeyvalue: Arc<Mutex<Vec<u8>>>,
+    /// Cache for account bytecodes, keyed by the bytecode hash.
+    /// Note that we don't remove entries on account code changes, since
+    /// those changes already affect the code hash stored in the account, and only
+    /// may result in this cache having useless data.
+    account_code_cache: Arc<Mutex<CodeCache>>,
+    account_code_cache_size: AtomicU64,
+}
+
+impl Clone for Store {
+    fn clone(&self) -> Self {
+        Self {
+            db: self.db.clone(),
+            trie_cache: self.trie_cache.clone(),
+            flatkeyvalue_control_tx: self.flatkeyvalue_control_tx.clone(),
+            trie_update_worker_tx: self.trie_update_worker_tx.clone(),
+            last_computed_flatkeyvalue: self.last_computed_flatkeyvalue.clone(),
+            account_code_cache: self.account_code_cache.clone(),
+            account_code_cache_size: AtomicU64::new(
+                self.account_code_cache_size
+                    .load(std::sync::atomic::Ordering::Relaxed),
+            ),
+        }
+    }
 }
 
 impl Store {
@@ -399,6 +429,10 @@ impl Store {
             flatkeyvalue_control_tx: fkv_tx,
             trie_update_worker_tx: trie_upd_tx,
             last_computed_flatkeyvalue: Arc::new(Mutex::new(last_written)),
+            account_code_cache: Arc::new(Mutex::new(LruCache::unbounded_with_hasher(
+                FxBuildHasher,
+            ))),
+            account_code_cache_size: AtomicU64::new(0),
         };
         let store_clone = store.clone();
         std::thread::spawn(move || {
@@ -1363,7 +1397,21 @@ impl StoreEngine for Store {
         .map_err(|e| StoreError::Custom(format!("Task panicked: {}", e)))?
     }
 
+    /// Get account code by its hash.
+    ///
+    /// Check if the code exists in the cache (attribute `account_code_cache`), if not,
+    /// reads the database, and if it exists, decodes and returns it.
     fn get_account_code(&self, code_hash: H256) -> Result<Option<Code>, StoreError> {
+        // check cache first
+        if let Some(code) = self
+            .account_code_cache
+            .lock()
+            .map_err(|_| StoreError::LockError)?
+            .get(&code_hash)
+        {
+            return Ok(Some(code.clone()));
+        }
+
         let cf = self.cf_handle(CF_ACCOUNT_CODES)?;
         let Some(bytes) = self
             .db
@@ -1378,6 +1426,39 @@ impl StoreEngine for Store {
             bytecode: Bytes::copy_from_slice(bytecode),
             jump_targets: <Vec<_>>::decode(targets)?,
         };
+
+        // insert into cache and evict if needed
+        {
+            let mut cache = self
+                .account_code_cache
+                .lock()
+                .map_err(|_| StoreError::LockError)?;
+            let code_size = code.size();
+            self.account_code_cache_size
+                .fetch_add(code_size as u64, std::sync::atomic::Ordering::SeqCst);
+            let cache_len = cache.len() + 1;
+            let mut current_size = self
+                .account_code_cache_size
+                .load(std::sync::atomic::Ordering::SeqCst);
+            debug!(
+                "[ACCOUNT CODE CACHE] cache elements (): {cache_len}, total size: {current_size} bytes"
+            );
+
+            while current_size > CODE_CACHE_MAX_SIZE {
+                if let Some((_, code)) = cache.pop_lru() {
+                    self.account_code_cache_size
+                        .fetch_sub(code.size() as u64, std::sync::atomic::Ordering::SeqCst);
+                    current_size = self
+                        .account_code_cache_size
+                        .load(std::sync::atomic::Ordering::SeqCst);
+                } else {
+                    break;
+                }
+            }
+
+            cache.get_or_insert(code_hash, || code.clone());
+        }
+
         Ok(Some(code))
     }