Skip to main content

reth_db/implementation/mdbx/
value_prefetch.rs

1//! Targeted read-ahead for mapped table values spanning multiple OS pages.
2
3use reth_libmdbx::{ffi, TableObject, TransactionKind};
4use std::borrow::Cow;
5
6/// Hints a mapped table value's pages before returning its bytes for decoding.
7pub(super) struct PrefetchValue<'a>(pub(super) Cow<'a, [u8]>);
8
9impl<'a> TableObject for PrefetchValue<'a> {
10    fn decode(_: &[u8]) -> reth_libmdbx::Result<Self> {
11        unreachable!("value prefetch requires the MDBX transaction context")
12    }
13
14    unsafe fn decode_val<K: TransactionKind>(
15        txn: *const ffi::MDBX_txn,
16        value: ffi::MDBX_val,
17    ) -> reth_libmdbx::Result<Self> {
18        #[cfg(test)]
19        tests::PREFETCH_READS.with(|count| count.set(count.get() + 1));
20        // SAFETY: MDBX invokes this decoder with a value from the live transaction.
21        if let Some((start, len)) = unsafe { prefetch_range::<K>(txn, value) } {
22            // MDBX disables general read-ahead. Hint only this value's pages to avoid
23            // serial faults while decoding large values, without enabling general read-ahead.
24            // SAFETY: The complete rounded range belongs to the live MDBX file mapping.
25            // Advice is best-effort and must not turn a successful lookup into an error.
26            let _ = unsafe { libc::madvise(start as *mut libc::c_void, len, libc::MADV_WILLNEED) };
27        }
28        // SAFETY: Preserve libmdbx's borrowing/copying behavior for the original value.
29        unsafe { <Cow<'a, [u8]> as TableObject>::decode_val::<K>(txn, value).map(Self) }
30    }
31}
32
33/// Returns the OS pages to prefetch, excluding small and dirty values.
34///
35/// # Safety
36/// `value` must be the unmodified value returned by MDBX for the live `txn` of kind `K`.
37unsafe fn prefetch_range<K: TransactionKind>(
38    txn: *const ffi::MDBX_txn,
39    value: ffi::MDBX_val,
40) -> Option<(usize, usize)> {
41    let page_size = page_size::get();
42    if value.iov_len <= page_size {
43        return None;
44    }
45    // Unchanged values can also be read through writable transactions.
46    // Dirty values can be heap-backed; never advise them, even with `return-borrowed`.
47    // SAFETY: This is the original value pointer, before any copy or transaction mutation.
48    if !K::IS_READ_ONLY && unsafe { ffi::mdbx_is_dirty(txn, value.iov_base) } != ffi::MDBX_SUCCESS {
49        return None;
50    }
51    page_range(value.iov_base as usize, value.iov_len, page_size)
52}
53
54/// Rounds a nonempty value's range to complete pages, rejecting invalid sizes and overflow.
55fn page_range(address: usize, len: usize, page_size: usize) -> Option<(usize, usize)> {
56    if len == 0 || !page_size.is_power_of_two() {
57        return None;
58    }
59    let start = address & !(page_size - 1);
60    let end = address.checked_add(len)?.checked_add(page_size - 1)? & !(page_size - 1);
61    Some((start, end.checked_sub(start)?))
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67    use crate::{
68        mdbx::DatabaseArguments, tables, test_utils::create_test_rw_db, DatabaseEnv,
69        DatabaseEnvKind,
70    };
71    use alloy_primitives::{Address, B256, U256};
72    use reth_db_api::{
73        cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO},
74        database::Database,
75        models::{ClientVersion, ShardedKey},
76        table::Encode,
77        transaction::{DbTx, DbTxMut},
78    };
79    use reth_primitives_traits::{Bytecode, StorageEntry};
80    use std::cell::Cell;
81
82    thread_local! {
83        pub(super) static PREFETCH_READS: Cell<usize> = const { Cell::new(0) };
84    }
85
86    /// Verifies decoder selection, including for small values whose advice is skipped.
87    fn with_prefetch_reads<R>(expected: usize, f: impl FnOnce() -> R) -> R {
88        let before = PREFETCH_READS.get();
89        let result = f();
90        assert_eq!(PREFETCH_READS.get() - before, expected);
91        result
92    }
93
94    #[test]
95    fn rounds_to_os_pages() {
96        assert_eq!(page_range(0x1014, 27_693, 4096), Some((0x1000, 7 * 4096)));
97        assert_eq!(page_range(0x1000, 4096, 4096), Some((0x1000, 4096)));
98        assert_eq!(page_range(0x1fff, 2, 4096), Some((0x1000, 8192)));
99        assert_eq!(page_range(0x1014, 0xfec, 4096), Some((0x1000, 4096)));
100        assert_eq!(page_range(0x1ffff, 2, 65536), Some((0x10000, 2 * 65536)));
101        assert_eq!(page_range(0x1000, 0, 4096), None);
102        assert_eq!(page_range(usize::MAX - 10, 20, 4096), None);
103        assert_eq!(page_range(usize::MAX - 10, 1, 4096), None);
104        assert_eq!(page_range(0x1000, 4096, 0), None);
105        assert_eq!(page_range(0x1000, 4096, 3), None);
106    }
107
108    /// Inspects the same prefetch eligibility decision used by the value decoder.
109    struct PrefetchRange(Option<(usize, usize)>);
110
111    impl TableObject for PrefetchRange {
112        fn decode(_: &[u8]) -> reth_libmdbx::Result<Self> {
113            unreachable!()
114        }
115
116        unsafe fn decode_val<K: TransactionKind>(
117            txn: *const ffi::MDBX_txn,
118            value: ffi::MDBX_val,
119        ) -> reth_libmdbx::Result<Self> {
120            // SAFETY: MDBX supplies the original value and its live transaction.
121            Ok(Self(unsafe { prefetch_range::<K>(txn, value) }))
122        }
123    }
124
125    #[test]
126    fn prefetches_clean_bytecode_in_read_only_and_writable_transactions() {
127        let dir = tempfile::tempdir().unwrap();
128        let mut db = DatabaseEnv::open(
129            dir.path(),
130            DatabaseEnvKind::RW,
131            DatabaseArguments::new(ClientVersion::default()),
132        )
133        .unwrap()
134        .with_metrics();
135        db.create_tables().unwrap();
136        // Larger than an OS page even on systems with 64 KiB pages.
137        let code = Bytecode::new_raw(vec![0x5b; 6 * page_size::get()].into());
138        let key = B256::repeat_byte(1);
139        let small_key = B256::repeat_byte(2);
140        let small_code = Bytecode::new_raw(vec![0].into());
141        let tx = db.tx_mut().unwrap();
142        let dbi = tx.get_dbi::<tables::Bytecodes>().unwrap();
143        DbTxMut::put::<tables::Bytecodes>(&tx, key, code.clone()).unwrap();
144        DbTxMut::put::<tables::Bytecodes>(&tx, small_key, small_code.clone()).unwrap();
145        assert!(tx
146            .inner()
147            .get::<PrefetchRange>(dbi, key.encode().as_ref())
148            .unwrap()
149            .unwrap()
150            .0
151            .is_none());
152        assert_eq!(tx.prefetch(true).get::<tables::Bytecodes>(key).unwrap(), Some(code.clone()));
153        assert_eq!(tx.get::<tables::Bytecodes>(key).unwrap(), Some(code.clone()));
154        tx.commit().unwrap();
155
156        // Catch-up execution reads clean, mapped values through a writable transaction.
157        let tx = db.tx_mut().unwrap();
158        let range = tx
159            .inner()
160            .get::<PrefetchRange>(dbi, key.encode().as_ref())
161            .unwrap()
162            .unwrap()
163            .0
164            .unwrap();
165        assert_eq!(range.0 % page_size::get(), 0);
166        assert!(range.1 >= code.original_byte_slice().len());
167        assert_eq!(
168            tx.prefetch(true).get_by_encoded_key::<tables::Bytecodes>(&key.encode()).unwrap(),
169            Some(code.clone())
170        );
171        // A value becomes ineligible again when overwritten in the same transaction.
172        let replacement = Bytecode::new_raw(vec![0; 6 * page_size::get()].into());
173        DbTxMut::put::<tables::Bytecodes>(&tx, key, replacement.clone()).unwrap();
174        assert!(tx
175            .inner()
176            .get::<PrefetchRange>(dbi, key.encode().as_ref())
177            .unwrap()
178            .unwrap()
179            .0
180            .is_none());
181        assert_eq!(tx.prefetch(true).get::<tables::Bytecodes>(key).unwrap(), Some(replacement));
182        tx.abort();
183
184        let tx = db.tx().unwrap();
185        assert!(tx
186            .inner()
187            .get::<PrefetchRange>(dbi, key.encode().as_ref())
188            .unwrap()
189            .unwrap()
190            .0
191            .is_some());
192        assert_eq!(tx.prefetch(true).get::<tables::Bytecodes>(key).unwrap(), Some(code));
193        assert!(tx
194            .inner()
195            .get::<PrefetchRange>(dbi, small_key.encode().as_ref())
196            .unwrap()
197            .unwrap()
198            .0
199            .is_none());
200        assert_eq!(
201            tx.prefetch(true).get::<tables::Bytecodes>(small_key).unwrap(),
202            Some(small_code)
203        );
204        assert!(tx
205            .inner()
206            .get::<PrefetchRange>(dbi, B256::ZERO.encode().as_ref())
207            .unwrap()
208            .is_none());
209        assert_eq!(tx.prefetch(true).get::<tables::Bytecodes>(B256::ZERO).unwrap(), None);
210    }
211
212    #[test]
213    fn prefetched_reads_preserve_decode_errors() {
214        let dir = tempfile::tempdir().unwrap();
215        let mut db = DatabaseEnv::open(
216            dir.path(),
217            DatabaseEnvKind::RW,
218            DatabaseArguments::new(ClientVersion::default()),
219        )
220        .unwrap();
221        db.create_tables().unwrap();
222        let key = ShardedKey::new(Address::ZERO, 1);
223        // One RoaringTreemap entry with an invalid bitmap header.
224        let mut invalid_history = vec![0; 6 * page_size::get()];
225        invalid_history[..8].copy_from_slice(&1u64.to_le_bytes());
226        let tx = db.tx_mut().unwrap();
227        tx.put::<tables::RawTable<tables::AccountsHistory>>(
228            tables::RawKey::new(key.clone()),
229            tables::RawValue::from_vec(invalid_history),
230        )
231        .unwrap();
232        tx.commit().unwrap();
233
234        let tx = db.tx().unwrap();
235        with_prefetch_reads(1, || {
236            assert_eq!(
237                tx.prefetch(true)
238                    .get::<tables::AccountsHistory>(key.clone())
239                    .unwrap_err()
240                    .to_string(),
241                tx.prefetch(false)
242                    .get::<tables::AccountsHistory>(key.clone())
243                    .unwrap_err()
244                    .to_string()
245            )
246        });
247        tx.prefetch(true);
248        with_prefetch_reads(2, || {
249            assert!(tx.get::<tables::AccountsHistory>(key.clone()).is_err());
250            assert!(tx.get::<tables::AccountsHistory>(key).is_err());
251        });
252    }
253
254    #[test]
255    fn transaction_and_cursor_prefetch_persist_until_changed() {
256        let db = create_test_rw_db();
257        let tx = db.tx_mut().unwrap();
258        let key = B256::repeat_byte(1);
259        let code = Bytecode::new_raw(vec![0x5b; 6 * page_size::get()].into());
260        let mut ordinary = tx.cursor_read::<tables::Bytecodes>().unwrap();
261        assert!(std::ptr::eq(tx.prefetch(true), &raw const tx));
262        // Writes, statistics, gets, misses, and cursor creation preserve the setting.
263        DbTxMut::put::<tables::Bytecodes>(&tx, key, code.clone()).unwrap();
264        assert_eq!(tx.entries::<tables::Bytecodes>().unwrap(), 1);
265        with_prefetch_reads(2, || {
266            assert_eq!(tx.get::<tables::Bytecodes>(key).unwrap(), Some(code.clone()));
267            assert_eq!(
268                tx.get_by_encoded_key::<tables::Bytecodes>(&key).unwrap(),
269                Some(code.clone())
270            );
271        });
272        with_prefetch_reads(1, || {
273            assert!(tx.get::<tables::Bytecodes>(B256::ZERO).unwrap().is_none());
274            tx.get::<tables::Bytecodes>(key).unwrap();
275        });
276        let mut prefetched = tx.cursor_write::<tables::Bytecodes>().unwrap();
277        let mut sibling = tx.cursor_read::<tables::Bytecodes>().unwrap();
278        with_prefetch_reads(0, || ordinary.first().unwrap());
279        with_prefetch_reads(3, || {
280            assert_eq!(prefetched.first().unwrap(), Some((key, code)));
281            sibling.first().unwrap();
282            tx.get::<tables::Bytecodes>(key).unwrap();
283        });
284
285        // A cursor can disable its inherited setting without changing its position or siblings.
286        with_prefetch_reads(0, || prefetched.prefetch(false).current().unwrap());
287        with_prefetch_reads(2, || {
288            sibling.current().unwrap();
289            tx.get::<tables::Bytecodes>(key).unwrap();
290        });
291        tx.prefetch(false);
292        with_prefetch_reads(0, || {
293            tx.get::<tables::Bytecodes>(key).unwrap();
294            tx.cursor_read::<tables::Bytecodes>().unwrap().first().unwrap();
295        });
296        with_prefetch_reads(1, || sibling.current().unwrap());
297
298        // Cursor overrides are independent of the transaction and persist through writes.
299        with_prefetch_reads(1, || prefetched.prefetch(true).current().unwrap());
300        prefetched.upsert(key, &Bytecode::new_raw(vec![0].into())).unwrap();
301        with_prefetch_reads(1, || prefetched.current().unwrap());
302        with_prefetch_reads(0, || {
303            ordinary.current().unwrap();
304            tx.get::<tables::Bytecodes>(key).unwrap();
305        });
306    }
307
308    #[test]
309    fn prefetch_covers_cursor_reads_and_walks() {
310        let db = create_test_rw_db();
311        let tx = db.tx_mut().unwrap();
312        let code = Bytecode::new_raw(vec![0x5b; 6 * page_size::get()].into());
313        let keys = [B256::repeat_byte(1), B256::repeat_byte(2), B256::repeat_byte(3)];
314        for key in keys {
315            DbTxMut::put::<tables::Bytecodes>(&tx, key, code.clone()).unwrap();
316        }
317        tx.commit().unwrap();
318        let tx = db.tx().unwrap();
319        let mut cursor = tx.prefetch(true).cursor_read::<tables::Bytecodes>().unwrap();
320        with_prefetch_reads(7, || {
321            assert_eq!(cursor.first().unwrap(), Some((keys[0], code.clone())));
322            assert_eq!(cursor.next().unwrap(), Some((keys[1], code.clone())));
323            assert_eq!(cursor.prev().unwrap(), Some((keys[0], code.clone())));
324            assert_eq!(cursor.last().unwrap(), Some((keys[2], code.clone())));
325            assert_eq!(cursor.current().unwrap(), Some((keys[2], code.clone())));
326            assert_eq!(cursor.seek(keys[1]).unwrap(), Some((keys[1], code.clone())));
327            assert_eq!(cursor.seek_exact(keys[0]).unwrap(), Some((keys[0], code.clone())));
328        });
329        let expected = keys.map(|key| (key, code.clone())).to_vec();
330        for start in [None, Some(keys[0])] {
331            with_prefetch_reads(3, || {
332                assert_eq!(
333                    cursor.walk(start).unwrap().collect::<Result<Vec<_>, _>>().unwrap(),
334                    expected
335                );
336            });
337        }
338        for start in [None, Some(keys[2])] {
339            with_prefetch_reads(3, || {
340                assert_eq!(
341                    cursor.walk_back(start).unwrap().collect::<Result<Vec<_>, _>>().unwrap(),
342                    expected.iter().rev().cloned().collect::<Vec<_>>()
343                );
344            });
345        }
346        with_prefetch_reads(3, || {
347            assert_eq!(
348                cursor
349                    .walk_range(keys[0]..=keys[2])
350                    .unwrap()
351                    .collect::<Result<Vec<_>, _>>()
352                    .unwrap(),
353                expected
354            );
355        });
356        with_prefetch_reads(3, || {
357            assert_eq!(
358                cursor.walk_range(..).unwrap().collect::<Result<Vec<_>, _>>().unwrap(),
359                expected
360            );
361        });
362        with_prefetch_reads(0, || {
363            assert_eq!(
364                cursor.prefetch(false).walk(None).unwrap().collect::<Result<Vec<_>, _>>().unwrap(),
365                expected
366            );
367        });
368        with_prefetch_reads(3, || {
369            assert_eq!(
370                cursor.prefetch(true).walk(None).unwrap().collect::<Result<Vec<_>, _>>().unwrap(),
371                expected
372            );
373        });
374    }
375
376    #[test]
377    fn prefetch_covers_duplicate_cursor_reads() {
378        let db = create_test_rw_db();
379        let tx = db.tx_mut().unwrap();
380        let address = Address::repeat_byte(1);
381        let next_address = Address::repeat_byte(2);
382        let entries =
383            [1, 2, 3].map(|i| StorageEntry { key: B256::repeat_byte(i), value: U256::from(i) });
384        for entry in entries {
385            DbTxMut::put::<tables::PlainStorageState>(&tx, address, entry).unwrap();
386        }
387        DbTxMut::put::<tables::PlainStorageState>(&tx, next_address, entries[0]).unwrap();
388        tx.commit().unwrap();
389        let tx = db.tx_mut().unwrap();
390        // Both read-only and writable duplicate cursor factories inherit the flag.
391        let mut cursor = tx.prefetch(true).cursor_dup_read::<tables::PlainStorageState>().unwrap();
392        with_prefetch_reads(7, || {
393            assert_eq!(cursor.first().unwrap(), Some((address, entries[0])));
394            assert_eq!(cursor.next_dup().unwrap(), Some((address, entries[1])));
395            assert_eq!(cursor.prev_dup().unwrap(), Some((address, entries[0])));
396            assert_eq!(cursor.next_dup_val().unwrap(), Some(entries[1]));
397            assert_eq!(cursor.last_dup().unwrap(), Some(entries[2]));
398            assert_eq!(cursor.next_no_dup().unwrap(), Some((next_address, entries[0])));
399            assert_eq!(
400                cursor.seek_by_key_subkey(address, entries[1].key).unwrap(),
401                Some(entries[1])
402            );
403        });
404        let mut cursor = tx.prefetch(true).cursor_dup_write::<tables::PlainStorageState>().unwrap();
405        for (key, subkey, reads) in [
406            (Some(address), Some(entries[0].key), 3),
407            (Some(address), None, 3),
408            (None, Some(entries[0].key), 4),
409            (None, None, 3),
410        ] {
411            with_prefetch_reads(reads, || {
412                assert_eq!(
413                    cursor.walk_dup(key, subkey).unwrap().collect::<Result<Vec<_>, _>>().unwrap(),
414                    entries.map(|entry| (address, entry)).to_vec()
415                );
416            });
417        }
418        with_prefetch_reads(0, || {
419            assert_eq!(
420                cursor.prefetch(false).seek_by_key_subkey(address, entries[0].key).unwrap(),
421                Some(entries[0])
422            );
423            assert_eq!(cursor.next_dup_val().unwrap(), Some(entries[1]));
424        });
425        with_prefetch_reads(2, || {
426            assert_eq!(
427                cursor.prefetch(true).seek_by_key_subkey(address, entries[0].key).unwrap(),
428                Some(entries[0])
429            );
430            assert_eq!(cursor.next_dup_val().unwrap(), Some(entries[1]));
431        });
432    }
433}