NextGraph
/
rust-rocksdb
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Make set_iterate_upper_bound method safe (#377)

Browse Source
master
wqfish 4 years ago committed by GitHub
parent 2ba70d304d
commit 81aa0163b8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 73 additions and 37 deletions
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Show Stats Download Patch File Download Diff File
  1. 55
      src/db.rs
  2. 16
      src/db_iterator.rs
  3. 25
      src/db_options.rs
  4. 8
      src/snapshot.rs

55
src/db.rs
Unescape View File

@ -416,15 +416,15 @@ impl DB {
pub fn iterator<'a: 'b, 'b>(&'a self, mode: IteratorMode) -> DBIterator<'b> { pub fn iterator<'a: 'b, 'b>(&'a self, mode: IteratorMode) -> DBIterator<'b> {
let readopts = ReadOptions::default(); let readopts = ReadOptions::default();
self.iterator_opt(mode, &readopts) self.iterator_opt(mode, readopts)
} }
pub fn iterator_opt<'a: 'b, 'b>( pub fn iterator_opt<'a: 'b, 'b>(
&'a self, &'a self,
mode: IteratorMode, mode: IteratorMode,
readopts: &ReadOptions, readopts: ReadOptions,
) -> DBIterator<'b> { ) -> DBIterator<'b> {
DBIterator::new(self, &readopts, mode) DBIterator::new(self, readopts, mode)
} }
/// Opens an iterator using the provided ReadOptions. /// Opens an iterator using the provided ReadOptions.
@ -432,10 +432,10 @@ impl DB {
pub fn iterator_cf_opt<'a: 'b, 'b>( pub fn iterator_cf_opt<'a: 'b, 'b>(
&'a self, &'a self,
cf_handle: &ColumnFamily, cf_handle: &ColumnFamily,
readopts: &ReadOptions, readopts: ReadOptions,
mode: IteratorMode, mode: IteratorMode,
) -> DBIterator<'b> { ) -> 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. /// 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> { pub fn full_iterator<'a: 'b, 'b>(&'a self, mode: IteratorMode) -> DBIterator<'b> {
let mut opts = ReadOptions::default(); let mut opts = ReadOptions::default();
opts.set_total_order_seek(true); 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> { 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); opts.set_prefix_same_as_start(true);
DBIterator::new( DBIterator::new(
self, self,
&opts, opts,
IteratorMode::From(prefix.as_ref(), Direction::Forward), IteratorMode::From(prefix.as_ref(), Direction::Forward),
) )
} }
@ -463,7 +463,7 @@ impl DB {
mode: IteratorMode, mode: IteratorMode,
) -> DBIterator<'b> { ) -> DBIterator<'b> {
let opts = ReadOptions::default(); 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>( pub fn full_iterator_cf<'a: 'b, 'b>(
@ -473,7 +473,7 @@ impl DB {
) -> DBIterator<'b> { ) -> DBIterator<'b> {
let mut opts = ReadOptions::default(); let mut opts = ReadOptions::default();
opts.set_total_order_seek(true); 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]>>( pub fn prefix_iterator_cf<'a: 'b, 'b, P: AsRef<[u8]>>(
@ -486,7 +486,7 @@ impl DB {
DBIterator::new_cf( DBIterator::new_cf(
self, self,
cf_handle, cf_handle,
&opts, opts,
IteratorMode::From(prefix.as_ref(), Direction::Forward), 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 /// Opens a raw iterator over the database, using the default read options
pub fn raw_iterator<'a: 'b, 'b>(&'a self) -> DBRawIterator<'b> { pub fn raw_iterator<'a: 'b, 'b>(&'a self) -> DBRawIterator<'b> {
let opts = ReadOptions::default(); 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 /// 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> { pub fn raw_iterator_cf<'a: 'b, 'b>(&'a self, cf_handle: &ColumnFamily) -> DBRawIterator<'b> {
let opts = ReadOptions::default(); 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 /// 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) DBRawIterator::new(self, readopts)
} }
@ -512,7 +512,7 @@ impl DB {
pub fn raw_iterator_cf_opt<'a: 'b, 'b>( pub fn raw_iterator_cf_opt<'a: 'b, 'b>(
&'a self, &'a self,
cf_handle: &ColumnFamily, cf_handle: &ColumnFamily,
readopts: &ReadOptions, readopts: ReadOptions,
) -> DBRawIterator<'b> { ) -> DBRawIterator<'b> {
DBRawIterator::new_cf(self, cf_handle, readopts) DBRawIterator::new_cf(self, cf_handle, readopts)
} }
@ -1056,6 +1056,31 @@ fn iterator_test_past_end() {
DB::destroy(&opts, path).unwrap(); 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::<Vec<_>>());
}
let opts = Options::default();
DB::destroy(&opts, path).unwrap();
}
#[test] #[test]
fn iterator_test_tailing() { fn iterator_test_tailing() {
let path = "_rust_rocksdb_iteratortest_tailing"; let path = "_rust_rocksdb_iteratortest_tailing";
@ -1070,7 +1095,7 @@ fn iterator_test_tailing() {
let r = db.put(k, v); let r = db.put(k, v);
assert!(r.is_ok()); 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 { for (k, v) in data_iter {
let r = db.put(k, v); let r = db.put(k, v);
assert!(r.is_ok()); assert!(r.is_ok());

16
src/db_iterator.rs
Unescape View File

@ -67,14 +67,21 @@ use std::slice;
/// ``` /// ```
pub struct DBRawIterator<'a> { pub struct DBRawIterator<'a> {
inner: *mut ffi::rocksdb_iterator_t, 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>, db: PhantomData<&'a DB>,
} }
impl<'a> DBRawIterator<'a> { 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 { unsafe {
DBRawIterator { DBRawIterator {
inner: ffi::rocksdb_create_iterator(db.inner, readopts.inner), inner: ffi::rocksdb_create_iterator(db.inner, readopts.inner),
_readopts: readopts,
db: PhantomData, db: PhantomData,
} }
} }
@ -83,11 +90,12 @@ impl<'a> DBRawIterator<'a> {
pub(crate) fn new_cf( pub(crate) fn new_cf(
db: &DB, db: &DB,
cf_handle: &ColumnFamily, cf_handle: &ColumnFamily,
readopts: &ReadOptions, readopts: ReadOptions,
) -> DBRawIterator<'a> { ) -> DBRawIterator<'a> {
unsafe { unsafe {
DBRawIterator { DBRawIterator {
inner: ffi::rocksdb_create_iterator_cf(db.inner, readopts.inner, cf_handle.inner), inner: ffi::rocksdb_create_iterator_cf(db.inner, readopts.inner, cf_handle.inner),
_readopts: readopts,
db: PhantomData, db: PhantomData,
} }
} }
@ -381,7 +389,7 @@ pub enum IteratorMode<'a> {
} }
impl<'a> DBIterator<'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 { let mut rv = DBIterator {
raw: DBRawIterator::new(db, readopts), raw: DBRawIterator::new(db, readopts),
direction: Direction::Forward, // blown away by set_mode() direction: Direction::Forward, // blown away by set_mode()
@ -394,7 +402,7 @@ impl<'a> DBIterator<'a> {
pub(crate) fn new_cf( pub(crate) fn new_cf(
db: &DB, db: &DB,
cf_handle: &ColumnFamily, cf_handle: &ColumnFamily,
readopts: &ReadOptions, readopts: ReadOptions,
mode: IteratorMode, mode: IteratorMode,
) -> DBIterator<'a> { ) -> DBIterator<'a> {
let mut rv = DBIterator { let mut rv = DBIterator {

25
src/db_options.rs
Unescape View File

@ -130,6 +130,7 @@ pub struct BlockBasedOptions {
pub struct ReadOptions { pub struct ReadOptions {
pub(crate) inner: *mut ffi::rocksdb_readoptions_t, pub(crate) inner: *mut ffi::rocksdb_readoptions_t,
iterate_upper_bound: Option<Vec<u8>>,
} }
// Safety note: auto-implementing Send on most db-related types is prevented by the inner FFI // Safety note: auto-implementing Send on most db-related types is prevented by the inner FFI
@ -1635,22 +1636,23 @@ 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. /// The upper bound itself is not included on the iteration result.
/// pub fn set_iterate_upper_bound<K: Into<Vec<u8>>>(&mut self, key: K) {
/// # Safety self.iterate_upper_bound = Some(key.into());
/// let upper_bound = self
/// This function will store a clone of key and will give a raw pointer of it to the .iterate_upper_bound
/// underlying C++ API, therefore, when given to any other [`DB`] method you must ensure .as_ref()
/// that this [`ReadOptions`] value does not leave the scope too early (e.g. `DB::iterator_cf_opt`). .expect("iterate_upper_bound must exist.");
pub unsafe fn set_iterate_upper_bound<K: AsRef<[u8]>>(&mut self, key: K) {
let key = key.as_ref(); unsafe {
ffi::rocksdb_readoptions_set_iterate_upper_bound( ffi::rocksdb_readoptions_set_iterate_upper_bound(
self.inner, self.inner,
key.as_ptr() as *const c_char, upper_bound.as_ptr() as *const c_char,
key.len() as size_t, upper_bound.len() as size_t,
); );
} }
}
pub fn set_prefix_same_as_start(&mut self, v: bool) { pub fn set_prefix_same_as_start(&mut self, v: bool) {
unsafe { ffi::rocksdb_readoptions_set_prefix_same_as_start(self.inner, v as c_uchar) } unsafe { ffi::rocksdb_readoptions_set_prefix_same_as_start(self.inner, v as c_uchar) }
@ -1702,6 +1704,7 @@ impl Default for ReadOptions {
unsafe { unsafe {
ReadOptions { ReadOptions {
inner: ffi::rocksdb_readoptions_create(), inner: ffi::rocksdb_readoptions_create(),
iterate_upper_bound: None,
} }
} }
} }

8
src/snapshot.rs
Unescape View File

@ -54,7 +54,7 @@ impl<'a> Snapshot<'a> {
pub fn iterator_opt(&self, mode: IteratorMode, mut readopts: ReadOptions) -> DBIterator<'a> { pub fn iterator_opt(&self, mode: IteratorMode, mut readopts: ReadOptions) -> DBIterator<'a> {
readopts.set_snapshot(self); readopts.set_snapshot(self);
DBIterator::new(self.db, &readopts, mode) DBIterator::new(self.db, readopts, mode)
} }
pub fn iterator_cf_opt( pub fn iterator_cf_opt(
@ -64,7 +64,7 @@ impl<'a> Snapshot<'a> {
mode: IteratorMode, mode: IteratorMode,
) -> DBIterator { ) -> DBIterator {
readopts.set_snapshot(self); 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. /// 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. /// 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 { pub fn raw_iterator_opt(&self, mut readopts: ReadOptions) -> DBRawIterator {
readopts.set_snapshot(self); 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. /// 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, mut readopts: ReadOptions,
) -> DBRawIterator { ) -> DBRawIterator {
readopts.set_snapshot(self); readopts.set_snapshot(self);
DBRawIterator::new_cf(self.db, cf_handle, &readopts) DBRawIterator::new_cf(self.db, cf_handle, readopts)
} }
pub fn get<K: AsRef<[u8]>>(&self, key: K) -> Result<Option<Vec<u8>>, Error> { pub fn get<K: AsRef<[u8]>>(&self, key: K) -> Result<Option<Vec<u8>>, Error> {

Write Preview
Loading…
Cancel
Save