Merge pull request #118 from fjall-rs/fix/lockless-ranges

2.7.2

Author: marvin-j97

Cargo.toml

@@ -2,7 +2,7 @@
 name = "lsm-tree"
 description = "A K.I.S.S. implementation of log-structured merge trees (LSM-trees/LSMTs)"
 license = "MIT OR Apache-2.0"
-version = "2.7.1"
+version = "2.7.2"
 edition = "2021"
 rust-version = "1.75.0"
 readme = "README.md"

src/abstract.rs

@@ -47,7 +47,10 @@ pub trait AbstractTree {
     fn register_segments(&self, segments: &[Segment]) -> crate::Result<()>;
 
     /// Write-locks the active memtable for exclusive access
-    fn lock_active_memtable(&self) -> RwLockWriteGuard<'_, Memtable>;
+    fn lock_active_memtable(&self) -> RwLockWriteGuard<'_, Arc<Memtable>>;
+
+    /// Clears the active memtable atomically.
+    fn clear_active_memtable(&self);
 
     /// Sets the active memtable.
     ///

src/blob_tree/gc/reader.rs

@@ -3,24 +3,24 @@
 // (found in the LICENSE-* files in the repository)
 
 use crate::{blob_tree::value::MaybeInlineValue, coding::Decode, Memtable};
-use std::{io::Cursor, sync::RwLockWriteGuard};
+use std::io::Cursor;
 use value_log::ValueHandle;
 
 #[allow(clippy::module_name_repetitions)]
 pub struct GcReader<'a> {
     tree: &'a crate::Tree,
-    memtable: &'a RwLockWriteGuard<'a, Memtable>,
+    memtable: &'a Memtable,
 }
 
 impl<'a> GcReader<'a> {
-    pub fn new(tree: &'a crate::Tree, memtable: &'a RwLockWriteGuard<'a, Memtable>) -> Self {
+    pub fn new(tree: &'a crate::Tree, memtable: &'a Memtable) -> Self {
         Self { tree, memtable }
     }
 
     fn get_internal(&self, key: &[u8]) -> crate::Result<Option<MaybeInlineValue>> {
         let Some(item) = self
             .tree
-            .get_internal_entry_with_lock(self.memtable, key, None)?
+            .get_internal_entry_with_memtable(self.memtable, key, None)?
             .map(|x| x.value)
         else {
             return Ok(None);

src/blob_tree/gc/writer.rs

@@ -6,18 +6,17 @@ use crate::{
     blob_tree::value::MaybeInlineValue, coding::Encode, value::InternalValue, Memtable, SeqNo,
     UserKey,
 };
-use std::sync::RwLockWriteGuard;
 use value_log::ValueHandle;
 
 #[allow(clippy::module_name_repetitions)]
 pub struct GcWriter<'a> {
     seqno: SeqNo,
     buffer: Vec<(UserKey, ValueHandle, u32)>,
-    memtable: &'a RwLockWriteGuard<'a, Memtable>,
+    memtable: &'a Memtable,
 }
 
 impl<'a> GcWriter<'a> {
-    pub fn new(seqno: SeqNo, memtable: &'a RwLockWriteGuard<'a, Memtable>) -> Self {
+    pub fn new(seqno: SeqNo, memtable: &'a Memtable) -> Self {
         Self {
             seqno,
             memtable,

src/blob_tree/mod.rs

@@ -230,6 +230,10 @@ impl BlobTree {
 }
 
 impl AbstractTree for BlobTree {
+    fn clear_active_memtable(&self) {
+        self.index.clear_active_memtable();
+    }
+
     fn l0_run_count(&self) -> usize {
         self.index.l0_run_count()
     }
@@ -411,7 +415,7 @@ impl AbstractTree for BlobTree {
         Ok(())
     }
 
-    fn lock_active_memtable(&self) -> std::sync::RwLockWriteGuard<'_, Memtable> {
+    fn lock_active_memtable(&self) -> std::sync::RwLockWriteGuard<'_, Arc<Memtable>> {
         self.index.lock_active_memtable()
     }
 

src/compaction/fifo.rs

@@ -133,7 +133,7 @@ mod tests {
         time::unix_timestamp,
         HashSet, KeyRange,
     };
-    use std::sync::Arc;
+    use std::sync::{atomic::AtomicBool, Arc};
     use test_log::test;
 
     #[allow(clippy::expect_used)]
@@ -180,6 +180,9 @@ mod tests {
             block_cache,
 
             bloom_filter: Some(crate::bloom::BloomFilter::with_fp_rate(1, 0.1)),
+
+            path: "a".into(),
+            is_deleted: AtomicBool::default(),
         }
         .into()
     }

src/compaction/leveled.rs

@@ -397,7 +397,10 @@ mod tests {
         time::unix_timestamp,
         Config, HashSet, KeyRange,
     };
-    use std::{path::Path, sync::Arc};
+    use std::{
+        path::Path,
+        sync::{atomic::AtomicBool, Arc},
+    };
     use test_log::test;
 
     fn string_key_range(a: &str, b: &str) -> KeyRange {
@@ -456,6 +459,9 @@ mod tests {
             block_cache,
 
             bloom_filter: Some(crate::bloom::BloomFilter::with_fp_rate(1, 0.1)),
+
+            path: "a".into(),
+            is_deleted: AtomicBool::default(),
         }
         .into()
     }

src/compaction/maintenance.rs

@@ -97,7 +97,7 @@ mod tests {
         },
         KeyRange,
     };
-    use std::sync::Arc;
+    use std::sync::{atomic::AtomicBool, Arc};
     use test_log::test;
 
     #[allow(clippy::expect_used)]
@@ -143,6 +143,9 @@ mod tests {
             block_cache,
 
             bloom_filter: Some(crate::bloom::BloomFilter::with_fp_rate(1, 0.1)),
+
+            path: "a".into(),
+            is_deleted: AtomicBool::default(),
         }
         .into()
     }

src/compaction/tiered.rs

@@ -168,7 +168,7 @@ mod tests {
         },
         HashSet, KeyRange, SeqNo,
     };
-    use std::sync::Arc;
+    use std::sync::{atomic::AtomicBool, Arc};
     use test_log::test;
 
     #[allow(clippy::expect_used)]
@@ -214,6 +214,9 @@ mod tests {
             block_cache,
 
             bloom_filter: Some(BloomFilter::with_fp_rate(1, 0.1)),
+
+            path: "a".into(),
+            is_deleted: AtomicBool::default(),
         }
         .into()
     }

src/compaction/worker.rs

@@ -24,7 +24,10 @@ use crate::{
 };
 use std::{
     path::Path,
-    sync::{atomic::AtomicU64, Arc, RwLock, RwLockWriteGuard},
+    sync::{
+        atomic::{AtomicBool, AtomicU64},
+        Arc, RwLock, RwLockWriteGuard,
+    },
     time::Instant,
 };
 
@@ -211,6 +214,19 @@ fn merge_segments(
         return Ok(());
     }
 
+    let Some(segments) = payload
+        .segment_ids
+        .iter()
+        .map(|&id| levels.get_segment(id))
+        .collect::<Option<Vec<_>>>()
+    else {
+        log::warn!(
+            "Compaction task created by {:?} contained segments not referenced in the level manifest",
+            opts.strategy.get_name(),
+        );
+        return Ok(());
+    };
+
     let segments_base_folder = opts.config.path.join(SEGMENTS_FOLDER);
 
     log::debug!(
@@ -381,6 +397,8 @@ fn merge_segments(
             let block_index = Arc::new(block_index);
 
             Ok(SegmentInner {
+                path: segment_file_path.to_path_buf(),
+
                 tree_id: opts.tree_id,
 
                 descriptor_table: opts.config.descriptor_table.clone(),
@@ -393,6 +411,8 @@ fn merge_segments(
                 block_index,
 
                 bloom_filter: Segment::load_bloom(&segment_file_path, trailer.offsets.bloom_ptr)?,
+
+                is_deleted: AtomicBool::default(),
             }
             .into())
         })
@@ -412,11 +432,7 @@ fn merge_segments(
     // NOTE: Mind lock order L -> M -> S
     log::trace!("compactor: acquiring levels manifest write lock");
     let mut levels = opts.levels.write().expect("lock is poisoned");
-
-    // IMPORTANT: Write lock memtable(s), otherwise segments may get deleted while a range read is happening
-    // NOTE: Mind lock order L -> M -> S
-    log::trace!("compactor: acquiring sealed memtables write lock");
-    let sealed_memtables_guard = opts.sealed_memtables.write().expect("lock is poisoned");
+    log::trace!("compactor: acquired levels manifest write lock");
 
     let swap_result = levels.atomic_swap(|recipe| {
         for segment in created_segments.iter().cloned() {
@@ -451,19 +467,11 @@ fn merge_segments(
             .insert(&segment_file_path, segment.global_id());
     }
 
-    // NOTE: Segments are registered, we can unlock the memtable(s) safely
-    drop(sealed_memtables_guard);
-
     // NOTE: If the application were to crash >here< it's fine
     // The segments are not referenced anymore, and will be
     // cleaned up upon recovery
-    for segment_id in &payload.segment_ids {
-        let segment_file_path = segments_base_folder.join(segment_id.to_string());
-        log::trace!("Removing old segment at {segment_file_path:?}");
-
-        if let Err(e) = std::fs::remove_file(segment_file_path) {
-            log::error!("Failed to cleanup file of deleted segment: {e:?}");
-        }
+    for segment in segments {
+        segment.mark_as_deleted();
     }
 
     // NOTE: Unlock level manifest before clearing old file descriptors
@@ -476,17 +484,6 @@ fn merge_segments(
     levels.show_segments(payload.segment_ids.iter().copied());
     drop(levels);
 
-    log::trace!(
-        "Closing file handles for old segment files: {:?}",
-        payload.segment_ids
-    );
-
-    for segment_id in &payload.segment_ids {
-        opts.config
-            .descriptor_table
-            .remove((opts.tree_id, *segment_id).into());
-    }
-
     log::trace!("Compaction successful");
 
     Ok(())
@@ -506,11 +503,17 @@ fn drop_segments(
         return Ok(());
     }
 
-    let segments_base_folder = opts.config.path.join(SEGMENTS_FOLDER);
-
-    // IMPORTANT: Write lock memtable, otherwise segments may get deleted while a range read is happening
-    log::trace!("Acquiring sealed memtables write lock");
-    let memtable_lock = opts.sealed_memtables.write().expect("lock is poisoned");
+    let Some(segments) = segment_ids
+        .iter()
+        .map(|id| levels.get_segment(id.segment_id()))
+        .collect::<Option<Vec<_>>>()
+    else {
+        log::warn!(
+        "Compaction task created by {:?} contained segments not referenced in the level manifest",
+        opts.strategy.get_name(),
+    );
+        return Ok(());
+    };
 
     // IMPORTANT: Write the segment with the removed segments first
     // Otherwise the folder is deleted, but the segment is still referenced!
@@ -525,26 +528,13 @@ fn drop_segments(
         }
     })?;
 
-    drop(memtable_lock);
     drop(levels);
 
     // NOTE: If the application were to crash >here< it's fine
     // The segments are not referenced anymore, and will be
     // cleaned up upon recovery
-    for key in segment_ids {
-        let segment_id = key.segment_id();
-
-        let segment_file_path = segments_base_folder.join(segment_id.to_string());
-        log::trace!("Removing old segment at {segment_file_path:?}");
-
-        if let Err(e) = std::fs::remove_file(segment_file_path) {
-            log::error!("Failed to cleanup file of deleted segment: {e:?}");
-        }
-    }
-
-    for key in segment_ids {
-        log::trace!("Closing file handles for segment data file");
-        opts.config.descriptor_table.remove(*key);
+    for segment in segments {
+        segment.mark_as_deleted();
     }
 
     log::trace!("Dropped {} segments", segment_ids.len());

src/level_manifest/level.rs

@@ -257,7 +257,7 @@ mod tests {
         },
         AbstractTree, KeyRange, Slice,
     };
-    use std::sync::Arc;
+    use std::sync::{atomic::AtomicBool, Arc};
     use test_log::test;
 
     #[allow(clippy::expect_used)]
@@ -303,6 +303,9 @@ mod tests {
             block_cache,
 
             bloom_filter: Some(crate::bloom::BloomFilter::with_fp_rate(1, 0.1)),
+
+            path: "a".into(),
+            is_deleted: AtomicBool::default(),
         }
         .into()
     }

src/level_manifest/mod.rs

@@ -378,6 +378,15 @@ impl LevelManifest {
         output
     }
 
+    pub(crate) fn get_segment(&self, id: SegmentId) -> Option<Segment> {
+        for level in &self.levels {
+            if let Some(segment) = level.segments.iter().find(|x| x.id() == id).cloned() {
+                return Some(segment);
+            }
+        }
+        None
+    }
+
     /// Returns a view into the levels, hiding all segments that currently are being compacted
     #[must_use]
     pub fn resolved_view(&self) -> Vec<Level> {