add compaction tests and fix bugs in compaction

Signed-off-by: Alex Chi <iskyzh@gmail.com>

Author: skyzh

.config/nextest.toml

@@ -0,0 +1,2 @@
+[profile.default]
+slow-timeout = { period = "10s", terminate-after = 3 }

mini-lsm-book/src/week2-01-compaction.md

@@ -8,6 +8,13 @@ In this chapter, you will:
 * Implement the logic to update the LSM states and manage SST files on the filesystem.
 * Update LSM read path to incorporate the LSM levels.
 
+To copy the test cases into the starter code and run them,
+
+```
+cargo x copy-test --week 2 --day 1
+cargo x scheck
+```
+
 ## Task 1: Compaction Implementation
 
 In this task, you will implement the core logic of doing a compaction -- merge sort a set of SST files into a sorted run. You will need to modify:

mini-lsm-book/src/week2-02-simple.md

@@ -7,6 +7,13 @@ In this chapter, you will:
 * Implement a simple leveled compaction strategy and simulate it on the compaction simulator.
 * Start compaction as a background task and implement a compaction trigger in the system.
 
+To copy the test cases into the starter code and run them,
+
+```
+cargo x copy-test --week 2 --day 2
+cargo x scheck
+```
+
 ## Task 1: Simple Leveled Compaction
 
 In this chapter, we are going to implement our first compaction strategy -- simple leveled compaction. In this task, you will need to modify:

mini-lsm-book/src/week2-03-tiered.md

@@ -9,6 +9,13 @@ In this chapter, you will:
 
 The tiered compaction we talk about in this chapter is the same as RocksDB's universal compaction. We will use these two terminologies interchangeably.
 
+To copy the test cases into the starter code and run them,
+
+```
+cargo x copy-test --week 2 --day 3
+cargo x scheck
+```
+
 ## Task 1: Universal Compaction
 
 In this chapter, you will implement RocksDB's universal compaction, which is of the tiered compaction family compaction strategies. Similar to the simple leveled compaction strategy, we only use number of files as the indicator in this compaction strategy. And when we trigger the compaction jobs, we always include a full sorted run (tier) in the compaction job.

mini-lsm-book/src/week2-04-leveled.md

@@ -7,6 +7,13 @@ In this chapter, you will:
 * Implement a leveled compaction strategy and simulate it on the compaction simulator.
 * Incorporate leveled compaction strategy into the system.
 
+To copy the test cases into the starter code and run them,
+
+```
+cargo x copy-test --week 2 --day 4
+cargo x scheck
+```
+
 ## Task 1: Leveled Compaction
 
 In chapter 2 day 2, you have implemented the simple leveled compaction strategies. However, the implementation has a few problems:

mini-lsm-mvcc/src/compact.rs

@@ -98,6 +98,7 @@ impl CompactionController {
     }
 }
 
