Add Snapshot PinnableSlice based API (rust-rocksdb#606)

schmidek · Jul 26, 2022 · f6e0797 · f6e0797
1 parent 7ce2768
commit f6e0797
Show file tree

Hide file tree

Showing 3 changed files with 192 additions and 2 deletions.
diff --git a/src/db.rs b/src/db.rs
@@ -145,6 +145,38 @@ pub trait DBAccess {
         key: K,
         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,
+        readopts: &ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = K>;
+
+    fn multi_get_cf_opt<'b, K, I, W>(
+        &self,
+        keys_cf: I,
+        readopts: &ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = (&'b W, K)>,
+        W: AsColumnFamilyRef + 'b;
 }
 
 impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> {
@@ -168,6 +200,48 @@ impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> {
     ) -> Result<Option<Vec<u8>>, Error> {
         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,
+        readopts: &ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = K>,
+    {
+        self.multi_get_opt(keys, readopts)
+    }
+
+    fn multi_get_cf_opt<'b, K, I, W>(
+        &self,
+        keys_cf: I,
+        readopts: &ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = (&'b W, K)>,
+        W: AsColumnFamilyRef + 'b,
+    {
+        self.multi_get_cf_opt(keys_cf, readopts)
+    }
 }
 
 /// A type alias to DB instance type with the single-threaded column family

diff --git 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
@@ -160,6 +160,99 @@ impl<'a, D: DBAccess> SnapshotWithThreadMode<'a, D> {
         readopts.set_snapshot(self);
         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
+        I: IntoIterator<Item = K>,
+    {
+        let readopts = ReadOptions::default();
+        self.multi_get_opt(keys, readopts)
+    }
+
+    /// Returns the bytes associated with the given key values and default read options.
+    pub fn multi_get_cf<'b, K, I, W>(&self, keys_cf: I) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = (&'b W, K)>,
+        W: AsColumnFamilyRef + 'b,
+    {
+        let readopts = ReadOptions::default();
+        self.multi_get_cf_opt(keys_cf, readopts)
+    }
+
+    /// Returns the bytes associated with the given key values and given read options.
+    pub fn multi_get_opt<K, I>(
+        &self,
+        keys: I,
+        mut readopts: ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = K>,
+    {
+        readopts.set_snapshot(self);
+        self.db.multi_get_opt(keys, &readopts)
+    }
+
+    /// Returns the bytes associated with the given key values, given column family and read options.
+    pub fn multi_get_cf_opt<'b, K, I, W>(
+        &self,
+        keys_cf: I,
+        mut readopts: ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = (&'b W, K)>,
+        W: AsColumnFamilyRef + 'b,
+    {
+        readopts.set_snapshot(self);
+        self.db.multi_get_cf_opt(keys_cf, &readopts)
+    }
 }
 
 impl<'a, D: DBAccess> Drop for SnapshotWithThreadMode<'a, D> {

diff --git 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]);
+}