From 037a057027565695d99e9ae9d87ba6a8f579b274 Mon Sep 17 00:00:00 2001
From: marvin-j97 <marvin.janke.97@gmail.com>
Date: Sat, 11 May 2024 16:09:04 +0200
Subject: [PATCH] refactor: sealed memtables

---
 src/prefix.rs     |  2 +-
 src/range.rs      |  2 +-
 src/tree.rs       |  8 ++++----
 src/tree_inner.rs | 19 ++++++++++++++++---
 4 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/src/prefix.rs b/src/prefix.rs
index aea2e59..40ee7f9 100644
--- a/src/prefix.rs
+++ b/src/prefix.rs
@@ -76,7 +76,7 @@ impl<'a> PrefixIterator<'a> {
 
         let mut iters: Vec<BoxedIterator<'a>> = segment_iters;
 
-        for memtable in lock.guard.sealed.values() {
+        for (_, memtable) in lock.guard.sealed.iter() {
             iters.push(Box::new(
                 memtable
                     .items
diff --git a/src/range.rs b/src/range.rs
index 0749ce1..56556f9 100644
--- a/src/range.rs
+++ b/src/range.rs
@@ -130,7 +130,7 @@ impl<'a> RangeIterator<'a> {
 
         let mut iters: Vec<BoxedIterator<'a>> = segment_iters;
 
-        for memtable in lock.guard.sealed.values() {
+        for (_, memtable) in lock.guard.sealed.iter() {
             iters.push(Box::new(memtable.items.range(range.clone()).map(|entry| {
                 Ok(Value::from((entry.key().clone(), entry.value().clone())))
             })));
diff --git a/src/tree.rs b/src/tree.rs
index 7515394..ff9452b 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -145,7 +145,7 @@ impl Tree {
         let mut memtable_lock = self.sealed_memtables.write().expect("lock is poisoned");
 
         for segment in segments {
-            memtable_lock.remove(&segment.metadata.id);
+            memtable_lock.remove(segment.metadata.id);
         }
 
         // NOTE: Segments are registered, we can unlock the memtable(s) safely
@@ -276,7 +276,7 @@ impl Tree {
         let yanked_memtable = Arc::new(yanked_memtable);
 
         let tmp_memtable_id = self.get_next_segment_id();
-        sealed_memtables.insert(tmp_memtable_id, yanked_memtable.clone());
+        sealed_memtables.add(tmp_memtable_id, yanked_memtable.clone());
 
         Some((tmp_memtable_id, yanked_memtable))
     }
@@ -295,7 +295,7 @@ impl Tree {
     /// May be used to restore the LSM-tree's in-memory state from some journals.
     pub fn add_sealed_memtable(&self, id: MemtableId, memtable: Arc<MemTable>) {
         let mut memtable_lock = self.sealed_memtables.write().expect("lock is poisoned");
-        memtable_lock.insert(id, memtable);
+        memtable_lock.add(id, memtable);
     }
 
     /// Scans the entire tree, returning the amount of items.
@@ -385,7 +385,7 @@ impl Tree {
 
         // Now look in sealed memtables
         let memtable_lock = self.sealed_memtables.read().expect("lock is poisoned");
-        for memtable in memtable_lock.values().rev() {
+        for (_, memtable) in memtable_lock.iter().rev() {
             if let Some(item) = memtable.get(&key, seqno) {
                 if evict_tombstone {
                     return Ok(ignore_tombstone_value(item));
diff --git a/src/tree_inner.rs b/src/tree_inner.rs
index 1ee1c65..6be552a 100644
--- a/src/tree_inner.rs
+++ b/src/tree_inner.rs
@@ -10,7 +10,6 @@ use crate::{
     BlockCache,
 };
 use std::{
-    collections::BTreeMap,
     path::PathBuf,
     sync::{atomic::AtomicU64, Arc, RwLock},
 };
@@ -20,8 +19,22 @@ pub type TreeId = u64;
 
 pub type MemtableId = u64;
 
-// TODO: Vec may be enough
-pub type SealedMemtables = BTreeMap<MemtableId, Arc<MemTable>>;
+#[derive(Default)]
+pub struct SealedMemtables(Vec<(MemtableId, Arc<MemTable>)>);
+
+impl SealedMemtables {
+    pub fn add(&mut self, id: MemtableId, memtable: Arc<MemTable>) {
+        self.0.push((id, memtable));
+    }
+
+    pub fn remove(&mut self, id_to_remove: MemtableId) {
+        self.0.retain(|(id, _)| *id != id_to_remove);
+    }
+
+    pub fn iter(&self) -> impl DoubleEndedIterator<Item = &(MemtableId, Arc<MemTable>)> {
+        self.0.iter()
+    }
+}
 
 pub fn get_next_tree_id() -> TreeId {
     static TREE_ID_COUNTER: AtomicU64 = AtomicU64::new(0);