From a0ad4fdcfea52f0a1b10d019c54c494e3f9488a5 Mon Sep 17 00:00:00 2001
From: mikhailOK <mikhail.kever@gmail.com>
Date: Mon, 24 May 2021 23:26:15 -0700
Subject: [PATCH] Fix multi_get (#511)

---
 src/db.rs        | 45 +++++++++++++++++++++++++++++----------------
 tests/test_db.rs | 22 ++++++++++++++--------
 2 files changed, 43 insertions(+), 24 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index a785de3..4172ee0 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -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()));
     }
 }