Fix multi_get (#511)

master
mikhailOK 4 years ago committed by GitHub
parent c0cb8a77c2
commit a0ad4fdcfe
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 41
      src/db.rs
  2. 22
      tests/test_db.rs

@ -747,7 +747,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
} }
/// Return the values associated with the given keys. /// Return the values associated with the given keys.
pub fn multi_get<K, I>(&self, keys: I) -> Result<Vec<Vec<u8>>, Error> pub fn multi_get<K, I>(&self, keys: I) -> Vec<Result<Option<Vec<u8>>, Error>>
where where
K: AsRef<[u8]>, K: AsRef<[u8]>,
I: IntoIterator<Item = K>, I: IntoIterator<Item = K>,
@ -760,7 +760,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
&self, &self,
keys: I, keys: I,
readopts: &ReadOptions, readopts: &ReadOptions,
) -> Result<Vec<Vec<u8>>, Error> ) -> Vec<Result<Option<Vec<u8>>, Error>>
where where
K: AsRef<[u8]>, K: AsRef<[u8]>,
I: IntoIterator<Item = K>, I: IntoIterator<Item = K>,
@ -773,8 +773,9 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
let mut values = vec![ptr::null_mut(); keys.len()]; let mut values = vec![ptr::null_mut(); keys.len()];
let mut values_sizes = vec![0_usize; keys.len()]; let mut values_sizes = vec![0_usize; keys.len()];
let mut errors = vec![ptr::null_mut(); keys.len()];
unsafe { unsafe {
ffi_try!(ffi::rocksdb_multi_get( ffi::rocksdb_multi_get(
self.inner, self.inner,
readopts.inner, readopts.inner,
ptr_keys.len(), ptr_keys.len(),
@ -782,14 +783,15 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
keys_sizes.as_ptr(), keys_sizes.as_ptr(),
values.as_mut_ptr(), values.as_mut_ptr(),
values_sizes.as_mut_ptr(), values_sizes.as_mut_ptr(),
)); errors.as_mut_ptr(),
);
} }
Ok(convert_values(values, values_sizes)) convert_values(values, values_sizes, errors)
} }
/// Return the values associated with the given keys and column families. /// Return the values associated with the given keys and column families.
pub fn multi_get_cf<K, I, W>(&self, keys: I) -> Result<Vec<Vec<u8>>, Error> pub fn multi_get_cf<K, I, W>(&self, keys: I) -> Vec<Result<Option<Vec<u8>>, Error>>
where where
K: AsRef<[u8]>, K: AsRef<[u8]>,
I: IntoIterator<Item = (W, K)>, I: IntoIterator<Item = (W, K)>,
@ -803,7 +805,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
&self, &self,
keys: I, keys: I,
readopts: &ReadOptions, readopts: &ReadOptions,
) -> Result<Vec<Vec<u8>>, Error> ) -> Vec<Result<Option<Vec<u8>>, Error>>
where where
K: AsRef<[u8]>, K: AsRef<[u8]>,
I: IntoIterator<Item = (W, K)>, I: IntoIterator<Item = (W, K)>,
@ -828,8 +830,9 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
let mut values = vec![ptr::null_mut(); boxed_keys.len()]; let mut values = vec![ptr::null_mut(); boxed_keys.len()];
let mut values_sizes = vec![0_usize; boxed_keys.len()]; let mut values_sizes = vec![0_usize; boxed_keys.len()];
let mut errors = vec![ptr::null_mut(); boxed_keys.len()];
unsafe { unsafe {
ffi_try!(ffi::rocksdb_multi_get_cf( ffi::rocksdb_multi_get_cf(
self.inner, self.inner,
readopts.inner, readopts.inner,
ptr_cfs.as_ptr(), ptr_cfs.as_ptr(),
@ -838,10 +841,11 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
keys_sizes.as_ptr(), keys_sizes.as_ptr(),
values.as_mut_ptr(), values.as_mut_ptr(),
values_sizes.as_mut_ptr(), values_sizes.as_mut_ptr(),
)); errors.as_mut_ptr(),
);
} }
Ok(convert_values(values, values_sizes)) convert_values(values, values_sizes, errors)
} }
/// Returns `false` if the given key definitely doesn't exist in the database, otherwise returns /// Returns `false` if the given key definitely doesn't exist in the database, otherwise returns
@ -1835,16 +1839,25 @@ fn convert_options(opts: &[(&str, &str)]) -> Result<Vec<(CString, CString)>, Err
.collect() .collect()
} }
fn convert_values(values: Vec<*mut c_char>, values_sizes: Vec<usize>) -> Vec<Vec<u8>> { fn convert_values(
values: Vec<*mut c_char>,
values_sizes: Vec<usize>,
errors: Vec<*mut c_char>,
) -> Vec<Result<Option<Vec<u8>>, Error>> {
values values
.into_iter() .into_iter()
.zip(values_sizes.into_iter()) .zip(values_sizes.into_iter())
.map(|(v, s)| { .zip(errors.into_iter())
let value = unsafe { slice::from_raw_parts(v as *const u8, s) }.into(); .map(|((v, s), e)| {
if e.is_null() {
let value = unsafe { crate::ffi_util::raw_data(v, s) };
unsafe { unsafe {
ffi::rocksdb_free(v as *mut c_void); ffi::rocksdb_free(v as *mut c_void);
} }
value Ok(value)
} else {
Err(Error::new(crate::ffi_util::error_message(e)))
}
}) })
.collect() .collect()
} }

@ -928,13 +928,17 @@ fn multi_get() {
db.put(b"k1", b"v1").unwrap(); db.put(b"k1", b"v1").unwrap();
db.put(b"k2", b"v2").unwrap(); db.put(b"k2", b"v2").unwrap();
let _ = db.multi_get(&[b"k0"; 40]);
let values = db let values = db
.multi_get(&[b"k0", b"k1", b"k2"]) .multi_get(&[b"k0", b"k1", b"k2"])
.expect("multi_get failed"); .into_iter()
.map(Result::unwrap)
.collect::<Vec<_>>();
assert_eq!(3, values.len()); assert_eq!(3, values.len());
assert!(values[0].is_empty()); assert_eq!(values[0], None);
assert_eq!(values[1], b"v1"); assert_eq!(values[1], Some(b"v1".to_vec()));
assert_eq!(values[2], b"v2"); assert_eq!(values[2], Some(b"v2".to_vec()));
} }
} }
@ -958,11 +962,13 @@ fn multi_get_cf() {
let values = db let values = db
.multi_get_cf(vec![(cf0, b"k0"), (cf1, b"k1"), (cf2, b"k2")]) .multi_get_cf(vec![(cf0, b"k0"), (cf1, b"k1"), (cf2, b"k2")])
.expect("multi_get failed"); .into_iter()
.map(Result::unwrap)
.collect::<Vec<_>>();
assert_eq!(3, values.len()); assert_eq!(3, values.len());
assert!(values[0].is_empty()); assert_eq!(values[0], None);
assert_eq!(values[1], b"v1"); assert_eq!(values[1], Some(b"v1".to_vec()));
assert_eq!(values[2], b"v2"); assert_eq!(values[2], Some(b"v2".to_vec()));
} }
} }

Loading…
Cancel
Save