store: introduce Database::get_with_rc_stripped method (#7483)

Add get_with_rc_stripped method to the Database trait which strips the
reference count from the returned value.

Previously, Store::get method called get_raw_bytes and then (for
reference counted columns) refcount::get_with_rc_logic to strip the
reference count.  With this change, it now calls either get_raw_bytes
or get_with_rc_stripped as needed.

The reason for this addition is to help with cold storage.  We don’t
want to keep refcounts in cold storage which means that we would need
to append dummy refcount to get_raw_bytes calls against cold storage.
By having a separate method which returns values without refcount, we
won’t need to worry about this dummy count.

As a side benefit, this removes a memory allocation when reading
a reference counted column.  Previosly get_payload method was used
to strip the refcount:

    fn get_payload(value: Vec<u8>) -> Option<Vec<u8>> {
        decode_value_with_rc(&value).0.map(|v| v.to_vec())
    }

The issue with it is that it takes a Vec as argument, takes a slice of
it and creates a new Vec from it.  Now the vector is simply truncated
by strip_refcount function.

core/store/src/db.rs

@@ -91,6 +91,14 @@ pub trait Database: Sync + Send {
     /// properly handle reference-counted columns.
     fn get_raw_bytes(&self, col: DBCol, key: &[u8]) -> io::Result<Option<Vec<u8>>>;
 
+    /// Returns value for given `key` forcing a reference count decoding.
+    ///
+    /// **Panics** if the column is not reference counted.
+    fn get_with_rc_stripped(&self, col: DBCol, key: &[u8]) -> io::Result<Option<Vec<u8>>> {
+        assert!(col.is_rc());
+        Ok(self.get_raw_bytes(col, key)?.and_then(crate::db::refcount::strip_refcount))
+    }
+
     /// Iterate over all items in given column in lexicographical order sorted
     /// by the key.
     ///

core/store/src/db/refcount.rs

@@ -51,6 +51,23 @@ pub fn decode_value_with_rc(bytes: &[u8]) -> (Option<&[u8]>, i64) {
     }
 }
 
+/// Strips refcount from an owned buffer.
+///
+/// Works like [`decode_value_with_rc`] but operates on an owned vector thus
+/// potentially avoiding memory allocations.  Returns None if the refcount on
+/// the argument is non-positive.
+pub(crate) fn strip_refcount(mut bytes: Vec<u8>) -> Option<Vec<u8>> {
+    if let Some(len) = bytes.len().checked_sub(8) {
+        if i64::from_le_bytes(bytes[len..].try_into().unwrap()) > 0 {
+            bytes.truncate(len);
+            return Some(bytes);
+        }
+    } else {
+        debug_assert!(bytes.is_empty());
+    }
+    None
+}
+
 /// Encode a positive reference count into the value.
 pub(crate) fn add_positive_refcount(data: &[u8], rc: std::num::NonZeroU32) -> Vec<u8> {
     [data, &i64::from(rc.get()).to_le_bytes()].concat()
@@ -95,39 +112,17 @@ pub(crate) fn refcount_merge<'a>(
     }
 }
 
-/// Returns value with reference count stripped if column is refcounted.
-///
-/// If the column is not refcounted, returns the value unchanged.
-///
-/// If the column is refcounted, extracts the reference count from it and
-/// returns value based on that.  If reference count is non-positive, returns
-/// `None`; otherwise returns the value with reference count stripped.  Empty
-/// values are treated as values with reference count zero.
-pub(crate) fn get_with_rc_logic(column: DBCol, value: Option<Vec<u8>>) -> Option<Vec<u8>> {
-    if column.is_rc() {
-        value.and_then(get_payload)
-    } else {
-        value
-    }
-}
-
-fn get_payload(value: Vec<u8>) -> Option<Vec<u8>> {
-    decode_value_with_rc(&value).0.map(|v| v.to_vec())
-}
-
 /// Iterator treats empty value as no value and strips refcount
 pub(crate) fn iter_with_rc_logic<'a>(
     column: DBCol,
     iterator: impl Iterator<Item = io::Result<(Box<[u8]>, Box<[u8]>)>> + 'a,
 ) -> crate::db::DBIterator<'a> {
     if column.is_rc() {
-        Box::new(iterator.filter_map(|item| {
-            let item = if let Ok((key, value)) = item {
-                Ok((key, get_payload(value.into_vec())?.into_boxed_slice()))
-            } else {
-                item
-            };
-            Some(item)
+        Box::new(iterator.filter_map(|item| match item {
+            Err(err) => Some(Err(err)),
+            Ok((key, value)) => {
+                strip_refcount(value.into_vec()).map(|value| Ok((key, value.into_boxed_slice())))
+            }
         }))
     } else {
         Box::new(iterator)
@@ -188,6 +183,9 @@ mod test {
         fn test(want_value: Option<&[u8]>, want_rc: i64, bytes: &[u8]) {
             let got = super::decode_value_with_rc(bytes);
             assert_eq!((want_value, want_rc), got);
+
+            let got = super::strip_refcount(bytes.to_vec());
+            assert_eq!(want_value, got.as_deref());
         }
 
         test(None, -2, MINUS_TWO);
@@ -273,44 +271,6 @@ mod test {
         test(Decision::Keep, ZERO);
     }
 
-    #[test]
-    fn get_with_rc_logic() {
-        fn get(col: DBCol, value: &[u8]) -> Option<Vec<u8>> {
-            assert_eq!(None, super::get_with_rc_logic(col, None));
-            super::get_with_rc_logic(col, Some(value.to_vec()))
-        }
-
-        fn test(want: Option<&[u8]>, col: DBCol, value: &[u8]) {
-            assert_eq!(want, get(col, value).as_ref().map(Vec::as_slice));
-        }
-
-        // Column without reference counting.  Values are returned as is.
-        assert!(!DBCol::Block.is_rc());
-        for value in [&b""[..], &b"foo"[..], MINUS_ONE, ZERO, PLUS_ONE] {
-            test(Some(value), DBCol::Block, value);
-        }
-
-        // Column with reference counting.  Count is extracted.
-        const RC_COL: DBCol = DBCol::State;
-        assert!(RC_COL.is_rc());
-
-        test(None, RC_COL, MINUS_ONE);
-        test(None, RC_COL, b"foo\xff\xff\xff\xff\xff\xff\xff\xff");
-        test(None, RC_COL, b"");
-        test(None, RC_COL, ZERO);
-        test(None, RC_COL, b"foo\x00\0\0\0\0\0\0\0");
-        test(Some(b""), RC_COL, PLUS_ONE);
-        test(Some(b"foo"), RC_COL, b"foo\x01\0\0\0\0\0\0\0");
-
-        check_debug_assert_or(
-            || {
-                let value = Some(b"short".to_vec());
-                super::get_with_rc_logic(RC_COL, value)
-            },
-            |got| assert_eq!(None, got),
-        );
-    }
-
     #[test]
     fn iter_with_rc_logic() {
         fn into_box(data: &[u8]) -> Box<[u8]> {

core/store/src/lib.rs

@@ -83,10 +83,11 @@ impl Store {
     }
 
     pub fn get(&self, column: DBCol, key: &[u8]) -> io::Result<Option<Vec<u8>>> {
-        let value = self
-            .storage
-            .get_raw_bytes(column, key)
-            .map(|result| refcount::get_with_rc_logic(column, result))?;
+        let value = if column.is_rc() {
+            self.storage.get_with_rc_stripped(column, key)
+        } else {
+            self.storage.get_raw_bytes(column, key)
+        }?;
         tracing::trace!(
             target: "store",
             db_op = "get",