+#[derive(Debug, Clone)]
 pub enum CompactionOptions {
     /// Leveled compaction with partial compaction + dynamic level support (= RocksDB's Leveled
     /// Compaction)
@@ -309,25 +310,25 @@ impl LsmStorageInner {
         };
         println!("running compaction task: {:?}", task);
         let sstables = self.compact(&task)?;
-        let files_added = sstables.len();
         let output = sstables.iter().map(|x| x.sst_id()).collect::<Vec<_>>();
         let ssts_to_remove = {
             let state_lock = self.state_lock.lock();
+            let mut snapshot = self.state.read().as_ref().clone();
+            let mut new_sst_ids = Vec::new();
+            for file_to_add in sstables {
+                new_sst_ids.push(file_to_add.sst_id());
+                let result = snapshot.sstables.insert(file_to_add.sst_id(), file_to_add);
+                assert!(result.is_none());
+            }
             let (mut snapshot, files_to_remove) = self
                 .compaction_controller
-                .apply_compaction_result(&self.state.read(), &task, &output);
+                .apply_compaction_result(&snapshot, &task, &output);
             let mut ssts_to_remove = Vec::with_capacity(files_to_remove.len());
             for file_to_remove in &files_to_remove {
                 let result = snapshot.sstables.remove(file_to_remove);
                 assert!(result.is_some(), "cannot remove {}.sst", file_to_remove);
                 ssts_to_remove.push(result.unwrap());
             }
-            let mut new_sst_ids = Vec::new();
-            for file_to_add in sstables {
-                new_sst_ids.push(file_to_add.sst_id());
-                let result = snapshot.sstables.insert(file_to_add.sst_id(), file_to_add);
-                assert!(result.is_none());
-            }
             let mut state = self.state.write();
             *state = Arc::new(snapshot);
             drop(state);
@@ -339,9 +340,10 @@ impl LsmStorageInner {
             ssts_to_remove
         };
         println!(
-            "compaction finished: {} files removed, {} files added",
+            "compaction finished: {} files removed, {} files added, output={:?}",
             ssts_to_remove.len(),
-            files_added
+            output.len(),
+            output
         );
         for sst in ssts_to_remove {
             std::fs::remove_file(self.path_of_sst(sst.sst_id()))?;

mini-lsm-mvcc/src/compact/leveled.rs

@@ -124,11 +124,11 @@ impl LeveledCompactionController {
                 "target level sizes: {:?}, real level sizes: {:?}, base_level: {}",
                 target_level_size
                     .iter()
-                    .map(|x| format!("{}MB", x / 1024 / 1024))
+                    .map(|x| format!("{:.3}MB", *x as f64 / 1024.0 / 1024.0))
                     .collect::<Vec<_>>(),
                 real_level_size
                     .iter()
-                    .map(|x| format!("{}MB", x / 1024 / 1024))
+                    .map(|x| format!("{:.3}MB", *x as f64 / 1024.0 / 1024.0))
                     .collect::<Vec<_>>(),
                 base_level,
             );

mini-lsm-mvcc/src/compact/simple_leveled.rs

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use serde::{Deserialize, Serialize};
 
 use crate::lsm_storage::LsmStorageState;
@@ -95,12 +97,20 @@ impl SimpleLeveledCompactionController {
             files_to_remove.extend(&snapshot.levels[upper_level - 1].1);
             snapshot.levels[upper_level - 1].1.clear();
         } else {
-            assert_eq!(
-                task.upper_level_sst_ids, snapshot.l0_sstables,
-                "sst mismatched"
-            );
-            files_to_remove.extend(&snapshot.l0_sstables);
-            snapshot.l0_sstables.clear();
+            files_to_remove.extend(&task.upper_level_sst_ids);
+            let mut l0_ssts_compacted = task
+                .upper_level_sst_ids
+                .iter()
+                .copied()
+                .collect::<HashSet<_>>();
+            let new_l0_sstables = snapshot
+                .l0_sstables
+                .iter()
+                .copied()
+                .filter(|x| !l0_ssts_compacted.remove(x))
+                .collect::<Vec<_>>();
+            assert!(l0_ssts_compacted.is_empty());
+            snapshot.l0_sstables = new_l0_sstables;
         }
         assert_eq!(
             task.lower_level_sst_ids,

mini-lsm-mvcc/src/lsm_storage.rs

@@ -94,6 +94,16 @@ impl LsmStorageOptions {
             num_memtable_limit: 2,
         }
     }
+
+    pub fn default_for_week2_test(compaction_options: CompactionOptions) -> Self {
+        Self {
+            block_size: 4096,
+            target_sst_size: 1 << 20, // 1MB
+            compaction_options,
+            enable_wal: false,
+            num_memtable_limit: 2,
+        }
+    }
 }
 
 fn range_overlap(

mini-lsm-mvcc/src/tests.rs

@@ -7,3 +7,6 @@ mod week1_day5;
 mod week1_day6;
 mod week1_day7;
 mod week2_day1;
+mod week2_day2;
+mod week2_day3;
+mod week2_day4;

mini-lsm-mvcc/src/tests/week2_day2.rs

@@ -0,0 +1 @@
+../../../mini-lsm/src/tests/week2_day2.rs

mini-lsm-mvcc/src/tests/week2_day3.rs

@@ -0,0 +1 @@
+../../../mini-lsm/src/tests/week2_day3.rs

mini-lsm-mvcc/src/tests/week2_day4.rs

@@ -0,0 +1 @@
+../../../mini-lsm/src/tests/week2_day4.rs

mini-lsm-starter/src/lsm_storage.rs

@@ -90,6 +90,16 @@ impl LsmStorageOptions {
             num_memtable_limit: 2,
         }
     }
+
+    pub fn default_for_week2_test(compaction_options: CompactionOptions) -> Self {
+        Self {
+            block_size: 4096,
+            target_sst_size: 1 << 20, // 1MB
+            compaction_options,
+            enable_wal: false,
+            num_memtable_limit: 2,
+        }
+    }
 }
 
 /// The storage interface of the LSM tree.

mini-lsm/src/compact.rs

@@ -98,6 +98,7 @@ impl CompactionController {
     }
 }
 
