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 DBCommon { } } + /// 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)`. + /// + /// 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>( + &self, + cf: &impl AsColumnFamilyRef, + key: K, + readopts: &ReadOptions, + ) -> (bool, Option) { + 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(value: S) -> Result { 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();