Skip to content

R
rust-rocksdb
Commit ed089679 authored by Jay's avatar Jay
Browse files
Options

refactor iterator (#12)

parent 1f605a95
Show whitespace changes
Inline Side-by-side
Showing with 152 additions and 170 deletions
src/lib.rs
View file @ ed089679
...@@ -27,7 +27,6 @@ pub mod comparator; ...@@ -27,7 +27,6 @@ pub mod comparator;
pub use librocksdb_sys::{DBCompactionStyle, DBComparator, DBCompressionType, pub use librocksdb_sys::{DBCompactionStyle, DBComparator, DBCompressionType,
new_bloom_filter, self as rocksdb_ffi}; new_bloom_filter, self as rocksdb_ffi};
pub use rocksdb::{DB, DBIterator, DBVector, Direction, IteratorMode, Kv, pub use rocksdb::{DB, DBIterator, DBVector, Kv, SeekKey, Writable, WriteBatch};
Writable, WriteBatch};
pub use rocksdb_options::{BlockBasedOptions, Options, WriteOptions}; pub use rocksdb_options::{BlockBasedOptions, Options, WriteOptions};
pub use merge_operator::MergeOperands; pub use merge_operator::MergeOperands;
src/rocksdb.rs
View file @ ed089679
...@@ -53,71 +53,22 @@ pub struct Snapshot<'a> { ...@@ -53,71 +53,22 @@ pub struct Snapshot<'a> {
pub struct DBIterator<'a> { pub struct DBIterator<'a> {
db: &'a DB, db: &'a DB,
inner: rocksdb_ffi::DBIterator, inner: rocksdb_ffi::DBIterator,
direction: Direction,
just_seeked: bool,
} }
pub enum Direction { pub enum SeekKey<'a> {
Forward,
Reverse,
}
pub type Kv<'a> = (&'a [u8], &'a [u8]);
impl<'a> Iterator for DBIterator<'a> {
type Item = Kv<'a>;
fn next(&mut self) -> Option<Kv<'a>> {
let native_iter = self.inner;
if self.just_seeked {
self.just_seeked = false;
} else {
match self.direction {
Direction::Forward => unsafe {
rocksdb_ffi::rocksdb_iter_next(native_iter)
},
Direction::Reverse => unsafe {
rocksdb_ffi::rocksdb_iter_prev(native_iter)
},
}
}
if unsafe { rocksdb_ffi::rocksdb_iter_valid(native_iter) } {
let mut key_len: size_t = 0;
let key_len_ptr: *mut size_t = &mut key_len;
let mut val_len: size_t = 0;
let val_len_ptr: *mut size_t = &mut val_len;
let key_ptr = unsafe {
rocksdb_ffi::rocksdb_iter_key(native_iter, key_len_ptr)
};
let key = unsafe {
slice::from_raw_parts(key_ptr, key_len as usize)
};
let val_ptr = unsafe {
rocksdb_ffi::rocksdb_iter_value(native_iter, val_len_ptr)
};
let val = unsafe {
slice::from_raw_parts(val_ptr, val_len as usize)
};
Some((key, val))
} else {
None
}
}
}
pub enum IteratorMode<'a> {
Start, Start,
End, End,
From(&'a [u8], Direction), Key(&'a [u8]),
} }
impl<'a> From<&'a [u8]> for SeekKey<'a> {
fn from(bs: &'a [u8]) -> SeekKey {
SeekKey::Key(bs)
}
}
impl<'a> DBIterator<'a> { impl<'a> DBIterator<'a> {
fn new(db: &'a DB, fn new(db: &'a DB, readopts: &ReadOptions, key: SeekKey) -> DBIterator<'a> {
readopts: &ReadOptions,
mode: IteratorMode)
-> DBIterator<'a> {
unsafe { unsafe {
let iterator = rocksdb_ffi::rocksdb_create_iterator(db.inner, let iterator = rocksdb_ffi::rocksdb_create_iterator(db.inner,
readopts.inner); readopts.inner);
...@@ -125,35 +76,74 @@ impl<'a> DBIterator<'a> { ...@@ -125,35 +76,74 @@ impl<'a> DBIterator<'a> {
let mut rv = DBIterator { let mut rv = DBIterator {
db: db, db: db,
inner: iterator, inner: iterator,
direction: Direction::Forward, // blown away by set_mode()
just_seeked: false,
}; };
rv.set_mode(mode); rv.seek(key);
rv rv
} }
} }
pub fn set_mode(&mut self, mode: IteratorMode) { pub fn seek(&mut self, key: SeekKey) -> bool {
unsafe { unsafe {
match mode { match key {
IteratorMode::Start => { SeekKey::Start => {
rocksdb_ffi::rocksdb_iter_seek_to_first(self.inner); rocksdb_ffi::rocksdb_iter_seek_to_first(self.inner)
self.direction = Direction::Forward;
} }
IteratorMode::End => { SeekKey::End => {
rocksdb_ffi::rocksdb_iter_seek_to_last(self.inner); rocksdb_ffi::rocksdb_iter_seek_to_last(self.inner)
self.direction = Direction::Reverse;
} }
IteratorMode::From(key, dir) => { SeekKey::Key(key) => {
rocksdb_ffi::rocksdb_iter_seek(self.inner, rocksdb_ffi::rocksdb_iter_seek(self.inner,
key.as_ptr(), key.as_ptr(),
key.len() as size_t); key.len() as size_t)
self.direction = dir; }
} }
}; }
self.just_seeked = true; self.valid()
}
pub fn prev(&mut self) -> bool {
unsafe {
rocksdb_ffi::rocksdb_iter_prev(self.inner);
}
self.valid()
}
pub fn next(&mut self) -> bool {
unsafe {
rocksdb_ffi::rocksdb_iter_next(self.inner);
}
self.valid()
}
pub fn key(&self) -> &[u8] {
assert!(self.valid());
let mut key_len: size_t = 0;
let key_len_ptr: *mut size_t = &mut key_len;
unsafe {
let key_ptr = rocksdb_ffi::rocksdb_iter_key(self.inner,
key_len_ptr);
slice::from_raw_parts(key_ptr, key_len as usize)
}
}
pub fn value(&self) -> &[u8] {
assert!(self.valid());
let mut val_len: size_t = 0;
let val_len_ptr: *mut size_t = &mut val_len;
unsafe {
let val_ptr = rocksdb_ffi::rocksdb_iter_value(self.inner,
val_len_ptr);
slice::from_raw_parts(val_ptr, val_len as usize)
}
}
pub fn kv(&self) -> Option<(Vec<u8>, Vec<u8>)> {
if self.valid() {
Some((self.key().to_vec(), self.value().to_vec()))
} else {
None
} }
} }
...@@ -164,7 +154,7 @@ impl<'a> DBIterator<'a> { ...@@ -164,7 +154,7 @@ impl<'a> DBIterator<'a> {
fn new_cf(db: &'a DB, fn new_cf(db: &'a DB,
cf_handle: DBCFHandle, cf_handle: DBCFHandle,
readopts: &ReadOptions, readopts: &ReadOptions,
mode: IteratorMode) key: SeekKey)
-> Result<DBIterator<'a>, String> { -> Result<DBIterator<'a>, String> {
unsafe { unsafe {
let iterator = let iterator =
...@@ -175,17 +165,29 @@ impl<'a> DBIterator<'a> { ...@@ -175,17 +165,29 @@ impl<'a> DBIterator<'a> {
let mut rv = DBIterator { let mut rv = DBIterator {
db: db, db: db,
inner: iterator, inner: iterator,
direction: Direction::Forward, // blown away by set_mode()
just_seeked: false,
}; };
rv.set_mode(mode); rv.seek(key);
Ok(rv) Ok(rv)
} }
} }
} }
pub type Kv = (Vec<u8>, Vec<u8>);
impl<'a> Iterator for DBIterator<'a> {
type Item = Kv;
fn next(&mut self) -> Option<Kv> {
let kv = self.kv();
if kv.is_some() {
DBIterator::next(self);
}
kv
}
}
impl<'a> Drop for DBIterator<'a> { impl<'a> Drop for DBIterator<'a> {
fn drop(&mut self) { fn drop(&mut self) {
unsafe { unsafe {
...@@ -196,19 +198,18 @@ impl<'a> Drop for DBIterator<'a> { ...@@ -196,19 +198,18 @@ impl<'a> Drop for DBIterator<'a> {
impl<'a> Snapshot<'a> { impl<'a> Snapshot<'a> {
pub fn new(db: &DB) -> Snapshot { pub fn new(db: &DB) -> Snapshot {
let snapshot = unsafe { let snapshot =
rocksdb_ffi::rocksdb_create_snapshot(db.inner) unsafe { rocksdb_ffi::rocksdb_create_snapshot(db.inner) };
};
Snapshot { Snapshot {
db: db, db: db,
inner: snapshot, inner: snapshot,
} }
} }
pub fn iterator(&self, mode: IteratorMode) -> DBIterator { pub fn iter(&self, key: SeekKey) -> DBIterator {
let mut readopts = ReadOptions::new(); let mut readopts = ReadOptions::new();
readopts.set_snapshot(self); readopts.set_snapshot(self);
DBIterator::new(self.db, &readopts, mode) DBIterator::new(self.db, &readopts, key)
} }
pub fn get(&self, key: &[u8]) -> Result<Option<DBVector>, String> { pub fn get(&self, key: &[u8]) -> Result<Option<DBVector>, String> {
...@@ -326,10 +327,7 @@ impl DB { ...@@ -326,10 +327,7 @@ impl DB {
// We need to store our CStrings in an intermediate vector // We need to store our CStrings in an intermediate vector
// so that their pointers remain valid. // so that their pointers remain valid.
let c_cfs: Vec<CString> = cfs_v.iter() let c_cfs: Vec<CString> = cfs_v.iter()
.map(|cf| { .map(|cf| CString::new(cf.as_bytes()).unwrap())
CString::new(cf.as_bytes())
.unwrap()
})
.collect(); .collect();
let cfnames: Vec<*const _> = c_cfs.iter() let cfnames: Vec<*const _> = c_cfs.iter()
...@@ -337,14 +335,12 @@ impl DB { ...@@ -337,14 +335,12 @@ impl DB {
.collect(); .collect();
// These handles will be populated by DB. // These handles will be populated by DB.
let cfhandles: Vec<rocksdb_ffi::DBCFHandle> = let cfhandles: Vec<rocksdb_ffi::DBCFHandle> = cfs_v.iter()
cfs_v.iter()
.map(|_| rocksdb_ffi::DBCFHandle(0 as *mut c_void)) .map(|_| rocksdb_ffi::DBCFHandle(0 as *mut c_void))
.collect(); .collect();
// TODO(tyler) allow options to be passed in. // TODO(tyler) allow options to be passed in.
let cfopts: Vec<rocksdb_ffi::DBOptions> = let cfopts: Vec<rocksdb_ffi::DBOptions> = cfs_v.iter()
cfs_v.iter()
.map(|_| unsafe { rocksdb_ffi::rocksdb_options_create() }) .map(|_| unsafe { rocksdb_ffi::rocksdb_options_create() })
.collect(); .collect();
...@@ -583,17 +579,17 @@ impl DB { ...@@ -583,17 +579,17 @@ impl DB {
self.cfs.get(name) self.cfs.get(name)
} }
pub fn iterator(&self, mode: IteratorMode) -> DBIterator { pub fn iter(&self, key: SeekKey) -> DBIterator {
let opts = ReadOptions::new(); let opts = ReadOptions::new();
DBIterator::new(&self, &opts, mode) DBIterator::new(&self, &opts, key)
} }
pub fn iterator_cf(&self, pub fn iter_cf(&self,
cf_handle: DBCFHandle, cf_handle: DBCFHandle,
mode: IteratorMode) key: SeekKey)
-> Result<DBIterator, String> { -> Result<DBIterator, String> {
let opts = ReadOptions::new(); let opts = ReadOptions::new();
DBIterator::new_cf(&self, cf_handle, &opts, mode) DBIterator::new_cf(&self, cf_handle, &opts, key)
} }
pub fn snapshot(&self) -> Snapshot { pub fn snapshot(&self) -> Snapshot {
...@@ -1091,7 +1087,7 @@ mod test { ...@@ -1091,7 +1087,7 @@ mod test {
db.put(b"k1", b"v1111").expect(""); db.put(b"k1", b"v1111").expect("");
db.put(b"k2", b"v2222").expect(""); db.put(b"k2", b"v2222").expect("");
db.put(b"k3", b"v3333").expect(""); db.put(b"k3", b"v3333").expect("");
let iter = db.iterator(IteratorMode::Start); let iter = db.iter(SeekKey::Start);
for (k, v) in iter { for (k, v) in iter {
println!("Hello {}: {}", println!("Hello {}: {}",
str::from_utf8(&*k).unwrap(), str::from_utf8(&*k).unwrap(),
......
src/rocksdb_options.rs
View file @ ed089679
...@@ -60,9 +60,8 @@ impl Drop for WriteOptions { ...@@ -60,9 +60,8 @@ impl Drop for WriteOptions {
impl Default for BlockBasedOptions { impl Default for BlockBasedOptions {
fn default() -> BlockBasedOptions { fn default() -> BlockBasedOptions {
let block_opts = unsafe { let block_opts =
rocksdb_ffi::rocksdb_block_based_options_create() unsafe { rocksdb_ffi::rocksdb_block_based_options_create() };
};
let rocksdb_ffi::DBBlockBasedTableOptions(opt_ptr) = block_opts; let rocksdb_ffi::DBBlockBasedTableOptions(opt_ptr) = block_opts;
if opt_ptr.is_null() { if opt_ptr.is_null() {
panic!("Could not create rocksdb block based options".to_string()); panic!("Could not create rocksdb block based options".to_string());
......
test/test_iterator.rs
View file @ ed089679
use rocksdb::{DB, Direction, IteratorMode, Writable, Kv}; use rocksdb::{DB, Writable, SeekKey, DBIterator, Kv};
use tempdir::TempDir; use tempdir::TempDir;
fn collect<'a, T: Iterator<Item=Kv<'a>>>(iter: T) -> Vec<(Vec<u8>, Vec<u8>)> { fn prev_collect<'a>(mut iter: DBIterator<'a>) -> Vec<Kv> {
iter.map(|(k, v)| (k.to_vec(), v.to_vec())).collect() let mut buf = vec![];
while iter.valid() {
buf.push(iter.kv().unwrap());
iter.prev();
}
buf
} }
#[test] #[test]
...@@ -28,96 +33,79 @@ pub fn test_iterator() { ...@@ -28,96 +33,79 @@ pub fn test_iterator() {
(k2.to_vec(), v2.to_vec()), (k2.to_vec(), v2.to_vec()),
(k3.to_vec(), v3.to_vec())]; (k3.to_vec(), v3.to_vec())];
let iterator1 = db.iterator(IteratorMode::Start); let mut iter = db.iter(SeekKey::Start);
assert_eq!(collect(iterator1), expected); assert_eq!(iter.collect::<Vec<_>>(), expected);
// Test that it's idempotent // Test that it's idempotent
let iterator1 = db.iterator(IteratorMode::Start); iter = db.iter(SeekKey::Start);
assert_eq!(collect(iterator1), expected); assert_eq!(iter.collect::<Vec<_>>(), expected);
let iterator1 = db.iterator(IteratorMode::Start);
assert_eq!(collect(iterator1), expected);
let iterator1 = db.iterator(IteratorMode::Start);
assert_eq!(collect(iterator1), expected);
// Test it in reverse a few times // Test it in reverse a few times
let iterator1 = db.iterator(IteratorMode::End); iter = db.iter(SeekKey::End);
let mut tmp_vec = collect(iterator1); let mut tmp_vec = prev_collect(iter);
tmp_vec.reverse(); tmp_vec.reverse();
assert_eq!(tmp_vec, expected); assert_eq!(tmp_vec, expected);
let iterator1 = db.iterator(IteratorMode::End); iter = db.iter(SeekKey::End);
let mut tmp_vec = collect(iterator1); let mut tmp_vec = prev_collect(iter);
tmp_vec.reverse();
assert_eq!(tmp_vec, expected);
let iterator1 = db.iterator(IteratorMode::End);
let mut tmp_vec = collect(iterator1);
tmp_vec.reverse();
assert_eq!(tmp_vec, expected);
let iterator1 = db.iterator(IteratorMode::End);
let mut tmp_vec = collect(iterator1);
tmp_vec.reverse();
assert_eq!(tmp_vec, expected);
let iterator1 = db.iterator(IteratorMode::End);
let mut tmp_vec = collect(iterator1);
tmp_vec.reverse(); tmp_vec.reverse();
assert_eq!(tmp_vec, expected); assert_eq!(tmp_vec, expected);
// Try it forward again // Try it forward again
let iterator1 = db.iterator(IteratorMode::Start); iter = db.iter(SeekKey::Start);
assert_eq!(collect(iterator1), expected); assert_eq!(iter.collect::<Vec<_>>(), expected);
let iterator1 = db.iterator(IteratorMode::Start); iter = db.iter(SeekKey::Start);
assert_eq!(collect(iterator1), expected); assert_eq!(iter.collect::<Vec<_>>(), expected);
let old_iterator = db.iterator(IteratorMode::Start); let old_iterator = db.iter(SeekKey::Start);
let p = db.put(&*k4, &*v4); let p = db.put(&*k4, &*v4);
assert!(p.is_ok()); assert!(p.is_ok());
let expected2 = vec![(k1.to_vec(), v1.to_vec()), let expected2 = vec![(k1.to_vec(), v1.to_vec()),
(k2.to_vec(), v2.to_vec()), (k2.to_vec(), v2.to_vec()),
(k3.to_vec(), v3.to_vec()), (k3.to_vec(), v3.to_vec()),
(k4.to_vec(), v4.to_vec())]; (k4.to_vec(), v4.to_vec())];
assert_eq!(collect(old_iterator), expected); assert_eq!(old_iterator.collect::<Vec<_>>(), expected);
let iterator1 = db.iterator(IteratorMode::Start); iter = db.iter(SeekKey::Start);
assert_eq!(collect(iterator1), expected2); assert_eq!(iter.collect::<Vec<_>>(), expected2);
let iterator1 = db.iterator(IteratorMode::From(b"k2", iter = db.iter(SeekKey::Key(k2));
Direction::Forward));
let expected = vec![(k2.to_vec(), v2.to_vec()), let expected = vec![(k2.to_vec(), v2.to_vec()),
(k3.to_vec(), v3.to_vec()), (k3.to_vec(), v3.to_vec()),
(k4.to_vec(), v4.to_vec())]; (k4.to_vec(), v4.to_vec())];
assert_eq!(collect(iterator1), expected); assert_eq!(iter.collect::<Vec<_>>(), expected);
let iterator1 = db.iterator(IteratorMode::From(b"k2", iter = db.iter(SeekKey::Key(k2));
Direction::Reverse));
let expected = vec![(k2.to_vec(), v2.to_vec()), (k1.to_vec(), v1.to_vec())]; let expected = vec![(k2.to_vec(), v2.to_vec()), (k1.to_vec(), v1.to_vec())];
assert_eq!(collect(iterator1), expected); assert_eq!(prev_collect(iter), expected);
let iterator1 = db.iterator(IteratorMode::From(b"k0", Direction::Forward)); iter = db.iter(SeekKey::Key(b"k0"));
assert!(iterator1.valid()); assert!(iter.valid());
let iterator2 = db.iterator(IteratorMode::From(b"k1", Direction::Forward)); iter.seek(SeekKey::Key(b"k1"));
assert!(iterator2.valid()); assert!(iter.valid());
let iterator3 = db.iterator(IteratorMode::From(b"k11", Direction::Forward)); iter.seek(SeekKey::Key(b"k11"));
assert!(iterator3.valid()); assert!(iter.valid());
let iterator4 = db.iterator(IteratorMode::From(b"k5", Direction::Forward)); iter.seek(SeekKey::Key(b"k5"));
assert!(!iterator4.valid()); assert!(!iter.valid());
let iterator5 = db.iterator(IteratorMode::From(b"k0", Direction::Reverse)); iter.seek(SeekKey::Key(b"k0"));
assert!(iterator5.valid()); assert!(iter.valid());
let iterator6 = db.iterator(IteratorMode::From(b"k1", Direction::Reverse)); iter.seek(SeekKey::Key(b"k1"));
assert!(iterator6.valid()); assert!(iter.valid());
let iterator7 = db.iterator(IteratorMode::From(b"k11", Direction::Reverse)); iter.seek(SeekKey::Key(b"k11"));
assert!(iterator7.valid()); assert!(iter.valid());
let iterator8 = db.iterator(IteratorMode::From(b"k5", Direction::Reverse)); iter.seek(SeekKey::Key(b"k5"));
assert!(!iterator8.valid()); assert!(!iter.valid());
let mut iterator1 = db.iterator(IteratorMode::From(b"k4", Direction::Forward)); iter.seek(SeekKey::Key(b"k4"));
iterator1.next(); assert!(iter.valid());
assert!(iterator1.valid()); iter.prev();
iterator1.next(); assert!(iter.valid());
assert!(!iterator1.valid()); iter.next();
assert!(iter.valid());
iter.next();
assert!(!iter.valid());
// Once iterator is invalid, it can't be reverted.
iter.prev();
assert!(!iter.valid());
} }
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment