From 4b8cb1848419f049ba9819def591ae416daccfbb Mon Sep 17 00:00:00 2001
From: pedrinfx <canaldopedrinho233@gmail.com>
Date: Sun, 19 Feb 2023 12:27:32 -0300
Subject: [PATCH] fix: update index on flush

---
 src/storage/lsm/index/mod.rs |  1 +
 src/storage/lsm/mod.rs       | 94 +++++++++++++++++-------------------
 2 files changed, 45 insertions(+), 50 deletions(-)

diff --git a/src/storage/lsm/index/mod.rs b/src/storage/lsm/index/mod.rs
index a8128fe..0dcefb9 100644
--- a/src/storage/lsm/index/mod.rs
+++ b/src/storage/lsm/index/mod.rs
@@ -23,6 +23,7 @@ pub fn write_index(path: &path::Path, index: &HashMap<String, String>) {
     let mut file = fs::File::create(_path).unwrap();
     file.write_all(&doc).unwrap();
 
+    file.sync_data().unwrap();
     file.flush().unwrap();
 }
 
diff --git a/src/storage/lsm/mod.rs b/src/storage/lsm/mod.rs
index d58696a..b59dfd7 100644
--- a/src/storage/lsm/mod.rs
+++ b/src/storage/lsm/mod.rs
@@ -116,18 +116,14 @@ impl Lsm {
             });
         }
 
-        let mut memtable = self.memtable.write().unwrap();
-
-        if memtable.contains_key(&key.to_string()) {
-            memtable.remove(&key.to_string());
+        let value = self.memtable.write().unwrap().remove(&key.to_string());
+        self.dense_index.write().unwrap().remove(&key.to_string());
+        self.bloom_filter.write().unwrap().delete(key);
 
-            drop(memtable);
-        } else {
-            self.dense_index.write().unwrap().remove(&key.to_string());
+        if let Some(value) = value {
+            self.memtable_size += mem::size_of_val(&value);
         }
 
-        self.bloom_filter.write().unwrap().delete(key);
-
         Ok(())
     }
 
@@ -160,38 +156,37 @@ impl Lsm {
 
     pub fn flush(&mut self) -> Result<()> {
         let memtable = self.get_memtable();
-
-        if memtable.is_empty() {
-            return Ok(());
-        }
-
         let mut dense_index = self.dense_index.write().unwrap();
 
-        let segments = sstable::Segment::from_tree(&memtable, &self.lsm_config.sstable_path);
-
-        for token in segments.1 {
-            dense_index.insert(token.0, token.1);
-        }
+        if memtable.is_empty() {
+            index::write_index(&self.lsm_config.sstable_path, dense_index.deref());
+            filter::write_filter(
+                &self.lsm_config.sstable_path,
+                self.bloom_filter.read().unwrap().deref(),
+            );
 
-        index::write_index(&self.lsm_config.sstable_path, dense_index.deref());
+            Ok(())
+        } else {
+            let segments = sstable::Segment::from_tree(&memtable, &self.lsm_config.sstable_path);
 
-        let mut keys = Vec::new();
+            for token in segments.1 {
+                dense_index.insert(token.0, token.1);
+            }
 
-        for segment in dense_index.deref() {
-            keys.push(segment.0.clone());
-        }
+            index::write_index(&self.lsm_config.sstable_path, dense_index.deref());
 
-        drop(dense_index);
+            drop(dense_index);
 
-        filter::write_filter(
-            &self.lsm_config.sstable_path,
-            self.bloom_filter.read().unwrap().deref(),
-        );
+            filter::write_filter(
+                &self.lsm_config.sstable_path,
+                self.bloom_filter.read().unwrap().deref(),
+            );
 
-        self.memtable.write().unwrap().clear();
-        self.memtable_size = 0;
+            self.memtable.write().unwrap().clear();
+            self.memtable_size = 0;
 
-        Ok(())
+            Ok(())
+        }
     }
 
     pub fn get_memtable(&self) -> BTreeMap<String, bson::Bson> {
@@ -243,27 +238,26 @@ impl Drop for Lsm {
         let memtable = self.memtable.read().unwrap();
         let mut dense_index = self.dense_index.write().unwrap();
 
-        if memtable.is_empty() {
-            return;
-        }
-
-        let segments = sstable::Segment::from_tree(memtable.deref(), &self.lsm_config.sstable_path);
+        if memtable.len() == 0 {
+            index::write_index(&self.lsm_config.sstable_path, dense_index.deref());
+            filter::write_filter(
+                &self.lsm_config.sstable_path,
+                self.bloom_filter.read().unwrap().deref(),
+            );
+        } else {
+            let segments =
+                sstable::Segment::from_tree(memtable.deref(), &self.lsm_config.sstable_path);
 
-        for token in segments.1 {
-            dense_index.insert(token.0, token.1);
-        }
+            for token in segments.1 {
+                dense_index.insert(token.0, token.1);
+            }
 
-        let mut keys = Vec::new();
+            index::write_index(&self.lsm_config.sstable_path, dense_index.deref());
 
-        for segment in dense_index.deref() {
-            keys.push(segment.0.clone());
+            filter::write_filter(
+                &self.lsm_config.sstable_path,
+                self.bloom_filter.read().unwrap().deref(),
+            );
         }
-
-        index::write_index(&self.lsm_config.sstable_path, dense_index.deref());
-
-        filter::write_filter(
-            &self.lsm_config.sstable_path,
-            self.bloom_filter.read().unwrap().deref(),
-        );
     }
 }