temporarily add memtable_size_sum

Author: marvin-j97

src/abstract.rs

@@ -25,6 +25,12 @@ pub type RangeItem = crate::Result<KvPair>;
 /// Generic Tree API
 #[enum_dispatch]
 pub trait AbstractTree {
+    // TODO: remove
+    #[doc(hidden)]
+    fn version_memtable_size_sum(&self) -> u64 {
+        self.get_version_history_lock().memtable_size_sum()
+    }
+
     #[doc(hidden)]
     fn next_table_id(&self) -> TableId;
 
@@ -79,16 +85,16 @@ pub trait AbstractTree {
         let sealed_ids = latest
             .sealed_memtables
             .iter()
-            .map(|mt| mt.0)
+            .map(|mt| mt.id)
             .collect::<Vec<_>>();
 
-        let flushed_size = latest.sealed_memtables.iter().map(|(_, x)| x.size()).sum();
+        let flushed_size = latest.sealed_memtables.iter().map(|mt| mt.size()).sum();
 
         let merger = Merger::new(
             latest
                 .sealed_memtables
                 .iter()
-                .map(|(_, mt)| mt.iter().map(Ok))
+                .map(|mt| mt.iter().map(Ok))
                 .collect::<Vec<_>>(),
         );
         let stream = CompactionStream::new(merger, seqno_threshold);
@@ -258,7 +264,7 @@ pub trait AbstractTree {
     /// Adds a sealed memtables.
     ///
     /// May be used to restore the LSM-tree's in-memory state from some journals.
-    fn add_sealed_memtable(&self, id: MemtableId, memtable: Arc<Memtable>);
+    fn add_sealed_memtable(&self, memtable: Arc<Memtable>);
 
     /// Performs compaction on the tree's levels, blocking the caller until it's done.
     ///

src/blob_tree/mod.rs

@@ -610,8 +610,8 @@ impl AbstractTree for BlobTree {
         self.index.set_active_memtable(memtable);
     }
 
-    fn add_sealed_memtable(&self, id: MemtableId, memtable: Arc<Memtable>) {
-        self.index.add_sealed_memtable(id, memtable);
+    fn add_sealed_memtable(&self, memtable: Arc<Memtable>) {
+        self.index.add_sealed_memtable(memtable);
     }
 
     fn compact(

src/memtable/mod.rs

@@ -3,6 +3,7 @@
 // (found in the LICENSE-* files in the repository)
 
 use crate::key::InternalKey;
+use crate::tree::inner::MemtableId;
 use crate::{
     value::{InternalValue, SeqNo, UserValue},
     ValueType,
@@ -14,8 +15,10 @@ use std::sync::atomic::AtomicU64;
 /// The memtable serves as an intermediary, ephemeral, sorted storage for new items
 ///
 /// When the Memtable exceeds some size, it should be flushed to a table.
-#[derive(Default)]
 pub struct Memtable {
+    #[doc(hidden)]
+    pub id: MemtableId,
+
     /// The actual content, stored in a lock-free skiplist.
     #[doc(hidden)]
     pub items: SkipMap<InternalKey, UserValue>,
@@ -32,6 +35,17 @@ pub struct Memtable {
 }
 
 impl Memtable {
+    #[doc(hidden)]
+    #[must_use]
+    pub fn new(id: MemtableId) -> Self {
+        Self {
+            id,
+            items: SkipMap::default(),
+            approximate_size: AtomicU64::default(),
+            highest_seqno: AtomicU64::default(),
+        }
+    }
+
     /// Clears the memtable.
     pub fn clear(&mut self) {
         self.items.clear();
@@ -161,7 +175,7 @@ mod tests {
     #[test]
     #[expect(clippy::unwrap_used)]
     fn memtable_mvcc_point_read() {
-        let memtable = Memtable::default();
+        let memtable = Memtable::new(0);
 
         memtable.insert(InternalValue::from_components(
             *b"hello-key-999991",
@@ -204,7 +218,7 @@ mod tests {
 
     #[test]
     fn memtable_get() {
-        let memtable = Memtable::default();
+        let memtable = Memtable::new(0);
 
         let value =
             InternalValue::from_components(b"abc".to_vec(), b"abc".to_vec(), 0, ValueType::Value);
@@ -216,7 +230,7 @@ mod tests {
 
     #[test]
     fn memtable_get_highest_seqno() {
-        let memtable = Memtable::default();
+        let memtable = Memtable::new(0);
 
         memtable.insert(InternalValue::from_components(
             b"abc".to_vec(),
@@ -262,7 +276,7 @@ mod tests {
 
     #[test]
     fn memtable_get_prefix() {
-        let memtable = Memtable::default();
+        let memtable = Memtable::new(0);
 
         memtable.insert(InternalValue::from_components(
             b"abc0".to_vec(),
@@ -300,7 +314,7 @@ mod tests {
 
     #[test]
     fn memtable_get_old_version() {
-        let memtable = Memtable::default();
+        let memtable = Memtable::new(0);
 
         memtable.insert(InternalValue::from_components(
             b"abc".to_vec(),

src/range.rs

@@ -192,7 +192,7 @@ impl TreeIter {
             }
 
             // Sealed memtables
-            for (_, memtable) in lock.version.sealed_memtables.iter() {
+            for memtable in lock.version.sealed_memtables.iter() {
                 let iter = memtable.range(range.clone());
 
                 iters.push(Box::new(

src/tree/mod.rs

@@ -455,7 +455,7 @@ impl AbstractTree for Tree {
             .write()
             .expect("lock is poisoned")
             .latest_version()
-            .active_memtable = Arc::new(Memtable::default());
+            .active_memtable = Arc::new(Memtable::new(self.memtable_id_counter.next()));
     }
 
     fn set_active_memtable(&self, memtable: Memtable) {
@@ -466,9 +466,9 @@ impl AbstractTree for Tree {
             .active_memtable = Arc::new(memtable);
     }
 
-    fn add_sealed_memtable(&self, id: MemtableId, memtable: Arc<Memtable>) {
+    fn add_sealed_memtable(&self, memtable: Arc<Memtable>) {
         let mut version_lock = self.version_history.write().expect("lock is poisoned");
-        version_lock.append_sealed_memtable(id, memtable);
+        version_lock.append_sealed_memtable(memtable);
     }
 
     fn compact(
@@ -522,20 +522,19 @@ impl AbstractTree for Tree {
         }
 
         let yanked_memtable = super_version.active_memtable;
-        let memtable_id = self.0.memtable_id_counter.next();
 
         let mut copy = version_history_lock.latest_version();
         copy.seqno = self.config.seqno.next();
-        copy.active_memtable = Arc::new(Memtable::default());
-        copy.sealed_memtables = Arc::new(
-            super_version
-                .sealed_memtables
-                .add(memtable_id, yanked_memtable.clone()),
-        );
+        copy.active_memtable = Arc::new(Memtable::new(self.memtable_id_counter.next()));
+        copy.sealed_memtables =
+            Arc::new(super_version.sealed_memtables.add(yanked_memtable.clone()));
 
         version_history_lock.append_version(copy);
 
-        log::trace!("rotate: added memtable id={memtable_id} to sealed memtables");
+        log::trace!(
+            "rotate: added memtable id={} to sealed memtables",
+            yanked_memtable.id,
+        );
 
         Some(yanked_memtable)
     }
@@ -566,7 +565,7 @@ impl AbstractTree for Tree {
         let sealed_count = super_version
             .sealed_memtables
             .iter()
-            .map(|(_, mt)| mt.len())
+            .map(|mt| mt.len())
             .sum::<usize>() as u64;
 
         (memtable_count + sealed_count + tables_item_count)
@@ -594,7 +593,7 @@ impl AbstractTree for Tree {
         let sealed = version
             .sealed_memtables
             .iter()
-            .map(|(_, table)| table.get_highest_seqno())
+            .map(|mt| mt.get_highest_seqno())
             .max()
             .flatten();
 
@@ -729,8 +728,8 @@ impl Tree {
         key: &[u8],
         seqno: SeqNo,
     ) -> Option<InternalValue> {
-        for (_, memtable) in super_version.sealed_memtables.iter().rev() {
-            if let Some(entry) = memtable.get(key, seqno) {
+        for mt in super_version.sealed_memtables.iter().rev() {
+            if let Some(entry) = mt.get(key, seqno) {
                 return Some(entry);
             }
         }

src/tree/sealed.rs

@@ -10,24 +10,24 @@ use std::sync::Arc;
 /// Memtable IDs are monotonically increasing, so we don't really
 /// need a search tree; also there are only a handful of them at most.
 #[derive(Clone, Default)]
-pub struct SealedMemtables(Vec<(MemtableId, Arc<Memtable>)>);
+pub struct SealedMemtables(Vec<Arc<Memtable>>);
 
 impl SealedMemtables {
     /// Copy-and-writes a new list with additional Memtable.
-    pub fn add(&self, id: MemtableId, memtable: Arc<Memtable>) -> Self {
+    pub fn add(&self, memtable: Arc<Memtable>) -> Self {
         let mut copy = self.clone();
-        copy.0.push((id, memtable));
+        copy.0.push(memtable);
         copy
     }
 
     /// Copy-and-writes a new list with the specified Memtable removed.
     pub fn remove(&self, id_to_remove: MemtableId) -> Self {
         let mut copy = self.clone();
-        copy.0.retain(|(id, _)| *id != id_to_remove);
+        copy.0.retain(|mt| mt.id != id_to_remove);
         copy
     }
 
-    pub fn iter(&self) -> impl DoubleEndedIterator<Item = &(MemtableId, Arc<Memtable>)> {
+    pub fn iter(&self) -> impl DoubleEndedIterator<Item = &Arc<Memtable>> {
         self.0.iter()
     }
 

src/version/super_version.rs

@@ -4,7 +4,7 @@
 
 use crate::{
     memtable::Memtable,
-    tree::{inner::MemtableId, sealed::SealedMemtables},
+    tree::sealed::SealedMemtables,
     version::{persist_version, Version},
     SeqNo, SequenceNumberCounter,
 };
@@ -31,7 +31,7 @@ impl SuperVersions {
     pub fn new(version: Version) -> Self {
         Self(
             vec![SuperVersion {
-                active_memtable: Arc::default(),
+                active_memtable: Arc::new(Memtable::new(0)),
                 sealed_memtables: Arc::default(),
                 version,
                 seqno: 0,
@@ -40,6 +40,24 @@ impl SuperVersions {
         )
     }
 
+    pub fn memtable_size_sum(&self) -> u64 {
+        let mut set = crate::HashMap::default();
+
+        for super_version in &self.0 {
+            set.entry(super_version.active_memtable.id)
+                .and_modify(|bytes| *bytes += super_version.active_memtable.size())
+                .or_insert_with(|| super_version.active_memtable.size());
+
+            for sealed in super_version.sealed_memtables.iter() {
+                set.entry(sealed.id)
+                    .and_modify(|bytes| *bytes += sealed.size())
+                    .or_insert_with(|| sealed.size());
+            }
+        }
+
+        set.into_values().sum()
+    }
+
     pub fn free_list_len(&self) -> usize {
         self.0.len().saturating_sub(1)
     }
@@ -140,9 +158,9 @@ impl SuperVersions {
         version.expect("should always find a SuperVersion")
     }
 
-    pub fn append_sealed_memtable(&mut self, id: MemtableId, memtable: Arc<Memtable>) {
+    pub fn append_sealed_memtable(&mut self, memtable: Arc<Memtable>) {
         let mut copy = self.latest_version();
-        copy.sealed_memtables = Arc::new(copy.sealed_memtables.add(id, memtable));
+        copy.sealed_memtables = Arc::new(copy.sealed_memtables.add(memtable));
         self.0.push_back(copy);
     }
 }