More scope guards!!

kernel/src/memory/Heap.cpp

@@ -7,6 +7,7 @@
 #include <luna/Alignment.h>
 #include <luna/LinkedList.h>
 #include <luna/SafeArithmetic.h>
+#include <luna/ScopeGuard.h>
 #include <luna/String.h>
 #include <luna/SystemError.h>
 
@@ -36,7 +37,9 @@ static DoublyLinkedList<HeapBlock> heap;
 static Result<HeapBlock*> allocate_pages(usize count)
 {
     u64 virt = TRY(KernelVM::alloc_several_pages(count));
+    auto vm_guard = make_scope_guard([&] { KernelVM::free_several_pages(virt, count).value(); });
     void* const ptr = (void*)TRY(MemoryManager::alloc_at(virt, count, MMU::ReadWrite | MMU::NoExecute));
+    vm_guard.deactivate();
     return (HeapBlock*)ptr;
 }
 

kernel/src/thread/Scheduler.cpp

@@ -4,6 +4,7 @@
 #include "arch/MMU.h"
 #include "memory/KernelVM.h"
 #include "memory/MemoryManager.h"
+#include <luna/ScopeGuard.h>
 #include <luna/Stack.h>
 
 static Thread g_idle;
@@ -44,12 +45,19 @@ namespace Scheduler
         return &g_idle;
     }
 
-    Result<void> kernel_thread_alloc_stack_and_append_impl(Thread* thread)
+    Result<void> new_kernel_thread_impl(Thread* thread)
     {
-        // FIXME: We will leak the thread if VM allocation or alloc_at fail.
+        // If anything fails, make sure to clean up.
+        auto thread_guard = make_scope_guard([&] { delete thread; });
+
         u64 thread_stack_vm = TRY(KernelVM::alloc_several_pages(4));
+        auto vm_guard = make_scope_guard([&] { KernelVM::free_several_pages(thread_stack_vm, 4).value(); });
+
         TRY(MemoryManager::alloc_at(thread_stack_vm, 4, MMU::NoExecute | MMU::ReadWrite));
 
+        thread_guard.deactivate();
+        vm_guard.deactivate();
+
         Stack thread_stack{thread_stack_vm, ARCH_PAGE_SIZE * 4};
         thread->set_sp(thread_stack.top());
 
@@ -66,7 +74,7 @@ namespace Scheduler
         thread->init_regs_kernel();
         thread->set_ip(address);
 
-        return kernel_thread_alloc_stack_and_append_impl(thread);
+        return new_kernel_thread_impl(thread);
     }
 
     Result<void> new_kernel_thread(void (*func)(void))
@@ -75,7 +83,7 @@ namespace Scheduler
         thread->init_regs_kernel();
         thread->set_ip((u64)func);
 
-        return kernel_thread_alloc_stack_and_append_impl(thread);
+        return new_kernel_thread_impl(thread);
     }
 
     Result<void> new_kernel_thread(void (*func)(void*), void* arg)
@@ -85,7 +93,7 @@ namespace Scheduler
         thread->set_ip((u64)func);
         thread->set_arguments((u64)arg, 0, 0, 0);
 
-        return kernel_thread_alloc_stack_and_append_impl(thread);
+        return new_kernel_thread_impl(thread);
     }
 
     Thread* pick_task()

luna/include/luna/ScopeGuard.h

@@ -0,0 +1,28 @@
+#pragma once
+
+template <typename Callback> class ScopeGuard
+{
+  public:
+    ScopeGuard(const Callback& callback) : m_callback(callback)
+    {
+    }
+
+    void deactivate()
+    {
+        m_activated = false;
+    }
+
+    ~ScopeGuard()
+    {
+        if (m_activated) m_callback();
+    }
+
+  private:
+    bool m_activated{true};
+    Callback m_callback;
+};
+
+template <typename Callback> [[nodiscard]] ScopeGuard<Callback> make_scope_guard(const Callback& callback)
+{
+    return {callback};
+}