fix compaction split bug and add 3.1 test

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

mini-lsm-mvcc/src/compact.rs

@@ -153,19 +153,22 @@ impl LsmStorageInner {
             }
 
             let builder_inner = builder.as_mut().unwrap();
-            builder_inner.add(iter.key(), iter.value());
 
             if builder_inner.estimated_size() >= self.options.target_sst_size && !same_as_last_key {
-                let sst_id = self.next_sst_id();
-                let builder = builder.take().unwrap();
-                let sst = Arc::new(builder.build(
+                let sst_id: usize = self.next_sst_id();
+                let old_builder = builder.take().unwrap();
+                let sst = Arc::new(old_builder.build(
                     sst_id,
                     Some(self.block_cache.clone()),
                     self.path_of_sst(sst_id),
                 )?);
                 new_sst.push(sst);
+                builder = Some(SsTableBuilder::new(self.options.block_size));
             }
 
+            let builder_inner = builder.as_mut().unwrap();
+            builder_inner.add(iter.key(), iter.value());
+
             if !same_as_last_key {
                 last_key.clear();
                 last_key.extend(iter.key().key_ref());

mini-lsm-mvcc/src/key.rs

@@ -35,6 +35,10 @@ impl<T: AsRef<[u8]>> Key<T> {
     pub fn is_empty(&self) -> bool {
         self.0.as_ref().is_empty()
     }
+
+    pub fn for_testing_ts(self) -> u64 {
+        self.1
+    }
 }
 
 impl Key<Vec<u8>> {
@@ -149,6 +153,10 @@ impl<'a> Key<&'a [u8]> {
     pub fn for_testing_from_slice_no_ts(slice: &'a [u8]) -> Self {
         Self(slice, TS_DEFAULT)
     }
+
+    pub fn for_testing_from_slice_with_ts(slice: &'a [u8], ts: u64) -> Self {
+        Self(slice, ts)
+    }
 }
 
 impl<T: AsRef<[u8]> + Debug> Debug for Key<T> {

mini-lsm-mvcc/src/tests.rs

@@ -12,3 +12,5 @@ mod week2_day3;
 mod week2_day4;
 // mod week2_day5;
 // mod week2_day6;
+mod week3_day1;
+mod week3_day2;

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

@@ -0,0 +1,54 @@
+use std::sync::Arc;
+
+use bytes::Bytes;
+use tempfile::tempdir;
+
+use crate::key::KeySlice;
+use crate::table::{FileObject, SsTable, SsTableBuilder, SsTableIterator};
+
+use super::harness::{check_iter_result_by_key_and_ts, generate_sst_with_ts};
+
+#[test]
+fn test_sst_build_multi_version_simple() {
+    let mut builder = SsTableBuilder::new(16);
+    builder.add(
+        KeySlice::for_testing_from_slice_with_ts(b"233", 233),
+        b"233333",
+    );
+    builder.add(
+        KeySlice::for_testing_from_slice_with_ts(b"233", 0),
+        b"2333333",
+    );
+    let dir = tempdir().unwrap();
+    builder.build_for_test(dir.path().join("1.sst")).unwrap();
+}
+
+fn generate_test_data() -> Vec<((Bytes, u64), Bytes)> {
+    (0..100)
+        .map(|id| {
+            (
+                (Bytes::from(format!("key{:05}", id / 5)), 5 - (id % 5)),
+                Bytes::from(format!("value{:05}", id)),
+            )
+        })
+        .collect()
+}
+
+#[test]
+fn test_sst_build_multi_version_hard() {
+    let dir = tempdir().unwrap();
+    let data = generate_test_data();
+    generate_sst_with_ts(1, dir.path().join("1.sst"), data.clone(), None);
+    let sst = Arc::new(
+        SsTable::open(
+            1,
+            None,
+            FileObject::open(&dir.path().join("1.sst")).unwrap(),
+        )
+        .unwrap(),
+    );
+    check_iter_result_by_key_and_ts(
+        &mut SsTableIterator::create_and_seek_to_first(sst).unwrap(),
+        data,
+    );
+}

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

@@ -0,0 +1,61 @@
+use std::time::Duration;
+
+use tempfile::tempdir;
+
+use crate::{
+    compact::CompactionOptions,
+    lsm_storage::{LsmStorageOptions, MiniLsm},
+    tests::harness::dump_files_in_dir,
+};
+
+#[test]
+fn test_task_1_2_integration() {
+    let dir = tempdir().unwrap();
+    let mut options = LsmStorageOptions::default_for_week2_test(CompactionOptions::NoCompaction);
+    options.enable_wal = true;
+    let storage = MiniLsm::open(&dir, options.clone()).unwrap();
+    let _txn = storage.new_txn().unwrap();
+    for i in 0..=20000 {
+        storage
+            .put(b"0", format!("{:02000}", i).as_bytes())
+            .unwrap();
+    }
+    std::thread::sleep(Duration::from_secs(1)); // wait until all memtables flush
+    while {
+        let snapshot = storage.inner.state.read();
+        !snapshot.imm_memtables.is_empty()
+    } {
+        storage.inner.force_flush_next_imm_memtable().unwrap();
+    }
+    assert!(storage.inner.state.read().l0_sstables.len() > 1);
+    storage.force_full_compaction().unwrap();
+    storage.dump_structure();
+    dump_files_in_dir(&dir);
+    assert!(storage.inner.state.read().l0_sstables.is_empty());
+    assert_eq!(storage.inner.state.read().levels.len(), 1);
+    // same key in the same SST
+    assert_eq!(storage.inner.state.read().levels[0].1.len(), 1);
+    for i in 0..=100 {
+        storage
+            .put(b"1", format!("{:02000}", i).as_bytes())
+            .unwrap();
+    }
+    storage
+        .inner
+        .force_freeze_memtable(&storage.inner.state_lock.lock())
+        .unwrap();
+    std::thread::sleep(Duration::from_secs(1)); // wait until all memtables flush
+    while {
+        let snapshot = storage.inner.state.read();
+        !snapshot.imm_memtables.is_empty()
+    } {
+        storage.inner.force_flush_next_imm_memtable().unwrap();
+    }
+    storage.force_full_compaction().unwrap();
+    storage.dump_structure();
+    dump_files_in_dir(&dir);
+    assert!(storage.inner.state.read().l0_sstables.is_empty());
+    assert_eq!(storage.inner.state.read().levels.len(), 1);
+    // same key in the same SST, now we should split two
+    assert_eq!(storage.inner.state.read().levels[0].1.len(), 2);
+}