Skip to content

Commit

Permalink
Add DBRawIteratorWithThreadMode::item method returning (key, value) (…
Browse files Browse the repository at this point in the history 
…rust-rocksdb#647)
  • Loading branch information
@mina86
mina86 authored Jun 30, 2022
1 parent 653d37f commit 0af1f8e
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 70 deletions.
63 changes: 37 additions & 26 deletions src/db_iterator.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -299,16 +299,7 @@ impl<'a, D: DBAccess> DBRawIteratorWithThreadMode<'a, D> {
/// Returns a slice of the current key.
pub fn key(&self) -> Option<&[u8]> {
if self.valid() {
// Safety Note: This is safe as all methods that may invalidate the buffer returned
// take `&mut self`, so borrow checker will prevent use of buffer after seek.
unsafe {
let mut key_len: size_t = 0;
let key_len_ptr: *mut size_t = &mut key_len;
let key_ptr =
ffi::rocksdb_iter_key(self.inner.as_ptr(), key_len_ptr) as *const c_uchar;

Some(slice::from_raw_parts(key_ptr, key_len as usize))
}
Some(self.key_impl())
} else {
None
}
Expand All @@ -317,20 +308,44 @@ impl<'a, D: DBAccess> DBRawIteratorWithThreadMode<'a, D> {
/// Returns a slice of the current value.
pub fn value(&self) -> Option<&[u8]> {
if self.valid() {
// Safety Note: This is safe as all methods that may invalidate the buffer returned
// take `&mut self`, so borrow checker will prevent use of buffer after seek.
unsafe {
let mut val_len: size_t = 0;
let val_len_ptr: *mut size_t = &mut val_len;
let val_ptr =
ffi::rocksdb_iter_value(self.inner.as_ptr(), val_len_ptr) as *const c_uchar;

Some(slice::from_raw_parts(val_ptr, val_len as usize))
}
Some(self.value_impl())
} else {
None
}
}

/// Returns pair with slice of the current key and current value.
pub fn item(&self) -> Option<(&[u8], &[u8])> {
if self.valid() {
Some((self.key_impl(), self.value_impl()))
} else {
None
}
}

/// Returns a slice of the current key; assumes the iterator is valid.
fn key_impl(&self) -> &[u8] {
// Safety Note: This is safe as all methods that may invalidate the buffer returned
// take `&mut self`, so borrow checker will prevent use of buffer after seek.
unsafe {
let mut key_len: size_t = 0;
let key_len_ptr: *mut size_t = &mut key_len;
let key_ptr = ffi::rocksdb_iter_key(self.inner.as_ptr(), key_len_ptr);
slice::from_raw_parts(key_ptr as *const c_uchar, key_len as usize)
}
}

/// Returns a slice of the current value; assumes the iterator is valid.
fn value_impl(&self) -> &[u8] {
// Safety Note: This is safe as all methods that may invalidate the buffer returned
// take `&mut self`, so borrow checker will prevent use of buffer after seek.
unsafe {
let mut val_len: size_t = 0;
let val_len_ptr: *mut size_t = &mut val_len;
let val_ptr = ffi::rocksdb_iter_value(self.inner.as_ptr(), val_len_ptr);
slice::from_raw_parts(val_ptr as *const c_uchar, val_len as usize)
}
}
}

impl<'a, D: DBAccess> Drop for DBRawIteratorWithThreadMode<'a, D> {
Expand Down Expand Up @@ -476,12 +491,8 @@ impl<'a, D: DBAccess> Iterator for DBIteratorWithThreadMode<'a, D> {
}
}

if self.raw.valid() {
// .key() and .value() only ever return None if valid == false, which we've just checked
Some((
Box::from(self.raw.key().unwrap()),
Box::from(self.raw.value().unwrap()),
))
if let Some((key, value)) = self.raw.item() {
Some((Box::from(key), Box::from(value)))
} else {
None
}
Expand Down
72 changes: 28 additions & 44 deletions tests/test_raw_iterator.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,23 @@ mod util;

use pretty_assertions::assert_eq;

use rocksdb::DB;
use rocksdb::{DBAccess, DBRawIteratorWithThreadMode, DB};
use util::DBPath;

fn assert_item<D: DBAccess>(iter: &DBRawIteratorWithThreadMode<'_, D>, key: &[u8], value: &[u8]) {
assert!(iter.valid());
assert_eq!(iter.key(), Some(key));
assert_eq!(iter.value(), Some(value));
assert_eq!(iter.item(), Some((key, value)));
}

fn assert_no_item<D: DBAccess>(iter: &DBRawIteratorWithThreadMode<'_, D>) {
assert!(!iter.valid());
assert_eq!(iter.key(), None);
assert_eq!(iter.value(), None);
assert_eq!(iter.item(), None);
}

#[test]
pub fn test_forwards_iteration() {
let n = DBPath::new("forwards_iteration");
Expand All @@ -31,24 +45,16 @@ pub fn test_forwards_iteration() {

let mut iter = db.raw_iterator();
iter.seek_to_first();

assert!(iter.valid());
assert_eq!(iter.key(), Some(b"k1".as_ref()));
assert_eq!(iter.value(), Some(b"v1".as_ref()));
assert_item(&iter, b"k1", b"v1");

iter.next();

assert!(iter.valid());
assert_eq!(iter.key(), Some(b"k2".as_ref()));
assert_eq!(iter.value(), Some(b"v2".as_ref()));
assert_item(&iter, b"k2", b"v2");

iter.next(); // k3
iter.next(); // k4
iter.next(); // invalid!

assert!(!iter.valid());
assert_eq!(iter.key(), None);
assert_eq!(iter.value(), None);
iter.next(); // invalid!
assert_no_item(&iter);
}
}

Expand All @@ -64,24 +70,16 @@ pub fn test_seek_last() {

let mut iter = db.raw_iterator();
iter.seek_to_last();

assert!(iter.valid());
assert_eq!(iter.key(), Some(b"k4".as_ref()));
assert_eq!(iter.value(), Some(b"v4".as_ref()));
assert_item(&iter, b"k4", b"v4");

iter.prev();

assert!(iter.valid());
assert_eq!(iter.key(), Some(b"k3".as_ref()));
assert_eq!(iter.value(), Some(b"v3".as_ref()));
assert_item(&iter, b"k3", b"v3");

iter.prev(); // k2
iter.prev(); // k1
iter.prev(); // invalid!

assert!(!iter.valid());
assert_eq!(iter.key(), None);
assert_eq!(iter.value(), None);
iter.prev(); // invalid!
assert_no_item(&iter);
}
}

Expand All @@ -97,16 +95,11 @@ pub fn test_seek() {
let mut iter = db.raw_iterator();
iter.seek(b"k2");

assert!(iter.valid());
assert_eq!(iter.key(), Some(b"k2".as_ref()));
assert_eq!(iter.value(), Some(b"v2".as_ref()));
assert_item(&iter, b"k2", b"v2");

// Check it gets the next key when the key doesn't exist
iter.seek(b"k3");

assert!(iter.valid());
assert_eq!(iter.key(), Some(b"k4".as_ref()));
assert_eq!(iter.value(), Some(b"v4".as_ref()));
assert_item(&iter, b"k4", b"v4");
}
}

Expand All @@ -121,10 +114,7 @@ pub fn test_seek_to_nonexistant() {

let mut iter = db.raw_iterator();
iter.seek(b"k2");

assert!(iter.valid());
assert_eq!(iter.key(), Some(b"k3".as_ref()));
assert_eq!(iter.value(), Some(b"v3".as_ref()));
assert_item(&iter, b"k3", b"v3");
}
}

Expand All @@ -139,16 +129,10 @@ pub fn test_seek_for_prev() {

let mut iter = db.raw_iterator();
iter.seek(b"k2");

assert!(iter.valid());
assert_eq!(iter.key(), Some(b"k2".as_ref()));
assert_eq!(iter.value(), Some(b"v2".as_ref()));
assert_item(&iter, b"k2", b"v2");

// Check it gets the previous key when the key doesn't exist
iter.seek_for_prev(b"k3");

assert!(iter.valid());
assert_eq!(iter.key(), Some(b"k2".as_ref()));
assert_eq!(iter.value(), Some(b"v2".as_ref()));
assert_item(&iter, b"k2", b"v2");
}
}

0 comments on commit 0af1f8e

Please sign in to comment.