From 3257d087dce78af0940e793158cf7175318886f4 Mon Sep 17 00:00:00 2001
From: Congyu <52687642+Congyuwang@users.noreply.github.com>
Date: Wed, 26 Apr 2023 16:50:23 +0800
Subject: [PATCH] Retrieve Value from `KeyMayExist` if value found in Cache or
 Memory (#759)

---
 src/db.rs        | 44 ++++++++++++++++++++++++++++++++++++++++++++
 src/ffi_util.rs  | 36 +++++++++++++++++++++++++++++++++++-
 tests/test_db.rs | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 112 insertions(+), 1 deletion(-)

diff --git a/src/db.rs b/src/db.rs
index 7542b7b..87becae 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -26,6 +26,7 @@ use crate::{
     WriteBatch, WriteOptions, DEFAULT_COLUMN_FAMILY_NAME,
 };
 
+use crate::ffi_util::CSlice;
 use libc::{self, c_char, c_int, c_uchar, c_void, size_t};
 use std::collections::BTreeMap;
 use std::ffi::{CStr, CString};
@@ -1272,6 +1273,49 @@ impl<T: ThreadMode, D: DBInner> DBCommon<T, D> {
         }
     }
 
+    /// If the key definitely does not exist in the database, then this method
+    /// returns `(false, None)`, else `(true, None)` if it may.
+    /// If the key is found in memory, then it returns `(true, Some<CSlice>)`.
+    ///
+    /// This check is potentially lighter-weight than calling `get()`. One way
+    /// to make this lighter weight is to avoid doing any IOs.
+    pub fn key_may_exist_cf_opt_value<K: AsRef<[u8]>>(
+        &self,
+        cf: &impl AsColumnFamilyRef,
+        key: K,
+        readopts: &ReadOptions,
+    ) -> (bool, Option<CSlice>) {
+        let key = key.as_ref();
+        let mut val: *mut c_char = ptr::null_mut();
+        let mut val_len: usize = 0;
+        let mut value_found: c_uchar = 0;
+        let may_exists = 0
+            != unsafe {
+                ffi::rocksdb_key_may_exist_cf(
+                    self.inner.inner(),
+                    readopts.inner,
+                    cf.inner(),
+                    key.as_ptr() as *const c_char,
+                    key.len() as size_t,
+                    &mut val,         /*value*/
+                    &mut val_len,     /*val_len*/
+                    ptr::null(),      /*timestamp*/
+                    0,                /*timestamp_len*/
+                    &mut value_found, /*value_found*/
+                )
+            };
+        // The value is only allocated (using malloc) and returned if it is found and
+        // value_found isn't NULL. In that case the user is responsible for freeing it.
+        if may_exists && value_found != 0 {
+            (
+                may_exists,
+                Some(unsafe { CSlice::from_raw_parts(val, val_len) }),
+            )
+        } else {
+            (may_exists, None)
+        }
+    }
+
     fn create_inner_cf_handle(
         &self,
         name: impl CStrLike,
diff --git a/src/ffi_util.rs b/src/ffi_util.rs
index 9c35102..19a9b11 100644
--- a/src/ffi_util.rs
+++ b/src/ffi_util.rs
@@ -14,7 +14,7 @@
 //
 
 use crate::Error;
-use libc::{self, c_char, c_void};
+use libc::{self, c_char, c_void, size_t};
 use std::ffi::{CStr, CString};
 use std::path::Path;
 use std::ptr;
@@ -195,6 +195,40 @@ impl<'a> CStrLike for &'a CString {
     }
 }
 
+/// Owned malloc-allocated memory slice.
+/// Do not derive `Clone` for this because it will cause double-free.
+pub struct CSlice {
+    data: *const c_char,
+    len: size_t,
+}
+
+impl CSlice {
+    /// Constructing such a slice may be unsafe.
+    ///
+    /// # Safety
+    /// The caller must ensure that the pointer and length are valid.
+    /// Moreover, `CSlice` takes the ownership of the memory and will free it using `libc::free`.
+    /// The caller must ensure that the memory is allocated by `libc::malloc`
+    /// and will not be freed by any other means.
+    pub(crate) unsafe fn from_raw_parts(data: *const c_char, len: size_t) -> Self {
+        Self { data, len }
+    }
+}
+
+impl AsRef<[u8]> for CSlice {
+    fn as_ref(&self) -> &[u8] {
+        unsafe { std::slice::from_raw_parts(self.data as *const u8, self.len) }
+    }
+}
+
+impl Drop for CSlice {
+    fn drop(&mut self) {
+        unsafe {
+            libc::free(self.data as *mut c_void);
+        }
+    }
+}
+
 #[test]
 fn test_c_str_like_bake() {
     fn test<S: CStrLike>(value: S) -> Result<usize, S::Error> {
diff --git a/tests/test_db.rs b/tests/test_db.rs
index 1a91ebd..336670b 100644
--- a/tests/test_db.rs
+++ b/tests/test_db.rs
@@ -14,6 +14,7 @@
 
 mod util;
 
+use std::convert::TryInto;
 use std::{mem, sync::Arc, thread, time::Duration};
 
 use pretty_assertions::assert_eq;
@@ -1296,6 +1297,38 @@ fn key_may_exist_cf() {
     }
 }
 
+#[test]
+fn key_may_exist_cf_value() {
+    let path = DBPath::new("_rust_key_may_exist_cf_value");
+
+    {
+        let mut opts = Options::default();
+        opts.create_if_missing(true);
+        opts.create_missing_column_families(true);
+        let db = DB::open_cf(&opts, &path, ["cf"]).unwrap();
+        let cf = db.cf_handle("cf").unwrap();
+
+        // put some entry into db
+        for i in 0..10000i32 {
+            let _ = db.put_cf(&cf, i.to_le_bytes(), i.to_le_bytes());
+        }
+
+        // call `key_may_exist_cf_opt_value`
+        for i in 0..10000i32 {
+            let (may_exist, value) =
+                db.key_may_exist_cf_opt_value(&cf, i.to_le_bytes(), &ReadOptions::default());
+
+            // all these numbers may exist
+            assert!(may_exist);
+
+            // check value correctness
+            if let Some(value) = value {
+                assert_eq!(i32::from_le_bytes(value.as_ref().try_into().unwrap()), i);
+            }
+        }
+    }
+}
+
 #[test]
 fn test_snapshot_outlive_db() {
     let t = trybuild::TestCases::new();