+#[derive(Debug, Clone)]
 pub enum CompactionOptions {
     /// Leveled compaction with partial compaction + dynamic level support (= RocksDB's Leveled
     /// Compaction)
@@ -309,25 +310,25 @@ impl LsmStorageInner {
         };
         println!("running compaction task: {:?}", task);
         let sstables = self.compact(&task)?;
-        let files_added = sstables.len();
         let output = sstables.iter().map(|x| x.sst_id()).collect::<Vec<_>>();
         let ssts_to_remove = {
             let state_lock = self.state_lock.lock();
+            let mut snapshot = self.state.read().as_ref().clone();
+            let mut new_sst_ids = Vec::new();
+            for file_to_add in sstables {
+                new_sst_ids.push(file_to_add.sst_id());
+                let result = snapshot.sstables.insert(file_to_add.sst_id(), file_to_add);
+                assert!(result.is_none());
+            }
             let (mut snapshot, files_to_remove) = self
                 .compaction_controller
-                .apply_compaction_result(&self.state.read(), &task, &output);
+                .apply_compaction_result(&snapshot, &task, &output);
             let mut ssts_to_remove = Vec::with_capacity(files_to_remove.len());
             for file_to_remove in &files_to_remove {
                 let result = snapshot.sstables.remove(file_to_remove);
                 assert!(result.is_some(), "cannot remove {}.sst", file_to_remove);
                 ssts_to_remove.push(result.unwrap());
             }
-            let mut new_sst_ids = Vec::new();
-            for file_to_add in sstables {
-                new_sst_ids.push(file_to_add.sst_id());
-                let result = snapshot.sstables.insert(file_to_add.sst_id(), file_to_add);
-                assert!(result.is_none());
-            }
             let mut state = self.state.write();
             *state = Arc::new(snapshot);
             drop(state);
@@ -339,9 +340,10 @@ impl LsmStorageInner {
             ssts_to_remove
         };
         println!(
-            "compaction finished: {} files removed, {} files added",
+            "compaction finished: {} files removed, {} files added, output={:?}",
             ssts_to_remove.len(),
-            files_added
+            output.len(),
+            output
         );
         for sst in ssts_to_remove {
             std::fs::remove_file(self.path_of_sst(sst.sst_id()))?;

mini-lsm/src/compact/leveled.rs

@@ -124,11 +124,11 @@ impl LeveledCompactionController {
                 "target level sizes: {:?}, real level sizes: {:?}, base_level: {}",
                 target_level_size
                     .iter()
-                    .map(|x| format!("{}MB", x / 1024 / 1024))
+                    .map(|x| format!("{:.3}MB", *x as f64 / 1024.0 / 1024.0))
                     .collect::<Vec<_>>(),
                 real_level_size
                     .iter()
-                    .map(|x| format!("{}MB", x / 1024 / 1024))
+                    .map(|x| format!("{:.3}MB", *x as f64 / 1024.0 / 1024.0))
                     .collect::<Vec<_>>(),
                 base_level,
             );

mini-lsm/src/compact/simple_leveled.rs

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use serde::{Deserialize, Serialize};
 
 use crate::lsm_storage::LsmStorageState;
@@ -95,12 +97,20 @@ impl SimpleLeveledCompactionController {
             files_to_remove.extend(&snapshot.levels[upper_level - 1].1);
             snapshot.levels[upper_level - 1].1.clear();
         } else {
-            assert_eq!(
-                task.upper_level_sst_ids, snapshot.l0_sstables,
-                "sst mismatched"
-            );
-            files_to_remove.extend(&snapshot.l0_sstables);
-            snapshot.l0_sstables.clear();
+            files_to_remove.extend(&task.upper_level_sst_ids);
+            let mut l0_ssts_compacted = task
+                .upper_level_sst_ids
+                .iter()
+                .copied()
+                .collect::<HashSet<_>>();
+            let new_l0_sstables = snapshot
+                .l0_sstables
+                .iter()
+                .copied()
+                .filter(|x| !l0_ssts_compacted.remove(x))
+                .collect::<Vec<_>>();
+            assert!(l0_ssts_compacted.is_empty());
+            snapshot.l0_sstables = new_l0_sstables;
         }
         assert_eq!(
             task.lower_level_sst_ids,

mini-lsm/src/lsm_storage.rs

@@ -94,6 +94,16 @@ impl LsmStorageOptions {
             num_memtable_limit: 2,
         }
     }
+
+    pub fn default_for_week2_test(compaction_options: CompactionOptions) -> Self {
+        Self {
+            block_size: 4096,
+            target_sst_size: 1 << 20, // 1MB
+            compaction_options,
+            enable_wal: false,
+            num_memtable_limit: 2,
+        }
+    }
 }
 
 fn range_overlap(

mini-lsm/src/tests.rs

@@ -7,3 +7,6 @@ mod week1_day5;
 mod week1_day6;
 mod week1_day7;
 mod week2_day1;
+mod week2_day2;
+mod week2_day3;
+mod week2_day4;