From 16e4eaf1c7b6a90e08761cfbc50375a6ba29b1ea Mon Sep 17 00:00:00 2001
From: Andrey Zheleznov <zheland.net@gmail.com>
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<Option<Vec<u8>>, Error>;
 
+    fn get_pinned_opt<K: AsRef<[u8]>>(
+        &self,
+        key: K,
+        readopts: &ReadOptions,
+    ) -> Result<Option<DBPinnableSlice>, Error>;
+
+    fn get_pinned_cf_opt<K: AsRef<[u8]>>(
+        &self,
+        cf: &impl AsColumnFamilyRef,
+        key: K,
+        readopts: &ReadOptions,
+    ) -> Result<Option<DBPinnableSlice>, Error>;
+
     fn multi_get_opt<K, I>(
         &self,
         keys: I,
@@ -188,6 +201,23 @@ impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> {
         self.get_cf_opt(cf, key, readopts)
     }
 
+    fn get_pinned_opt<K: AsRef<[u8]>>(
+        &self,
+        key: K,
+        readopts: &ReadOptions,
+    ) -> Result<Option<DBPinnableSlice>, Error> {
+        self.get_pinned_opt(key, readopts)
+    }
+
+    fn get_pinned_cf_opt<K: AsRef<[u8]>>(
+        &self,
+        cf: &impl AsColumnFamilyRef,
+        key: K,
+        readopts: &ReadOptions,
+    ) -> Result<Option<DBPinnableSlice>, Error> {
+        self.get_pinned_cf_opt(cf, key, readopts)
+    }
+
     fn multi_get_opt<K, I>(
         &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<K: AsRef<[u8]>>(&self, key: K) -> Result<Option<DBPinnableSlice>, 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<K: AsRef<[u8]>>(
+        &self,
+        cf: &impl AsColumnFamilyRef,
+        key: K,
+    ) -> Result<Option<DBPinnableSlice>, 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<K: AsRef<[u8]>>(
+        &self,
+        key: K,
+        mut readopts: ReadOptions,
+    ) -> Result<Option<DBPinnableSlice>, 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<K: AsRef<[u8]>>(
+        &self,
+        cf: &impl AsColumnFamilyRef,
+        key: K,
+        mut readopts: ReadOptions,
+    ) -> Result<Option<DBPinnableSlice>, 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<K: AsRef<[u8]>, I>(&self, keys: I) -> Vec<Result<Option<Vec<u8>>, 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]);
+}