perf: remove heap allocation in snapshot point read path

+ also not using MvccStream because we just need to terminate on first matching item
fjall-rs · Dec 14, 2024 · ab28434 · ab28434
1 parent 26847b3
commit ab28434
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 35 deletions.
diff --git a/src/segment/mod.rs b/src/segment/mod.rs
@@ -334,7 +334,6 @@ impl Segment {
         key: K,
         seqno: Option<SeqNo>,
     ) -> crate::Result<Option<InternalValue>> {
-        use crate::mvcc_stream::MvccStream;
         use block_index::BlockIndex;
         use value_block::{CachePolicy, ValueBlock};
         use value_block_consumer::ValueBlockConsumer;
@@ -379,11 +378,7 @@ impl Segment {
             None,
         );
         reader.lo_block_size = block.header.data_length.into();
-        reader.lo_block_items = Some(ValueBlockConsumer::with_bounds(
-            block,
-            &Some(key.into()), // TODO: this may cause a heap alloc
-            &None,
-        ));
+        reader.lo_block_items = Some(ValueBlockConsumer::with_bounds(block, Some(key), None));
         reader.lo_initialized = true;
 
         // NOTE: For finding a specific seqno,
@@ -403,7 +398,7 @@ impl Segment {
         // unfortunately is in the next block
         //
         // Also because of weak tombstones, we may have to look further than the first item we encounter
-        let reader = reader.filter(|x| {
+        let mut reader = reader.filter(|x| {
             match x {
                 Ok(entry) => {
                     // Check for seqno if needed
@@ -417,7 +412,7 @@ impl Segment {
             }
         });
 
-        let Some(entry) = MvccStream::new(reader).next().transpose()? else {
+        let Some(entry) = reader.next().transpose()? else {
             return Ok(None);
         };
 

diff --git a/src/segment/reader.rs b/src/segment/reader.rs
@@ -103,7 +103,11 @@ impl Reader {
             Ok(Some((
                 block.header.data_length.into(),
                 block.header.previous_block_offset,
-                ValueBlockConsumer::with_bounds(block, &self.start_key, &self.end_key),
+                ValueBlockConsumer::with_bounds(
+                    block,
+                    self.start_key.as_deref(),
+                    self.end_key.as_deref(),
+                ),
             )))
         })
     }

diff --git a/src/segment/value_block_consumer.rs b/src/segment/value_block_consumer.rs
@@ -3,7 +3,7 @@
 // (found in the LICENSE-* files in the repository)
 
 use super::value_block::ValueBlock;
-use crate::{value::InternalValue, UserKey};
+use crate::value::InternalValue;
 use std::sync::Arc;
 
 pub struct ValueBlockConsumer {
@@ -15,14 +15,14 @@ pub struct ValueBlockConsumer {
 impl ValueBlockConsumer {
     #[must_use]
     pub fn new(inner: Arc<ValueBlock>) -> Self {
-        Self::with_bounds(inner, &None, &None)
+        Self::with_bounds(inner, None, None)
     }
 
     #[must_use]
     pub fn with_bounds(
         inner: Arc<ValueBlock>,
-        start_key: &Option<UserKey>,
-        end_key: &Option<UserKey>,
+        start_key: Option<&[u8]>,
+        end_key: Option<&[u8]>,
     ) -> Self {
         let mut lo = start_key.as_ref().map_or(0, |key| {
             inner.items.partition_point(|x| &*x.key.user_key < *key)
@@ -90,12 +90,9 @@ impl DoubleEndedIterator for ValueBlockConsumer {
 #[allow(clippy::expect_used)]
 mod tests {
     use super::*;
-    use crate::{
-        segment::{
-            block::{checksum::Checksum, header::Header},
-            value_block::BlockOffset,
-        },
-        Slice,
+    use crate::segment::{
+        block::{checksum::Checksum, header::Header},
+        value_block::BlockOffset,
     };
     use test_log::test;
 
@@ -227,15 +224,13 @@ mod tests {
             InternalValue::from_components(*b"e", vec![], 0, crate::ValueType::Value),
         ]);
 
-        let mut iter =
-            ValueBlockConsumer::with_bounds(block.clone().into(), &Some(Slice::from(*b"c")), &None);
+        let mut iter = ValueBlockConsumer::with_bounds(block.clone().into(), Some(b"c"), None);
         assert_eq!(*b"c", &*iter.next().expect("should exist").key.user_key);
         assert_eq!(*b"d", &*iter.next().expect("should exist").key.user_key);
         assert_eq!(*b"e", &*iter.next().expect("should exist").key.user_key);
         iter_closed!(iter);
 
-        let mut iter =
-            ValueBlockConsumer::with_bounds(block.into(), &Some(Slice::from(*b"c")), &None);
+        let mut iter = ValueBlockConsumer::with_bounds(block.into(), Some(b"c"), None);
         assert_eq!(
             *b"e",
             &*iter.next_back().expect("should exist").key.user_key
@@ -261,15 +256,13 @@ mod tests {
             InternalValue::from_components(*b"e", vec![], 0, crate::ValueType::Value),
         ]);
 
-        let mut iter =
-            ValueBlockConsumer::with_bounds(block.clone().into(), &None, &Some(Slice::from(*b"c")));
+        let mut iter = ValueBlockConsumer::with_bounds(block.clone().into(), None, Some(b"c"));
         assert_eq!(*b"a", &*iter.next().expect("should exist").key.user_key);
         assert_eq!(*b"b", &*iter.next().expect("should exist").key.user_key);
         assert_eq!(*b"c", &*iter.next().expect("should exist").key.user_key);
         iter_closed!(iter);
 
-        let mut iter =
-            ValueBlockConsumer::with_bounds(block.into(), &None, &Some(Slice::from(*b"c")));
+        let mut iter = ValueBlockConsumer::with_bounds(block.into(), None, Some(b"c"));
         assert_eq!(
             *b"c",
             &*iter.next_back().expect("should exist").key.user_key
@@ -294,12 +287,10 @@ mod tests {
             InternalValue::from_components(*b"e", vec![], 0, crate::ValueType::Value),
         ]);
 
-        let mut iter =
-            ValueBlockConsumer::with_bounds(block.clone().into(), &None, &Some(Slice::from(*b"a")));
+        let mut iter = ValueBlockConsumer::with_bounds(block.clone().into(), None, Some(b"a"));
         iter_closed!(iter);
 
-        let mut iter =
-            ValueBlockConsumer::with_bounds(block.into(), &None, &Some(Slice::from(*b"a"))).rev();
+        let mut iter = ValueBlockConsumer::with_bounds(block.into(), None, Some(b"a")).rev();
         iter_closed!(iter);
     }
 
@@ -313,12 +304,10 @@ mod tests {
             InternalValue::from_components(*b"e", vec![], 0, crate::ValueType::Value),
         ]);
 
-        let mut iter =
-            ValueBlockConsumer::with_bounds(block.clone().into(), &Some(Slice::from(*b"f")), &None);
+        let mut iter = ValueBlockConsumer::with_bounds(block.clone().into(), Some(b"f"), None);
         iter_closed!(iter);
 
-        let mut iter =
-            ValueBlockConsumer::with_bounds(block.into(), &Some(Slice::from(*b"f")), &None).rev();
+        let mut iter = ValueBlockConsumer::with_bounds(block.into(), Some(b"f"), None).rev();
         iter_closed!(iter);
     }
 }