3 files changed, 43 insertions, 16 deletions

diff --git a/arch/x86_64/include/arch/memory/heap/bump_allocator.hpp b/arch/x86_64/include/arch/memory/heap/bump_allocator.hpp
index d31783d..595eeea 100644
--- a/arch/x86_64/include/arch/memory/heap/bump_allocator.hpp
+++ b/arch/x86_64/include/arch/memory/heap/bump_allocator.hpp
@@ -1,6 +1,7 @@
 #ifndef TEACHOS_ARCH_X86_64_MEMORY_HEAP_BUMP_ALLOCATOR_HPP
 #define TEACHOS_ARCH_X86_64_MEMORY_HEAP_BUMP_ALLOCATOR_HPP
 
+#include <atomic>
 #include <cstdint>
 
 namespace teachos::arch::memory::heap
@@ -44,9 +45,9 @@ namespace teachos::arch::memory::heap
     auto deallocate(uint8_t * pointer, std::size_t size) -> void;
 
   private:
-    std::size_t heap_start;  ///< Start of the allocatable heap area
-    std::size_t heap_end;    ///< End of the allocatable heap area
-    std::size_t next;        ///< Current address, which is the start of still unused allocatable heap area
+    std::size_t heap_start;     ///< Start of the allocatable heap area
+    std::size_t heap_end;       ///< End of the allocatable heap area
+    std::atomic_uint64_t next;  ///< Current address, which is the start of still unused allocatable heap area
   };
 
 }  // namespace teachos::arch::memory::heap
diff --git a/arch/x86_64/src/kernel/main.cpp b/arch/x86_64/src/kernel/main.cpp
index 13526f4..023327e 100644
--- a/arch/x86_64/src/kernel/main.cpp
+++ b/arch/x86_64/src/kernel/main.cpp
@@ -20,9 +20,10 @@ namespace teachos::arch::kernel
                                                 memory::heap::HEAP_START + memory::heap::HEAP_SIZE};
     auto test = heap_allocator.allocate(1024);
     auto test2 = new (test) memory::multiboot::memory_information{};
-    test2->kernel_end = 5000;
-    auto test3 = test2->kernel_end;
-    if (test || test2 || test3)
+    auto test3 = *test2;
+    test3.kernel_end = 5000;
+    auto test4 = test3.kernel_end;
+    if (test || test2 || test3.kernel_end || test4)
     {
       video::vga::text::write("Kernel remapping successfull", video::vga::text::common_attributes::green_on_black);
     }
@@ -31,7 +32,5 @@ namespace teachos::arch::kernel
     // allocator.allocate_frame still works?
     // TODO: Fix unmapping old level 4 page table and turn it into guard page, use Stack Probes for stack allocation if
     // possible.
-
-    // TODO: Align up and down for the bump allocator. https://os.phil-opp.com/kernel-heap/#a-bump-allocator
   }
 }  // namespace teachos::arch::kernel
diff --git a/arch/x86_64/src/memory/heap/bump_allocator.cpp b/arch/x86_64/src/memory/heap/bump_allocator.cpp
index 1ab8ea9..19ced47 100644
--- a/arch/x86_64/src/memory/heap/bump_allocator.cpp
+++ b/arch/x86_64/src/memory/heap/bump_allocator.cpp
@@ -2,17 +2,43 @@
 
 #include "arch/exception_handling/assert.hpp"
 
+#include <limits>
+#include <type_traits>
+
 namespace teachos::arch::memory::heap
 {
-  auto bump_allocator::allocate(std::size_t size) -> void *
+  namespace
   {
-    auto alloc_start = next;
-    auto alloc_end = next + size;
-
-    arch::exception_handling::assert(alloc_end <= heap_end, "[Heap Allocator] Out of memory!");
+    template<typename T>
+    auto saturating_add(T x, T y) -> T
+      requires std::is_unsigned_v<T>
+    {
+      if (x > std::numeric_limits<T>::max() - y)
+      {
+        return std::numeric_limits<T>::max();
+      }
+      T result = x + y;
+      return result;
+    }
+  }  // namespace
 
-    next = alloc_end;
-    return reinterpret_cast<void *>(alloc_start);
+  auto bump_allocator::allocate(std::size_t size) -> void *
+  {
+    // Repeat allocation until it succeeds, has to be done, because another allocator could overtake it at any time
+    // causing the value to differ and the calculation to have to be redone.
+    for (;;)
+    {
+      auto alloc_start = next.load(std::memory_order::relaxed);
+      auto const alloc_end = saturating_add(alloc_start, size);
+      arch::exception_handling::assert(alloc_end <= heap_end, "[Heap Allocator] Out of memory");
+      // Check if the atomic value is still the one initally loaded, if it isn't we have been overtaken by another
+      // thread and need to redo the calculation.
+      auto const updated = next.compare_exchange_strong(alloc_start, alloc_end, std::memory_order::relaxed);
+      if (updated)
+      {
+        return reinterpret_cast<void *>(alloc_start);
+      }
+    }
   }
 
   auto bump_allocator::deallocate(uint8_t * pointer, std::size_t size) -> void
@@ -21,4 +47,5 @@ namespace teachos::arch::memory::heap
     {
     }
   }
-}  // namespace teachos::arch::memory::heap
\ No newline at end of file
+
+}  // namespace teachos::arch::memory::heap