From 1a739b5d6ea4257e8e2b607e43152047016346e4 Mon Sep 17 00:00:00 2001
From: Hendiadyoin1 <leon2002.la@gmail.com>
Date: Mon, 14 Feb 2022 13:40:14 +0100
Subject: [PATCH] Kernel: Try to reuse empty slabheaps before expanding the
 kmalloc-heap

---
 Kernel/Heap/kmalloc.cpp | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/Kernel/Heap/kmalloc.cpp b/Kernel/Heap/kmalloc.cpp
index eeeb4c4bc0..d7b62e40f2 100644
--- a/Kernel/Heap/kmalloc.cpp
+++ b/Kernel/Heap/kmalloc.cpp
@@ -168,6 +168,31 @@ public:
         return total;
     }
 
+    bool try_purge()
+    {
+        bool did_purge = false;
+
+        // Note: We cannot remove children from the list when using a structured loop,
+        //       Because we need to advance the iterator before we delete the underlying
+        //       value, so we have to iterate manually
+
+        auto block = m_usable_blocks.begin();
+        while (block != m_usable_blocks.end()) {
+            if (block->allocated_bytes() != 0) {
+                ++block;
+                continue;
+            }
+            auto& block_to_remove = *block;
+            ++block;
+            block_to_remove.list_node.remove();
+            block_to_remove.~KmallocSlabBlock();
+            kfree_aligned(&block_to_remove);
+
+            did_purge = true;
+        }
+        return did_purge;
+    }
+
 private:
     size_t m_slab_size { 0 };
 
@@ -205,6 +230,22 @@ struct KmallocGlobalData {
                 return ptr;
         }
 
+        // NOTE: This size calculation is a mirror of kmalloc_aligned(KmallocSlabBlock)
+        if (size <= KmallocSlabBlock::block_size * 2 + sizeof(ptrdiff_t) + sizeof(size_t)) {
+            // FIXME: We should propagate a freed pointer, to find the specific subheap it belonged to
+            //        This would save us iterating over them in the next step and remove a recursion
+            bool did_purge = false;
+            for (auto& slabheap : slabheaps) {
+                if (slabheap.try_purge()) {
+                    dbgln_if(KMALLOC_DEBUG, "Kmalloc purged block(s) from slabheap of size {} to avoid expansion", slabheap.slab_size());
+                    did_purge = true;
+                    break;
+                }
+            }
+            if (did_purge)
+                return allocate(size);
+        }
+
         if (!try_expand(size)) {
             PANIC("OOM when trying to expand kmalloc heap.");
         }