feat(storage): add in-memory BAL retention (#23873)

Co-authored-by: Matthias Seitz <[email protected]>

Author: 0xKarl98

Cargo.lock

@@ -449,7 +449,7 @@ alloy-sol-types = { version = "1.5.6", default-features = false }
 
 alloy-chains = { version = "0.2.33", default-features = false }
 alloy-eip2124 = { version = "0.2.0", default-features = false }
-alloy-eip7928 = { version = "0.3.4", default-features = false }
+alloy-eip7928 = { version = "0.3.5", default-features = false }
 alloy-evm = { version = "0.34.0", default-features = false }
 alloy-rlp = { version = "0.3.13", default-features = false, features = ["core-net"] }
 alloy-trie = { version = "0.9.4", default-features = false }

Cargo.toml

@@ -35,6 +35,7 @@ reth-static-file-types = { workspace = true, features = ["std"] }
 reth-fs-util.workspace = true
 
 # ethereum
+alloy-eip7928.workspace = true
 alloy-eips.workspace = true
 alloy-genesis.workspace = true
 alloy-primitives.workspace = true

crates/storage/provider/Cargo.toml

@@ -1,32 +1,129 @@
+use alloy_eip7928::BAL_RETENTION_PERIOD_SLOTS;
 use alloy_primitives::{BlockHash, BlockNumber, Bytes};
 use parking_lot::RwLock;
+use reth_prune_types::PruneMode;
 use reth_storage_api::{BalStore, GetBlockAccessListLimit};
 use reth_storage_errors::provider::ProviderResult;
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::{BTreeMap, HashMap},
+    sync::Arc,
+};
 
 /// Basic in-memory BAL store keyed by block hash.
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone)]
 pub struct InMemoryBalStore {
-    entries: Arc<RwLock<HashMap<BlockHash, Bytes>>>,
+    config: BalConfig,
+    inner: Arc<RwLock<InMemoryBalStoreInner>>,
+}
+
+impl InMemoryBalStore {
+    /// Creates a new in-memory BAL store with the given config.
+    pub fn new(config: BalConfig) -> Self {
+        Self { config, inner: Arc::new(RwLock::new(InMemoryBalStoreInner::default())) }
+    }
+}
+
+impl Default for InMemoryBalStore {
+    fn default() -> Self {
+        Self::new(BalConfig::default())
+    }
+}
+
+/// Configuration for BAL storage.
+#[derive(Debug, Clone, Copy, Eq, PartialEq)]
+pub struct BalConfig {
+    /// Retention policy for BALs kept in memory.
+    in_memory_retention: Option<PruneMode>,
+}
+
+impl BalConfig {
+    /// Returns a config with no in-memory BAL retention limit.
+    pub const fn unbounded() -> Self {
+        Self { in_memory_retention: None }
+    }
+
+    /// Returns a config with the given in-memory BAL retention policy.
+    pub const fn with_in_memory_retention(in_memory_retention: PruneMode) -> Self {
+        Self { in_memory_retention: Some(in_memory_retention) }
+    }
+}
+
+impl Default for BalConfig {
+    fn default() -> Self {
+        Self::with_in_memory_retention(PruneMode::Distance(BAL_RETENTION_PERIOD_SLOTS))
+    }
+}
+
+#[derive(Debug, Default)]
+struct InMemoryBalStoreInner {
+    entries: HashMap<BlockHash, BalEntry>,
+    hashes_by_number: BTreeMap<BlockNumber, Vec<BlockHash>>,
+    highest_block_number: Option<BlockNumber>,
+}
+
+impl InMemoryBalStoreInner {
+    // Inserts a BAL and keeps the block-number index in sync.
+    fn insert(&mut self, block_hash: BlockHash, block_number: BlockNumber, bal: Bytes) {
+        let empty_block_number =
+            self.entries.insert(block_hash, BalEntry { block_number, bal }).and_then(|entry| {
+                let hashes = self.hashes_by_number.get_mut(&entry.block_number)?;
+                hashes.retain(|hash| *hash != block_hash);
+                hashes.is_empty().then_some(entry.block_number)
+            });
+
+        if let Some(block_number) = empty_block_number {
+            self.hashes_by_number.remove(&block_number);
+        }
+
+        self.hashes_by_number.entry(block_number).or_default().push(block_hash);
+        self.highest_block_number = Some(
+            self.highest_block_number.map_or(block_number, |highest| highest.max(block_number)),
+        );
+    }
+
+    // Removes BALs outside the configured retention window.
+    fn prune(&mut self, prune_mode: Option<PruneMode>) {
+        let Some(prune_mode) = prune_mode else { return };
+        let Some(tip) = self.highest_block_number else { return };
+
+        while let Some((&block_number, _)) = self.hashes_by_number.first_key_value() {
+            if !prune_mode.should_prune(block_number, tip) {
+                break
+            }
+
+            let Some((_, hashes)) = self.hashes_by_number.pop_first() else { break };
+            for hash in hashes {
+                self.entries.remove(&hash);
+            }
+        }
+    }
+}
+
+#[derive(Debug)]
+struct BalEntry {
+    block_number: BlockNumber,
+    bal: Bytes,
 }
 
 impl BalStore for InMemoryBalStore {
     fn insert(
         &self,
         block_hash: BlockHash,
-        _block_number: BlockNumber,
+        block_number: BlockNumber,
         bal: Bytes,
     ) -> ProviderResult<()> {
-        self.entries.write().insert(block_hash, bal);
+        let mut inner = self.inner.write();
+        inner.insert(block_hash, block_number, bal);
+        inner.prune(self.config.in_memory_retention);
         Ok(())
     }
 
     fn get_by_hashes(&self, block_hashes: &[BlockHash]) -> ProviderResult<Vec<Option<Bytes>>> {
-        let entries = self.entries.read();
+        let inner = self.inner.read();
         let mut result = Vec::with_capacity(block_hashes.len());
 
         for hash in block_hashes {
-            result.push(entries.get(hash).cloned());
+            result.push(inner.entries.get(hash).map(|entry| entry.bal.clone()));
         }
 
         Ok(result)
@@ -38,11 +135,15 @@ impl BalStore for InMemoryBalStore {
         limit: GetBlockAccessListLimit,
         out: &mut Vec<Bytes>,
     ) -> ProviderResult<()> {
-        let entries = self.entries.read();
+        let inner = self.inner.read();
         let mut size = 0;
 
         for hash in block_hashes {
-            let bal = entries.get(hash).cloned().unwrap_or_else(|| Bytes::from_static(&[0xc0]));
+            let bal = inner
+                .entries
+                .get(hash)
+                .map(|entry| entry.bal.clone())
+                .unwrap_or_else(|| Bytes::from_static(&[0xc0]));
             size += bal.len();
             out.push(bal);
 
@@ -106,4 +207,54 @@ mod tests {
 
         assert_eq!(limited, vec![bal0, bal1]);
     }
+
+    #[test]
+    fn default_retention_prunes_old_bals() {
+        let store = InMemoryBalStore::default();
+        let old_hash = B256::random();
+        let retained_hash = B256::random();
+        let tip_hash = B256::random();
+        let old_bal = Bytes::from_static(b"old");
+        let retained_bal = Bytes::from_static(b"retained");
+        let tip_bal = Bytes::from_static(b"tip");
+
+        store.insert(old_hash, 1, old_bal).unwrap();
+        store.insert(retained_hash, BAL_RETENTION_PERIOD_SLOTS, retained_bal.clone()).unwrap();
+        store.insert(tip_hash, BAL_RETENTION_PERIOD_SLOTS + 2, tip_bal.clone()).unwrap();
+
+        assert_eq!(
+            store.get_by_hashes(&[old_hash, retained_hash, tip_hash]).unwrap(),
+            vec![None, Some(retained_bal), Some(tip_bal)]
+        );
+    }
+
+    #[test]
+    fn unbounded_retention_keeps_old_bals() {
+        let store = InMemoryBalStore::new(BalConfig::unbounded());
+        let old_hash = B256::random();
+        let tip_hash = B256::random();
+        let old_bal = Bytes::from_static(b"old");
+        let tip_bal = Bytes::from_static(b"tip");
+
+        store.insert(old_hash, 1, old_bal.clone()).unwrap();
+        store.insert(tip_hash, BAL_RETENTION_PERIOD_SLOTS + 1, tip_bal.clone()).unwrap();
+
+        assert_eq!(
+            store.get_by_hashes(&[old_hash, tip_hash]).unwrap(),
+            vec![Some(old_bal), Some(tip_bal)]
+        );
+    }
+
+    #[test]
+    fn reinserting_hash_updates_number_index() {
+        let store =
+            InMemoryBalStore::new(BalConfig::with_in_memory_retention(PruneMode::Before(2)));
+        let hash = B256::random();
+        let bal = Bytes::from_static(b"bal");
+
+        store.insert(hash, 1, Bytes::from_static(b"old")).unwrap();
+        store.insert(hash, 2, bal.clone()).unwrap();
+
+        assert_eq!(store.get_by_hashes(&[hash]).unwrap(), vec![Some(bal)]);
+    }
 }

crates/storage/provider/src/bal.rs

@@ -39,7 +39,7 @@ pub mod either_writer;
 pub use either_writer::*;
 
 mod bal;
-pub use bal::InMemoryBalStore;
+pub use bal::{BalConfig, InMemoryBalStore};
 
 pub use reth_chain_state::{
     CanonStateNotification, CanonStateNotificationSender, CanonStateNotificationStream,