From dbd2ca6e4f5bd4a268cddacace8a7d4132ee955d Mon Sep 17 00:00:00 2001 From: Xuejie Xiao Date: Sun, 27 Jan 2019 09:13:06 +0800 Subject: [PATCH 1/5] 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, 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, Error> { + self.get_pinned_opt(key, &ReadOptions::default()) + } + + pub fn get_pinned_cf_opt( + &self, + cf: ColumnFamily, + key: &[u8], + 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(), + )); + } + + 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, 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 +1665,40 @@ fn to_cpath>(path: P) -> Result { } } +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]); +} From e19dad01419e8fdc5604fe8e913dfb33ddb04959 Mon Sep 17 00:00:00 2001 From: Xuejie Xiao Date: Fri, 1 Feb 2019 09:39:51 +0800 Subject: [PATCH 2/5] Add lifetime to DBPinnableSlice to make sure it cannot outlive DB --- src/db.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/src/db.rs b/src/db.rs index 40d7c13..dffb520 100644 --- a/src/db.rs +++ b/src/db.rs @@ -1665,11 +1665,12 @@ fn to_cpath>(path: P) -> Result { } } -pub struct DBPinnableSlice { +pub struct DBPinnableSlice<'a> { ptr: *mut ffi::rocksdb_pinnableslice_t, + db: PhantomData<&'a DB> } -impl Deref for DBPinnableSlice { +impl<'a> Deref for DBPinnableSlice<'a> { type Target = [u8]; fn deref(&self) -> &[u8] { @@ -1681,7 +1682,7 @@ impl Deref for DBPinnableSlice { } } -impl Drop for DBPinnableSlice { +impl<'a> Drop for DBPinnableSlice<'a> { fn drop(&mut self) { unsafe { ffi::rocksdb_pinnableslice_destroy(self.ptr,); @@ -1689,13 +1690,16 @@ impl Drop for DBPinnableSlice { } } -impl DBPinnableSlice { +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 { - DBPinnableSlice { ptr, } + pub unsafe fn from_c(ptr: *mut ffi::rocksdb_pinnableslice_t) -> DBPinnableSlice<'a> { + DBPinnableSlice { + ptr, + db: PhantomData, + } } } From 81091a05ba74721a9cb1e20a290428bf296d399e Mon Sep 17 00:00:00 2001 From: Xuejie Xiao Date: Sat, 2 Feb 2019 07:32:15 +0800 Subject: [PATCH 3/5] Ergonomic changes via AsRef --- src/db.rs | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/src/db.rs b/src/db.rs index dffb520..430609b 100644 --- a/src/db.rs +++ b/src/db.rs @@ -938,7 +938,7 @@ impl DB { self.get_cf_opt(cf, key.as_ref(), &ReadOptions::default()) } - pub fn get_pinned_opt(&self, key: &[u8], readopts: &ReadOptions) -> Result, Error> { + pub fn get_pinned_opt>(&self, key: T, readopts: &ReadOptions) -> Result, Error> { if readopts.inner.is_null() { return Err(Error::new( "Unable to create RocksDB read options. \ @@ -950,6 +950,7 @@ impl DB { )); } + let key = key.as_ref(); unsafe { let val = ffi_try!(ffi::rocksdb_get_pinned( self.inner, @@ -965,14 +966,14 @@ impl DB { } } - pub fn get_pinned(&self, key: &[u8]) -> Result, Error> { + pub fn get_pinned>(&self, key: T) -> Result, Error> { self.get_pinned_opt(key, &ReadOptions::default()) } - pub fn get_pinned_cf_opt( + pub fn get_pinned_cf_opt>( &self, cf: ColumnFamily, - key: &[u8], + key: T, readopts: &ReadOptions, ) -> Result, Error> { if readopts.inner.is_null() { @@ -986,6 +987,7 @@ impl DB { )); } + let key = key.as_ref(); unsafe { let val = ffi_try!(ffi::rocksdb_get_pinned_cf( self.inner, @@ -1002,7 +1004,7 @@ impl DB { } } - pub fn get_pinned_cf(&self, cf: ColumnFamily, key: &[u8]) -> Result, Error> { + pub fn get_pinned_cf>(&self, cf: ColumnFamily, key: T) -> Result, Error> { self.get_pinned_cf_opt(cf, key, &ReadOptions::default()) } @@ -1670,6 +1672,13 @@ pub struct DBPinnableSlice<'a> { 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]; From a651b19aa52bbfacda26cedc55a367f241e342f5 Mon Sep 17 00:00:00 2001 From: Xuejie Xiao Date: Fri, 8 Feb 2019 11:58:23 +0800 Subject: [PATCH 4/5] Code review fixes --- src/db.rs | 58 +++++++++++++++++++++++++----------- src/lib.rs | 4 +-- tests/test_pinnable_slice.rs | 10 +++---- 3 files changed, 48 insertions(+), 24 deletions(-) diff --git a/src/db.rs b/src/db.rs index 430609b..d56e459 100644 --- a/src/db.rs +++ b/src/db.rs @@ -938,14 +938,20 @@ impl DB { self.get_cf_opt(cf, key.as_ref(), &ReadOptions::default()) } - pub fn get_pinned_opt>(&self, key: T, readopts: &ReadOptions) -> Result, Error> { + /// 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." + This is a fairly trivial call, and its \ + failure may be indicative of a \ + mis-compiled or mis-loaded RocksDB \ + library." .to_owned(), )); } @@ -966,23 +972,29 @@ impl DB { } } - pub fn get_pinned>(&self, key: T) -> Result, Error> { + /// 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()) } - pub fn get_pinned_cf_opt>( + /// 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: T, + 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." + This is a fairly trivial call, and its \ + failure may be indicative of a \ + mis-compiled or mis-loaded RocksDB \ + library." .to_owned(), )); } @@ -1004,7 +1016,14 @@ impl DB { } } - pub fn get_pinned_cf>(&self, cf: ColumnFamily, key: T) -> Result, Error> { + /// 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()) } @@ -1667,9 +1686,14 @@ 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> + db: PhantomData<&'a DB>, } impl<'a> AsRef<[u8]> for DBPinnableSlice<'a> { @@ -1685,8 +1709,8 @@ impl<'a> Deref for DBPinnableSlice<'a> { 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,) + let val = ffi::rocksdb_pinnableslice_value(self.ptr, &mut val_len) as *mut u8; + slice::from_raw_parts(val, val_len) } } } @@ -1694,7 +1718,7 @@ impl<'a> Deref for DBPinnableSlice<'a> { impl<'a> Drop for DBPinnableSlice<'a> { fn drop(&mut self) { unsafe { - ffi::rocksdb_pinnableslice_destroy(self.ptr,); + ffi::rocksdb_pinnableslice_destroy(self.ptr); } } } diff --git a/src/lib.rs b/src/lib.rs index ec60c7c..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::{ - new_bloom_filter, DBCompactionStyle, DBCompressionType, DBIterator, DBRawIterator, DBPinnableSlice, 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 index 35e0c63..be3db7a 100644 --- a/tests/test_pinnable_slice.rs +++ b/tests/test_pinnable_slice.rs @@ -3,16 +3,16 @@ extern crate rocksdb; use rocksdb::{Options, DB}; #[test] -pub fn test_pinnable_slice() { +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(); + opts.create_if_missing(true); + let db = DB::open(&opts, path).unwrap(); - db.put(b"k1", b"value12345",).unwrap(); + db.put(b"k1", b"value12345").unwrap(); - let result = db.get_pinned(b"k1",); + let result = db.get_pinned(b"k1"); assert!(result.is_ok()); let value = result.unwrap(); From aa2c6959418392759bce37e25dcd933d26be8984 Mon Sep 17 00:00:00 2001 From: Xuejie Xiao Date: Fri, 8 Feb 2019 20:26:00 +0800 Subject: [PATCH 5/5] Use DBPath in PinnableSlice test --- tests/test_pinnable_slice.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_pinnable_slice.rs b/tests/test_pinnable_slice.rs index be3db7a..1ccb875 100644 --- a/tests/test_pinnable_slice.rs +++ b/tests/test_pinnable_slice.rs @@ -1,14 +1,16 @@ extern crate rocksdb; +mod util; use rocksdb::{Options, DB}; +use util::DBPath; #[test] fn test_pinnable_slice() { - let path = "_rust_rocksdb_pinnable_slice_test"; + 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(); + let db = DB::open(&opts, &path).unwrap(); db.put(b"k1", b"value12345").unwrap();