fix: prevent dropping of new blob files

Author: marvin-j97

Cargo.toml

@@ -2,14 +2,14 @@
 name = "value-log"
 description = "Value log implementation for key-value separated LSM storage"
 license = "MIT OR Apache-2.0"
-version = "1.0.0-pre.2"
+version = "1.0.0-pre.3"
 edition = "2021"
 rust-version = "1.74.0"
 readme = "README.md"
 include = ["src/**/*", "LICENSE-APACHE", "LICENSE-MIT", "README.md"]
 repository = "https://github.com/fjall-rs/value-log"
 homepage = "https://github.com/fjall-rs/value-log"
-keywords = ["database", "lsmt", "lsm", "wisckey", "key-value"]
+keywords = ["database", "blobdb", "lsm", "wisckey", "key-value"]
 categories = ["data-structures", "database-implementations", "algorithms"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

src/lib.rs

@@ -94,6 +94,10 @@ mod key_range;
 mod manifest;
 mod mock;
 mod path;
+
+#[doc(hidden)]
+pub mod scanner;
+
 mod segment;
 mod serde;
 mod value;

src/scanner.rs

@@ -0,0 +1,58 @@
+use crate::{id::SegmentId, ValueHandle, ValueLog};
+use std::{collections::BTreeMap, sync::MutexGuard};
+
+#[derive(Debug, Default)]
+pub struct SegmentCounter {
+    pub size: u64,
+    pub item_count: u64,
+}
+
+pub type SizeMap = BTreeMap<SegmentId, SegmentCounter>;
+
+pub struct Scanner<'a, I: Iterator<Item = std::io::Result<(ValueHandle, u32)>>> {
+    iter: I,
+
+    #[allow(unused)]
+    lock_guard: MutexGuard<'a, ()>,
+
+    size_map: SizeMap,
+}
+
+impl<'a, I: Iterator<Item = std::io::Result<(ValueHandle, u32)>>> Scanner<'a, I> {
+    pub fn new(vlog: &'a ValueLog, iter: I) -> Self {
+        Self {
+            iter,
+            lock_guard: vlog.rollover_guard.lock().expect("lock is poisoned"),
+            size_map: BTreeMap::default(),
+        }
+    }
+
+    pub fn finish(self) -> SizeMap {
+        self.size_map
+    }
+
+    pub fn scan(&mut self) -> crate::Result<()> {
+        for handle in self.iter.by_ref() {
+            let (handle, size) = handle.map_err(|_| {
+                crate::Error::Io(std::io::Error::new(
+                    std::io::ErrorKind::Other,
+                    "Index returned error",
+                ))
+            })?;
+            let size = u64::from(size);
+
+            self.size_map
+                .entry(handle.segment_id)
+                .and_modify(|x| {
+                    x.item_count += 1;
+                    x.size += size;
+                })
+                .or_insert_with(|| SegmentCounter {
+                    size,
+                    item_count: 1,
+                });
+        }
+
+        Ok(())
+    }
+}

src/value_log.rs

