diff --git a/src/db.rs b/src/db.rs index 5c1bc5e..9a92a04 100644 --- a/src/db.rs +++ b/src/db.rs @@ -416,15 +416,15 @@ impl DB { pub fn iterator<'a: 'b, 'b>(&'a self, mode: IteratorMode) -> DBIterator<'b> { let readopts = ReadOptions::default(); - self.iterator_opt(mode, &readopts) + self.iterator_opt(mode, readopts) } pub fn iterator_opt<'a: 'b, 'b>( &'a self, mode: IteratorMode, - readopts: &ReadOptions, + readopts: ReadOptions, ) -> DBIterator<'b> { - DBIterator::new(self, &readopts, mode) + DBIterator::new(self, readopts, mode) } /// Opens an iterator using the provided ReadOptions. @@ -432,10 +432,10 @@ impl DB { pub fn iterator_cf_opt<'a: 'b, 'b>( &'a self, cf_handle: &ColumnFamily, - readopts: &ReadOptions, + readopts: ReadOptions, mode: IteratorMode, ) -> DBIterator<'b> { - DBIterator::new_cf(self, cf_handle, &readopts, mode) + DBIterator::new_cf(self, cf_handle, readopts, mode) } /// Opens an iterator with `set_total_order_seek` enabled. @@ -444,7 +444,7 @@ impl DB { pub fn full_iterator<'a: 'b, 'b>(&'a self, mode: IteratorMode) -> DBIterator<'b> { let mut opts = ReadOptions::default(); opts.set_total_order_seek(true); - DBIterator::new(self, &opts, mode) + DBIterator::new(self, opts, mode) } pub fn prefix_iterator<'a: 'b, 'b, P: AsRef<[u8]>>(&'a self, prefix: P) -> DBIterator<'b> { @@ -452,7 +452,7 @@ impl DB { opts.set_prefix_same_as_start(true); DBIterator::new( self, - &opts, + opts, IteratorMode::From(prefix.as_ref(), Direction::Forward), ) } @@ -463,7 +463,7 @@ impl DB { mode: IteratorMode, ) -> DBIterator<'b> { let opts = ReadOptions::default(); - DBIterator::new_cf(self, cf_handle, &opts, mode) + DBIterator::new_cf(self, cf_handle, opts, mode) } pub fn full_iterator_cf<'a: 'b, 'b>( @@ -473,7 +473,7 @@ impl DB { ) -> DBIterator<'b> { let mut opts = ReadOptions::default(); opts.set_total_order_seek(true); - DBIterator::new_cf(self, cf_handle, &opts, mode) + DBIterator::new_cf(self, cf_handle, opts, mode) } pub fn prefix_iterator_cf<'a: 'b, 'b, P: AsRef<[u8]>>( @@ -486,7 +486,7 @@ impl DB { DBIterator::new_cf( self, cf_handle, - &opts, + opts, IteratorMode::From(prefix.as_ref(), Direction::Forward), ) } @@ -494,17 +494,17 @@ impl DB { /// Opens a raw iterator over the database, using the default read options pub fn raw_iterator<'a: 'b, 'b>(&'a self) -> DBRawIterator<'b> { let opts = ReadOptions::default(); - DBRawIterator::new(self, &opts) + DBRawIterator::new(self, opts) } /// Opens a raw iterator over the given column family, using the default read options pub fn raw_iterator_cf<'a: 'b, 'b>(&'a self, cf_handle: &ColumnFamily) -> DBRawIterator<'b> { let opts = ReadOptions::default(); - DBRawIterator::new_cf(self, cf_handle, &opts) + DBRawIterator::new_cf(self, cf_handle, opts) } /// Opens a raw iterator over the database, using the given read options - pub fn raw_iterator_opt<'a: 'b, 'b>(&'a self, readopts: &ReadOptions) -> DBRawIterator<'b> { + pub fn raw_iterator_opt<'a: 'b, 'b>(&'a self, readopts: ReadOptions) -> DBRawIterator<'b> { DBRawIterator::new(self, readopts) } @@ -512,7 +512,7 @@ impl DB { pub fn raw_iterator_cf_opt<'a: 'b, 'b>( &'a self, cf_handle: &ColumnFamily, - readopts: &ReadOptions, + readopts: ReadOptions, ) -> DBRawIterator<'b> { DBRawIterator::new_cf(self, cf_handle, readopts) } @@ -1056,6 +1056,31 @@ fn iterator_test_past_end() { DB::destroy(&opts, path).unwrap(); } +#[test] +fn iterator_test_upper_bound() { + let path = "_rust_rocksdb_iteratortest_upper_bound"; + { + let db = DB::open_default(path).unwrap(); + db.put(b"k1", b"v1").unwrap(); + db.put(b"k2", b"v2").unwrap(); + db.put(b"k3", b"v3").unwrap(); + db.put(b"k4", b"v4").unwrap(); + db.put(b"k5", b"v5").unwrap(); + + let mut readopts = ReadOptions::default(); + readopts.set_iterate_upper_bound(b"k4".to_vec()); + + let iter = db.iterator_opt(IteratorMode::Start, readopts); + let expected: Vec<_> = vec![(b"k1", b"v1"), (b"k2", b"v2"), (b"k3", b"v3")] + .into_iter() + .map(|(k, v)| (k.to_vec().into_boxed_slice(), v.to_vec().into_boxed_slice())) + .collect(); + assert_eq!(expected, iter.collect::>()); + } + let opts = Options::default(); + DB::destroy(&opts, path).unwrap(); +} + #[test] fn iterator_test_tailing() { let path = "_rust_rocksdb_iteratortest_tailing"; @@ -1070,7 +1095,7 @@ fn iterator_test_tailing() { let r = db.put(k, v); assert!(r.is_ok()); - let tail_iter = db.iterator_opt(IteratorMode::Start, &ro); + let tail_iter = db.iterator_opt(IteratorMode::Start, ro); for (k, v) in data_iter { let r = db.put(k, v); assert!(r.is_ok()); diff --git a/src/db_iterator.rs b/src/db_iterator.rs index f539e78..576da45 100644 --- a/src/db_iterator.rs +++ b/src/db_iterator.rs @@ -67,14 +67,21 @@ use std::slice; /// ``` pub struct DBRawIterator<'a> { inner: *mut ffi::rocksdb_iterator_t, + + /// When iterate_upper_bound is set, the inner C iterator keeps a pointer to the upper bound + /// inside `_readopts`. Storing this makes sure the upper bound is always alive when the + /// iterator is being used. + _readopts: ReadOptions, + db: PhantomData<&'a DB>, } impl<'a> DBRawIterator<'a> { - pub(crate) fn new(db: &DB, readopts: &ReadOptions) -> DBRawIterator<'a> { + pub(crate) fn new(db: &DB, readopts: ReadOptions) -> DBRawIterator<'a> { unsafe { DBRawIterator { inner: ffi::rocksdb_create_iterator(db.inner, readopts.inner), + _readopts: readopts, db: PhantomData, } } @@ -83,11 +90,12 @@ impl<'a> DBRawIterator<'a> { pub(crate) fn new_cf( db: &DB, cf_handle: &ColumnFamily, - readopts: &ReadOptions, + readopts: ReadOptions, ) -> DBRawIterator<'a> { unsafe { DBRawIterator { inner: ffi::rocksdb_create_iterator_cf(db.inner, readopts.inner, cf_handle.inner), + _readopts: readopts, db: PhantomData, } } @@ -381,7 +389,7 @@ pub enum IteratorMode<'a> { } impl<'a> DBIterator<'a> { - pub(crate) fn new(db: &DB, readopts: &ReadOptions, mode: IteratorMode) -> DBIterator<'a> { + pub(crate) fn new(db: &DB, readopts: ReadOptions, mode: IteratorMode) -> DBIterator<'a> { let mut rv = DBIterator { raw: DBRawIterator::new(db, readopts), direction: Direction::Forward, // blown away by set_mode() @@ -394,7 +402,7 @@ impl<'a> DBIterator<'a> { pub(crate) fn new_cf( db: &DB, cf_handle: &ColumnFamily, - readopts: &ReadOptions, + readopts: ReadOptions, mode: IteratorMode, ) -> DBIterator<'a> { let mut rv = DBIterator { diff --git a/src/db_options.rs b/src/db_options.rs index 67d2b34..4c80e32 100644 --- a/src/db_options.rs +++ b/src/db_options.rs @@ -130,6 +130,7 @@ pub struct BlockBasedOptions { pub struct ReadOptions { pub(crate) inner: *mut ffi::rocksdb_readoptions_t, + iterate_upper_bound: Option>, } // Safety note: auto-implementing Send on most db-related types is prevented by the inner FFI @@ -1635,21 +1636,22 @@ impl ReadOptions { } } - /// Set the upper bound for an iterator. + /// Sets the upper bound for an iterator. /// The upper bound itself is not included on the iteration result. - /// - /// # Safety - /// - /// This function will store a clone of key and will give a raw pointer of it to the - /// underlying C++ API, therefore, when given to any other [`DB`] method you must ensure - /// that this [`ReadOptions`] value does not leave the scope too early (e.g. `DB::iterator_cf_opt`). - pub unsafe fn set_iterate_upper_bound>(&mut self, key: K) { - let key = key.as_ref(); - ffi::rocksdb_readoptions_set_iterate_upper_bound( - self.inner, - key.as_ptr() as *const c_char, - key.len() as size_t, - ); + pub fn set_iterate_upper_bound>>(&mut self, key: K) { + self.iterate_upper_bound = Some(key.into()); + let upper_bound = self + .iterate_upper_bound + .as_ref() + .expect("iterate_upper_bound must exist."); + + unsafe { + ffi::rocksdb_readoptions_set_iterate_upper_bound( + self.inner, + upper_bound.as_ptr() as *const c_char, + upper_bound.len() as size_t, + ); + } } pub fn set_prefix_same_as_start(&mut self, v: bool) { @@ -1702,6 +1704,7 @@ impl Default for ReadOptions { unsafe { ReadOptions { inner: ffi::rocksdb_readoptions_create(), + iterate_upper_bound: None, } } } diff --git a/src/snapshot.rs b/src/snapshot.rs index a73fc77..2d5df31 100644 --- a/src/snapshot.rs +++ b/src/snapshot.rs @@ -54,7 +54,7 @@ impl<'a> Snapshot<'a> { pub fn iterator_opt(&self, mode: IteratorMode, mut readopts: ReadOptions) -> DBIterator<'a> { readopts.set_snapshot(self); - DBIterator::new(self.db, &readopts, mode) + DBIterator::new(self.db, readopts, mode) } pub fn iterator_cf_opt( @@ -64,7 +64,7 @@ impl<'a> Snapshot<'a> { mode: IteratorMode, ) -> DBIterator { readopts.set_snapshot(self); - DBIterator::new_cf(self.db, cf_handle, &readopts, mode) + DBIterator::new_cf(self.db, cf_handle, readopts, mode) } /// Opens a raw iterator over the data in this snapshot, using the default read options. @@ -82,7 +82,7 @@ impl<'a> Snapshot<'a> { /// Opens a raw iterator over the data in this snapshot, using the given read options. pub fn raw_iterator_opt(&self, mut readopts: ReadOptions) -> DBRawIterator { readopts.set_snapshot(self); - DBRawIterator::new(self.db, &readopts) + DBRawIterator::new(self.db, readopts) } /// Opens a raw iterator over the data in this snapshot under the given column family, using the given read options. @@ -92,7 +92,7 @@ impl<'a> Snapshot<'a> { mut readopts: ReadOptions, ) -> DBRawIterator { readopts.set_snapshot(self); - DBRawIterator::new_cf(self.db, cf_handle, &readopts) + DBRawIterator::new_cf(self.db, cf_handle, readopts) } pub fn get>(&self, key: K) -> Result>, Error> {