From 3b157a65733b30127fb3743132fe90a4bf79bf4b Mon Sep 17 00:00:00 2001
From: pjht <pterpstra@pterpstra.com>
Date: Sun, 7 Jul 2024 08:01:09 -0500
Subject: [PATCH] Count allocated and freed bytes in the kernel heap

---
 src/kernel_heap.rs | 69 ++++++++++++++++++++++++++++++++++++----------
 src/qemu_exit.rs   |  3 ++
 2 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/src/kernel_heap.rs b/src/kernel_heap.rs
index 3a8d56c..a409f37 100644
--- a/src/kernel_heap.rs
+++ b/src/kernel_heap.rs
@@ -1,50 +1,89 @@
-use crate::virtual_memory::KERNEL_SPACE;
+use crate::{println, virtual_memory::KERNEL_SPACE};
 
 use core::{
     alloc::{GlobalAlloc, Layout},
     ptr::NonNull,
+    sync::atomic::{AtomicUsize, Ordering},
 };
 use linked_list_allocator::hole::HoleList;
 use spin::Mutex;
 use x86_64::structures::paging::PageTableFlags;
 
-struct Heap(Mutex<HoleList>);
+pub struct Heap {
+    heap: Mutex<HoleList>,
+    bytes_alloced: AtomicUsize,
+    bytes_freed: AtomicUsize,
+    pmem_size: AtomicUsize,
+}
+
+fn format_byte_count(count: usize) -> (f64, &'static str) {
+    let mut count = count as f64;
+    let mut prefix = 0;
+    while count >= 1024.0 {
+        count /= 1024.0;
+        prefix += 1;
+    }
+    let prefix = ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei"][prefix];
+    (count, prefix)
+}
+
+impl Heap {
+    pub fn print_stats(&self) {
+        let (fmtd_alloced, alloced_pfx) =
+            format_byte_count(self.bytes_alloced.load(Ordering::Relaxed));
+        let (fmtd_freed, freed_pfx) = format_byte_count(self.bytes_freed.load(Ordering::Relaxed));
+        let (fmtd_total, total_pfx) = format_byte_count(
+            self.bytes_alloced.load(Ordering::Relaxed) - self.bytes_freed.load(Ordering::Relaxed),
+        );
+        let (fmtd_pmem, pmem_pfx) = format_byte_count(self.pmem_size.load(Ordering::Relaxed));
+        println!(
+            "[HEAP] {:2} {}B allocated, {:2} {}B freed ({:2} {}B total)",
+            fmtd_alloced, alloced_pfx, fmtd_freed, freed_pfx, fmtd_total, total_pfx
+        );
+        println!("[HEAP] {:2} {}B physical memory used for heap", fmtd_pmem, pmem_pfx);
+    }
+}
 
 unsafe impl Send for Heap {}
 unsafe impl Sync for Heap {}
 
 unsafe impl GlobalAlloc for Heap {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        let mut locked_self = self.0.lock();
-        let ptr = locked_self
+        let mut locked_self = self.heap.lock();
+        let (ptr, true_layout) = locked_self
             .allocate_first_fit(layout)
-            .map(|(allocation, _)| allocation.as_ptr())
+            .map(|(allocation, true_layout)| (allocation.as_ptr(), true_layout))
             .unwrap_or_else(|_| {
                 drop(locked_self);
                 let num_pages = layout.size().div_ceil(4096) * 2;
-                unsafe {
-                    self.dealloc(
-                        KERNEL_SPACE.lock().map_free(num_pages, PageTableFlags::empty()).unwrap(),
-                        Layout::from_size_align(num_pages * 4096, 4096).unwrap(),
-                    );
-                }
-                self.0
+                self.pmem_size.fetch_add(num_pages * 4096, Ordering::Relaxed);
+                let ptr = KERNEL_SPACE.lock().map_free(num_pages, PageTableFlags::empty()).unwrap();
+                let layout = Layout::from_size_align(num_pages * 4096, 4096).unwrap();
+                unsafe { self.heap.lock().deallocate(NonNull::new_unchecked(ptr), layout) };
+                self.heap
                     .lock()
                     .allocate_first_fit(layout)
-                    .map(|(allocation, _)| allocation.as_ptr())
+                    .map(|(allocation, true_layout)| (allocation.as_ptr(), true_layout))
                     .unwrap()
             });
         assert!((ptr as usize & (layout.align() - 1)) == 0);
+        self.bytes_alloced.fetch_add(true_layout.size(), Ordering::Relaxed);
         ptr
     }
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        unsafe { self.0.lock().deallocate(NonNull::new_unchecked(ptr), layout) };
+        self.bytes_freed.fetch_add(layout.size(), Ordering::Relaxed);
+        unsafe { self.heap.lock().deallocate(NonNull::new_unchecked(ptr), layout) };
     }
 }
 
 #[global_allocator]
-static HEAP: Heap = Heap(Mutex::new(HoleList::empty()));
+pub static HEAP: Heap = Heap {
+    heap: Mutex::new(HoleList::empty()),
+    bytes_alloced: AtomicUsize::new(0),
+    bytes_freed: AtomicUsize::new(0),
+    pmem_size: AtomicUsize::new(0),
+};
 
 #[alloc_error_handler]
 fn alloc_error_handler(layout: Layout) -> ! {
diff --git a/src/qemu_exit.rs b/src/qemu_exit.rs
index 4404ea3..e152882 100644
--- a/src/qemu_exit.rs
+++ b/src/qemu_exit.rs
@@ -1,6 +1,9 @@
 use x86_64::instructions::port::Port;
 
+use crate::kernel_heap::HEAP;
+
 pub fn exit_qemu() -> ! {
+    HEAP.print_stats();
     unsafe {
         Port::new(0xf4).write(0u32);
     }