From 235ca01fddd9696341339da7e1c16d05bdda3f9f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Coss=C3=ADo?= <luis.cossio@outlook.com>
Date: Fri, 4 Oct 2024 05:03:40 -0300
Subject: [PATCH] expose storage-wide stats (#12)

---
 src/lib.rs             |  1 +
 src/payload_storage.rs |  8 ++++----
 src/slotted_page.rs    |  4 ++--
 src/stats.rs           | 45 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 52 insertions(+), 6 deletions(-)
 create mode 100644 src/stats.rs

diff --git a/src/lib.rs b/src/lib.rs
index 700a2f3..b9e505b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,6 +3,7 @@ mod page_tracker;
 pub mod payload;
 mod payload_storage;
 mod slotted_page;
+mod stats;
 mod utils_copied;
 
 pub use payload_storage::PayloadStorage;
diff --git a/src/payload_storage.rs b/src/payload_storage.rs
index 14e298c..bd40cb8 100644
--- a/src/payload_storage.rs
+++ b/src/payload_storage.rs
@@ -10,10 +10,10 @@ use std::path::PathBuf;
 
 pub struct PayloadStorage {
     page_tracker: PageTracker,
-    new_page_size: usize, // page size in bytes when creating new pages
-    pages: HashMap<u32, SlottedPageMmap>, // page_id -> mmap page
+    pub(super) new_page_size: usize, // page size in bytes when creating new pages
+    pub(super) pages: HashMap<u32, SlottedPageMmap>, // page_id -> mmap page
     max_page_id: u32,
-    page_emptiness: PriorityQueue<PageId, usize>,
+    pub(super) page_emptiness: PriorityQueue<PageId, usize>,
     base_path: PathBuf,
 }
 
@@ -261,7 +261,7 @@ impl PayloadStorage {
 
                 // check if we should defrag this page
                 let frag_threshold =
-                    SlottedPageMmap::FRAGMENTATION_THRESHOLD_RATIO * page.page_size() as f32;
+                    SlottedPageMmap::FRAGMENTATION_THRESHOLD_RATIO * page.size() as f32;
                 if frag_space < frag_threshold.ceil() as usize {
                     // page is not fragmented enough, skip
                     return None;
diff --git a/src/slotted_page.rs b/src/slotted_page.rs
index 1e9bd29..aec5b16 100644
--- a/src/slotted_page.rs
+++ b/src/slotted_page.rs
@@ -274,7 +274,7 @@ impl SlottedPageMmap {
             .expect("this should never overflow, otherwise the page is corrupted")
     }
 
-    pub fn page_size(&self) -> usize {
+    pub fn size(&self) -> usize {
         self.header.page_size()
     }
 
@@ -289,7 +289,7 @@ impl SlottedPageMmap {
         let mut fragmented_space = 0;
 
         let mut slot_id = 0;
-        let mut last_offset = self.page_size() as u64;
+        let mut last_offset = self.size() as u64;
         while let Some(slot) = self.get_slot_ref(&slot_id) {
             // if the slot is deleted, we can consider it empty space
             if slot.deleted {
diff --git a/src/stats.rs b/src/stats.rs
new file mode 100644
index 0000000..9ca2321
--- /dev/null
+++ b/src/stats.rs
@@ -0,0 +1,45 @@
+use crate::PayloadStorage;
+
+pub struct StorageStats {
+    /// The number of pages in the storage
+    pub pages_count: usize,
+
+    /// The default size of pages in bytes
+    pub default_page_bytes: usize,
+
+    /// The number of bytes unused in between values
+    pub fragmented_bytes: usize,
+
+    /// The number of bytes available in all pages.
+    pub available_bytes: usize,
+
+    /// The total size of the storage in bytes
+    pub total_size_bytes: usize,
+}
+
+impl PayloadStorage {
+    /// Storage-wide statistics
+    pub fn get_stats(&self) -> StorageStats {
+        let pages_count = self.pages.len();
+        let default_page_bytes = self.new_page_size;
+        let available_bytes = self
+            .page_emptiness
+            .iter()
+            .map(|(_, &free_space)| free_space)
+            .sum();
+
+        let mut fragmented_bytes = 0;
+        let mut total_size_bytes = 0;
+        for page in self.pages.values() {
+            fragmented_bytes += page.fragmented_space();
+            total_size_bytes += page.size();
+        }
+        StorageStats {
+            pages_count,
+            default_page_bytes,
+            fragmented_bytes,
+            total_size_bytes,
+            available_bytes,
+        }
+    }
+}