@@ -4,14 +4,14 @@ use crate::{
     index::Writer as IndexWriter,
     manifest::{SegmentManifest, SEGMENTS_FOLDER, VLOG_MARKER},
     path::absolute_path,
+    scanner::{Scanner, SizeMap},
     segment::merge::MergeReader,
     value::UserValue,
     version::Version,
     Config, IndexReader, SegmentWriter, ValueHandle,
 };
 use byteorder::{BigEndian, ReadBytesExt};
 use std::{
-    collections::BTreeMap,
     fs::File,
     io::{BufReader, Read, Seek},
     path::PathBuf,
@@ -63,7 +63,7 @@ pub struct ValueLogInner {
 
     /// Guards the rollover (compaction) process to only
     /// allow one to happen at a time
-    rollover_guard: Mutex<()>,
+    pub(crate) rollover_guard: Mutex<()>,
 }
 
 impl ValueLog {
@@ -174,6 +174,8 @@ impl ValueLog {
     ///
     /// Will return `Err` if an IO error occurs.
     pub fn register_writer(&self, writer: SegmentWriter) -> crate::Result<()> {
+        let _lock = self.rollover_guard.lock().expect("lock is poisoned");
+
         self.manifest.register(writer)?;
         Ok(())
     }
@@ -349,51 +351,9 @@ impl ValueLog {
         self.manifest.space_amp()
     }
 
-    /// Scans the given index and collecting GC statistics.
-    ///
-    /// # Errors
-    ///
-    /// Will return `Err` if an IO error occurs.
-    #[allow(clippy::result_unit_err)]
-    pub fn scan_for_stats(
-        &self,
-        iter: impl Iterator<Item = std::io::Result<(ValueHandle, u32)>>,
-    ) -> crate::Result<()> {
-        #[derive(Debug)]
-        struct SegmentCounter {
-            size: u64,
-            item_count: u64,
-        }
-
-        // IMPORTANT: Only allow 1 rollover or GC at any given time
-        let _guard = self.rollover_guard.lock().expect("lock is poisoned");
-
-        log::info!("--- GC report for vLog @ {:?} ---", self.path);
-
-        let mut size_map = BTreeMap::<SegmentId, SegmentCounter>::new();
-
-        for handle in iter {
-            let (handle, size) = handle.map_err(|_| {
-                crate::Error::Io(std::io::Error::new(
-                    std::io::ErrorKind::Other,
-                    "Index returned error",
-                ))
-            })?;
-            let size = u64::from(size);
-
-            size_map
-                .entry(handle.segment_id)
-                .and_modify(|x| {
-                    x.item_count += 1;
-                    x.size += size;
-                })
-                .or_insert_with(|| SegmentCounter {
-                    size,
-                    item_count: 1,
-                });
-        }
-
-        for (&id, counter) in &size_map {
+    #[doc(hidden)]
+    pub fn consume_scan_result(&self, segment_ids: &[u64], size_map: &SizeMap) {
+        for (&id, counter) in size_map {
             let used_size = counter.size;
             let alive_item_count = counter.item_count;
 
@@ -412,20 +372,14 @@ impl ValueLog {
                 "Blob file #{id} has {}/{} stale MiB ({:.1}% stale, {stale_items}/{total_items} items) - space amp: {space_amp})",
                 stale_bytes / 1_024 / 1_024,
                 total_bytes / 1_024 / 1_024,
-                stale_ratio * 100.0
+                stale_ratio * 100.0,
             );
 
             segment.gc_stats.set_stale_bytes(stale_bytes);
             segment.gc_stats.set_stale_items(stale_items);
         }
 
-        for id in self
-            .manifest
-            .segments
-            .read()
-            .expect("lock is poisoned")
-            .keys()
-        {
+        for id in segment_ids {
             let segment = self
                 .manifest
                 .get_segment(*id)
@@ -440,13 +394,36 @@ impl ValueLog {
                 self.mark_as_stale(&[*id]);
             }
         }
+    }
 
+    /// Scans the given index and collecting GC statistics.
+    ///
+    /// # Errors
+    ///
+    /// Will return `Err` if an IO error occurs.
+    #[allow(clippy::result_unit_err)]
+    pub fn scan_for_stats(
+        &self,
+        iter: impl Iterator<Item = std::io::Result<(ValueHandle, u32)>>,
+    ) -> crate::Result<()> {
+        let mut scanner = Scanner::new(self, iter);
+
+        let segment_ids = self.manifest.list_segment_ids();
+
+        scanner.scan()?;
+        let size_map = scanner.finish();
+        self.consume_scan_result(&segment_ids, &size_map);
+
+        Ok(())
+    }
+
+    /// Prints GC report.
+    pub fn print_gc_report(&self) {
+        log::info!("--- GC report for vLog @ {:?} ---", self.path);
         log::info!("Total bytes: {}", self.manifest.total_bytes());
         log::info!("Stale bytes: {}", self.manifest.stale_bytes());
         log::info!("Space amp: {}", self.space_amp());
         log::info!("--- GC report done ---");
-
-        Ok(())
     }
 
     #[doc(hidden)]

tests/accidental_drop_rc.rs

@@ -0,0 +1,76 @@
+// This test is to prevent a race condition that may cause new blob files
+// to be accidentally dropped during GC.
+//
+// When a blob file is registered is after a `scan_for_stats`, it has an reference
+// count of 0. Then it would be dropped even though it was just created.
+
+use test_log::test;
+use value_log::{Config, IndexWriter, MockIndex, MockIndexWriter, ValueLog};
+
+#[test]
+fn accidental_drop_rc() -> value_log::Result<()> {
+    let folder = tempfile::tempdir()?;
+    let vl_path = folder.path();
+
+    let index = MockIndex::default();
+
+    let value_log = ValueLog::open(vl_path, Config::default())?;
+
+    for key in ["a", "b"] {
+        let value = &key;
+
+        let mut index_writer = MockIndexWriter(index.clone());
+        let mut writer = value_log.get_writer()?;
+
+        let handle = writer.get_next_value_handle(key.as_bytes());
+        index_writer.insert_indirect(key.as_bytes(), handle, value.len() as u32)?;
+
+        writer.write(key.as_bytes(), value.as_bytes())?;
+        value_log.register_writer(writer)?;
+    }
+
+    assert_eq!(2, value_log.segment_count());
+
+    value_log.scan_for_stats(index.read().unwrap().values().cloned().map(Ok))?;
+    value_log.drop_stale_segments()?;
+    assert_eq!(2, value_log.segment_count());
+
+    let segment_ids = value_log.manifest.list_segment_ids();
+
+    // NOTE: Now start a new GC scan
+    let index_lock = index.read().unwrap();
+    let mut scanner =
+        value_log::scanner::Scanner::new(&value_log, index_lock.values().cloned().map(Ok));
+    scanner.scan()?;
+    let scan_result = scanner.finish();
+    drop(index_lock);
+
+    // NOTE: Now, we create a new blob file, that won't be referenced in the GC report
+    {
+        let key = "c";
+        let value = &key;
+
+        let mut index_writer = MockIndexWriter(index.clone());
+        let mut writer = value_log.get_writer()?;
+
+        let handle = writer.get_next_value_handle(key.as_bytes());
+        index_writer.insert_indirect(key.as_bytes(), handle, value.len() as u32)?;
+
+        writer.write(key.as_bytes(), value.as_bytes())?;
+        value_log.register_writer(writer)?;
+    }
+    assert_eq!(3, value_log.segment_count());
+
+    // NOTE: Now we can consume the scan result, which, in a bad implementation
+    // would cause the new blob file to be marked as stale
+    //
+    // But we are forced to pass the list of segment IDs we saw before starting the
+    // scan, which prevents marking ones as stale that were created later
+    value_log.consume_scan_result(&segment_ids, &scan_result);
+
+    // IMPORTANT: The new blob file should not be dropped
+    value_log.drop_stale_segments()?;
+    assert_eq!(3, value_log.segment_count());
+
+    Ok(())
+}