From e04c6220ffea579e51903528ccb9f82c98300255 Mon Sep 17 00:00:00 2001
From: Troels Bjerre Lund <troels@google.com>
Date: Fri, 6 Jan 2023 12:43:46 +0000
Subject: [PATCH] [K/N] custom-alloc: free largepages on sweep ^KT-55364

The default behavior on a sweep is to push empty pages into a separate
stack, which will be freed if not used before the next GC starts. This
serves two purposes: it reduces the number of system allocations, and it
avoids a race condition inside AtomicStack::Pop.

Neither of these are relevant for LargePages, since LargePages are never
reused and it is only the GC thread that calls
AtomicStack<LargePage>::Pop. The change is to free LargePages
immediately instead of waiting for the next GC cycle.


Co-authored-by: Troels Lund <troels@google.com>

Merge-request: KOTLIN-MR-598
Merged-by: Alexander Shabalin <alexander.shabalin@jetbrains.com>
---
 .../runtime/src/custom_alloc/cpp/Heap.cpp     |  2 +-
 .../src/custom_alloc/cpp/PageStore.hpp        | 38 ++++++++++++-------
 2 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/kotlin-native/runtime/src/custom_alloc/cpp/Heap.cpp b/kotlin-native/runtime/src/custom_alloc/cpp/Heap.cpp
index 2a865ba6ca3..a39e02a1905 100644
--- a/kotlin-native/runtime/src/custom_alloc/cpp/Heap.cpp
+++ b/kotlin-native/runtime/src/custom_alloc/cpp/Heap.cpp
@@ -40,7 +40,7 @@ void Heap::Sweep() noexcept {
         smallPages_[blockSize].Sweep();
     }
     mediumPages_.Sweep();
-    largePages_.Sweep();
+    largePages_.SweepAndFree();
 }
 
 MediumPage* Heap::GetMediumPage(uint32_t cellCount) noexcept {
diff --git a/kotlin-native/runtime/src/custom_alloc/cpp/PageStore.hpp b/kotlin-native/runtime/src/custom_alloc/cpp/PageStore.hpp
index 0ba8665a8ef..27808c580e5 100644
--- a/kotlin-native/runtime/src/custom_alloc/cpp/PageStore.hpp
+++ b/kotlin-native/runtime/src/custom_alloc/cpp/PageStore.hpp
@@ -22,26 +22,24 @@ public:
         while ((page = empty_.Pop())) page->Destroy();
     }
 
-    T* SweepAndFreeEmpty(AtomicStack<T>& from, AtomicStack<T>& to) noexcept {
-        T* page;
-        while ((page = from.Pop())) {
-            if (!page->Sweep()) {
-                empty_.Push(page);
-            } else {
-                to.Push(page);
-                return page;
-            }
-        }
-        return nullptr;
+    void Sweep() noexcept {
+        while (SweepSingle(unswept_, ready_)) {}
     }
 
-    void Sweep() noexcept {
-        while (SweepAndFreeEmpty(unswept_, ready_)) {}
+    void SweepAndFree() noexcept {
+        T* page;
+        while ((page = unswept_.Pop())) {
+            if (page->Sweep()) {
+                ready_.Push(page);
+            } else {
+                page->Destroy();
+            }
+        }
     }
 
     T* GetPage(uint32_t cellCount) noexcept {
         T* page;
-        if ((page = SweepAndFreeEmpty(unswept_, used_))) {
+        if ((page = SweepSingle(unswept_, used_))) {
             return page;
         }
         if ((page = ready_.Pop())) {
@@ -70,6 +68,18 @@ public:
     }
 
 private:
+    T* SweepSingle(AtomicStack<T>& from, AtomicStack<T>& to) noexcept {
+        T* page;
+        while ((page = from.Pop())) {
+            if (page->Sweep()) {
+                to.Push(page);
+                return page;
+            }
+            empty_.Push(page);
+        }
+        return nullptr;
+    }
+
     AtomicStack<T> empty_;
     AtomicStack<T> ready_;
     AtomicStack<T> used_;