Refactoring: rely on Range/RangeInclusive to harden the API. (#68)

Author: fulmicoton

Cargo.toml

@@ -1,7 +1,7 @@
 workspace = { members = ["mrecordlog_cli"] }
 [package]
 name = "mrecordlog"
-version = "0.4.0"
+version = "0.5.0"
 edition = "2021"
 license = "MIT"
 description = "Quickwit's shared record log."

src/block_read_write.rs

@@ -48,7 +48,7 @@ impl<'a> From<&'a [u8]> for ArrayReader<'a> {
     }
 }
 
-impl<'a> BlockRead for ArrayReader<'a> {
+impl BlockRead for ArrayReader<'_> {
     fn next_block(&mut self) -> io::Result<bool> {
         if self.data.len() < BLOCK_NUM_BYTES {
             return Ok(false);

src/mem/mod.rs

@@ -1,5 +1,6 @@
 mod queue;
 mod queues;
+mod rolling_buffer;
 mod summary;
 
 pub(crate) use self::queue::MemQueue;

src/mem/queue.rs

@@ -1,94 +1,11 @@
-use std::borrow::Cow;
-use std::collections::VecDeque;
-use std::ops::{Bound, RangeBounds};
+use std::ops::{Bound, RangeBounds, RangeToInclusive};
 
+use super::rolling_buffer::RollingBuffer;
 use crate::error::AppendError;
 use crate::mem::QueueSummary;
 use crate::rolling::FileNumber;
 use crate::Record;
 
-#[derive(Default)]
-struct RollingBuffer {
-    buffer: VecDeque<u8>,
-}
-
-impl RollingBuffer {
-    fn new() -> Self {
-        RollingBuffer {
-            buffer: VecDeque::new(),
-        }
-    }
-
-    fn len(&self) -> usize {
-        self.buffer.len()
-    }
-
-    fn capacity(&self) -> usize {
-        self.buffer.capacity()
-    }
-
-    fn clear(&mut self) {
-        self.buffer.clear();
-        self.buffer.shrink_to_fit();
-    }
-
-    fn drain_start(&mut self, pos: usize) {
-        let target_capacity = self.len() * 9 / 8;
-        self.buffer.drain(..pos);
-        // In order to avoid leaking memory we shrink the buffer.
-        // The last maximum length (= the length before drain)
-        // is a good estimate of what we will need in the future.
-        //
-        // We add 1/8 to that in order to make sure that we don't end up
-        // shrinking  / allocating for small variations.
-
-        if self.buffer.capacity() > target_capacity {
-            self.buffer.shrink_to(target_capacity);
-        }
-    }
-
-    fn extend(&mut self, slice: &[u8]) {
-        self.buffer.extend(slice.iter().copied());
-    }
-
-    fn get_range(&self, bounds: impl RangeBounds<usize>) -> Cow<[u8]> {
-        let start = match bounds.start_bound() {
-            Bound::Included(pos) => *pos,
-            Bound::Excluded(pos) => pos + 1,
-            Bound::Unbounded => 0,
-        };
-
-        let end = match bounds.end_bound() {
-            Bound::Included(pos) => pos + 1,
-            Bound::Excluded(pos) => *pos,
-            Bound::Unbounded => self.len(),
-        };
-
-        let (left_part_of_queue, right_part_of_queue) = self.buffer.as_slices();
-
-        if end < left_part_of_queue.len() {
-            Cow::Borrowed(&left_part_of_queue[start..end])
-        } else if start >= left_part_of_queue.len() {
-            let start = start - left_part_of_queue.len();
-            let end = end - left_part_of_queue.len();
-
-            Cow::Borrowed(&right_part_of_queue[start..end])
-        } else {
-            // VecDeque is a rolling buffer. As a result, we do not have
-            // access to a continuous buffer.
-            //
-            // Here the requested slice cross the boundary and we need to allocate and copy the data
-            // in a new buffer.
-            let mut res = Vec::with_capacity(end - start);
-            res.extend_from_slice(&left_part_of_queue[start..]);
-            let end = end - left_part_of_queue.len();
-            res.extend_from_slice(&right_part_of_queue[..end]);
-
-            Cow::Owned(res)
-        }
-    }
-}
-
 #[derive(Clone)]
 struct RecordMeta {
     start_offset: usize,
@@ -247,7 +164,8 @@ impl MemQueue {
     ///
     /// If truncating to a future position, make the queue go forward to that position.
     /// Return the number of record removed.
-    pub fn truncate(&mut self, truncate_up_to_pos: u64) -> usize {
+    pub fn truncate_head(&mut self, truncate_range: RangeToInclusive<u64>) -> usize {
+        let truncate_up_to_pos = truncate_range.end;
         if self.start_position > truncate_up_to_pos {
             return 0;
         }
@@ -267,7 +185,8 @@ impl MemQueue {
         for record_meta in &mut self.record_metas {
             record_meta.start_offset -= start_offset_to_keep;
         }
-        self.concatenated_records.drain_start(start_offset_to_keep);
+        self.concatenated_records
+            .truncate_head(..start_offset_to_keep);
         self.start_position = truncate_up_to_pos + 1;
         first_record_to_keep
     }

src/mem/queues.rs

@@ -1,5 +1,5 @@
 use std::collections::HashMap;
-use std::ops::RangeBounds;
+use std::ops::{RangeBounds, RangeToInclusive};
 
 use tracing::{info, warn};
 
@@ -154,9 +154,9 @@ impl MemQueues {
     ///
     /// If there are no records `<= position`, the method will
     /// not do anything.
-    pub fn truncate(&mut self, queue: &str, position: u64) -> Option<usize> {
+    pub fn truncate(&mut self, queue: &str, position: RangeToInclusive<u64>) -> Option<usize> {
         if let Ok(queue) = self.get_queue_mut(queue) {
-            Some(queue.truncate(position))
+            Some(queue.truncate_head(position))
         } else {
             None
         }

src/mem/rolling_buffer.rs

@@ -0,0 +1,89 @@
+use std::borrow::Cow;
+use std::collections::VecDeque;
+use std::ops::{Bound, RangeBounds, RangeTo};
+
+#[derive(Default)]
+pub struct RollingBuffer {
+    buffer: VecDeque<u8>,
+}
+
+impl RollingBuffer {
+    pub fn new() -> Self {
+        RollingBuffer {
+            buffer: VecDeque::new(),
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        self.buffer.len()
+    }
+
+    pub fn capacity(&self) -> usize {
+        self.buffer.capacity()
+    }
+
+    pub fn clear(&mut self) {
+        self.buffer.clear();
+        self.buffer.shrink_to_fit();
+    }
+
+    // Removes all of the data up to pos byte excluded (meaning pos is kept).
+    //
+    // If we notice that the rolling buffer was very large, this function may shrink
+    // it.
+    pub fn truncate_head(&mut self, first_pos_to_keep: RangeTo<usize>) {
+        let target_capacity = self.len() * 9 / 8;
+        self.buffer.drain(first_pos_to_keep);
+        // In order to avoid leaking memory we shrink the buffer.
+        // The last maximum length (= the length before drain)
+        // is a good estimate of what we will need in the future.
+        //
+        // We add 1/8 to that in order to make sure that we don't end up
+        // shrinking  / allocating for small variations.
+
+        if self.buffer.capacity() > target_capacity {
+            self.buffer.shrink_to(target_capacity);
+        }
+    }
+
+    pub fn extend(&mut self, slice: &[u8]) {
+        self.buffer.extend(slice.iter().copied());
+    }
+
+    pub fn get_range(&self, bounds: impl RangeBounds<usize>) -> Cow<[u8]> {
+        let start = match bounds.start_bound() {
+            Bound::Included(pos) => *pos,
+            Bound::Excluded(pos) => pos + 1,
+            Bound::Unbounded => 0,
+        };
+
+        let end = match bounds.end_bound() {
+            Bound::Included(pos) => pos + 1,
+            Bound::Excluded(pos) => *pos,
+            Bound::Unbounded => self.len(),
+        };
+
+        let (left_part_of_queue, right_part_of_queue) = self.buffer.as_slices();
+
+        if end < left_part_of_queue.len() {
+            Cow::Borrowed(&left_part_of_queue[start..end])
+        } else if start >= left_part_of_queue.len() {
+            let start = start - left_part_of_queue.len();
+            let end = end - left_part_of_queue.len();
+
+            Cow::Borrowed(&right_part_of_queue[start..end])
+        } else {
+            // VecDeque is a rolling buffer. As a result, we do not have
+            // access to a continuous buffer.
+            //
+            // Here the requested slice cross the boundary and we need to allocate and copy the data
+            // in a new buffer.
+            let mut res = Vec::with_capacity(end - start);
+            res.extend_from_slice(&left_part_of_queue[start..]);
+            let end = end - left_part_of_queue.len();
+            res.extend_from_slice(&right_part_of_queue[..end]);
+
+            Cow::Owned(res)
+        }
+    }
+}

src/mem/tests.rs

@@ -85,7 +85,7 @@ fn test_mem_queues_truncate() {
             .append_record("droopy", &1.into(), 5, b"payer")
             .unwrap();
     }
-    mem_queues.truncate("droopy", 3);
+    mem_queues.truncate("droopy", ..=3);
     let droopy: Vec<Record> = mem_queues.range("droopy", 0..).unwrap().collect();
     assert_eq!(
         &droopy[..],
@@ -210,7 +210,7 @@ fn test_mem_queues_keep_filenum() {
     assert!(!files[0].can_be_deleted());
     assert!(!files[1].can_be_deleted());
 
-    mem_queues.truncate("droopy", 1);
+    mem_queues.truncate("droopy", ..=1);
 
     assert!(!files[0].can_be_deleted());
     assert!(!files[1].can_be_deleted());
@@ -223,13 +223,13 @@ fn test_mem_queues_keep_filenum() {
     assert!(!files[1].can_be_deleted());
     assert!(!files[2].can_be_deleted());
 
-    mem_queues.truncate("droopy", 3);
+    mem_queues.truncate("droopy", ..=3);
 
     assert!(files[0].can_be_deleted());
     assert!(files[1].can_be_deleted());
     assert!(!files[2].can_be_deleted());
 
-    mem_queues.truncate("droopy", 4);
+    mem_queues.truncate("droopy", ..=4);
 
     let empty_queues = mem_queues.empty_queues().collect::<Vec<_>>();
     assert_eq!(empty_queues.len(), 1);

src/multi_record_log.rs

@@ -1,5 +1,5 @@
 use std::io;
-use std::ops::RangeBounds;
+use std::ops::{RangeBounds, RangeToInclusive};
 use std::path::Path;
 
 use bytes::Buf;
@@ -72,8 +72,11 @@ impl MultiRecordLog {
                                 .map_err(|_| ReadRecordError::Corruption)?;
                         }
                     }
-                    MultiPlexedRecord::Truncate { position, queue } => {
-                        in_mem_queues.truncate(queue, position);
+                    MultiPlexedRecord::Truncate {
+                        truncate_range,
+                        queue,
+                    } => {
+                        in_mem_queues.truncate(queue, truncate_range);
                     }
                     MultiPlexedRecord::RecordPosition { queue, position } => {
                         in_mem_queues.ack_position(queue, position);
@@ -227,20 +230,30 @@ impl MultiRecordLog {
         Ok(())
     }
 
-    /// Truncates the queue up to `position`, included. This method immediately truncates the
-    /// underlying in-memory queue whereas the backing log files are deleted asynchronously when
-    /// they become exclusively composed of deleted records.
+    /// Truncates the queue up to a given `position`, included. This method immediately
+    /// truncates the underlying in-memory queue whereas the backing log files are deleted
+    /// asynchronously when they become exclusively composed of deleted records.
     ///
     /// This method will always truncate the record log and release the associated memory.
     /// It returns the number of records deleted.
-    pub fn truncate(&mut self, queue: &str, position: u64) -> Result<usize, TruncateError> {
-        info!(position = position, queue = queue, "truncate queue");
+    pub fn truncate(
+        &mut self,
+        queue: &str,
+        truncate_range: RangeToInclusive<u64>,
+    ) -> Result<usize, TruncateError> {
+        info!(range=?truncate_range, queue = queue, "truncate queue");
         if !self.queue_exists(queue) {
             return Err(TruncateError::MissingQueue(queue.to_string()));
         }
         self.record_log_writer
-            .write_record(MultiPlexedRecord::Truncate { position, queue })?;
-        let removed_count = self.in_mem_queues.truncate(queue, position).unwrap_or(0);
+            .write_record(MultiPlexedRecord::Truncate {
+                truncate_range,
+                queue,
+            })?;
+        let removed_count = self
+            .in_mem_queues
+            .truncate(queue, truncate_range)
+            .unwrap_or(0);
         self.run_gc_if_necessary()?;
         self.persist_on_policy()?;
         Ok(removed_count)

src/persist_policy.rs

@@ -18,15 +18,16 @@ impl PersistAction {
 /// We have two type of operations on the mrecordlog.
 ///
 /// Critical records are relatively rare and really need to be persisted:
-/// - RecordPosition { queue: &'a str, position: u64 },
-/// - DeleteQueue.
+///   - RecordPosition { queue: &'a str, position: u64 },
+///   - DeleteQueue.
 ///
 /// For these operations, we want to always flush and fsync.
 ///
 /// On the other hand,
-/// - Truncate
-/// - AppendRecords
-/// are considered are more frequent and one might want to sacrifice
+///   - Truncate
+///   - AppendRecords
+///
+/// are both considered are more frequent and one might want to sacrifice
 /// persistence guarantees for performance.
 ///
 /// The `PersistPolicy` defines the trade-off applied for the second kind of