From dbd2ca6e4f5bd4a268cddacace8a7d4132ee955d Mon Sep 17 00:00:00 2001
From: Xuejie Xiao <xxuejie@gmail.com>
Date: Sun, 27 Jan 2019 09:13:06 +0800
Subject: [PATCH] Implement PinnableSlice based API

---
 src/db.rs                    | 102 +++++++++++++++++++++++++++++++++++
 src/lib.rs                   |   2 +-
 tests/test_pinnable_slice.rs |  24 +++++++++
 3 files changed, 127 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_pinnable_slice.rs

diff --git a/src/db.rs b/src/db.rs
index 8cdf53a..40d7c13 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -938,6 +938,74 @@ impl DB {
         self.get_cf_opt(cf, key.as_ref(), &ReadOptions::default())
     }
 
+    pub fn get_pinned_opt(&self, key: &[u8], readopts: &ReadOptions) -> Result<Option<DBPinnableSlice>, Error> {
+        if readopts.inner.is_null() {
+            return Err(Error::new(
+                "Unable to create RocksDB read options. \
+                                   This is a fairly trivial call, and its \
+                                   failure may be indicative of a \
+                                   mis-compiled or mis-loaded RocksDB \
+                                   library."
+                    .to_owned(),
+            ));
+        }
+
+        unsafe {
+            let val = ffi_try!(ffi::rocksdb_get_pinned(
+                self.inner,
+                readopts.inner,
+                key.as_ptr() as *const c_char,
+                key.len() as size_t,
+            ));
+            if val.is_null() {
+                Ok(None)
+            } else {
+                Ok(Some(DBPinnableSlice::from_c(val)))
+            }
+        }
+    }
+
+    pub fn get_pinned(&self, key: &[u8]) -> Result<Option<DBPinnableSlice>, Error> {
+        self.get_pinned_opt(key, &ReadOptions::default())
+    }
+
+    pub fn get_pinned_cf_opt(
+        &self,
+        cf: ColumnFamily,
+        key: &[u8],
+        readopts: &ReadOptions,
+    ) -> Result<Option<DBPinnableSlice>, Error> {
+        if readopts.inner.is_null() {
+            return Err(Error::new(
+                "Unable to create RocksDB read options. \
+                                   This is a fairly trivial call, and its \
+                                   failure may be indicative of a \
+                                   mis-compiled or mis-loaded RocksDB \
+                                   library."
+                    .to_owned(),
+            ));
+        }
+
+        unsafe {
+            let val = ffi_try!(ffi::rocksdb_get_pinned_cf(
+                self.inner,
+                readopts.inner,
+                cf.inner,
+                key.as_ptr() as *const c_char,
+                key.len() as size_t,
+            ));
+            if val.is_null() {
+                Ok(None)
+            } else {
+                Ok(Some(DBPinnableSlice::from_c(val)))
+            }
+        }
+    }
+
+    pub fn get_pinned_cf(&self, cf: ColumnFamily, key: &[u8]) -> Result<Option<DBPinnableSlice>, Error> {
+        self.get_pinned_cf_opt(cf, key, &ReadOptions::default())
+    }
+
     pub fn create_cf(&self, name: &str, opts: &Options) -> Result<ColumnFamily, Error> {
         let cname = match CString::new(name.as_bytes()) {
             Ok(c) => c,
@@ -1597,6 +1665,40 @@ fn to_cpath<P: AsRef<Path>>(path: P) -> Result<CString, Error> {
     }
 }
 
+pub struct DBPinnableSlice {
+    ptr: *mut ffi::rocksdb_pinnableslice_t,
+}
+
+impl Deref for DBPinnableSlice {
+    type Target = [u8];
+
+    fn deref(&self) -> &[u8] {
+        unsafe {
+            let mut val_len: size_t = 0;
+            let val = ffi::rocksdb_pinnableslice_value(self.ptr, &mut val_len,) as *mut u8;
+            slice::from_raw_parts(val, val_len,)
+        }
+    }
+}
+
+impl Drop for DBPinnableSlice {
+    fn drop(&mut self) {
+        unsafe {
+            ffi::rocksdb_pinnableslice_destroy(self.ptr,);
+        }
+    }
+}
+
+impl DBPinnableSlice {
+    /// Used to wrap a PinnableSlice from rocksdb to avoid unnecessary memcpy
+    ///
+    /// # Unsafe
+    /// Requires that the pointer must be generated by rocksdb_get_pinned
+    pub unsafe fn from_c(ptr: *mut ffi::rocksdb_pinnableslice_t) -> DBPinnableSlice {
+        DBPinnableSlice { ptr, }
+    }
+}
+
 #[test]
 fn test_db_vector() {
     use std::mem;
diff --git a/src/lib.rs b/src/lib.rs
index fe94c43..ec60c7c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -71,7 +71,7 @@ mod slice_transform;
 
 pub use compaction_filter::Decision as CompactionDecision;
 pub use db::{
-    DBCompactionStyle, DBCompressionType, DBIterator, DBRawIterator, DBRecoveryMode, DBVector,
+    new_bloom_filter, DBCompactionStyle, DBCompressionType, DBIterator, DBRawIterator, DBPinnableSlice, DBRecoveryMode, DBVector,
     Direction, IteratorMode, ReadOptions, Snapshot, WriteBatch,
 };
 
diff --git a/tests/test_pinnable_slice.rs b/tests/test_pinnable_slice.rs
new file mode 100644
index 0000000..35e0c63
--- /dev/null
+++ b/tests/test_pinnable_slice.rs
@@ -0,0 +1,24 @@
+extern crate rocksdb;
+
+use rocksdb::{Options, DB};
+
+#[test]
+pub fn test_pinnable_slice() {
+    let path = "_rust_rocksdb_pinnable_slice_test";
+
+    let mut opts = Options::default();
+    opts.create_if_missing(true,);
+    let db = DB::open(&opts, path,).unwrap();
+
+    db.put(b"k1", b"value12345",).unwrap();
+
+    let result = db.get_pinned(b"k1",);
+    assert!(result.is_ok());
+
+    let value = result.unwrap();
+    assert!(value.is_some());
+
+    let pinnable_slice = value.unwrap();
+
+    assert_eq!(b"12345", &pinnable_slice[5..10]);
+}