redb/
table.rs

1use crate::db::TransactionGuard;
2use crate::sealed::Sealed;
3use crate::tree_store::{
4    AccessGuardMut, Btree, BtreeExtractIf, BtreeHeader, BtreeMut, BtreeRangeIter, MAX_PAIR_LENGTH,
5    MAX_VALUE_LENGTH, PageHint, PageNumber, PageTrackerPolicy, RawBtree, TransactionalMemory,
6};
7use crate::types::{Key, MutInPlaceValue, Value};
8use crate::{AccessGuard, StorageError, WriteTransaction};
9use crate::{Result, TableHandle};
10use std::borrow::Borrow;
11use std::fmt::{Debug, Formatter};
12use std::marker::PhantomData;
13use std::ops::RangeBounds;
14use std::sync::{Arc, Mutex};
15
16/// Informational storage stats about a table
17#[derive(Debug)]
18pub struct TableStats {
19    pub(crate) tree_height: u32,
20    pub(crate) leaf_pages: u64,
21    pub(crate) branch_pages: u64,
22    pub(crate) stored_leaf_bytes: u64,
23    pub(crate) metadata_bytes: u64,
24    pub(crate) fragmented_bytes: u64,
25}
26
27impl TableStats {
28    /// Maximum traversal distance to reach the deepest (key, value) pair in the table
29    pub fn tree_height(&self) -> u32 {
30        self.tree_height
31    }
32
33    /// Number of leaf pages that store user data
34    pub fn leaf_pages(&self) -> u64 {
35        self.leaf_pages
36    }
37
38    /// Number of branch pages in the btree that store user data
39    pub fn branch_pages(&self) -> u64 {
40        self.branch_pages
41    }
42
43    /// Number of bytes consumed by keys and values that have been inserted.
44    /// Does not include indexing overhead
45    pub fn stored_bytes(&self) -> u64 {
46        self.stored_leaf_bytes
47    }
48
49    /// Number of bytes consumed by keys in internal branch pages, plus other metadata
50    pub fn metadata_bytes(&self) -> u64 {
51        self.metadata_bytes
52    }
53
54    /// Number of bytes consumed by fragmentation, both in data pages and internal metadata tables
55    pub fn fragmented_bytes(&self) -> u64 {
56        self.fragmented_bytes
57    }
58}
59
60/// A table containing key-value mappings
61pub struct Table<'txn, K: Key + 'static, V: Value + 'static> {
62    name: String,
63    transaction: &'txn WriteTransaction,
64    tree: BtreeMut<'txn, K, V>,
65}
66
67impl<K: Key + 'static, V: Value + 'static> TableHandle for Table<'_, K, V> {
68    fn name(&self) -> &str {
69        &self.name
70    }
71}
72
73impl<'txn, K: Key + 'static, V: Value + 'static> Table<'txn, K, V> {
74    pub(crate) fn new(
75        name: &str,
76        table_root: Option<BtreeHeader>,
77        freed_pages: Arc<Mutex<Vec<PageNumber>>>,
78        allocated_pages: Arc<Mutex<PageTrackerPolicy>>,
79        mem: Arc<TransactionalMemory>,
80        transaction: &'txn WriteTransaction,
81    ) -> Table<'txn, K, V> {
82        Table {
83            name: name.to_string(),
84            transaction,
85            tree: BtreeMut::new(
86                table_root,
87                transaction.transaction_guard(),
88                mem,
89                freed_pages,
90                allocated_pages,
91            ),
92        }
93    }
94
95    #[allow(dead_code)]
96    pub(crate) fn print_debug(&self, include_values: bool) -> Result {
97        self.tree.print_debug(include_values)
98    }
99
100    /// Removes and returns the first key-value pair in the table
101    pub fn pop_first(&mut self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
102        // TODO: optimize this
103        let first = self
104            .iter()?
105            .next()
106            .map(|x| x.map(|(key, _)| K::as_bytes(&key.value()).as_ref().to_vec()));
107        if let Some(owned_key) = first {
108            let owned_key = owned_key?;
109            let key = K::from_bytes(&owned_key);
110            let value = self.remove(&key)?.unwrap();
111            drop(key);
112            Ok(Some((AccessGuard::with_owned_value(owned_key), value)))
113        } else {
114            Ok(None)
115        }
116    }
117
118    /// Removes and returns the last key-value pair in the table
119    pub fn pop_last(&mut self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
120        // TODO: optimize this
121        let last = self
122            .iter()?
123            .next_back()
124            .map(|x| x.map(|(key, _)| K::as_bytes(&key.value()).as_ref().to_vec()));
125        if let Some(owned_key) = last {
126            let owned_key = owned_key?;
127            let key = K::from_bytes(&owned_key);
128            let value = self.remove(&key)?.unwrap();
129            drop(key);
130            Ok(Some((AccessGuard::with_owned_value(owned_key), value)))
131        } else {
132            Ok(None)
133        }
134    }
135
136    /// Applies `predicate` to all key-value pairs. All entries for which
137    /// `predicate` evaluates to `true` are returned in an iterator, and those which are read from the iterator are removed
138    ///
139    /// Note: values not read from the iterator will not be removed
140    pub fn extract_if<F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
141        &mut self,
142        predicate: F,
143    ) -> Result<ExtractIf<K, V, F>> {
144        self.extract_from_if::<K::SelfType<'_>, F>(.., predicate)
145    }
146
147    /// Applies `predicate` to all key-value pairs in the specified range. All entries for which
148    /// `predicate` evaluates to `true` are returned in an iterator, and those which are read from the iterator are removed
149    ///
150    /// Note: values not read from the iterator will not be removed
151    pub fn extract_from_if<'a, KR, F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
152        &mut self,
153        range: impl RangeBounds<KR> + 'a,
154        predicate: F,
155    ) -> Result<ExtractIf<K, V, F>>
156    where
157        KR: Borrow<K::SelfType<'a>> + 'a,
158    {
159        self.tree
160            .extract_from_if(&range, predicate)
161            .map(ExtractIf::new)
162    }
163
164    /// Applies `predicate` to all key-value pairs. All entries for which
165    /// `predicate` evaluates to `false` are removed.
166    ///
167    pub fn retain<F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
168        &mut self,
169        predicate: F,
170    ) -> Result {
171        self.tree.retain_in::<K::SelfType<'_>, F>(predicate, ..)
172    }
173
174    /// Applies `predicate` to all key-value pairs in the range `start..end`. All entries for which
175    /// `predicate` evaluates to `false` are removed.
176    ///
177    pub fn retain_in<'a, KR, F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
178        &mut self,
179        range: impl RangeBounds<KR> + 'a,
180        predicate: F,
181    ) -> Result
182    where
183        KR: Borrow<K::SelfType<'a>> + 'a,
184    {
185        self.tree.retain_in(predicate, range)
186    }
187
188    /// Insert mapping of the given key to the given value
189    ///
190    /// If key is already present it is replaced
191    ///
192    /// Returns the old value, if the key was present in the table, otherwise None is returned
193    pub fn insert<'k, 'v>(
194        &mut self,
195        key: impl Borrow<K::SelfType<'k>>,
196        value: impl Borrow<V::SelfType<'v>>,
197    ) -> Result<Option<AccessGuard<V>>> {
198        let value_len = V::as_bytes(value.borrow()).as_ref().len();
199        if value_len > MAX_VALUE_LENGTH {
200            return Err(StorageError::ValueTooLarge(value_len));
201        }
202        let key_len = K::as_bytes(key.borrow()).as_ref().len();
203        if key_len > MAX_VALUE_LENGTH {
204            return Err(StorageError::ValueTooLarge(key_len));
205        }
206        if value_len + key_len > MAX_PAIR_LENGTH {
207            return Err(StorageError::ValueTooLarge(value_len + key_len));
208        }
209        self.tree.insert(key.borrow(), value.borrow())
210    }
211
212    /// Removes the given key
213    ///
214    /// Returns the old value, if the key was present in the table
215    pub fn remove<'a>(
216        &mut self,
217        key: impl Borrow<K::SelfType<'a>>,
218    ) -> Result<Option<AccessGuard<V>>> {
219        self.tree.remove(key.borrow())
220    }
221}
222
223impl<K: Key + 'static, V: MutInPlaceValue + 'static> Table<'_, K, V> {
224    /// Reserve space to insert a key-value pair
225    ///
226    /// If key is already present it is replaced
227    ///
228    /// The returned reference will have length equal to `value_length`
229    pub fn insert_reserve<'a>(
230        &mut self,
231        key: impl Borrow<K::SelfType<'a>>,
232        value_length: u32,
233    ) -> Result<AccessGuardMut<V>> {
234        if value_length as usize > MAX_VALUE_LENGTH {
235            return Err(StorageError::ValueTooLarge(value_length as usize));
236        }
237        let key_len = K::as_bytes(key.borrow()).as_ref().len();
238        if key_len > MAX_VALUE_LENGTH {
239            return Err(StorageError::ValueTooLarge(key_len));
240        }
241        if value_length as usize + key_len > MAX_PAIR_LENGTH {
242            return Err(StorageError::ValueTooLarge(value_length as usize + key_len));
243        }
244        self.tree.insert_reserve(key.borrow(), value_length)
245    }
246}
247
248impl<K: Key + 'static, V: Value + 'static> ReadableTableMetadata for Table<'_, K, V> {
249    fn stats(&self) -> Result<TableStats> {
250        let tree_stats = self.tree.stats()?;
251
252        Ok(TableStats {
253            tree_height: tree_stats.tree_height,
254            leaf_pages: tree_stats.leaf_pages,
255            branch_pages: tree_stats.branch_pages,
256            stored_leaf_bytes: tree_stats.stored_leaf_bytes,
257            metadata_bytes: tree_stats.metadata_bytes,
258            fragmented_bytes: tree_stats.fragmented_bytes,
259        })
260    }
261
262    fn len(&self) -> Result<u64> {
263        self.tree.len()
264    }
265}
266
267impl<K: Key + 'static, V: Value + 'static> ReadableTable<K, V> for Table<'_, K, V> {
268    fn get<'a>(&self, key: impl Borrow<K::SelfType<'a>>) -> Result<Option<AccessGuard<V>>> {
269        self.tree.get(key.borrow())
270    }
271
272    fn range<'a, KR>(&self, range: impl RangeBounds<KR> + 'a) -> Result<Range<K, V>>
273    where
274        KR: Borrow<K::SelfType<'a>> + 'a,
275    {
276        self.tree
277            .range(&range)
278            .map(|x| Range::new(x, self.transaction.transaction_guard()))
279    }
280
281    fn first(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
282        self.tree.first()
283    }
284
285    fn last(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
286        self.tree.last()
287    }
288}
289
290impl<K: Key, V: Value> Sealed for Table<'_, K, V> {}
291
292impl<K: Key + 'static, V: Value + 'static> Drop for Table<'_, K, V> {
293    fn drop(&mut self) {
294        self.transaction.close_table(
295            &self.name,
296            &self.tree,
297            self.tree.get_root().map(|x| x.length).unwrap_or_default(),
298        );
299    }
300}
301
302fn debug_helper<K: Key + 'static, V: Value + 'static>(
303    f: &mut Formatter<'_>,
304    name: &str,
305    len: Result<u64>,
306    first: Result<Option<(AccessGuard<K>, AccessGuard<V>)>>,
307    last: Result<Option<(AccessGuard<K>, AccessGuard<V>)>>,
308) -> std::fmt::Result {
309    write!(f, "Table [ name: \"{name}\", ")?;
310    if let Ok(len) = len {
311        if len == 0 {
312            write!(f, "No entries")?;
313        } else if len == 1 {
314            if let Ok(first) = first {
315                let (key, value) = first.as_ref().unwrap();
316                write!(f, "One key-value: {:?} = {:?}", key.value(), value.value())?;
317            } else {
318                write!(f, "I/O Error accessing table!")?;
319            }
320        } else {
321            if let Ok(first) = first {
322                let (key, value) = first.as_ref().unwrap();
323                write!(f, "first: {:?} = {:?}, ", key.value(), value.value())?;
324            } else {
325                write!(f, "I/O Error accessing table!")?;
326            }
327            if len > 2 {
328                write!(f, "...{} more entries..., ", len - 2)?;
329            }
330            if let Ok(last) = last {
331                let (key, value) = last.as_ref().unwrap();
332                write!(f, "last: {:?} = {:?}", key.value(), value.value())?;
333            } else {
334                write!(f, "I/O Error accessing table!")?;
335            }
336        }
337    } else {
338        write!(f, "I/O Error accessing table!")?;
339    }
340    write!(f, " ]")?;
341
342    Ok(())
343}
344
345impl<K: Key + 'static, V: Value + 'static> Debug for Table<'_, K, V> {
346    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
347        debug_helper(f, &self.name, self.len(), self.first(), self.last())
348    }
349}
350
351pub trait ReadableTableMetadata {
352    /// Retrieves information about storage usage for the table
353    fn stats(&self) -> Result<TableStats>;
354
355    /// Returns the number of entries in the table
356    fn len(&self) -> Result<u64>;
357
358    /// Returns `true` if the table is empty
359    fn is_empty(&self) -> Result<bool> {
360        Ok(self.len()? == 0)
361    }
362}
363
364pub trait ReadableTable<K: Key + 'static, V: Value + 'static>: ReadableTableMetadata {
365    /// Returns the value corresponding to the given key
366    fn get<'a>(&self, key: impl Borrow<K::SelfType<'a>>) -> Result<Option<AccessGuard<V>>>;
367
368    /// Returns a double-ended iterator over a range of elements in the table
369    ///
370    /// # Examples
371    ///
372    /// Usage:
373    /// ```rust
374    /// use redb::*;
375    /// # use tempfile::NamedTempFile;
376    /// const TABLE: TableDefinition<&str, u64> = TableDefinition::new("my_data");
377    ///
378    /// # fn main() -> Result<(), Error> {
379    /// # let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
380    /// # let filename = tmpfile.path();
381    /// let db = Database::create(filename)?;
382    /// let write_txn = db.begin_write()?;
383    /// {
384    ///     let mut table = write_txn.open_table(TABLE)?;
385    ///     table.insert("a", &0)?;
386    ///     table.insert("b", &1)?;
387    ///     table.insert("c", &2)?;
388    /// }
389    /// write_txn.commit()?;
390    ///
391    /// let read_txn = db.begin_read()?;
392    /// let table = read_txn.open_table(TABLE)?;
393    /// let mut iter = table.range("a".."c")?;
394    /// let (key, value) = iter.next().unwrap()?;
395    /// assert_eq!("a", key.value());
396    /// assert_eq!(0, value.value());
397    /// # Ok(())
398    /// # }
399    /// ```
400    fn range<'a, KR>(&self, range: impl RangeBounds<KR> + 'a) -> Result<Range<K, V>>
401    where
402        KR: Borrow<K::SelfType<'a>> + 'a;
403
404    /// Returns the first key-value pair in the table, if it exists
405    fn first(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>>;
406
407    /// Returns the last key-value pair in the table, if it exists
408    fn last(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>>;
409
410    /// Returns a double-ended iterator over all elements in the table
411    fn iter(&self) -> Result<Range<K, V>> {
412        self.range::<K::SelfType<'_>>(..)
413    }
414}
415
416/// A read-only untyped table
417pub struct ReadOnlyUntypedTable {
418    tree: RawBtree,
419}
420
421impl Sealed for ReadOnlyUntypedTable {}
422
423impl ReadableTableMetadata for ReadOnlyUntypedTable {
424    /// Retrieves information about storage usage for the table
425    fn stats(&self) -> Result<TableStats> {
426        let tree_stats = self.tree.stats()?;
427
428        Ok(TableStats {
429            tree_height: tree_stats.tree_height,
430            leaf_pages: tree_stats.leaf_pages,
431            branch_pages: tree_stats.branch_pages,
432            stored_leaf_bytes: tree_stats.stored_leaf_bytes,
433            metadata_bytes: tree_stats.metadata_bytes,
434            fragmented_bytes: tree_stats.fragmented_bytes,
435        })
436    }
437
438    fn len(&self) -> Result<u64> {
439        self.tree.len()
440    }
441}
442
443impl ReadOnlyUntypedTable {
444    pub(crate) fn new(
445        root_page: Option<BtreeHeader>,
446        fixed_key_size: Option<usize>,
447        fixed_value_size: Option<usize>,
448        mem: Arc<TransactionalMemory>,
449    ) -> Self {
450        Self {
451            tree: RawBtree::new(root_page, fixed_key_size, fixed_value_size, mem),
452        }
453    }
454}
455
456/// A read-only table
457pub struct ReadOnlyTable<K: Key + 'static, V: Value + 'static> {
458    name: String,
459    tree: Btree<K, V>,
460    transaction_guard: Arc<TransactionGuard>,
461}
462
463impl<K: Key + 'static, V: Value + 'static> TableHandle for ReadOnlyTable<K, V> {
464    fn name(&self) -> &str {
465        &self.name
466    }
467}
468
469impl<K: Key + 'static, V: Value + 'static> ReadOnlyTable<K, V> {
470    pub(crate) fn new(
471        name: String,
472        root_page: Option<BtreeHeader>,
473        hint: PageHint,
474        guard: Arc<TransactionGuard>,
475        mem: Arc<TransactionalMemory>,
476    ) -> Result<ReadOnlyTable<K, V>> {
477        Ok(ReadOnlyTable {
478            name,
479            tree: Btree::new(root_page, hint, guard.clone(), mem)?,
480            transaction_guard: guard,
481        })
482    }
483
484    pub fn get<'a>(
485        &self,
486        key: impl Borrow<K::SelfType<'a>>,
487    ) -> Result<Option<AccessGuard<'static, V>>> {
488        self.tree.get(key.borrow())
489    }
490
491    /// This method is like [`ReadableTable::range()`], but the iterator is reference counted and keeps the transaction
492    /// alive until it is dropped.
493    pub fn range<'a, KR>(&self, range: impl RangeBounds<KR>) -> Result<Range<'static, K, V>>
494    where
495        KR: Borrow<K::SelfType<'a>>,
496    {
497        self.tree
498            .range(&range)
499            .map(|x| Range::new(x, self.transaction_guard.clone()))
500    }
501}
502
503impl<K: Key + 'static, V: Value + 'static> ReadableTableMetadata for ReadOnlyTable<K, V> {
504    fn stats(&self) -> Result<TableStats> {
505        let tree_stats = self.tree.stats()?;
506
507        Ok(TableStats {
508            tree_height: tree_stats.tree_height,
509            leaf_pages: tree_stats.leaf_pages,
510            branch_pages: tree_stats.branch_pages,
511            stored_leaf_bytes: tree_stats.stored_leaf_bytes,
512            metadata_bytes: tree_stats.metadata_bytes,
513            fragmented_bytes: tree_stats.fragmented_bytes,
514        })
515    }
516
517    fn len(&self) -> Result<u64> {
518        self.tree.len()
519    }
520}
521
522impl<K: Key + 'static, V: Value + 'static> ReadableTable<K, V> for ReadOnlyTable<K, V> {
523    fn get<'a>(&self, key: impl Borrow<K::SelfType<'a>>) -> Result<Option<AccessGuard<V>>> {
524        self.tree.get(key.borrow())
525    }
526
527    fn range<'a, KR>(&self, range: impl RangeBounds<KR> + 'a) -> Result<Range<K, V>>
528    where
529        KR: Borrow<K::SelfType<'a>> + 'a,
530    {
531        self.tree
532            .range(&range)
533            .map(|x| Range::new(x, self.transaction_guard.clone()))
534    }
535
536    fn first(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
537        self.tree.first()
538    }
539
540    fn last(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
541        self.tree.last()
542    }
543}
544
545impl<K: Key, V: Value> Sealed for ReadOnlyTable<K, V> {}
546
547impl<K: Key + 'static, V: Value + 'static> Debug for ReadOnlyTable<K, V> {
548    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
549        debug_helper(f, &self.name, self.len(), self.first(), self.last())
550    }
551}
552
553pub struct ExtractIf<
554    'a,
555    K: Key + 'static,
556    V: Value + 'static,
557    F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
558> {
559    inner: BtreeExtractIf<'a, K, V, F>,
560}
561
562impl<
563    'a,
564    K: Key + 'static,
565    V: Value + 'static,
566    F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
567> ExtractIf<'a, K, V, F>
568{
569    pub(crate) fn new(inner: BtreeExtractIf<'a, K, V, F>) -> Self {
570        Self { inner }
571    }
572}
573
574impl<
575    'a,
576    K: Key + 'static,
577    V: Value + 'static,
578    F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
579> Iterator for ExtractIf<'a, K, V, F>
580{
581    type Item = Result<(AccessGuard<'a, K>, AccessGuard<'a, V>)>;
582
583    fn next(&mut self) -> Option<Self::Item> {
584        let entry = self.inner.next()?;
585        Some(entry.map(|entry| {
586            let (page, key_range, value_range) = entry.into_raw();
587            let key = AccessGuard::with_page(page.clone(), key_range);
588            let value = AccessGuard::with_page(page, value_range);
589            (key, value)
590        }))
591    }
592}
593
594impl<
595    K: Key + 'static,
596    V: Value + 'static,
597    F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
598> DoubleEndedIterator for ExtractIf<'_, K, V, F>
599{
600    fn next_back(&mut self) -> Option<Self::Item> {
601        let entry = self.inner.next_back()?;
602        Some(entry.map(|entry| {
603            let (page, key_range, value_range) = entry.into_raw();
604            let key = AccessGuard::with_page(page.clone(), key_range);
605            let value = AccessGuard::with_page(page, value_range);
606            (key, value)
607        }))
608    }
609}
610
611#[derive(Clone)]
612pub struct Range<'a, K: Key + 'static, V: Value + 'static> {
613    inner: BtreeRangeIter<K, V>,
614    _transaction_guard: Arc<TransactionGuard>,
615    // This lifetime is here so that `&` can be held on `Table` preventing concurrent mutation
616    _lifetime: PhantomData<&'a ()>,
617}
618
619impl<K: Key + 'static, V: Value + 'static> Range<'_, K, V> {
620    pub(super) fn new(inner: BtreeRangeIter<K, V>, guard: Arc<TransactionGuard>) -> Self {
621        Self {
622            inner,
623            _transaction_guard: guard,
624            _lifetime: Default::default(),
625        }
626    }
627}
628
629impl<'a, K: Key + 'static, V: Value + 'static> Iterator for Range<'a, K, V> {
630    type Item = Result<(AccessGuard<'a, K>, AccessGuard<'a, V>)>;
631
632    fn next(&mut self) -> Option<Self::Item> {
633        self.inner.next().map(|x| {
634            x.map(|entry| {
635                let (page, key_range, value_range) = entry.into_raw();
636                let key = AccessGuard::with_page(page.clone(), key_range);
637                let value = AccessGuard::with_page(page, value_range);
638                (key, value)
639            })
640        })
641    }
642}
643
644impl<K: Key + 'static, V: Value + 'static> DoubleEndedIterator for Range<'_, K, V> {
645    fn next_back(&mut self) -> Option<Self::Item> {
646        self.inner.next_back().map(|x| {
647            x.map(|entry| {
648                let (page, key_range, value_range) = entry.into_raw();
649                let key = AccessGuard::with_page(page.clone(), key_range);
650                let value = AccessGuard::with_page(page, value_range);
651                (key, value)
652            })
653        })
654    }
655}