reth_provider/
bal.rs

1use alloy_eip7928::BAL_RETENTION_PERIOD_SLOTS;
2use alloy_eips::NumHash;
3use alloy_primitives::{BlockHash, BlockNumber, Bytes};
4use parking_lot::RwLock;
5use reth_prune_types::PruneMode;
6use reth_storage_api::{
7    BalNotification, BalNotificationStream, BalStore, GetBlockAccessListLimit, RawBal,
8};
9use reth_storage_errors::provider::ProviderResult;
10use reth_tokio_util::EventSender;
11use std::{
12    collections::{BTreeMap, HashMap},
13    sync::Arc,
14};
15
16mod rocksdb;
17pub use rocksdb::RocksDBBalStore;
18
19/// Basic in-memory BAL store keyed by block hash.
20#[derive(Debug, Clone)]
21pub struct InMemoryBalStore {
22    config: BalConfig,
23    inner: Arc<RwLock<InMemoryBalStoreInner>>,
24    notifications: EventSender<BalNotification>,
25}
26
27impl InMemoryBalStore {
28    /// Creates a new in-memory BAL store with the given config.
29    pub fn new(config: BalConfig) -> Self {
30        let notifications = EventSender::new(DEFAULT_BAL_NOTIFICATION_CHANNEL_SIZE);
31        Self {
32            config,
33            inner: Arc::new(RwLock::new(InMemoryBalStoreInner::default())),
34            notifications,
35        }
36    }
37}
38
39// Match the canonical state broadcast buffer so BAL subscriptions behave like the existing
40// in-memory notification path. This is a bounded best-effort channel, not a durability boundary.
41const DEFAULT_BAL_NOTIFICATION_CHANNEL_SIZE: usize = 256;
42
43impl Default for InMemoryBalStore {
44    fn default() -> Self {
45        Self::new(BalConfig::default())
46    }
47}
48
49/// Configuration for BAL storage.
50#[derive(Debug, Clone, Copy, Eq, PartialEq)]
51pub struct BalConfig {
52    /// Retention policy for BALs kept in memory.
53    in_memory_retention: Option<PruneMode>,
54}
55
56impl BalConfig {
57    /// Default block distance for BALs kept in memory.
58    pub const DEFAULT_IN_MEMORY_RETENTION_DISTANCE: u64 = BAL_RETENTION_PERIOD_SLOTS;
59
60    /// Returns a config with no in-memory BAL retention limit.
61    pub const fn unbounded() -> Self {
62        Self { in_memory_retention: None }
63    }
64
65    /// Returns a config that keeps BALs within the given block distance in memory.
66    pub const fn with_in_memory_retention_distance(blocks: u64) -> Self {
67        Self::with_in_memory_retention(PruneMode::Distance(blocks))
68    }
69
70    /// Returns a config with the given in-memory BAL retention policy.
71    pub const fn with_in_memory_retention(in_memory_retention: PruneMode) -> Self {
72        Self { in_memory_retention: Some(in_memory_retention) }
73    }
74}
75
76impl Default for BalConfig {
77    fn default() -> Self {
78        Self::with_in_memory_retention_distance(Self::DEFAULT_IN_MEMORY_RETENTION_DISTANCE)
79    }
80}
81
82#[derive(Debug, Default)]
83struct InMemoryBalStoreInner {
84    entries: HashMap<BlockHash, BalEntry>,
85    hashes_by_number: BTreeMap<BlockNumber, Vec<BlockHash>>,
86    highest_block_number: Option<BlockNumber>,
87}
88
89impl InMemoryBalStoreInner {
90    // Inserts a BAL and keeps the block-number index in sync.
91    fn insert(&mut self, block_hash: BlockHash, block_number: BlockNumber, bal: Bytes) {
92        let empty_block_number =
93            self.entries.insert(block_hash, BalEntry { block_number, bal }).and_then(|entry| {
94                let hashes = self.hashes_by_number.get_mut(&entry.block_number)?;
95                hashes.retain(|hash| *hash != block_hash);
96                hashes.is_empty().then_some(entry.block_number)
97            });
98
99        if let Some(block_number) = empty_block_number {
100            self.hashes_by_number.remove(&block_number);
101        }
102
103        self.hashes_by_number.entry(block_number).or_default().push(block_hash);
104        self.highest_block_number = Some(
105            self.highest_block_number.map_or(block_number, |highest| highest.max(block_number)),
106        );
107    }
108
109    // Removes BALs outside the configured retention window for the given chain tip.
110    fn prune(&mut self, prune_mode: Option<PruneMode>, tip: BlockNumber) -> usize {
111        let Some(prune_mode) = prune_mode else { return 0 };
112
113        let mut pruned = 0;
114        while let Some((&block_number, _)) = self.hashes_by_number.first_key_value() {
115            if !prune_mode.should_prune(block_number, tip) {
116                break
117            }
118
119            let Some((_, hashes)) = self.hashes_by_number.pop_first() else { break };
120            for hash in hashes {
121                pruned += usize::from(self.entries.remove(&hash).is_some());
122            }
123        }
124        pruned
125    }
126}
127
128#[derive(Debug)]
129struct BalEntry {
130    block_number: BlockNumber,
131    bal: Bytes,
132}
133
134impl BalStore for InMemoryBalStore {
135    fn insert(&self, num_hash: NumHash, bal: RawBal) -> ProviderResult<()> {
136        let mut inner = self.inner.write();
137        inner.insert(num_hash.hash, num_hash.number, bal.as_raw().clone());
138        if let Some(highest_block_number) = inner.highest_block_number {
139            // This preserves insert-time cleanup based on the highest inserted BAL block.
140            inner.prune(self.config.in_memory_retention, highest_block_number);
141        }
142        self.notifications.notify(BalNotification::new(num_hash, bal));
143        Ok(())
144    }
145
146    fn insert_many(&self, entries: Vec<(NumHash, RawBal)>) -> ProviderResult<()> {
147        if entries.is_empty() {
148            return Ok(())
149        }
150
151        let mut inner = self.inner.write();
152        inner.entries.reserve(entries.len());
153        for (num_hash, bal) in &entries {
154            inner.insert(num_hash.hash, num_hash.number, bal.as_raw().clone());
155        }
156        if let Some(highest_block_number) = inner.highest_block_number {
157            inner.prune(self.config.in_memory_retention, highest_block_number);
158        }
159        drop(inner);
160
161        for (num_hash, bal) in entries {
162            self.notifications.notify(BalNotification::new(num_hash, bal));
163        }
164        Ok(())
165    }
166
167    fn flush(&self, _blocks: &[NumHash]) -> ProviderResult<()> {
168        Ok(())
169    }
170
171    fn prune(&self, tip: BlockNumber) -> ProviderResult<usize> {
172        Ok(self.inner.write().prune(self.config.in_memory_retention, tip))
173    }
174
175    fn get_by_hashes(&self, block_hashes: &[BlockHash]) -> ProviderResult<Vec<Option<Bytes>>> {
176        let inner = self.inner.read();
177        let mut result = Vec::with_capacity(block_hashes.len());
178
179        for hash in block_hashes {
180            result.push(inner.entries.get(hash).map(|entry| entry.bal.clone()));
181        }
182
183        Ok(result)
184    }
185
186    fn append_by_hashes_with_limit(
187        &self,
188        block_hashes: &[BlockHash],
189        limit: GetBlockAccessListLimit,
190        out: &mut Vec<Option<Bytes>>,
191    ) -> ProviderResult<()> {
192        let inner = self.inner.read();
193        let mut size = 0;
194
195        for hash in block_hashes {
196            let bal = inner.entries.get(hash).map(|entry| entry.bal.clone());
197            size += bal.as_ref().map_or(1, |bytes| bytes.len());
198            out.push(bal);
199
200            if limit.exceeds(size) {
201                break
202            }
203        }
204
205        Ok(())
206    }
207
208    fn bal_stream(&self) -> BalNotificationStream {
209        self.notifications.new_listener()
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use alloy_primitives::B256;
217    use tokio_stream::StreamExt;
218
219    #[test]
220    fn insert_and_lookup_by_hash() {
221        let store = InMemoryBalStore::default();
222        let hash = B256::random();
223        let missing = B256::random();
224        let bal = Bytes::from_static(b"bal");
225
226        store.insert(NumHash::new(1, hash), RawBal::from(bal.clone())).unwrap();
227
228        assert_eq!(store.get_by_hashes(&[hash, missing]).unwrap(), vec![Some(bal), None]);
229    }
230
231    #[test]
232    fn insert_many_and_lookup_by_hash() {
233        let store = InMemoryBalStore::default();
234        let hash0 = B256::random();
235        let hash1 = B256::random();
236        let bal0 = RawBal::from(Bytes::from_static(b"bal0"));
237        let bal1 = RawBal::from(Bytes::from_static(b"bal1"));
238
239        store
240            .insert_many(vec![
241                (NumHash::new(1, hash0), bal0.clone()),
242                (NumHash::new(2, hash1), bal1),
243            ])
244            .unwrap();
245
246        assert_eq!(
247            store.get_by_hashes(&[hash0, hash1]).unwrap(),
248            vec![Some(bal0.as_raw().clone()), Some(Bytes::from_static(b"bal1"))]
249        );
250    }
251
252    #[test]
253    fn flush_is_noop() {
254        let store = InMemoryBalStore::default();
255
256        store.flush(&[]).unwrap();
257    }
258
259    #[test]
260    fn limited_lookup_returns_prefix() {
261        let store = InMemoryBalStore::default();
262        let hash0 = B256::random();
263        let hash1 = B256::random();
264        let hash2 = B256::random();
265        let bal0 = Bytes::from_static(&[0xc1, 0x01]);
266        let bal1 = Bytes::from_static(&[0xc1, 0x02]);
267        let bal2 = Bytes::from_static(&[0xc1, 0x03]);
268
269        store.insert(NumHash::new(1, hash0), RawBal::from(bal0.clone())).unwrap();
270        store.insert(NumHash::new(2, hash1), RawBal::from(bal1.clone())).unwrap();
271        store.insert(NumHash::new(3, hash2), RawBal::from(bal2)).unwrap();
272
273        let limited = store
274            .get_by_hashes_with_limit(
275                &[hash0, hash1, hash2],
276                GetBlockAccessListLimit::ResponseSizeSoftLimit(2),
277            )
278            .unwrap();
279
280        assert_eq!(limited, vec![Some(bal0), Some(bal1)]);
281    }
282
283    #[test]
284    fn default_retention_prunes_old_bals() {
285        let store = InMemoryBalStore::default();
286        let old_hash = B256::random();
287        let retained_hash = B256::random();
288        let tip_hash = B256::random();
289        let old_bal = Bytes::from_static(b"old");
290        let retained_bal = Bytes::from_static(b"retained");
291        let tip_bal = Bytes::from_static(b"tip");
292
293        store.insert(NumHash::new(1, old_hash), RawBal::from(old_bal)).unwrap();
294        store
295            .insert(
296                NumHash::new(BAL_RETENTION_PERIOD_SLOTS, retained_hash),
297                RawBal::from(retained_bal.clone()),
298            )
299            .unwrap();
300        store
301            .insert(
302                NumHash::new(BAL_RETENTION_PERIOD_SLOTS + 2, tip_hash),
303                RawBal::from(tip_bal.clone()),
304            )
305            .unwrap();
306
307        assert_eq!(
308            store.get_by_hashes(&[old_hash, retained_hash, tip_hash]).unwrap(),
309            vec![None, Some(retained_bal), Some(tip_bal)]
310        );
311    }
312
313    #[test]
314    fn prune_uses_chain_tip() {
315        let store =
316            InMemoryBalStore::new(BalConfig::with_in_memory_retention(PruneMode::Distance(2)));
317        let old_hash = B256::random();
318        let retained_hash = B256::random();
319        let old_bal = Bytes::from_static(b"old");
320        let retained_bal = Bytes::from_static(b"retained");
321
322        store.insert(NumHash::new(7, old_hash), RawBal::from(old_bal)).unwrap();
323        store.insert(NumHash::new(8, retained_hash), RawBal::from(retained_bal.clone())).unwrap();
324
325        assert_eq!(store.prune(10).unwrap(), 1);
326        assert_eq!(
327            store.get_by_hashes(&[old_hash, retained_hash]).unwrap(),
328            vec![None, Some(retained_bal)]
329        );
330    }
331
332    #[test]
333    fn insert_prunes_from_highest_inserted_block() {
334        let store =
335            InMemoryBalStore::new(BalConfig::with_in_memory_retention(PruneMode::Distance(2)));
336        let old_hash = B256::random();
337        let high_hash = B256::random();
338        let late_hash = B256::random();
339        let high_bal = Bytes::from_static(b"high");
340        let late_bal = Bytes::from_static(b"late");
341
342        store.insert(NumHash::new(7, old_hash), RawBal::from(Bytes::from_static(b"old"))).unwrap();
343        store.insert(NumHash::new(10, high_hash), RawBal::from(high_bal.clone())).unwrap();
344        store.insert(NumHash::new(8, late_hash), RawBal::from(late_bal.clone())).unwrap();
345
346        assert_eq!(
347            store.get_by_hashes(&[old_hash, high_hash, late_hash]).unwrap(),
348            vec![None, Some(high_bal), Some(late_bal)]
349        );
350    }
351
352    #[test]
353    fn unbounded_retention_keeps_old_bals() {
354        let store = InMemoryBalStore::new(BalConfig::unbounded());
355        let old_hash = B256::random();
356        let tip_hash = B256::random();
357        let old_bal = Bytes::from_static(b"old");
358        let tip_bal = Bytes::from_static(b"tip");
359
360        store.insert(NumHash::new(1, old_hash), RawBal::from(old_bal.clone())).unwrap();
361        store
362            .insert(
363                NumHash::new(BAL_RETENTION_PERIOD_SLOTS + 1, tip_hash),
364                RawBal::from(tip_bal.clone()),
365            )
366            .unwrap();
367
368        assert_eq!(
369            store.get_by_hashes(&[old_hash, tip_hash]).unwrap(),
370            vec![Some(old_bal), Some(tip_bal)]
371        );
372        assert_eq!(store.prune(BAL_RETENTION_PERIOD_SLOTS + 2).unwrap(), 0);
373    }
374
375    #[test]
376    fn in_memory_retention_distance_prunes_old_bals() {
377        let store = InMemoryBalStore::new(BalConfig::with_in_memory_retention_distance(2));
378        let old_hash = B256::random();
379        let retained_hash = B256::random();
380        let tip_hash = B256::random();
381        let old_bal = Bytes::from_static(b"old");
382        let retained_bal = Bytes::from_static(b"retained");
383        let tip_bal = Bytes::from_static(b"tip");
384
385        store.insert(NumHash::new(1, old_hash), RawBal::from(old_bal)).unwrap();
386        store.insert(NumHash::new(2, retained_hash), RawBal::from(retained_bal.clone())).unwrap();
387        store.insert(NumHash::new(4, tip_hash), RawBal::from(tip_bal.clone())).unwrap();
388
389        assert_eq!(
390            store.get_by_hashes(&[old_hash, retained_hash, tip_hash]).unwrap(),
391            vec![None, Some(retained_bal), Some(tip_bal)]
392        );
393    }
394
395    #[test]
396    fn reinserting_hash_updates_number_index() {
397        let store =
398            InMemoryBalStore::new(BalConfig::with_in_memory_retention(PruneMode::Before(2)));
399        let hash = B256::random();
400        let bal = Bytes::from_static(b"bal");
401
402        store.insert(NumHash::new(1, hash), RawBal::from(Bytes::from_static(b"old"))).unwrap();
403        store.insert(NumHash::new(2, hash), RawBal::from(bal.clone())).unwrap();
404
405        assert_eq!(store.get_by_hashes(&[hash]).unwrap(), vec![Some(bal)]);
406    }
407
408    #[tokio::test]
409    async fn insert_notifies_subscribers() {
410        let store = InMemoryBalStore::default();
411        let hash = B256::random();
412        let block_number = 7;
413        let bal = Bytes::from_static(b"bal");
414        let mut stream = store.bal_stream();
415
416        let raw_bal = RawBal::from(bal);
417
418        store.insert(NumHash::new(block_number, hash), raw_bal.clone()).unwrap();
419
420        assert_eq!(
421            stream.next().await.unwrap(),
422            BalNotification::new(NumHash::new(block_number, hash), raw_bal)
423        );
424    }
425
426    #[tokio::test]
427    async fn insert_many_notifies_subscribers() {
428        let store = InMemoryBalStore::default();
429        let mut stream = store.bal_stream();
430        let hash0 = B256::random();
431        let hash1 = B256::random();
432        let bal0 = RawBal::from(Bytes::from_static(b"bal0"));
433        let bal1 = RawBal::from(Bytes::from_static(b"bal1"));
434
435        store
436            .insert_many(vec![
437                (NumHash::new(1, hash0), bal0.clone()),
438                (NumHash::new(2, hash1), bal1.clone()),
439            ])
440            .unwrap();
441
442        assert_eq!(
443            stream.next().await.unwrap(),
444            BalNotification::new(NumHash::new(1, hash0), bal0)
445        );
446        assert_eq!(
447            stream.next().await.unwrap(),
448            BalNotification::new(NumHash::new(2, hash1), bal1)
449        );
450    }
451
452    #[test]
453    fn insert_without_subscribers_still_succeeds() {
454        let store = InMemoryBalStore::default();
455
456        assert!(store
457            .insert(NumHash::new(1, B256::random()), RawBal::from(Bytes::from_static(b"bal")))
458            .is_ok());
459    }
460
461    #[tokio::test]
462    async fn bal_stream_skips_lagged_notifications() {
463        let store = InMemoryBalStore::new(BalConfig::unbounded());
464        let mut stream = store.bal_stream();
465
466        for number in 0..=DEFAULT_BAL_NOTIFICATION_CHANNEL_SIZE as u64 {
467            store
468                .insert(
469                    NumHash::new(number, B256::random()),
470                    RawBal::from(Bytes::from(vec![number as u8])),
471                )
472                .unwrap();
473        }
474
475        let first = stream.next().await.unwrap();
476        let second = stream.next().await.unwrap();
477
478        assert_eq!(first.num_hash.number, 1);
479        assert_eq!(second.num_hash.number, 2);
480    }
481
482    #[tokio::test]
483    async fn cloned_store_shares_notification_channel() {
484        let store = InMemoryBalStore::default();
485        let clone = store.clone();
486        let hash = B256::random();
487        let block_number = 9;
488        let bal = Bytes::from_static(b"bal");
489        let mut stream = clone.bal_stream();
490
491        let raw_bal = RawBal::from(bal);
492
493        store.insert(NumHash::new(block_number, hash), raw_bal.clone()).unwrap();
494
495        assert_eq!(
496            stream.next().await.unwrap(),
497            BalNotification::new(NumHash::new(block_number, hash), raw_bal)
498        );
499    }
500}
reth_provider/bal.rs

reth_provider/
bal.rs
