fix: avoid leveled compaction crash when recovering from manifest (#63)

* Fix: Avoid leveled copaction crash when recovering from manifest

* Also sort SSTs in manifest recovery

* Add `in_recovery` flag to `apply_compaction_result`

- Don't sort the SSTs inside `apply_compaction_result` if in recovery

mini-lsm-mvcc/src/compact.rs

@@ -73,10 +73,11 @@ impl CompactionController {
         snapshot: &LsmStorageState,
         task: &CompactionTask,
         output: &[usize],
+        in_recovery: bool,
     ) -> (LsmStorageState, Vec<usize>) {
         match (self, task) {
             (CompactionController::Leveled(ctrl), CompactionTask::Leveled(task)) => {
-                ctrl.apply_compaction_result(snapshot, task, output)
+                ctrl.apply_compaction_result(snapshot, task, output, in_recovery)
             }
             (CompactionController::Simple(ctrl), CompactionTask::Simple(task)) => {
                 ctrl.apply_compaction_result(snapshot, task, output)
@@ -381,7 +382,8 @@ impl LsmStorageInner {
             }
             let (mut snapshot, files_to_remove) = self
                 .compaction_controller
-                .apply_compaction_result(&snapshot, &task, &output);
+                .apply_compaction_result(&snapshot, &task, &output, false);
+
             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);

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

@@ -160,6 +160,7 @@ impl LeveledCompactionController {
         snapshot: &LsmStorageState,
         task: &LeveledCompactionTask,
         output: &[usize],
+        in_recovery: bool,
     ) -> (LsmStorageState, Vec<usize>) {
         let mut snapshot = snapshot.clone();
         let mut files_to_remove = Vec::new();
@@ -216,14 +217,17 @@ impl LeveledCompactionController {
             .collect::<Vec<_>>();
         assert!(lower_level_sst_ids_set.is_empty());
         new_lower_level_ssts.extend(output);
-        new_lower_level_ssts.sort_by(|x, y| {
-            snapshot
-                .sstables
-                .get(x)
-                .unwrap()
-                .first_key()
-                .cmp(snapshot.sstables.get(y).unwrap().first_key())
-        });
+        // Don't sort the SST IDs during recovery because actual SSTs are not loaded at that point
+        if !in_recovery {
+            new_lower_level_ssts.sort_by(|x, y| {
+                snapshot
+                    .sstables
+                    .get(x)
+                    .unwrap()
+                    .first_key()
+                    .cmp(snapshot.sstables.get(y).unwrap().first_key())
+            });
+        }
         snapshot.levels[task.lower_level - 1].1 = new_lower_level_ssts;
         (snapshot, files_to_remove)
     }

mini-lsm-mvcc/src/lsm_storage.rs

@@ -368,8 +368,8 @@ impl LsmStorageInner {
                         memtables.insert(x);
                     }
                     ManifestRecord::Compaction(task, output) => {
-                        let (new_state, _) =
-                            compaction_controller.apply_compaction_result(&state, &task, &output);
+                        let (new_state, _) = compaction_controller
+                            .apply_compaction_result(&state, &task, &output, true);
                         // TODO: apply remove again
                         state = new_state;
                         next_sst_id =
@@ -400,6 +400,18 @@ impl LsmStorageInner {
 
             next_sst_id += 1;
 
+            // Sort SSTs on each level
+            for (_id, ssts) in &mut state.levels {
+                ssts.sort_by(|x, y| {
+                    state
+                        .sstables
+                        .get(x)
+                        .unwrap()
+                        .first_key()
+                        .cmp(state.sstables.get(y).unwrap().first_key())
+                })
+            }
+
             // recover memtables
             if options.enable_wal {
                 let mut wal_cnt = 0;