reth_storage_api/
database_provider.rs

1use alloc::vec::Vec;
2use core::ops::{Bound, RangeBounds};
3use reth_db_api::{
4    common::KeyValue,
5    cursor::DbCursorRO,
6    database::Database,
7    table::Table,
8    transaction::{DbTx, DbTxMut},
9    DatabaseError,
10};
11use reth_prune_types::PruneModes;
12use reth_storage_errors::provider::ProviderResult;
13
14/// Database provider.
15pub trait DBProvider: Sized {
16    /// Underlying database transaction held by the provider.
17    type Tx: DbTx;
18
19    /// Returns a reference to the underlying transaction.
20    fn tx_ref(&self) -> &Self::Tx;
21
22    /// Returns a mutable reference to the underlying transaction.
23    fn tx_mut(&mut self) -> &mut Self::Tx;
24
25    /// Consumes the provider and returns the underlying transaction.
26    fn into_tx(self) -> Self::Tx;
27
28    /// Disables long-lived read transaction safety guarantees for leaks prevention and
29    /// observability improvements.
30    ///
31    /// CAUTION: In most of the cases, you want the safety guarantees for long read transactions
32    /// enabled. Use this only if you're sure that no write transaction is open in parallel, meaning
33    /// that Reth as a node is offline and not progressing.
34    fn disable_long_read_transaction_safety(mut self) -> Self {
35        self.tx_mut().disable_long_read_transaction_safety();
36        self
37    }
38
39    /// Commit database transaction
40    fn commit(self) -> ProviderResult<bool>;
41
42    /// Returns a reference to prune modes.
43    fn prune_modes_ref(&self) -> &PruneModes;
44
45    /// Return full table as Vec
46    fn table<T: Table>(&self) -> Result<Vec<KeyValue<T>>, DatabaseError>
47    where
48        T::Key: Default + Ord,
49    {
50        self.tx_ref()
51            .cursor_read::<T>()?
52            .walk(Some(T::Key::default()))?
53            .collect::<Result<Vec<_>, DatabaseError>>()
54    }
55
56    /// Return a list of entries from the table, based on the given range.
57    #[inline]
58    fn get<T: Table>(
59        &self,
60        range: impl RangeBounds<T::Key>,
61    ) -> Result<Vec<KeyValue<T>>, DatabaseError> {
62        self.tx_ref().cursor_read::<T>()?.walk_range(range)?.collect::<Result<Vec<_>, _>>()
63    }
64
65    /// Iterates over read only values in the given table and collects them into a vector.
66    ///
67    /// Early-returns if the range is empty, without opening a cursor transaction.
68    fn cursor_read_collect<T: Table<Key = u64>>(
69        &self,
70        range: impl RangeBounds<T::Key>,
71    ) -> ProviderResult<Vec<T::Value>> {
72        let capacity = match range_size_hint(&range) {
73            Some(0) | None => return Ok(Vec::new()),
74            Some(capacity) => capacity,
75        };
76        let mut cursor = self.tx_ref().cursor_read::<T>()?;
77        self.cursor_collect_with_capacity(&mut cursor, range, capacity)
78    }
79
80    /// Iterates over read only values in the given table and collects them into a vector.
81    fn cursor_collect<T: Table<Key = u64>>(
82        &self,
83        cursor: &mut impl DbCursorRO<T>,
84        range: impl RangeBounds<T::Key>,
85    ) -> ProviderResult<Vec<T::Value>> {
86        let capacity = range_size_hint(&range).unwrap_or(0);
87        self.cursor_collect_with_capacity(cursor, range, capacity)
88    }
89
90    /// Iterates over read only values in the given table and collects them into a vector with
91    /// capacity.
92    fn cursor_collect_with_capacity<T: Table<Key = u64>>(
93        &self,
94        cursor: &mut impl DbCursorRO<T>,
95        range: impl RangeBounds<T::Key>,
96        capacity: usize,
97    ) -> ProviderResult<Vec<T::Value>> {
98        let mut items = Vec::with_capacity(capacity);
99        for entry in cursor.walk_range(range)? {
100            items.push(entry?.1);
101        }
102        Ok(items)
103    }
104
105    /// Remove list of entries from the table. Returns the number of entries removed.
106    #[inline]
107    fn remove<T: Table>(&self, range: impl RangeBounds<T::Key>) -> Result<usize, DatabaseError>
108    where
109        Self::Tx: DbTxMut,
110    {
111        let mut entries = 0;
112        let mut cursor_write = self.tx_ref().cursor_write::<T>()?;
113        let mut walker = cursor_write.walk_range(range)?;
114        while walker.next().transpose()?.is_some() {
115            walker.delete_current()?;
116            entries += 1;
117        }
118        Ok(entries)
119    }
120
121    /// Return a list of entries from the table, and remove them, based on the given range.
122    #[inline]
123    fn take<T: Table>(
124        &self,
125        range: impl RangeBounds<T::Key>,
126    ) -> Result<Vec<KeyValue<T>>, DatabaseError>
127    where
128        Self::Tx: DbTxMut,
129    {
130        let mut cursor_write = self.tx_ref().cursor_write::<T>()?;
131        let mut walker = cursor_write.walk_range(range)?;
132        let mut items = Vec::new();
133        while let Some(i) = walker.next().transpose()? {
134            walker.delete_current()?;
135            items.push(i)
136        }
137        Ok(items)
138    }
139}
140
141/// Database provider factory.
142#[auto_impl::auto_impl(&, Arc)]
143pub trait DatabaseProviderFactory: Send + Sync {
144    /// Database this factory produces providers for.
145    type DB: Database;
146
147    /// Provider type returned by the factory.
148    type Provider: DBProvider<Tx = <Self::DB as Database>::TX>;
149
150    /// Read-write provider type returned by the factory.
151    type ProviderRW: DBProvider<Tx = <Self::DB as Database>::TXMut>;
152
153    /// Create new read-only database provider.
154    fn database_provider_ro(&self) -> ProviderResult<Self::Provider>;
155
156    /// Create new read-write database provider.
157    fn database_provider_rw(&self) -> ProviderResult<Self::ProviderRW>;
158}
159
160/// Helper type alias to get the associated transaction type from a [`DatabaseProviderFactory`].
161pub type FactoryTx<F> = <<F as DatabaseProviderFactory>::DB as Database>::TX;
162
163fn range_size_hint(range: &impl RangeBounds<u64>) -> Option<usize> {
164    let start = match range.start_bound().cloned() {
165        Bound::Included(start) => start,
166        Bound::Excluded(start) => start.checked_add(1)?,
167        Bound::Unbounded => 0,
168    };
169    let end = match range.end_bound().cloned() {
170        Bound::Included(end) => end.saturating_add(1),
171        Bound::Excluded(end) => end,
172        Bound::Unbounded => return None,
173    };
174    end.checked_sub(start).map(|x| x as _)
175}