reth_db/implementation/mdbx/
cursor.rs

1//! Cursor wrapper for libmdbx-sys.
2
3use super::utils::*;
4use crate::{
5    metrics::{Operation, TableOperationMetrics},
6    DatabaseError,
7};
8use reth_db_api::{
9    common::{PairResult, ValueOnlyResult},
10    cursor::{
11        DbCursorRO, DbCursorRW, DbDupCursorRO, DbDupCursorRW, DupWalker, RangeWalker,
12        ReverseWalker, Walker,
13    },
14    table::{Compress, Decode, Decompress, DupSort, Encode, IntoVec, Table},
15};
16use reth_libmdbx::{Error as MDBXError, TransactionKind, WriteFlags, RO, RW};
17use reth_storage_errors::db::{DatabaseErrorInfo, DatabaseWriteError, DatabaseWriteOperation};
18use std::{borrow::Cow, collections::Bound, marker::PhantomData, ops::RangeBounds};
19
20/// Read only Cursor.
21pub type CursorRO<T> = Cursor<RO, T>;
22/// Read write cursor.
23pub type CursorRW<T> = Cursor<RW, T>;
24
25/// Cursor wrapper to access KV items.
26#[derive(Debug)]
27pub struct Cursor<K: TransactionKind, T: Table> {
28    /// Inner `libmdbx` cursor.
29    pub(crate) inner: reth_libmdbx::Cursor<K>,
30    /// Cache buffer that receives compressed values.
31    buf: Vec<u8>,
32    /// Per-table operation metrics. If `None`, metrics are not recorded.
33    metrics: Option<TableOperationMetrics>,
34    /// Phantom data to enforce encoding/decoding.
35    _dbi: PhantomData<T>,
36    /// Whether all value reads through this cursor should request read-ahead.
37    #[cfg(target_os = "linux")]
38    prefetch: bool,
39}
40
41impl<K: TransactionKind, T: Table> Cursor<K, T> {
42    pub(crate) const fn new_with_metrics(
43        inner: reth_libmdbx::Cursor<K>,
44        metrics: Option<TableOperationMetrics>,
45        #[cfg(target_os = "linux")] prefetch: bool,
46    ) -> Self {
47        Self {
48            inner,
49            buf: Vec::new(),
50            metrics,
51            _dbi: PhantomData,
52            #[cfg(target_os = "linux")]
53            prefetch,
54        }
55    }
56
57    /// If `self.metrics` is `Some(...)`, record a metric with the provided operation and value
58    /// size.
59    ///
60    /// Otherwise, just execute the closure.
61    fn execute_with_operation_metric<R>(
62        &mut self,
63        operation: Operation,
64        value_size: Option<usize>,
65        f: impl FnOnce(&mut Self) -> R,
66    ) -> R {
67        if let Some(metrics) = self.metrics.clone() {
68            metrics[operation.index()].record(value_size, || f(self))
69        } else {
70            f(self)
71        }
72    }
73}
74
75/// Decodes a `(key, value)` pair from the database.
76#[expect(clippy::type_complexity)]
77pub fn decode<T>(
78    res: Result<Option<(Cow<'_, [u8]>, Cow<'_, [u8]>)>, impl Into<DatabaseErrorInfo>>,
79) -> PairResult<T>
80where
81    T: Table,
82    T::Key: Decode,
83    T::Value: Decompress,
84{
85    res.map_err(|e| DatabaseError::Read(e.into()))?.map(decoder::<T>).transpose()
86}
87
88/// Some types don't support compression (eg. B256), and we don't want to be copying them to the
89/// allocated buffer when we can just use their reference.
90macro_rules! compress_to_buf_or_ref {
91    ($self:expr, $value:expr) => {
92        if let Some(value) = $value.uncompressable_ref() {
93            Some(value)
94        } else {
95            $self.buf.clear();
96            $value.compress_to_buf(&mut $self.buf);
97            None
98        }
99    };
100}
101
102/// Selects the value decoder inside the existing MDBX cursor lookup.
103macro_rules! read_pair {
104    ($self:ident, $method:ident($($arg:expr),* $(,)?)) => {{
105        #[cfg(target_os = "linux")]
106        if $self.prefetch {
107            $self.inner.$method::<Cow<'_, [u8]>, super::value_prefetch::PrefetchValue<'_>>($($arg),*)
108                .map(|entry| entry.map(|(key, value)| (key, value.0)))
109        } else {
110            $self.inner.$method($($arg),*)
111        }
112        #[cfg(not(target_os = "linux"))]
113        $self.inner.$method($($arg),*)
114    }};
115}
116
117/// Selects the decoder for cursor operations returning only a value.
118macro_rules! read_value {
119    ($self:ident, $method:ident($($arg:expr),* $(,)?)) => {{
120        #[cfg(target_os = "linux")]
121        if $self.prefetch {
122            $self.inner.$method::<super::value_prefetch::PrefetchValue<'_>>($($arg),*)
123                .map(|entry| entry.map(|value| value.0))
124        } else {
125            $self.inner.$method($($arg),*)
126        }
127        #[cfg(not(target_os = "linux"))]
128        $self.inner.$method($($arg),*)
129    }};
130}
131
132impl<K: TransactionKind, T: Table> DbCursorRO<T> for Cursor<K, T> {
133    fn first(&mut self) -> PairResult<T> {
134        decode::<T>(read_pair!(self, first()))
135    }
136
137    fn seek_exact(&mut self, key: <T as Table>::Key) -> PairResult<T> {
138        decode::<T>(read_pair!(self, set_key(key.encode().as_ref())))
139    }
140
141    fn seek(&mut self, key: <T as Table>::Key) -> PairResult<T> {
142        decode::<T>(read_pair!(self, set_range(key.encode().as_ref())))
143    }
144
145    fn next(&mut self) -> PairResult<T> {
146        decode::<T>(read_pair!(self, next()))
147    }
148
149    fn prev(&mut self) -> PairResult<T> {
150        decode::<T>(read_pair!(self, prev()))
151    }
152
153    fn last(&mut self) -> PairResult<T> {
154        decode::<T>(read_pair!(self, last()))
155    }
156
157    fn current(&mut self) -> PairResult<T> {
158        decode::<T>(read_pair!(self, get_current()))
159    }
160
161    fn walk(&mut self, start_key: Option<T::Key>) -> Result<Walker<'_, T, Self>, DatabaseError> {
162        let start = if let Some(start_key) = start_key {
163            self.seek(start_key).transpose()
164        } else {
165            self.first().transpose()
166        };
167
168        Ok(Walker::new(self, start))
169    }
170
171    fn walk_range(
172        &mut self,
173        range: impl RangeBounds<T::Key>,
174    ) -> Result<RangeWalker<'_, T, Self>, DatabaseError> {
175        let start = match range.start_bound().cloned() {
176            Bound::Included(key) => self.seek(key),
177            Bound::Excluded(_key) => {
178                unreachable!("Rust doesn't allow for Bound::Excluded in starting bounds");
179            }
180            Bound::Unbounded => self.first(),
181        };
182        let start = start.transpose();
183        Ok(RangeWalker::new(self, start, range.end_bound().cloned()))
184    }
185
186    fn walk_back(
187        &mut self,
188        start_key: Option<T::Key>,
189    ) -> Result<ReverseWalker<'_, T, Self>, DatabaseError> {
190        let start =
191            if let Some(start_key) = start_key { self.seek(start_key) } else { self.last() }
192                .transpose();
193
194        Ok(ReverseWalker::new(self, start))
195    }
196
197    #[cfg(target_os = "linux")]
198    fn prefetch(&mut self, enabled: bool) -> &mut Self {
199        self.prefetch = enabled;
200        self
201    }
202}
203
204impl<K: TransactionKind, T: DupSort> DbDupCursorRO<T> for Cursor<K, T> {
205    /// Returns the previous `(key, value)` pair of a DUPSORT table.
206    fn prev_dup(&mut self) -> PairResult<T> {
207        decode::<T>(read_pair!(self, prev_dup()))
208    }
209
210    /// Returns the next `(key, value)` pair of a DUPSORT table.
211    fn next_dup(&mut self) -> PairResult<T> {
212        decode::<T>(read_pair!(self, next_dup()))
213    }
214
215    /// Returns the last `value` of the current duplicate `key`.
216    fn last_dup(&mut self) -> ValueOnlyResult<T> {
217        read_value!(self, last_dup())
218            .map_err(|e| DatabaseError::Read(e.into()))?
219            .map(decode_one::<T>)
220            .transpose()
221    }
222
223    /// Returns the next `(key, value)` pair skipping the duplicates.
224    fn next_no_dup(&mut self) -> PairResult<T> {
225        decode::<T>(read_pair!(self, next_nodup()))
226    }
227
228    /// Returns the next `value` of a duplicate `key`.
229    fn next_dup_val(&mut self) -> ValueOnlyResult<T> {
230        read_pair!(self, next_dup())
231            .map_err(|e| DatabaseError::Read(e.into()))?
232            .map(decode_value::<T>)
233            .transpose()
234    }
235
236    fn seek_by_key_subkey(
237        &mut self,
238        key: <T as Table>::Key,
239        subkey: <T as DupSort>::SubKey,
240    ) -> ValueOnlyResult<T> {
241        read_value!(self, get_both_range(key.encode().as_ref(), subkey.encode().as_ref()))
242            .map_err(|e| DatabaseError::Read(e.into()))?
243            .map(decode_one::<T>)
244            .transpose()
245    }
246
247    /// Depending on its arguments, returns an iterator starting at:
248    /// - Some(key), Some(subkey): a `key` item whose data is >= than `subkey`
249    /// - Some(key), None: first item of a specified `key`
250    /// - None, Some(subkey): like first case, but in the first key
251    /// - None, None: first item in the table of a DUPSORT table.
252    fn walk_dup(
253        &mut self,
254        key: Option<T::Key>,
255        subkey: Option<T::SubKey>,
256    ) -> Result<DupWalker<'_, T, Self>, DatabaseError> {
257        let start = match (key, subkey) {
258            (Some(key), Some(subkey)) => {
259                let encoded_key = key.encode();
260                read_value!(self, get_both_range(encoded_key.as_ref(), subkey.encode().as_ref()))
261                    .map_err(|e| DatabaseError::Read(e.into()))?
262                    .map(|val| decoder::<T>((Cow::Borrowed(encoded_key.as_ref()), val)))
263            }
264            (Some(key), None) => {
265                let encoded_key = key.encode();
266                read_value!(self, set(encoded_key.as_ref()))
267                    .map_err(|e| DatabaseError::Read(e.into()))?
268                    .map(|val| decoder::<T>((Cow::Borrowed(encoded_key.as_ref()), val)))
269            }
270            (None, Some(subkey)) => {
271                if let Some((key, _)) = self.first()? {
272                    let encoded_key = key.encode();
273                    read_value!(
274                        self,
275                        get_both_range(encoded_key.as_ref(), subkey.encode().as_ref())
276                    )
277                    .map_err(|e| DatabaseError::Read(e.into()))?
278                    .map(|val| decoder::<T>((Cow::Borrowed(encoded_key.as_ref()), val)))
279                } else {
280                    Some(Err(DatabaseError::Read(MDBXError::NotFound.into())))
281                }
282            }
283            (None, None) => self.first().transpose(),
284        };
285
286        Ok(DupWalker::<'_, T, Self> { cursor: self, start })
287    }
288}
289
290impl<T: Table> DbCursorRW<T> for Cursor<RW, T> {
291    /// Database operation that will update an existing row if a specified value already
292    /// exists in a table, and insert a new row if the specified value doesn't already exist
293    ///
294    /// For a DUPSORT table, `upsert` will not actually update-or-insert. If the key already exists,
295    /// it will append the value to the subkey, even if the subkeys are the same. So if you want
296    /// to properly upsert, you'll need to `seek_exact` & `delete_current` if the key+subkey was
297    /// found, before calling `upsert`.
298    fn upsert(&mut self, key: T::Key, value: &T::Value) -> Result<(), DatabaseError> {
299        let key = key.encode();
300        let value = compress_to_buf_or_ref!(self, value);
301        self.execute_with_operation_metric(
302            Operation::CursorUpsert,
303            Some(value.unwrap_or(&self.buf).len()),
304            |this| {
305                this.inner
306                    .put(key.as_ref(), value.unwrap_or(&this.buf), WriteFlags::UPSERT)
307                    .map_err(|e| {
308                        DatabaseWriteError {
309                            info: e.into(),
310                            operation: DatabaseWriteOperation::CursorUpsert,
311                            table_name: T::NAME,
312                            key: key.into_vec(),
313                        }
314                        .into()
315                    })
316            },
317        )
318    }
319
320    fn insert(&mut self, key: T::Key, value: &T::Value) -> Result<(), DatabaseError> {
321        let key = key.encode();
322        let value = compress_to_buf_or_ref!(self, value);
323        self.execute_with_operation_metric(
324            Operation::CursorInsert,
325            Some(value.unwrap_or(&self.buf).len()),
326            |this| {
327                this.inner
328                    .put(key.as_ref(), value.unwrap_or(&this.buf), WriteFlags::NO_OVERWRITE)
329                    .map_err(|e| {
330                        DatabaseWriteError {
331                            info: e.into(),
332                            operation: DatabaseWriteOperation::CursorInsert,
333                            table_name: T::NAME,
334                            key: key.into_vec(),
335                        }
336                        .into()
337                    })
338            },
339        )
340    }
341
342    /// Appends the data to the end of the table. Consequently, the append operation
343    /// will fail if the inserted key is less than the last table key
344    fn append(&mut self, key: T::Key, value: &T::Value) -> Result<(), DatabaseError> {
345        let key = key.encode();
346        let value = compress_to_buf_or_ref!(self, value);
347        self.execute_with_operation_metric(
348            Operation::CursorAppend,
349            Some(value.unwrap_or(&self.buf).len()),
350            |this| {
351                this.inner
352                    .put(key.as_ref(), value.unwrap_or(&this.buf), WriteFlags::APPEND)
353                    .map_err(|e| {
354                        DatabaseWriteError {
355                            info: e.into(),
356                            operation: DatabaseWriteOperation::CursorAppend,
357                            table_name: T::NAME,
358                            key: key.into_vec(),
359                        }
360                        .into()
361                    })
362            },
363        )
364    }
365
366    fn delete_current(&mut self) -> Result<(), DatabaseError> {
367        self.execute_with_operation_metric(Operation::CursorDeleteCurrent, None, |this| {
368            this.inner.del(WriteFlags::CURRENT).map_err(|e| DatabaseError::Delete(e.into()))
369        })
370    }
371}
372
373impl<T: DupSort> DbDupCursorRW<T> for Cursor<RW, T> {
374    fn delete_current_duplicates(&mut self) -> Result<(), DatabaseError> {
375        self.execute_with_operation_metric(Operation::CursorDeleteCurrentDuplicates, None, |this| {
376            this.inner.del(WriteFlags::NO_DUP_DATA).map_err(|e| DatabaseError::Delete(e.into()))
377        })
378    }
379
380    fn append_dup(&mut self, key: T::Key, value: T::Value) -> Result<(), DatabaseError> {
381        let key = key.encode();
382        let value = compress_to_buf_or_ref!(self, value);
383        self.execute_with_operation_metric(
384            Operation::CursorAppendDup,
385            Some(value.unwrap_or(&self.buf).len()),
386            |this| {
387                this.inner
388                    .put(key.as_ref(), value.unwrap_or(&this.buf), WriteFlags::APPEND_DUP)
389                    .map_err(|e| {
390                        DatabaseWriteError {
391                            info: e.into(),
392                            operation: DatabaseWriteOperation::CursorAppendDup,
393                            table_name: T::NAME,
394                            key: key.into_vec(),
395                        }
396                        .into()
397                    })
398            },
399        )
400    }
401}
402
403#[cfg(test)]
404mod tests {
405
406    use crate::{
407        mdbx::{DatabaseArguments, DatabaseEnv, DatabaseEnvKind},
408        tables::StorageChangeSets,
409        Database,
410    };
411    use alloy_primitives::{Address, B256, U256};
412    use reth_db_api::{
413        cursor::{DbCursorRO, DbDupCursorRW},
414        models::{BlockNumberAddress, ClientVersion},
415        table::TableImporter,
416        transaction::{DbTx, DbTxMut},
417    };
418    use reth_primitives_traits::StorageEntry;
419    use tempfile::TempDir;
420
421    fn create_test_db() -> DatabaseEnv {
422        let path = TempDir::new().unwrap();
423        let mut db = DatabaseEnv::open(
424            path.path(),
425            DatabaseEnvKind::RW,
426            DatabaseArguments::new(ClientVersion::default()),
427        )
428        .unwrap();
429        db.create_tables().unwrap();
430        db
431    }
432
433    #[test]
434    fn test_import_table_with_range_works_on_dupsort() {
435        let addr1 = Address::with_last_byte(1);
436        let addr2 = Address::with_last_byte(2);
437        let addr3 = Address::with_last_byte(3);
438        let source_db = create_test_db();
439        let target_db = create_test_db();
440        let test_data = vec![
441            (
442                BlockNumberAddress((100, addr1)),
443                StorageEntry { key: B256::with_last_byte(1), value: U256::from(100) },
444            ),
445            (
446                BlockNumberAddress((100, addr1)),
447                StorageEntry { key: B256::with_last_byte(2), value: U256::from(200) },
448            ),
449            (
450                BlockNumberAddress((100, addr1)),
451                StorageEntry { key: B256::with_last_byte(3), value: U256::from(300) },
452            ),
453            (
454                BlockNumberAddress((101, addr1)),
455                StorageEntry { key: B256::with_last_byte(1), value: U256::from(400) },
456            ),
457            (
458                BlockNumberAddress((101, addr2)),
459                StorageEntry { key: B256::with_last_byte(1), value: U256::from(500) },
460            ),
461            (
462                BlockNumberAddress((101, addr2)),
463                StorageEntry { key: B256::with_last_byte(2), value: U256::from(600) },
464            ),
465            (
466                BlockNumberAddress((102, addr3)),
467                StorageEntry { key: B256::with_last_byte(1), value: U256::from(700) },
468            ),
469        ];
470
471        // setup data
472        let tx = source_db.tx_mut().unwrap();
473        {
474            let mut cursor = tx.cursor_dup_write::<StorageChangeSets>().unwrap();
475            for (key, value) in &test_data {
476                cursor.append_dup(*key, *value).unwrap();
477            }
478        }
479        tx.commit().unwrap();
480
481        // import data from source db to target
482        let source_tx = source_db.tx().unwrap();
483        let target_tx = target_db.tx_mut().unwrap();
484
485        target_tx
486            .import_table_with_range::<StorageChangeSets, _>(
487                &source_tx,
488                Some(BlockNumberAddress((100, Address::ZERO))),
489                BlockNumberAddress((102, Address::repeat_byte(0xff))),
490            )
491            .unwrap();
492        target_tx.commit().unwrap();
493
494        // fetch all data from target db
495        let verify_tx = target_db.tx().unwrap();
496        let mut cursor = verify_tx.cursor_dup_read::<StorageChangeSets>().unwrap();
497        let copied: Vec<_> = cursor.walk(None).unwrap().collect::<Result<Vec<_>, _>>().unwrap();
498
499        // verify each entry matches the test data
500        assert_eq!(copied.len(), test_data.len(), "Should copy all entries including duplicates");
501        for ((copied_key, copied_value), (expected_key, expected_value)) in
502            copied.iter().zip(test_data.iter())
503        {
504            assert_eq!(copied_key, expected_key);
505            assert_eq!(copied_value, expected_value);
506        }
507    }
508}
reth_db/implementation/mdbx/cursor.rs

reth_db/implementation/mdbx/
cursor.rs
