Merge pull request #246 from nervosnetwork/pinnable-slice

Implement PinnableSlice based API

src/db.rs

@@ -938,6 +938,95 @@ impl DB {
         self.get_cf_opt(cf, key.as_ref(), &ReadOptions::default())
     }
 
+    /// Return the value associated with a key using RocksDB's PinnableSlice
+    /// so as to avoid unnecessary memory copy.
+    pub fn get_pinned_opt<K: AsRef<[u8]>>(
+        &self,
+        key: K,
+        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(),
+            ));
+        }
+
+        let key = key.as_ref();
+        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)))
+            }
+        }
+    }
+
+    /// Return the value associated with a key using RocksDB's PinnableSlice
+    /// so as to avoid unnecessary memory copy. Similar to get_pinned_opt but
+    /// leverages default options.
+    pub fn get_pinned<K: AsRef<[u8]>>(&self, key: K) -> Result<Option<DBPinnableSlice>, Error> {
+        self.get_pinned_opt(key, &ReadOptions::default())
+    }
+
+    /// Return the value associated with a key using RocksDB's PinnableSlice
+    /// so as to avoid unnecessary memory copy. Similar to get_pinned_opt but
+    /// allows specifying ColumnFamily
+    pub fn get_pinned_cf_opt<K: AsRef<[u8]>>(
+        &self,
+        cf: ColumnFamily,
+        key: K,
+        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(),
+            ));
+        }
+
+        let key = key.as_ref();
+        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)))
+            }
+        }
+    }
+
+    /// Return the value associated with a key using RocksDB's PinnableSlice
+    /// so as to avoid unnecessary memory copy. Similar to get_pinned_cf_opt but
+    /// leverages default options.
+    pub fn get_pinned_cf<K: AsRef<[u8]>>(
+        &self,
+        cf: ColumnFamily,
+        key: K,
+    ) -> 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 +1686,56 @@ fn to_cpath<P: AsRef<Path>>(path: P) -> Result<CString, Error> {
     }
 }
 
+/// Wrapper around RocksDB PinnableSlice struct.
+///
+/// With a pinnable slice, we can directly leverage in-memory data within
+/// RocksDB toa void unnecessary memory copies. The struct here wraps the
+/// returned raw pointer and ensures proper finalization work.
+pub struct DBPinnableSlice<'a> {
+    ptr: *mut ffi::rocksdb_pinnableslice_t,
+    db: PhantomData<&'a DB>,
+}
+
+impl<'a> AsRef<[u8]> for DBPinnableSlice<'a> {
+    fn as_ref(&self) -> &[u8] {
+        // Implement this via Deref so as not to repeat ourselves
+        &*self
+    }
+}
+
+impl<'a> Deref for DBPinnableSlice<'a> {
+    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<'a> Drop for DBPinnableSlice<'a> {
+    fn drop(&mut self) {
+        unsafe {
+            ffi::rocksdb_pinnableslice_destroy(self.ptr);
+        }
+    }
+}
+
+impl<'a> DBPinnableSlice<'a> {
+    /// 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<'a> {
+        DBPinnableSlice {
+            ptr,
+            db: PhantomData,
+        }
+    }
+}
+
 #[test]
 fn test_db_vector() {
     use std::mem;

src/lib.rs

@@ -71,8 +71,8 @@ mod slice_transform;
 
 pub use compaction_filter::Decision as CompactionDecision;
 pub use db::{
-    DBCompactionStyle, DBCompressionType, DBIterator, DBRawIterator, DBRecoveryMode, DBVector,
-    Direction, IteratorMode, ReadOptions, Snapshot, WriteBatch,
+    DBCompactionStyle, DBCompressionType, DBIterator, DBPinnableSlice, DBRawIterator,
+    DBRecoveryMode, DBVector, Direction, IteratorMode, ReadOptions, Snapshot, WriteBatch,
 };
 
 pub use slice_transform::SliceTransform;

tests/test_pinnable_slice.rs

@@ -0,0 +1,26 @@
+extern crate rocksdb;
+mod util;
+
+use rocksdb::{Options, DB};
+use util::DBPath;
+
+#[test]
+fn test_pinnable_slice() {
+    let path = DBPath::new("_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]);
+}