From bac6c720f815038b82e1d3d92cb96aeca3668bc7 Mon Sep 17 00:00:00 2001
From: dennis zhuang <killme2008@gmail.com>
Date: Tue, 28 Jun 2022 15:11:04 +0800
Subject: [PATCH] feat: impl bytes_allocated for memtable (#55)

---
 src/storage/src/memtable.rs       |  5 +++++
 src/storage/src/memtable/btree.rs | 13 ++++++++++---
 src/storage/src/memtable/tests.rs |  3 +++
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/storage/src/memtable.rs b/src/storage/src/memtable.rs
index 6e77c221e2..99de24d363 100644
--- a/src/storage/src/memtable.rs
+++ b/src/storage/src/memtable.rs
@@ -122,6 +122,11 @@ impl KeyValues {
     pub fn is_empty(&self) -> bool {
         self.len() == 0
     }
+
+    pub fn estimated_memory_size(&self) -> usize {
+        self.keys.iter().fold(0, |acc, v| acc + v.memory_size())
+            + self.values.iter().fold(0, |acc, v| acc + v.memory_size())
+    }
 }
 
 pub struct DefaultMemtableBuilder {}
diff --git a/src/storage/src/memtable/btree.rs b/src/storage/src/memtable/btree.rs
index 7a2bcb76bd..e81ba48563 100644
--- a/src/storage/src/memtable/btree.rs
+++ b/src/storage/src/memtable/btree.rs
@@ -1,7 +1,10 @@
 use std::cmp::Ordering;
 use std::collections::{btree_map, BTreeMap};
 use std::ops::Bound;
-use std::sync::{Arc, RwLock};
+use std::sync::{
+    atomic::{AtomicUsize, Ordering as AtomicOrdering},
+    Arc, RwLock,
+};
 
 use datatypes::prelude::*;
 use datatypes::value::Value;
@@ -22,6 +25,7 @@ type RwLockMap = RwLock<BTreeMap<InnerKey, RowValue>>;
 pub struct BTreeMemtable {
     schema: MemtableSchema,
     map: Arc<RwLockMap>,
+    estimated_bytes: AtomicUsize,
 }
 
 impl BTreeMemtable {
@@ -29,6 +33,7 @@ impl BTreeMemtable {
         BTreeMemtable {
             schema,
             map: Arc::new(RwLock::new(BTreeMap::new())),
+            estimated_bytes: AtomicUsize::new(0),
         }
     }
 }
@@ -39,8 +44,10 @@ impl Memtable for BTreeMemtable {
     }
 
     fn write(&self, kvs: &KeyValues) -> Result<()> {
-        let mut map = self.map.write().unwrap();
+        self.estimated_bytes
+            .fetch_add(kvs.estimated_memory_size(), AtomicOrdering::Relaxed);
 
+        let mut map = self.map.write().unwrap();
         let iter_row = IterRow::new(kvs);
         for (inner_key, row_value) in iter_row {
             map.insert(inner_key, row_value);
@@ -58,7 +65,7 @@ impl Memtable for BTreeMemtable {
     }
 
     fn bytes_allocated(&self) -> usize {
-        unimplemented!()
+        self.estimated_bytes.load(AtomicOrdering::Relaxed)
     }
 }
 
diff --git a/src/storage/src/memtable/tests.rs b/src/storage/src/memtable/tests.rs
index 81a38947dc..564c68732a 100644
--- a/src/storage/src/memtable/tests.rs
+++ b/src/storage/src/memtable/tests.rs
@@ -175,6 +175,7 @@ fn write_iter_memtable_case(ctx: &TestContext) {
     // Test iterating an empty memtable.
     let mut iter = ctx.memtable.iter(IterContext::default()).unwrap();
     assert!(iter.next().unwrap().is_none());
+    assert_eq!(0, ctx.memtable.bytes_allocated());
 
     // Init test data.
     write_kvs(
@@ -199,6 +200,8 @@ fn write_iter_memtable_case(ctx: &TestContext) {
         &[None, Some(5), None],             // values
     );
 
+    assert_eq!(216, ctx.memtable.bytes_allocated());
+
     let batch_sizes = [1, 4, 8, consts::READ_BATCH_SIZE];
     for batch_size in batch_sizes {
         let iter_ctx = IterContext {