Skip to main content

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