From 16e4eaf1c7b6a90e08761cfbc50375a6ba29b1ea Mon Sep 17 00:00:00 2001 From: Andrey Zheleznov Date: Thu, 7 Apr 2022 12:46:21 +0300 Subject: [PATCH] Add Snapshot PinnableSlice based API (#606) --- src/db.rs | 30 ++++++++++++++++++++++ src/snapshot.rs | 48 ++++++++++++++++++++++++++++++++++-- tests/test_pinnable_slice.rs | 23 +++++++++++++++++ 3 files changed, 99 insertions(+), 2 deletions(-) diff --git a/src/db.rs b/src/db.rs index 3b3997f..bc7fadc 100644 --- a/src/db.rs +++ b/src/db.rs @@ -146,6 +146,19 @@ pub trait DBAccess { readopts: &ReadOptions, ) -> Result>, Error>; + fn get_pinned_opt>( + &self, + key: K, + readopts: &ReadOptions, + ) -> Result, Error>; + + fn get_pinned_cf_opt>( + &self, + cf: &impl AsColumnFamilyRef, + key: K, + readopts: &ReadOptions, + ) -> Result, Error>; + fn multi_get_opt( &self, keys: I, @@ -188,6 +201,23 @@ impl DBAccess for DBWithThreadMode { self.get_cf_opt(cf, key, readopts) } + fn get_pinned_opt>( + &self, + key: K, + readopts: &ReadOptions, + ) -> Result, Error> { + self.get_pinned_opt(key, readopts) + } + + fn get_pinned_cf_opt>( + &self, + cf: &impl AsColumnFamilyRef, + key: K, + readopts: &ReadOptions, + ) -> Result, Error> { + self.get_pinned_cf_opt(cf, key, readopts) + } + fn multi_get_opt( &self, keys: I, diff --git a/src/snapshot.rs b/src/snapshot.rs index 0a0ea06..c996c34 100644 --- a/src/snapshot.rs +++ b/src/snapshot.rs @@ -13,8 +13,8 @@ // limitations under the License. use crate::{ - db::DBAccess, ffi, AsColumnFamilyRef, DBIteratorWithThreadMode, DBRawIteratorWithThreadMode, - Error, IteratorMode, ReadOptions, DB, + db::DBAccess, ffi, AsColumnFamilyRef, DBIteratorWithThreadMode, DBPinnableSlice, + DBRawIteratorWithThreadMode, Error, IteratorMode, ReadOptions, DB, }; /// A type alias to keep compatibility. See [`SnapshotWithThreadMode`] for details @@ -161,6 +161,50 @@ impl<'a, D: DBAccess> SnapshotWithThreadMode<'a, D> { self.db.get_cf_opt(cf, key.as_ref(), &readopts) } + /// 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> { + let readopts = ReadOptions::default(); + self.get_pinned_opt(key, readopts) + } + + /// 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: &impl AsColumnFamilyRef, + key: K, + ) -> Result, Error> { + let readopts = ReadOptions::default(); + self.get_pinned_cf_opt(cf, key.as_ref(), readopts) + } + + /// 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, + mut readopts: ReadOptions, + ) -> Result, Error> { + readopts.set_snapshot(self); + self.db.get_pinned_opt(key.as_ref(), &readopts) + } + + /// 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: &impl AsColumnFamilyRef, + key: K, + mut readopts: ReadOptions, + ) -> Result, Error> { + readopts.set_snapshot(self); + self.db.get_pinned_cf_opt(cf, key.as_ref(), &readopts) + } + /// Returns the bytes associated with the given key values and default read options. pub fn multi_get, I>(&self, keys: I) -> Vec>, Error>> where diff --git a/tests/test_pinnable_slice.rs b/tests/test_pinnable_slice.rs index d117440..279150e 100644 --- a/tests/test_pinnable_slice.rs +++ b/tests/test_pinnable_slice.rs @@ -39,3 +39,26 @@ fn test_pinnable_slice() { assert_eq!(b"12345", &pinnable_slice[5..10]); } + +#[test] +fn test_snapshot_pinnable_slice() { + let path = DBPath::new("_rust_rocksdb_snapshot_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 snap = db.snapshot(); + assert!(db.put(b"k1", b"value23456").is_ok()); + + let result = snap.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]); +}