From 678bbd29cac72b819ed2d8304418f72146d7fb52 Mon Sep 17 00:00:00 2001
From: Tom <tomut@yahoo.com>
Date: Fri, 4 Sep 2020 15:31:56 -0600
Subject: [PATCH] Kernel: Fix heap expansion loop

By being a bit too greedy and only allocating how much we need for
the failing allocation, we can end up in an infinite loop trying
to expand the heap further. That's because there are other allocations
(e.g. logging, vmobjects, regions, ...) that happen before we finally
retry the failed allocation request.

Also fix allocating in page size increments, which lead to an assertion
when the heap had to grow more than the 1 MiB backup.
---
 Kernel/Heap/Heap.h      |  5 +++++
 Kernel/Heap/kmalloc.cpp | 14 +++++++++++++-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/Kernel/Heap/Heap.h b/Kernel/Heap/Heap.h
index 6c59082c2e..7e71c5fb56 100644
--- a/Kernel/Heap/Heap.h
+++ b/Kernel/Heap/Heap.h
@@ -257,6 +257,7 @@ public:
 
     void* allocate(size_t size)
     {
+        int attempt = 0;
         do {
             for (auto* subheap = &m_heaps; subheap; subheap = subheap->next) {
                 if (void* ptr = subheap->heap.allocate(size))
@@ -269,6 +270,10 @@ public:
             // This is especially true for the kmalloc heap, where adding memory
             // requires several other objects to be allocated just to be able to
             // expand the heap.
+
+            // To avoid an infinite expansion loop, limit to two attempts
+            if (attempt++ >= 2)
+                break;
         } while (expand_memory(size));
         return nullptr;
     }
diff --git a/Kernel/Heap/kmalloc.cpp b/Kernel/Heap/kmalloc.cpp
index 2e83579429..247bc55b88 100644
--- a/Kernel/Heap/kmalloc.cpp
+++ b/Kernel/Heap/kmalloc.cpp
@@ -59,22 +59,30 @@ struct KmallocGlobalHeap {
         {
         }
 
+        bool m_adding { false };
         bool add_memory(size_t allocation_request)
         {
             if (!MemoryManager::is_initialized()) {
                 klog() << "kmalloc(): Cannot expand heap before MM is initialized!";
                 return false;
             }
+            ASSERT(!m_adding);
+            TemporaryChange change(m_adding, true);
             // At this point we have very little memory left. Any attempt to
             // kmalloc() could fail, so use our backup memory first, so we
             // can't really reliably allocate even a new region of memory.
             // This is why we keep a backup region, which we can
             auto region = move(m_global_heap.m_backup_memory);
             if (!region) {
+                // Be careful to not log too much here. We don't want to trigger
+                // any further calls to kmalloc(). We're already out of memory
+                // and don't have any backup memory, either!
                 klog() << "kmalloc(): Cannot expand heap: no backup memory";
                 return false;
             }
 
+            // At this point we should have at least enough memory from the
+            // backup region to be able to log properly
             klog() << "kmalloc(): Adding memory to heap at " << region->vaddr() << ", bytes: " << region->size();
 
             auto& subheap = m_global_heap.m_heap.add_subheap(region->vaddr().as_ptr(), region->size());
@@ -91,7 +99,11 @@ struct KmallocGlobalHeap {
             // was big enough to likely satisfy the request
             if (subheap.free_bytes() < allocation_request) {
                 // Looks like we probably need more
-                size_t memory_size = max(decltype(m_global_heap.m_heap)::calculate_memory_for_bytes(allocation_request), (size_t)(1 * MiB));
+                size_t memory_size = PAGE_ROUND_UP(decltype(m_global_heap.m_heap)::calculate_memory_for_bytes(allocation_request));
+                // Add some more to the new heap. We're already using it for other 
+                // allocations not including the original allocation_request
+                // that triggered heap expansion. If we don't allocate
+                memory_size += 1 * MiB;
                 region = MM.allocate_kernel_region(memory_size, "kmalloc subheap", Region::Access::Read | Region::Access::Write);
                 if (region) {
                     klog() << "kmalloc(): Adding even more memory to heap at " << region->vaddr() << ", bytes: " << region->size();