diff options
author | mikhailOK <mikhail.kever@gmail.com> | 2021-05-24 23:26:15 -0700 |
---|---|---|
committer | GitHub <noreply@github.com> | 2021-05-25 09:26:15 +0300 |
commit | a0ad4fdcfea52f0a1b10d019c54c494e3f9488a5 (patch) | |
tree | 42922255b4f72821c01bdd06a0a02d66647cf820 | |
parent | c0cb8a77c2f73a757ab74ce658c93d75f6b74f2d (diff) |
Fix multi_get (#511)
-rw-r--r-- | src/db.rs | 45 | ||||
-rw-r--r-- | tests/test_db.rs | 22 |
2 files changed, 43 insertions, 24 deletions
@@ -747,7 +747,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> { } /// 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 K: AsRef<[u8]>, I: IntoIterator<Item = K>, @@ -760,7 +760,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> { &self, keys: I, readopts: &ReadOptions, - ) -> Result<Vec<Vec<u8>>, Error> + ) -> Vec<Result<Option<Vec<u8>>, Error>> where K: AsRef<[u8]>, 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_sizes = vec![0_usize; keys.len()]; + let mut errors = vec![ptr::null_mut(); keys.len()]; unsafe { - ffi_try!(ffi::rocksdb_multi_get( + ffi::rocksdb_multi_get( self.inner, readopts.inner, ptr_keys.len(), @@ -782,14 +783,15 @@ impl<T: ThreadMode> DBWithThreadMode<T> { keys_sizes.as_ptr(), values.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. - 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 K: AsRef<[u8]>, I: IntoIterator<Item = (W, K)>, @@ -803,7 +805,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> { &self, keys: I, readopts: &ReadOptions, - ) -> Result<Vec<Vec<u8>>, Error> + ) -> Vec<Result<Option<Vec<u8>>, Error>> where K: AsRef<[u8]>, 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_sizes = vec![0_usize; boxed_keys.len()]; + let mut errors = vec![ptr::null_mut(); boxed_keys.len()]; unsafe { - ffi_try!(ffi::rocksdb_multi_get_cf( + ffi::rocksdb_multi_get_cf( self.inner, readopts.inner, ptr_cfs.as_ptr(), @@ -838,10 +841,11 @@ impl<T: ThreadMode> DBWithThreadMode<T> { keys_sizes.as_ptr(), values.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 @@ -1835,16 +1839,25 @@ fn convert_options(opts: &[(&str, &str)]) -> Result<Vec<(CString, CString)>, Err .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 .into_iter() .zip(values_sizes.into_iter()) - .map(|(v, s)| { - let value = unsafe { slice::from_raw_parts(v as *const u8, s) }.into(); - unsafe { - ffi::rocksdb_free(v as *mut c_void); + .zip(errors.into_iter()) + .map(|((v, s), e)| { + if e.is_null() { + let value = unsafe { crate::ffi_util::raw_data(v, s) }; + unsafe { + ffi::rocksdb_free(v as *mut c_void); + } + Ok(value) + } else { + Err(Error::new(crate::ffi_util::error_message(e))) } - value }) .collect() } diff --git a/tests/test_db.rs b/tests/test_db.rs index 8aa682f..f43b28c 100644 --- a/tests/test_db.rs +++ b/tests/test_db.rs @@ -928,13 +928,17 @@ fn multi_get() { db.put(b"k1", b"v1").unwrap(); db.put(b"k2", b"v2").unwrap(); + let _ = db.multi_get(&[b"k0"; 40]); + let values = db .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!(values[0].is_empty()); - assert_eq!(values[1], b"v1"); - assert_eq!(values[2], b"v2"); + assert_eq!(values[0], None); + assert_eq!(values[1], Some(b"v1".to_vec())); + assert_eq!(values[2], Some(b"v2".to_vec())); } } @@ -958,11 +962,13 @@ fn multi_get_cf() { let values = db .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!(values[0].is_empty()); - assert_eq!(values[1], b"v1"); - assert_eq!(values[2], b"v2"); + assert_eq!(values[0], None); + assert_eq!(values[1], Some(b"v1".to_vec())); + assert_eq!(values[2], Some(b"v2".to_vec())); } } |