reth_provider/
bal.rs

1use alloy_eip7928::BAL_RETENTION_PERIOD_SLOTS;
2use alloy_eips::NumHash;
3use alloy_primitives::{map::B256Map, BlockHash, BlockNumber, Bytes};
4use parking_lot::RwLock;
5use reth_prune_types::PruneMode;
6use reth_storage_api::{BalStore, GetBlockAccessListLimit, RawBal};
7use reth_storage_errors::provider::ProviderResult;
8use std::{collections::BTreeMap, sync::Arc};
9
10mod rocksdb;
11pub use rocksdb::RocksDBBalStore;
12
13/// Basic in-memory BAL store keyed by block hash.
14#[derive(Debug, Clone)]
15pub struct InMemoryBalStore {
16    config: BalConfig,
17    inner: Arc<RwLock<InMemoryBalStoreInner>>,
18}
19
20impl InMemoryBalStore {
21    /// Creates a new in-memory BAL store with the given config.
22    pub fn new(config: BalConfig) -> Self {
23        Self { config, inner: Arc::new(RwLock::new(InMemoryBalStoreInner::default())) }
24    }
25}
26
27impl Default for InMemoryBalStore {
28    fn default() -> Self {
29        Self::new(BalConfig::default())
30    }
31}
32
33/// Configuration for BAL storage.
34#[derive(Debug, Clone, Copy, Eq, PartialEq)]
35pub struct BalConfig {
36    /// Retention policy for BALs kept in memory.
37    in_memory_retention: Option<PruneMode>,
38}
39
40impl BalConfig {
41    /// Default block distance for BALs kept in memory.
42    pub const DEFAULT_IN_MEMORY_RETENTION_DISTANCE: u64 = BAL_RETENTION_PERIOD_SLOTS;
43
44    /// Returns a config with no in-memory BAL retention limit.
45    pub const fn unbounded() -> Self {
46        Self { in_memory_retention: None }
47    }
48
49    /// Returns a config that keeps BALs within the given block distance in memory.
50    pub const fn with_in_memory_retention_distance(blocks: u64) -> Self {
51        Self::with_in_memory_retention(PruneMode::Distance(blocks))
52    }
53
54    /// Returns a config with the given in-memory BAL retention policy.
55    pub const fn with_in_memory_retention(in_memory_retention: PruneMode) -> Self {
56        Self { in_memory_retention: Some(in_memory_retention) }
57    }
58}
59
60impl Default for BalConfig {
61    fn default() -> Self {
62        Self::with_in_memory_retention_distance(Self::DEFAULT_IN_MEMORY_RETENTION_DISTANCE)
63    }
64}
65
66#[derive(Debug, Default)]
67struct InMemoryBalStoreInner {
68    entries: B256Map<BalEntry>,
69    hashes_by_number: BTreeMap<BlockNumber, Vec<BlockHash>>,
70    highest_block_number: Option<BlockNumber>,
71}
72
73impl InMemoryBalStoreInner {
74    // Inserts a BAL and keeps the block-number index in sync.
75    fn insert(&mut self, block_hash: BlockHash, block_number: BlockNumber, bal: Bytes) {
76        let empty_block_number =
77            self.entries.insert(block_hash, BalEntry { block_number, bal }).and_then(|entry| {
78                let hashes = self.hashes_by_number.get_mut(&entry.block_number)?;
79                hashes.retain(|hash| *hash != block_hash);
80                hashes.is_empty().then_some(entry.block_number)
81            });
82
83        if let Some(block_number) = empty_block_number {
84            self.hashes_by_number.remove(&block_number);
85        }
86
87        self.hashes_by_number.entry(block_number).or_default().push(block_hash);
88        self.highest_block_number = Some(
89            self.highest_block_number.map_or(block_number, |highest| highest.max(block_number)),
90        );
91    }
92
93    // Removes BALs outside the configured retention window for the given chain tip.
94    fn prune(&mut self, prune_mode: Option<PruneMode>, tip: BlockNumber) -> usize {
95        let Some(prune_mode) = prune_mode else { return 0 };
96
97        let mut pruned = 0;
98        while let Some((&block_number, _)) = self.hashes_by_number.first_key_value() {
99            if !prune_mode.should_prune(block_number, tip) {
100                break
101            }
102
103            let Some((_, hashes)) = self.hashes_by_number.pop_first() else { break };
104            for hash in hashes {
105                pruned += usize::from(self.entries.remove(&hash).is_some());
106            }
107        }
108        pruned
109    }
110}
111
112#[derive(Debug)]
113struct BalEntry {
114    block_number: BlockNumber,
115    bal: Bytes,
116}
117
118impl BalStore for InMemoryBalStore {
119    fn insert(&self, num_hash: NumHash, bal: RawBal) -> ProviderResult<()> {
120        let mut inner = self.inner.write();
121        inner.insert(num_hash.hash, num_hash.number, bal.as_raw().clone());
122        if let Some(highest_block_number) = inner.highest_block_number {
123            // This preserves insert-time cleanup based on the highest inserted BAL block.
124            inner.prune(self.config.in_memory_retention, highest_block_number);
125        }
126        Ok(())
127    }
128
129    fn insert_many(&self, entries: Vec<(NumHash, RawBal)>) -> ProviderResult<()> {
130        if entries.is_empty() {
131            return Ok(())
132        }
133
134        let mut inner = self.inner.write();
135        inner.entries.reserve(entries.len());
136        for (num_hash, bal) in &entries {
137            inner.insert(num_hash.hash, num_hash.number, bal.as_raw().clone());
138        }
139        if let Some(highest_block_number) = inner.highest_block_number {
140            inner.prune(self.config.in_memory_retention, highest_block_number);
141        }
142        Ok(())
143    }
144
145    fn flush(&self, _blocks: &[NumHash]) -> ProviderResult<()> {
146        Ok(())
147    }
148
149    fn prune(&self, tip: BlockNumber) -> ProviderResult<usize> {
150        Ok(self.inner.write().prune(self.config.in_memory_retention, tip))
151    }
152
153    fn get_by_hashes(&self, block_hashes: &[BlockHash]) -> ProviderResult<Vec<Option<Bytes>>> {
154        let inner = self.inner.read();
155        let mut result = Vec::with_capacity(block_hashes.len());
156
157        for hash in block_hashes {
158            result.push(inner.entries.get(hash).map(|entry| entry.bal.clone()));
159        }
160
161        Ok(result)
162    }
163
164    fn append_by_hashes_with_limit(
165        &self,
166        block_hashes: &[BlockHash],
167        limit: GetBlockAccessListLimit,
168        out: &mut Vec<Option<Bytes>>,
169    ) -> ProviderResult<()> {
170        let inner = self.inner.read();
171        let mut size = 0;
172
173        for hash in block_hashes {
174            let bal = inner.entries.get(hash).map(|entry| entry.bal.clone());
175            size += bal.as_ref().map_or(1, |bytes| bytes.len());
176            out.push(bal);
177
178            if limit.exceeds(size) {
179                break
180            }
181        }
182
183        Ok(())
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190    use alloy_primitives::B256;
191
192    #[test]
193    fn insert_and_lookup_by_hash() {
194        let store = InMemoryBalStore::default();
195        let hash = B256::random();
196        let missing = B256::random();
197        let bal = Bytes::from_static(b"bal");
198
199        store.insert(NumHash::new(1, hash), RawBal::from(bal.clone())).unwrap();
200
201        assert_eq!(store.get_by_hashes(&[hash, missing]).unwrap(), vec![Some(bal), None]);
202    }
203
204    #[test]
205    fn insert_many_and_lookup_by_hash() {
206        let store = InMemoryBalStore::default();
207        let hash0 = B256::random();
208        let hash1 = B256::random();
209        let bal0 = RawBal::from(Bytes::from_static(b"bal0"));
210        let bal1 = RawBal::from(Bytes::from_static(b"bal1"));
211
212        store
213            .insert_many(vec![
214                (NumHash::new(1, hash0), bal0.clone()),
215                (NumHash::new(2, hash1), bal1),
216            ])
217            .unwrap();
218
219        assert_eq!(
220            store.get_by_hashes(&[hash0, hash1]).unwrap(),
221            vec![Some(bal0.as_raw().clone()), Some(Bytes::from_static(b"bal1"))]
222        );
223    }
224
225    #[test]
226    fn flush_is_noop() {
227        let store = InMemoryBalStore::default();
228
229        store.flush(&[]).unwrap();
230    }
231
232    #[test]
233    fn limited_lookup_returns_prefix() {
234        let store = InMemoryBalStore::default();
235        let hash0 = B256::random();
236        let hash1 = B256::random();
237        let hash2 = B256::random();
238        let bal0 = Bytes::from_static(&[0xc1, 0x01]);
239        let bal1 = Bytes::from_static(&[0xc1, 0x02]);
240        let bal2 = Bytes::from_static(&[0xc1, 0x03]);
241
242        store.insert(NumHash::new(1, hash0), RawBal::from(bal0.clone())).unwrap();
243        store.insert(NumHash::new(2, hash1), RawBal::from(bal1.clone())).unwrap();
244        store.insert(NumHash::new(3, hash2), RawBal::from(bal2)).unwrap();
245
246        let limited = store
247            .get_by_hashes_with_limit(
248                &[hash0, hash1, hash2],
249                GetBlockAccessListLimit::ResponseSizeSoftLimit(2),
250            )
251            .unwrap();
252
253        assert_eq!(limited, vec![Some(bal0), Some(bal1)]);
254    }
255
256    #[test]
257    fn default_retention_prunes_old_bals() {
258        let store = InMemoryBalStore::default();
259        let old_hash = B256::random();
260        let retained_hash = B256::random();
261        let tip_hash = B256::random();
262        let old_bal = Bytes::from_static(b"old");
263        let retained_bal = Bytes::from_static(b"retained");
264        let tip_bal = Bytes::from_static(b"tip");
265
266        store.insert(NumHash::new(1, old_hash), RawBal::from(old_bal)).unwrap();
267        store
268            .insert(
269                NumHash::new(BAL_RETENTION_PERIOD_SLOTS, retained_hash),
270                RawBal::from(retained_bal.clone()),
271            )
272            .unwrap();
273        store
274            .insert(
275                NumHash::new(BAL_RETENTION_PERIOD_SLOTS + 2, tip_hash),
276                RawBal::from(tip_bal.clone()),
277            )
278            .unwrap();
279
280        assert_eq!(
281            store.get_by_hashes(&[old_hash, retained_hash, tip_hash]).unwrap(),
282            vec![None, Some(retained_bal), Some(tip_bal)]
283        );
284    }
285
286    #[test]
287    fn prune_uses_chain_tip() {
288        let store =
289            InMemoryBalStore::new(BalConfig::with_in_memory_retention(PruneMode::Distance(2)));
290        let old_hash = B256::random();
291        let retained_hash = B256::random();
292        let old_bal = Bytes::from_static(b"old");
293        let retained_bal = Bytes::from_static(b"retained");
294
295        store.insert(NumHash::new(7, old_hash), RawBal::from(old_bal)).unwrap();
296        store.insert(NumHash::new(8, retained_hash), RawBal::from(retained_bal.clone())).unwrap();
297
298        assert_eq!(store.prune(10).unwrap(), 1);
299        assert_eq!(
300            store.get_by_hashes(&[old_hash, retained_hash]).unwrap(),
301            vec![None, Some(retained_bal)]
302        );
303    }
304
305    #[test]
306    fn insert_prunes_from_highest_inserted_block() {
307        let store =
308            InMemoryBalStore::new(BalConfig::with_in_memory_retention(PruneMode::Distance(2)));
309        let old_hash = B256::random();
310        let high_hash = B256::random();
311        let late_hash = B256::random();
312        let high_bal = Bytes::from_static(b"high");
313        let late_bal = Bytes::from_static(b"late");
314
315        store.insert(NumHash::new(7, old_hash), RawBal::from(Bytes::from_static(b"old"))).unwrap();
316        store.insert(NumHash::new(10, high_hash), RawBal::from(high_bal.clone())).unwrap();
317        store.insert(NumHash::new(8, late_hash), RawBal::from(late_bal.clone())).unwrap();
318
319        assert_eq!(
320            store.get_by_hashes(&[old_hash, high_hash, late_hash]).unwrap(),
321            vec![None, Some(high_bal), Some(late_bal)]
322        );
323    }
324
325    #[test]
326    fn unbounded_retention_keeps_old_bals() {
327        let store = InMemoryBalStore::new(BalConfig::unbounded());
328        let old_hash = B256::random();
329        let tip_hash = B256::random();
330        let old_bal = Bytes::from_static(b"old");
331        let tip_bal = Bytes::from_static(b"tip");
332
333        store.insert(NumHash::new(1, old_hash), RawBal::from(old_bal.clone())).unwrap();
334        store
335            .insert(
336                NumHash::new(BAL_RETENTION_PERIOD_SLOTS + 1, tip_hash),
337                RawBal::from(tip_bal.clone()),
338            )
339            .unwrap();
340
341        assert_eq!(
342            store.get_by_hashes(&[old_hash, tip_hash]).unwrap(),
343            vec![Some(old_bal), Some(tip_bal)]
344        );
345        assert_eq!(store.prune(BAL_RETENTION_PERIOD_SLOTS + 2).unwrap(), 0);
346    }
347
348    #[test]
349    fn in_memory_retention_distance_prunes_old_bals() {
350        let store = InMemoryBalStore::new(BalConfig::with_in_memory_retention_distance(2));
351        let old_hash = B256::random();
352        let retained_hash = B256::random();
353        let tip_hash = B256::random();
354        let old_bal = Bytes::from_static(b"old");
355        let retained_bal = Bytes::from_static(b"retained");
356        let tip_bal = Bytes::from_static(b"tip");
357
358        store.insert(NumHash::new(1, old_hash), RawBal::from(old_bal)).unwrap();
359        store.insert(NumHash::new(2, retained_hash), RawBal::from(retained_bal.clone())).unwrap();
360        store.insert(NumHash::new(4, tip_hash), RawBal::from(tip_bal.clone())).unwrap();
361
362        assert_eq!(
363            store.get_by_hashes(&[old_hash, retained_hash, tip_hash]).unwrap(),
364            vec![None, Some(retained_bal), Some(tip_bal)]
365        );
366    }
367
368    #[test]
369    fn reinserting_hash_updates_number_index() {
370        let store =
371            InMemoryBalStore::new(BalConfig::with_in_memory_retention(PruneMode::Before(2)));
372        let hash = B256::random();
373        let bal = Bytes::from_static(b"bal");
374
375        store.insert(NumHash::new(1, hash), RawBal::from(Bytes::from_static(b"old"))).unwrap();
376        store.insert(NumHash::new(2, hash), RawBal::from(bal.clone())).unwrap();
377
378        assert_eq!(store.get_by_hashes(&[hash]).unwrap(), vec![Some(bal)]);
379    }
380}
reth_provider/bal.rs

reth_provider/
bal.rs
