diff --git a/core/store/src/db.rs b/core/store/src/db.rs index 9c1625938ac..9ce57f63c7d 100644 --- a/core/store/src/db.rs +++ b/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>>; + /// 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>> { + 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. /// diff --git a/core/store/src/db/refcount.rs b/core/store/src/db/refcount.rs index 7ca75c440fc..88400b51068 100644 --- a/core/store/src/db/refcount.rs +++ b/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) -> Option> { + 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 { [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>) -> Option> { - if column.is_rc() { - value.and_then(get_payload) - } else { - value - } -} - -fn get_payload(value: Vec) -> Option> { - 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, 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> { - 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]> { diff --git a/core/store/src/lib.rs b/core/store/src/lib.rs index 1d8874833e4..face16acea1 100644 --- a/core/store/src/lib.rs +++ b/core/store/src/lib.rs @@ -83,10 +83,11 @@ impl Store { } pub fn get(&self, column: DBCol, key: &[u8]) -> io::Result>> { - 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",