1 files changed, 35 insertions, 8 deletions

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