add week 2 day 5 + 6 tests

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

mini-lsm/src/tests/harness.rs

@@ -1,4 +1,4 @@
-use std::{collections::BTreeMap, path::Path, sync::Arc, time::Duration};
+use std::{collections::BTreeMap, ops::Bound, path::Path, sync::Arc, time::Duration};
 
 use anyhow::{bail, Result};
 use bytes::Bytes;
@@ -9,7 +9,7 @@ use crate::{
         TieredCompactionOptions,
     },
     iterators::StorageIterator,
-    key::KeySlice,
+    key::{KeySlice, TS_ENABLED},
     lsm_storage::{BlockCache, LsmStorageInner, MiniLsm},
     table::{SsTable, SsTableBuilder},
 };
@@ -171,11 +171,12 @@ pub fn compaction_bench(storage: Arc<MiniLsm>) {
     let gen_key = |i| format!("{:010}", i); // 10B
     let gen_value = |i| format!("{:0110}", i); // 110B
     let mut max_key = 0;
+    let overlaps = if TS_ENABLED { 10000 } else { 20000 };
     for iter in 0..10 {
         let range_begin = iter * 5000;
-        for i in range_begin..(range_begin + 10000) {
+        for i in range_begin..(range_begin + overlaps) {
             // 120B per key, 4MB data populated
-            let key = gen_key(i);
+            let key: String = gen_key(i);
             let version = key_map.get(&i).copied().unwrap_or_default() + 1;
             let value = gen_value(version);
             key_map.insert(i, version);
@@ -184,17 +185,24 @@ pub fn compaction_bench(storage: Arc<MiniLsm>) {
         }
     }
 
+    let mut expected_key_value_pairs = Vec::new();
     for i in 0..(max_key + 40000) {
         let key = gen_key(i);
         let value = storage.get(key.as_bytes()).unwrap();
         if let Some(val) = key_map.get(&i) {
             let expected_value = gen_value(*val);
-            assert_eq!(value, Some(Bytes::from(expected_value)));
+            assert_eq!(value, Some(Bytes::from(expected_value.clone())));
+            expected_key_value_pairs.push((Bytes::from(key), Bytes::from(expected_value)));
         } else {
             assert!(value.is_none());
         }
     }
 
+    check_lsm_iter_result_by_key(
+        &mut storage.scan(Bound::Unbounded, Bound::Unbounded).unwrap(),
+        expected_key_value_pairs,
+    );
+
     while {
         let snapshot = storage.inner.state.read();
         !snapshot.imm_memtables.is_empty()
@@ -324,3 +332,9 @@ pub fn check_compaction_ratio(storage: Arc<MiniLsm>) {
         }
     }
 }
+
+pub fn dump_files_in_dir(path: impl AsRef<Path>) {
+    for f in path.as_ref().read_dir().unwrap() {
+        println!("{}", f.unwrap().path().display())
+    }
+}

mini-lsm/src/tests/week2_day5.rs

@@ -0,0 +1,81 @@
+use tempfile::tempdir;
+
+use crate::{
+    compact::{
+        CompactionOptions, LeveledCompactionOptions, SimpleLeveledCompactionOptions,
+        TieredCompactionOptions,
+    },
+    lsm_storage::{LsmStorageOptions, MiniLsm},
+    tests::harness::dump_files_in_dir,
+};
+
+#[test]
+fn test_integration_leveled() {
+    test_integration(CompactionOptions::Leveled(LeveledCompactionOptions {
+        level_size_multiplier: 2,
+        level0_file_num_compaction_trigger: 2,
+        max_levels: 3,
+        base_level_size_mb: 1,
+    }))
+}
+
+#[test]
+fn test_integration_tiered() {
+    test_integration(CompactionOptions::Tiered(TieredCompactionOptions {
+        num_tiers: 3,
+        max_size_amplification_percent: 200,
+        size_ratio: 1,
+        min_merge_width: 3,
+    }))
+}
+
+#[test]
+fn test_integration_simple() {
+    test_integration(CompactionOptions::Simple(SimpleLeveledCompactionOptions {
+        size_ratio_percent: 200,
+        level0_file_num_compaction_trigger: 2,
+        max_levels: 3,
+    }));
+}
+
+fn test_integration(compaction_options: CompactionOptions) {
+    let dir = tempdir().unwrap();
+    let storage = MiniLsm::open(
+        &dir,
+        LsmStorageOptions::default_for_week2_test(compaction_options.clone()),
+    )
+    .unwrap();
+    for i in 0..=20 {
+        storage.put(b"0", format!("v{}", i).as_bytes()).unwrap();
+        if i % 2 == 0 {
+            storage.put(b"1", format!("v{}", i).as_bytes()).unwrap();
+        } else {
+            storage.delete(b"1").unwrap();
+        }
+        if i % 2 == 1 {
+            storage.put(b"2", format!("v{}", i).as_bytes()).unwrap();
+        } else {
+            storage.delete(b"2").unwrap();
+        }
+        storage
+            .inner
+            .force_freeze_memtable(&storage.inner.state_lock.lock())
+            .unwrap();
+    }
+    storage.close().unwrap();
+    // ensure all SSTs are flushed
+    assert!(storage.inner.state.read().memtable.is_empty());
+    assert!(storage.inner.state.read().imm_memtables.is_empty());
+    storage.dump_structure();
+    drop(storage);
+    dump_files_in_dir(&dir);
+
+    let storage = MiniLsm::open(
+        &dir,
+        LsmStorageOptions::default_for_week2_test(compaction_options.clone()),
+    )
+    .unwrap();
+    assert_eq!(&storage.get(b"0").unwrap().unwrap()[..], b"v20".as_slice());
+    assert_eq!(&storage.get(b"1").unwrap().unwrap()[..], b"v20".as_slice());
+    assert_eq!(storage.get(b"2").unwrap(), None);
+}

mini-lsm/src/tests/week2_day6.rs

@@ -0,0 +1,77 @@
+use tempfile::tempdir;
+
+use crate::{
+    compact::{
+        CompactionOptions, LeveledCompactionOptions, SimpleLeveledCompactionOptions,
+        TieredCompactionOptions,
+    },
+    lsm_storage::{LsmStorageOptions, MiniLsm},
+    tests::harness::dump_files_in_dir,
+};
+
+#[test]
+fn test_integration_leveled() {
+    test_integration(CompactionOptions::Leveled(LeveledCompactionOptions {
+        level_size_multiplier: 2,
+        level0_file_num_compaction_trigger: 2,
+        max_levels: 3,
+        base_level_size_mb: 1,
+    }))
+}
+
+#[test]
+fn test_integration_tiered() {
+    test_integration(CompactionOptions::Tiered(TieredCompactionOptions {
+        num_tiers: 3,
+        max_size_amplification_percent: 200,
+        size_ratio: 1,
+        min_merge_width: 3,
+    }))
+}
+
+#[test]
+fn test_integration_simple() {
+    test_integration(CompactionOptions::Simple(SimpleLeveledCompactionOptions {
+        size_ratio_percent: 200,
+        level0_file_num_compaction_trigger: 2,
+        max_levels: 3,
+    }));
+}
+
+fn test_integration(compaction_options: CompactionOptions) {
+    let dir = tempdir().unwrap();
+    let mut options = LsmStorageOptions::default_for_week2_test(compaction_options);
+    options.enable_wal = true;
+    let storage = MiniLsm::open(&dir, options.clone()).unwrap();
+    for i in 0..=20 {
+        storage.put(b"0", format!("v{}", i).as_bytes()).unwrap();
+        if i % 2 == 0 {
+            storage.put(b"1", format!("v{}", i).as_bytes()).unwrap();
+        } else {
+            storage.delete(b"1").unwrap();
+        }
+        if i % 2 == 1 {
+            storage.put(b"2", format!("v{}", i).as_bytes()).unwrap();
+        } else {
+            storage.delete(b"2").unwrap();
+        }
+        storage
+            .inner
+            .force_freeze_memtable(&storage.inner.state_lock.lock())
+            .unwrap();
+    }
+    storage.close().unwrap();
+    // ensure some SSTs are not flushed
+    assert!(
+        !storage.inner.state.read().memtable.is_empty()
+            || !storage.inner.state.read().imm_memtables.is_empty()
+    );
+    storage.dump_structure();
+    drop(storage);
+    dump_files_in_dir(&dir);
+
+    let storage = MiniLsm::open(&dir, options).unwrap();
+    assert_eq!(&storage.get(b"0").unwrap().unwrap()[..], b"v20".as_slice());
+    assert_eq!(&storage.get(b"1").unwrap().unwrap()[..], b"v20".as_slice());
+    assert_eq!(storage.get(b"2").unwrap(), None);
+}