diff --git a/src/db.rs b/src/db.rs index 8cdf53a..d56e459 100644 --- a/src/db.rs +++ b/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>( + &self, + key: K, + readopts: &ReadOptions, + ) -> Result, 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>(&self, key: K) -> Result, 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>( + &self, + cf: ColumnFamily, + key: K, + readopts: &ReadOptions, + ) -> Result, 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>( + &self, + cf: ColumnFamily, + key: K, + ) -> Result, Error> { + self.get_pinned_cf_opt(cf, key, &ReadOptions::default()) + } + pub fn create_cf(&self, name: &str, opts: &Options) -> Result { let cname = match CString::new(name.as_bytes()) { Ok(c) => c, @@ -1597,6 +1686,56 @@ fn to_cpath>(path: P) -> Result { } } +/// 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; diff --git a/src/lib.rs b/src/lib.rs index fe94c43..caf2d20 100644 --- a/src/lib.rs +++ b/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; diff --git a/tests/test_pinnable_slice.rs b/tests/test_pinnable_slice.rs new file mode 100644 index 0000000..1ccb875 --- /dev/null +++ b/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]); +}