aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorAlex Auvolat <alex@adnab.me>2022-06-03 16:18:56 +0200
committerAlex Auvolat <alex@adnab.me>2022-06-03 16:18:56 +0200
commitf7a1c70089bf453554f5ae787da4865caf0ee5c4 (patch)
treea71156b041bfe7445e884d50ecb3a665f017e61c
parentbd2997a4534a4d06bc0cc636a2c0b50fe793057a (diff)
downloadgarage-f7a1c70089bf453554f5ae787da4865caf0ee5c4.tar.gz
garage-f7a1c70089bf453554f5ae787da4865caf0ee5c4.zip
Implement iterator for LMDB
-rw-r--r--src/db/lmdb_adapter.rs106
1 files changed, 98 insertions, 8 deletions
diff --git a/src/db/lmdb_adapter.rs b/src/db/lmdb_adapter.rs
index ec1a5444..d568b5c5 100644
--- a/src/db/lmdb_adapter.rs
+++ b/src/db/lmdb_adapter.rs
@@ -3,11 +3,12 @@ use core::ops::Bound;
use core::pin::Pin;
use core::ptr::NonNull;
-use std::cell::RefCell;
use std::collections::HashMap;
+use std::convert::TryInto;
use std::sync::{Arc, RwLock};
-use heed::{Env, RoTxn, RwTxn, UntypedDatabase as Database};
+use heed::types::ByteSlice;
+use heed::{BytesDecode, Env, RoTxn, RwTxn, UntypedDatabase as Database};
use crate::{
Db, Error, IDb, ITx, ITxFn, IValue, Result, TxError, TxFnResult, TxResult, Value, ValueIter,
@@ -101,11 +102,17 @@ impl IDb for LmdbDb {
}
fn remove(&self, tree: usize, key: &[u8]) -> Result<bool> {
- unimplemented!()
+ let tree = self.get_tree(tree)?;
+ let mut tx = self.db.write_txn()?;
+ let deleted = tree.delete(&mut tx, &key)?;
+ tx.commit()?;
+ Ok(deleted)
}
fn len(&self, tree: usize) -> Result<usize> {
- unimplemented!()
+ let tree = self.get_tree(tree)?;
+ let tx = self.db.read_txn()?;
+ Ok(tree.len(&tx)?.try_into().unwrap())
}
fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<()> {
@@ -117,11 +124,15 @@ impl IDb for LmdbDb {
}
fn iter(&self, tree: usize) -> Result<ValueIter<'_>> {
- unimplemented!()
+ let tree = self.get_tree(tree)?;
+ let tx = self.db.read_txn()?;
+ TxAndIterator::make(tx, |tx| Ok(tree.iter(tx)?))
}
fn iter_rev(&self, tree: usize) -> Result<ValueIter<'_>> {
- unimplemented!()
+ let tree = self.get_tree(tree)?;
+ let tx = self.db.read_txn()?;
+ TxAndIterator::make(tx, |tx| Ok(tree.rev_iter(tx)?))
}
fn range<'r>(
@@ -130,7 +141,9 @@ impl IDb for LmdbDb {
low: Bound<&'r [u8]>,
high: Bound<&'r [u8]>,
) -> Result<ValueIter<'_>> {
- unimplemented!()
+ let tree = self.get_tree(tree)?;
+ let tx = self.db.read_txn()?;
+ TxAndIterator::make(tx, |tx| Ok(tree.range(tx, &(low, high))?))
}
fn range_rev<'r>(
&self,
@@ -138,7 +151,9 @@ impl IDb for LmdbDb {
low: Bound<&'r [u8]>,
high: Bound<&'r [u8]>,
) -> Result<ValueIter<'_>> {
- unimplemented!()
+ let tree = self.get_tree(tree)?;
+ let tx = self.db.read_txn()?;
+ TxAndIterator::make(tx, |tx| Ok(tree.rev_range(tx, &(low, high))?))
}
// ----
@@ -262,3 +277,78 @@ impl<'a> std::borrow::Borrow<[u8]> for TxAndValuePin<'a> {
self.as_ref()
}
}
+
+// ----
+
+type IteratorItem<'a> = heed::Result<(
+ <ByteSlice as BytesDecode<'a>>::DItem,
+ <ByteSlice as BytesDecode<'a>>::DItem,
+)>;
+
+struct TxAndIterator<'a, I>
+where
+ I: Iterator<Item = IteratorItem<'a>> + 'a,
+{
+ tx: RoTxn<'a>,
+ iter: Option<I>,
+ _pin: PhantomPinned,
+}
+
+impl<'a, I> TxAndIterator<'a, I>
+where
+ I: Iterator<Item = IteratorItem<'a>> + 'a,
+{
+ fn make<F>(tx: RoTxn<'a>, iterfun: F) -> Result<ValueIter<'a>>
+ where
+ F: FnOnce(&'a RoTxn<'a>) -> Result<I>,
+ {
+ let res = TxAndIterator {
+ tx,
+ iter: None,
+ _pin: PhantomPinned,
+ };
+ let mut boxed = Box::pin(res);
+
+ unsafe {
+ let tx = NonNull::from(&boxed.tx);
+ let iter = iterfun(tx.as_ref())?;
+
+ let mut_ref: Pin<&mut TxAndIterator<'a, I>> = Pin::as_mut(&mut boxed);
+ Pin::get_unchecked_mut(mut_ref).iter = Some(iter);
+ }
+
+ Ok(Box::new(TxAndIteratorPin(boxed)))
+ }
+}
+
+impl<'a, I> Drop for TxAndIterator<'a, I>
+where
+ I: Iterator<Item = IteratorItem<'a>> + 'a,
+{
+ fn drop(&mut self) {
+ drop(self.iter.take());
+ }
+}
+
+struct TxAndIteratorPin<'a, I: Iterator<Item = IteratorItem<'a>> + 'a>(
+ Pin<Box<TxAndIterator<'a, I>>>,
+);
+
+impl<'a, I> Iterator for TxAndIteratorPin<'a, I>
+where
+ I: Iterator<Item = IteratorItem<'a>> + 'a,
+{
+ type Item = Result<(Value<'a>, Value<'a>)>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let iter_ref = unsafe {
+ let mut_ref: Pin<&mut TxAndIterator<'a, I>> = Pin::as_mut(&mut self.0);
+ Pin::get_unchecked_mut(mut_ref).iter.as_mut()
+ };
+ match iter_ref.unwrap().next() {
+ None => None,
+ Some(Err(e)) => Some(Err(e.into())),
+ Some(Ok((k, v))) => Some(Ok((k.into(), v.into()))),
+ }
+ }
+}