From e35a5aa550c8520e7ba4b66e5f8ae979502cf87a Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Mon, 16 Aug 2021 13:38:22 +0300
Subject: [PATCH] WIP: track mem usage

---
 .../src/layered_repository/inmemory_layer.rs  | 14 ++++++++++++
 .../src/layered_repository/storage_layer.rs   | 22 +++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/pageserver/src/layered_repository/inmemory_layer.rs b/pageserver/src/layered_repository/inmemory_layer.rs
index 70a7b7216e..cc356efaed 100644
--- a/pageserver/src/layered_repository/inmemory_layer.rs
+++ b/pageserver/src/layered_repository/inmemory_layer.rs
@@ -51,6 +51,11 @@ pub struct InMemoryLayerInner {
     /// `segsizes` tracks the size of the segment at different points in time.
     ///
     segsizes: BTreeMap<Lsn, u32>,
+
+    ///
+    /// Memory usage
+    ///
+    mem_used: usize,
 }
 
 impl Layer for InMemoryLayer {
@@ -186,6 +191,7 @@ impl InMemoryLayer {
                 drop_lsn: None,
                 page_versions: BTreeMap::new(),
                 segsizes: BTreeMap::new(),
+                mem_used: 0,
             }),
         })
     }
@@ -228,6 +234,9 @@ impl InMemoryLayer {
             self.timelineid,
             lsn
         );
+
+        let mem_size = pv.get_mem_size();
+
         let mut inner = self.inner.lock().unwrap();
 
         let old = inner.page_versions.insert((blknum, lsn), pv);
@@ -238,6 +247,8 @@ impl InMemoryLayer {
                 "Page version of rel {} blk {} at {} already exists",
                 self.seg.rel, blknum, lsn
             );
+        } else {
+            inner.mem_used += mem_size;
         }
 
         // Also update the relation size, if this extended the relation.
@@ -313,6 +324,7 @@ impl InMemoryLayer {
         );
         let mut page_versions = BTreeMap::new();
         let mut segsizes = BTreeMap::new();
+        let mut mem_used = 0;
 
         let seg = src.get_seg_tag();
 
@@ -333,6 +345,7 @@ impl InMemoryLayer {
                 page_image: Some(img),
                 record: None,
             };
+            mem_used += pv.get_mem_size();
             page_versions.insert((blknum, lsn), pv);
         }
 
@@ -346,6 +359,7 @@ impl InMemoryLayer {
                 drop_lsn: None,
                 page_versions: page_versions,
                 segsizes: segsizes,
+                mem_used: mem_used,
             }),
         })
     }
diff --git a/pageserver/src/layered_repository/storage_layer.rs b/pageserver/src/layered_repository/storage_layer.rs
index 0a181a1aac..142bf76b96 100644
--- a/pageserver/src/layered_repository/storage_layer.rs
+++ b/pageserver/src/layered_repository/storage_layer.rs
@@ -69,6 +69,28 @@ pub struct PageVersion {
     pub record: Option<WALRecord>,
 }
 
+impl PageVersion {
+    pub fn get_mem_size(&self) -> usize {
+        let mut sz = 0;
+
+        // every page version has some fixed overhead.
+        sz += 16;
+
+        if let Some(img) = &self.page_image {
+            sz += img.len();
+        }
+
+        if let Some(rec) = &self.record {
+            sz += rec.rec.len();
+
+            // Some per-record overhead. Not very accurate, but close enough
+            sz += 32;
+        }
+
+        sz
+    }
+}
+
 ///
 /// Data needed to reconstruct a page version
 ///