[K/N] custom-alloc: handle finalizers

This commit adds finalizers to the custom allocator. Compared to the
existing solution:

* The finalizer queue is replaced by an AtomicStack<ExtraObjectData>.
* All objects with finalizers get meta objects attached. This was not
  previously the case for CleanerImpl, but is now needed to link objects
  together in the finalization queue.
* The finalizer queue is built during SweepExtraObjects, instead of
  during regular sweeping.
* Cleaners are executed by the finalizer thread, and no longer by a
  separate worker thread.

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

Merge-request: KOTLIN-MR-592
Merged-by: Alexander Shabalin <alexander.shabalin@jetbrains.com>

kotlin-native/runtime/build.gradle.kts

@@ -260,6 +260,18 @@ bitcode {
             }
         }
 
+        module("experimental_memory_manager_custom") {
+            srcRoot.set(layout.projectDirectory.dir("src/mm"))
+            headersDirs.from(files("src/gc/common/cpp", "src/main/cpp", "src/custom_alloc/cpp"))
+            sourceSets {
+                main {}
+                testFixtures {}
+                test {}
+            }
+            
+            compilerArgs.add("-DCUSTOM_ALLOCATOR")
+        }
+
         module("common_gc") {
             srcRoot.set(layout.projectDirectory.dir("src/gc/common"))
             headersDirs.from(files("src/mm/cpp", "src/main/cpp"))

kotlin-native/runtime/src/custom_alloc/cpp/AtomicStack.hpp

@@ -8,14 +8,26 @@
 
 #include <atomic>
 
-#include "CustomLogging.hpp"
 #include "KAssert.h"
+#include "Utils.hpp"
 
 namespace kotlin::alloc {
 
 template <class T>
-class AtomicStack {
+class AtomicStack: private kotlin::MoveOnly {
 public:
+    AtomicStack() noexcept = default;
+
+    AtomicStack(AtomicStack&& other) noexcept
+        : stack_(other.stack_.exchange(nullptr, std::memory_order_acq_rel)) {}
+
+    AtomicStack& operator=(AtomicStack&& other) noexcept {
+        // Not using swap idiom, because implementing swap of two atomics requires DCAS or locks.
+        auto newHead = other.stack_.exchange(nullptr, std::memory_order_acq_rel);
+        stack_.store(newHead, std::memory_order_acq_rel);
+        return *this;
+    }
+
     // Pop() is not fully thread-safe, in that the returned page must not be
     // immediately freed, if another thread might be simultaneously Popping
     // from the same stack. As of writing this comment, this is handled by only
@@ -23,7 +35,6 @@ public:
     T* Pop() noexcept {
         T* elm = stack_.load(std::memory_order_acquire);
         while (elm && !stack_.compare_exchange_weak(elm, elm->next_, std::memory_order_acq_rel)) {}
-        CustomAllocDebug("AtomicStack(%p)::Pop() = %p", this, elm);
         return elm;
     }
 
@@ -34,13 +45,12 @@ public:
         } while (!stack_.compare_exchange_weak(head, elm, std::memory_order_acq_rel));
     }
 
-    void TransferAllFrom(AtomicStack<T>& other) noexcept {
-        // Clear out the `other` stack.
-        T* otherHead = nullptr;
-        while (!other.stack_.compare_exchange_weak(otherHead, nullptr, std::memory_order_acq_rel)) {}
+    // This will put the contents of the other stack on top of this stack
+    void TransferAllFrom(AtomicStack<T> other) noexcept {
+        T* otherHead = other.stack_.exchange(nullptr, std::memory_order_relaxed);
         // If the `other` stack was empty, do nothing.
         if (!otherHead) return;
-        // If `this` stack is empty, just copy the `other` stack over 
+        // If `this` stack is empty, just copy the `other` stack over
         T* thisHead = nullptr;
         if (stack_.compare_exchange_strong(thisHead, otherHead, std::memory_order_acq_rel)) {
             return;
@@ -59,8 +69,17 @@ public:
 
     bool isEmpty() noexcept { return stack_.load(std::memory_order_relaxed) == nullptr; }
 
-    ~AtomicStack() noexcept {
-        RuntimeAssert(isEmpty(), "AtomicStack must be empty when destroyed");
+    // Not thread-safe. Named like this to make AtomicStack compatible with FinalizerQueue
+    size_t size() {
+        size_t size = 0;
+        for (T* elm = stack_.load(std::memory_order_relaxed); elm != nullptr; elm = elm->next_) {
+            ++size;
+        }
+        return size;
+    }
+
+    ~AtomicStack() {
+        RuntimeAssert(isEmpty(), "AtomicStack must be empty on destruction");
     }
 
 private:

kotlin-native/runtime/src/custom_alloc/cpp/CustomAllocConstants.hpp

@@ -11,6 +11,7 @@
 
 #include "SmallPage.hpp"
 #include "MediumPage.hpp"
+#include "ExtraObjectPage.hpp"
 
 inline constexpr const size_t KiB = 1024;
 
@@ -25,4 +26,8 @@ inline constexpr const size_t MEDIUM_PAGE_CELL_COUNT =
 
 inline constexpr const size_t LARGE_PAGE_SIZE_THRESHOLD = (MEDIUM_PAGE_CELL_COUNT - 1);
 
+inline constexpr const size_t EXTRA_OBJECT_PAGE_SIZE = 64 * KiB;
+inline constexpr const int EXTRA_OBJECT_COUNT =
+        (EXTRA_OBJECT_PAGE_SIZE - sizeof(kotlin::alloc::ExtraObjectPage)) / sizeof(kotlin::alloc::ExtraObjectCell);
+
 #endif

kotlin-native/runtime/src/custom_alloc/cpp/CustomAllocator.cpp

@@ -13,10 +13,16 @@
 
 #include "ConcurrentMarkAndSweep.hpp"
 #include "CustomLogging.hpp"
+#include "ExtraObjectData.hpp"
+#include "ExtraObjectPage.hpp"
 #include "GCScheduler.hpp"
 #include "LargePage.hpp"
 #include "MediumPage.hpp"
+#include "Memory.h"
 #include "SmallPage.hpp"
+#include "GCImpl.hpp"
+#include "TypeInfo.h"
+#include "Types.h"
 
 namespace kotlin::alloc {
 
@@ -48,7 +54,7 @@ uint64_t ArrayAllocatedDataSize(const TypeInfo* typeInfo, uint32_t count) noexce
 }
 
 CustomAllocator::CustomAllocator(Heap& heap, gc::GCSchedulerThreadData& gcScheduler) noexcept :
-    heap_(heap), gcScheduler_(gcScheduler), mediumPage_(nullptr) {
+    heap_(heap), gcScheduler_(gcScheduler), mediumPage_(nullptr), extraObjectPage_(nullptr) {
     CustomAllocInfo("CustomAllocator::CustomAllocator(heap)");
     memset(smallPages_, 0, sizeof(smallPages_));
 }
@@ -58,7 +64,12 @@ ObjHeader* CustomAllocator::CreateObject(const TypeInfo* typeInfo) noexcept {
     size_t allocSize = ObjectAllocatedDataSize(typeInfo);
     auto* heapObject = new (Allocate(allocSize)) HeapObjHeader();
     auto* object = &heapObject->object;
-    object->typeInfoOrMeta_ = const_cast<TypeInfo*>(typeInfo);
+    if (typeInfo->flags_ & TF_HAS_FINALIZER) {
+        auto* extraObject = CreateExtraObject();
+        object->typeInfoOrMeta_ = reinterpret_cast<TypeInfo*>(new (extraObject) mm::ExtraObjectData(object, typeInfo));
+    } else {
+        object->typeInfoOrMeta_ = const_cast<TypeInfo*>(typeInfo);
+    }
     return object;
 }
 
@@ -72,10 +83,40 @@ ArrayHeader* CustomAllocator::CreateArray(const TypeInfo* typeInfo, uint32_t cou
     return array;
 }
 
+mm::ExtraObjectData* CustomAllocator::CreateExtraObject() noexcept {
+    CustomAllocDebug("CustomAllocator::CreateExtraObject()");
+    ExtraObjectPage* page = extraObjectPage_;
+    if (page) {
+        mm::ExtraObjectData* block = page->TryAllocate();
+        if (block) {
+            memset(block, 0, sizeof(mm::ExtraObjectData));
+            return block;
+        }
+    }
+    CustomAllocDebug("Failed to allocate in current ExtraObjectPage");
+    while ((page = heap_.GetExtraObjectPage())) {
+        mm::ExtraObjectData* block = page->TryAllocate();
+        if (block) {
+            extraObjectPage_ = page;
+            memset(block, 0, sizeof(mm::ExtraObjectData));
+            return block;
+        }
+    }
+    return nullptr;
+}
+
+// static
+mm::ExtraObjectData& CustomAllocator::CreateExtraObjectDataForObject(
+        mm::ThreadData* threadData, ObjHeader* baseObject, const TypeInfo* info) noexcept {
+    mm::ExtraObjectData* extraObject = threadData->gc().impl().alloc().CreateExtraObject();
+    return *new (extraObject) mm::ExtraObjectData(baseObject, info);
+}
+
 void CustomAllocator::PrepareForGC() noexcept {
     CustomAllocInfo("CustomAllocator@%p::PrepareForGC()", this);
     mediumPage_ = nullptr;
     memset(smallPages_, 0, sizeof(smallPages_));
+    extraObjectPage_ = nullptr;
 }
 
 uint8_t* CustomAllocator::Allocate(uint64_t size) noexcept {

kotlin-native/runtime/src/custom_alloc/cpp/CustomAllocator.hpp

@@ -9,6 +9,8 @@
 #include <atomic>
 #include <cstring>
 
+#include "ExtraObjectData.hpp"
+#include "ExtraObjectPage.hpp"
 #include "GCScheduler.hpp"
 #include "Heap.hpp"
 #include "MediumPage.hpp"
@@ -25,6 +27,11 @@ public:
 
     ArrayHeader* CreateArray(const TypeInfo* typeInfo, uint32_t count) noexcept;
 
+    mm::ExtraObjectData* CreateExtraObject() noexcept;
+
+    static mm::ExtraObjectData& CreateExtraObjectDataForObject(
+            mm::ThreadData* threadData, ObjHeader* baseObject, const TypeInfo* info) noexcept;
+
     void PrepareForGC() noexcept;
 
 private:
@@ -37,6 +44,7 @@ private:
     gc::GCSchedulerThreadData& gcScheduler_;
     MediumPage* mediumPage_;
     SmallPage* smallPages_[SMALL_PAGE_MAX_BLOCK_SIZE + 1];
+    ExtraObjectPage* extraObjectPage_;
 };
 
 } // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/CustomAllocatorTest.cpp

@@ -11,7 +11,6 @@
 #include "Memory.h"
 #include "gtest/gtest.h"
 #include "Heap.hpp"
-#include "SmallPage.hpp"
 #include "TypeInfo.h"
 
 namespace {
@@ -19,7 +18,7 @@ namespace {
 using Heap = typename kotlin::alloc::Heap;
 using CustomAllocator = typename kotlin::alloc::CustomAllocator;
 
-#define MIN_BLOCK_SIZE 2
+inline constexpr int MIN_BLOCK_SIZE = 2;
 
 TEST(CustomAllocTest, SmallAllocNonNull) {
     const int N = 200;
@@ -59,7 +58,6 @@ TEST(CustomAllocTest, SmallAllocSameSmallPage) {
 TEST(CustomAllocTest, TwoAllocatorsDifferentPages) {
     for (int blocks = MIN_BLOCK_SIZE; blocks < 2000; ++blocks) {
         Heap heap;
-        kotlin::gc::GCScheduler scheduler;
         kotlin::gc::GCSchedulerConfig config;
         kotlin::gc::GCSchedulerThreadData schedulerData1(config, [](auto&) {});
         kotlin::gc::GCSchedulerThreadData schedulerData2(config, [](auto&) {});
@@ -73,5 +71,21 @@ TEST(CustomAllocTest, TwoAllocatorsDifferentPages) {
     }
 }
 
+using Data = typename kotlin::mm::ExtraObjectData;
+
+TEST(CustomAllocTest, AllocExtraObjectNonNullZeroed) {
+    Heap heap;
+    kotlin::gc::GCSchedulerConfig config;
+    kotlin::gc::GCSchedulerThreadData schedulerData(config, [](auto&) {});
+    CustomAllocator ca(heap, schedulerData);
+    for (int i = 1; i < 10; ++i) {
+        uint8_t* obj = reinterpret_cast<uint8_t*>(ca.CreateExtraObject());
+        EXPECT_TRUE(obj);
+        for (size_t j = 0; j < sizeof(Data); ++j) {
+            EXPECT_FALSE(obj[j]);
+        }
+    }
+}
+
 #undef MIN_BLOCK_SIZE
 } // namespace

kotlin-native/runtime/src/custom_alloc/cpp/CustomFinalizerProcessor.cpp

@@ -0,0 +1,103 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "CustomFinalizerProcessor.hpp"
+
+#include <cstdint>
+#include <mutex>
+#include <thread>
+
+#include "AtomicStack.hpp"
+#include "CustomLogging.hpp"
+#include "ExtraObjectData.hpp"
+#include "FinalizerHooks.hpp"
+#include "Memory.h"
+#include "Runtime.h"
+
+namespace kotlin::alloc {
+
+void CustomFinalizerProcessor::StartFinalizerThreadIfNone() noexcept {
+    CustomAllocDebug("CustomFinalizerProcessor::StartFinalizerThreadIfNone()");
+    std::unique_lock guard(threadCreatingMutex_);
+    if (finalizerThread_.joinable()) return;
+
+    finalizerThread_ = ScopedThread(ScopedThread::attributes().name("Custom finalizer processor"), [this] {
+        Kotlin_initRuntimeIfNeeded();
+        {
+            std::unique_lock guard(initializedMutex_);
+            initialized_ = true;
+        }
+        initializedCondVar_.notify_all();
+        while (true) {
+            std::unique_lock lock(finalizerQueueMutex_);
+            finalizerQueueCondVar_.wait(lock, [this] {
+                return finalizedEpoch_ != scheduledEpoch_ || shutdownFlag_;
+            });
+            if (finalizedEpoch_ == scheduledEpoch_) {
+                RuntimeAssert(shutdownFlag_, "Nothing to do, but no shutdownFlag_ is set on wakeup");
+                RuntimeAssert(finalizerQueue_.isEmpty(), "Finalizer queue should be empty when killing finalizer thread");
+                break;
+            }
+            int64_t queueEpoch = scheduledEpoch_;
+            Queue finalizerQueue = std::move(finalizerQueue_);
+            lock.unlock();
+            ThreadStateGuard guard(ThreadState::kRunnable);
+            ExtraObjectCell* cell;
+            while ((cell = finalizerQueue.Pop())) {
+                auto* extraObject = cell->Data();
+                auto* baseObject = extraObject->GetBaseObject();
+                RunFinalizers(baseObject);
+                extraObject->setFlag(mm::ExtraObjectData::FLAGS_FINALIZED);
+            }
+            while (finalizedEpoch_ != queueEpoch) {
+                epochDoneCallback_(++finalizedEpoch_);
+            }
+        }
+        {
+            std::unique_lock guard(initializedMutex_);
+            initialized_ = false;
+        }
+        initializedCondVar_.notify_all();
+    });
+}
+
+void CustomFinalizerProcessor::StopFinalizerThread() noexcept {
+    CustomAllocDebug("CustomFinalizerProcessor::StopFinalizerThread()");
+    {
+        std::unique_lock guard(finalizerQueueMutex_);
+        if (!finalizerThread_.joinable()) return;
+        shutdownFlag_ = true;
+        finalizerQueueCondVar_.notify_all();
+    }
+    finalizerThread_.join();
+    shutdownFlag_ = false;
+    RuntimeAssert(finalizerQueue_.isEmpty(), "Finalizer queue should be empty when killing finalizer thread");
+    std::unique_lock guard(finalizerQueueMutex_);
+    finalizerQueueCondVar_.notify_all();
+}
+
+void CustomFinalizerProcessor::ScheduleTasks(Queue&& tasks, int64_t epoch) noexcept {
+    std::unique_lock guard(finalizerQueueMutex_);
+    StartFinalizerThreadIfNone();
+    finalizerQueue_.TransferAllFrom(std::move(tasks));
+    scheduledEpoch_ = epoch;
+    finalizerQueueCondVar_.notify_all();
+}
+
+bool CustomFinalizerProcessor::IsRunning() noexcept {
+    return finalizerThread_.joinable();
+}
+
+void CustomFinalizerProcessor::WaitFinalizerThreadInitialized() noexcept {
+    CustomAllocDebug("CustomFinalizerProcessor::WaitFinalizerThreadInitialized()");
+    std::unique_lock guard(initializedMutex_);
+    initializedCondVar_.wait(guard, [this] { return initialized_; });
+}
+
+CustomFinalizerProcessor::~CustomFinalizerProcessor() {
+    StopFinalizerThread();
+}
+
+} // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/CustomFinalizerProcessor.hpp

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_CUSTOMFINALIZERPROCESSOR_HPP_
+#define CUSTOM_ALLOC_CPP_CUSTOMFINALIZERPROCESSOR_HPP_
+
+#include <cstdint>
+#include <condition_variable>
+
+#include "AtomicStack.hpp"
+#include "ExtraObjectPage.hpp"
+#include "ScopedThread.hpp"
+
+namespace kotlin::alloc {
+
+class CustomFinalizerProcessor : Pinned {
+public:
+    using Queue = typename kotlin::alloc::AtomicStack<kotlin::alloc::ExtraObjectCell>;
+    explicit CustomFinalizerProcessor(std::function<void(int64_t)> epochDoneCallback) : epochDoneCallback_(std::move(epochDoneCallback)) {}
+    void ScheduleTasks(Queue&& tasks, int64_t epoch) noexcept;
+    void StopFinalizerThread() noexcept;
+    bool IsRunning() noexcept;
+    void StartFinalizerThreadIfNone() noexcept;
+    void WaitFinalizerThreadInitialized() noexcept;
+    ~CustomFinalizerProcessor();
+
+private:
+    ScopedThread finalizerThread_;
+    Queue finalizerQueue_;
+    std::condition_variable finalizerQueueCondVar_;
+    std::mutex finalizerQueueMutex_;
+    std::function<void(int64_t)> epochDoneCallback_;
+    int64_t scheduledEpoch_ = 0;
+    int64_t finalizedEpoch_ = 0;
+    bool shutdownFlag_ = false;
+
+    std::mutex initializedMutex_;
+    std::condition_variable initializedCondVar_;
+    bool initialized_ = false;
+
+    std::mutex threadCreatingMutex_;
+};
+
+} // namespace kotlin::alloc
+
+#endif // CUSTOM_ALLOC_CPP_CUSTOMFINALIZERPROCESSOR_HPP_

kotlin-native/runtime/src/custom_alloc/cpp/ExtraObjectPage.cpp

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "ExtraObjectPage.hpp"
+
+#include <atomic>
+#include <cstring>
+#include <random>
+
+#include "AtomicStack.hpp"
+#include "CustomAllocConstants.hpp"
+#include "CustomLogging.hpp"
+#include "ExtraObjectData.hpp"
+#include "GCApi.hpp"
+
+namespace kotlin::alloc {
+
+ExtraObjectPage* ExtraObjectPage::Create() noexcept {
+    CustomAllocInfo("ExtraObjectPage::Create()");
+    return new (SafeAlloc(EXTRA_OBJECT_PAGE_SIZE)) ExtraObjectPage();
+}
+
+ExtraObjectPage::ExtraObjectPage() noexcept {
+    CustomAllocInfo("ExtraObjectPage(%p)::ExtraObjectPage()", this);
+    nextFree_ = cells_;
+    ExtraObjectCell* end = cells_ + EXTRA_OBJECT_COUNT;
+    for (ExtraObjectCell* cell = cells_; cell < end; cell = cell->next_) {
+        cell->next_ = cell + 1;
+    }
+}
+
+void ExtraObjectPage::Destroy() noexcept {
+    std_support::free(this);
+}
+
+mm::ExtraObjectData* ExtraObjectPage::TryAllocate() noexcept {
+    if (nextFree_ >= cells_ + EXTRA_OBJECT_COUNT) {
+        return nullptr;
+    }
+    ExtraObjectCell* freeBlock = nextFree_;
+    nextFree_ = freeBlock->next_;
+    CustomAllocDebug("ExtraObjectPage(%p)::TryAllocate() = %p", this, freeBlock);
+    return freeBlock->Data();
+}
+
+bool ExtraObjectPage::Sweep(AtomicStack<ExtraObjectCell>& finalizerQueue) noexcept {
+    CustomAllocInfo("ExtraObjectPage(%p)::Sweep()", this);
+    // `end` is after the last legal allocation of a block, but does not
+    // necessarily match an actual block starting point.
+    ExtraObjectCell* end = cells_ + EXTRA_OBJECT_COUNT;
+    bool alive = false;
+    ExtraObjectCell** nextFree = &nextFree_;
+    for (ExtraObjectCell* cell = cells_; cell < end; ++cell) {
+        // If the current cell is free, move on.
+        if (cell == *nextFree) {
+            nextFree = &cell->next_;
+            continue;
+        }
+        // If the current cell was marked, it's alive, and the whole page is alive.
+        if (!SweepExtraObject(cell, finalizerQueue)) {
+            alive = true;
+            continue;
+        }
+        // Free the current block and insert it into the free list.
+        cell->next_ = *nextFree;
+        *nextFree = cell;
+        nextFree = &cell->next_;
+    }
+    return alive;
+}
+
+} // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/ExtraObjectPage.hpp

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_EXTRA_OBJECTPAGE_HPP_
+#define CUSTOM_ALLOC_CPP_EXTRA_OBJECTPAGE_HPP_
+
+#include <atomic>
+#include <cstdint>
+
+#include "AtomicStack.hpp"
+#include "ExtraObjectData.hpp"
+
+namespace kotlin::alloc {
+
+struct ExtraObjectCell {
+    mm::ExtraObjectData* Data() { return reinterpret_cast<mm::ExtraObjectData*>(data_); }
+
+    // This is used to simultaneously build two lists: a free list and a finalizers queue.
+    // A cell cannot exist in both of them, but can be in neither when it's alive.
+    ExtraObjectCell* next_;
+    struct alignas(mm::ExtraObjectData) {
+        uint8_t data_[sizeof(mm::ExtraObjectData)];
+    };
+};
+
+class alignas(8) ExtraObjectPage {
+public:
+    static ExtraObjectPage* Create() noexcept;
+
+    void Destroy() noexcept;
+
+    // Tries to allocate in current page, returns null if no free block in page
+    mm::ExtraObjectData* TryAllocate() noexcept;
+
+    bool Sweep(AtomicStack<ExtraObjectCell>& finalizerQueue) noexcept;
+
+private:
+    friend class AtomicStack<ExtraObjectPage>;
+
+    ExtraObjectPage() noexcept;
+
+    // Used for linking pages together in `pages` queue or in `unswept` queue.
+    ExtraObjectPage* next_;
+    ExtraObjectCell* nextFree_;
+    ExtraObjectCell cells_[];
+};
+
+} // namespace kotlin::alloc
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/ExtraObjectPageTest.cpp

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include <cstddef>
+#include <cstdint>
+
+#include "AtomicStack.hpp"
+#include "CustomAllocConstants.hpp"
+#include "ExtraObjectData.hpp"
+#include "ExtraObjectPage.hpp"
+#include "gtest/gtest.h"
+#include "TypeInfo.h"
+
+namespace {
+
+using Data = typename kotlin::mm::ExtraObjectData;
+using Cell = typename kotlin::alloc::ExtraObjectCell;
+using Page = typename kotlin::alloc::ExtraObjectPage;
+using Queue = typename kotlin::alloc::AtomicStack<Cell>;
+
+Data* alloc(Page* page) {
+    Data* ptr = page->TryAllocate();
+    if (ptr) {
+        memset(ptr, 0, sizeof(Data));
+    }
+    return ptr;
+}
+
+TEST(CustomAllocTest, ExtraObjectPageConsequtiveAlloc) {
+    Page* page = Page::Create();
+    uint8_t* prev = reinterpret_cast<uint8_t*>(alloc(page));
+    uint8_t* cur;
+    while ((cur = reinterpret_cast<uint8_t*>(alloc(page)))) {
+        EXPECT_EQ(prev + sizeof(Cell), cur);
+        prev = cur;
+    }
+    free(page);
+}
+
+TEST(CustomAllocTest, ExtraObjectPageSweepEmptyPage) {
+    Page* page = Page::Create();
+    Queue finalizerQueue;
+    EXPECT_FALSE(page->Sweep(finalizerQueue));
+    EXPECT_EQ(finalizerQueue.size(), size_t(0));
+    free(page);
+}
+
+TEST(CustomAllocTest, ExtraObjectPageSweepFullFinalizedPage) {
+    Page* page = Page::Create();
+    int count = 0;
+    Data* ptr;
+    while ((ptr = alloc(page))) {
+        ptr->setFlag(Data::FLAGS_FINALIZED);
+        ++count;
+    }
+    EXPECT_EQ(count, EXTRA_OBJECT_COUNT);
+    Queue finalizerQueue;
+    EXPECT_FALSE(page->Sweep(finalizerQueue));
+    EXPECT_EQ(finalizerQueue.size(), size_t(0));
+    free(page);
+}
+
+} // namespace

kotlin-native/runtime/src/custom_alloc/cpp/GCApi.cpp

@@ -9,6 +9,8 @@
 
 #include "ConcurrentMarkAndSweep.hpp"
 #include "CustomLogging.hpp"
+#include "FinalizerHooks.hpp"
+#include "KAssert.h"
 #include "ObjectFactory.hpp"
 
 namespace kotlin::alloc {
@@ -19,14 +21,60 @@ bool TryResetMark(void* ptr) noexcept {
     Node& node = Node::FromData(ptr);
     NodeRef ref = NodeRef(node);
     auto& objectData = ref.ObjectData();
-    if (!objectData.tryResetMark()) {
-        auto* objHeader = ref.GetObjHeader();
-        if (HasFinalizers(objHeader)) {
-            CustomAllocWarning("FINALIZER IGNORED");
-        }
+    bool reset = objectData.tryResetMark();
+    CustomAllocDebug("TryResetMark(%p) = %d", ptr, reset);
+    return reset;
+}
+
+using ObjectFactory = mm::ObjectFactory<gc::ConcurrentMarkAndSweep>;
+using ExtraObjectsFactory = mm::ExtraObjectDataFactory;
+
+static void KeepAlive(ObjHeader* baseObject) noexcept {
+    auto& objectData = ObjectFactory::NodeRef::From(baseObject).ObjectData();
+    objectData.tryMark();
+}
+
+static bool IsAlive(ObjHeader* baseObject) noexcept {
+    auto& objectData = ObjectFactory::NodeRef::From(baseObject).ObjectData();
+    return objectData.marked();
+}
+
+bool SweepExtraObject(ExtraObjectCell* extraObjectCell, AtomicStack<ExtraObjectCell>& finalizerQueue) noexcept {
+    auto* extraObject = extraObjectCell->Data();
+    if (extraObject->getFlag(mm::ExtraObjectData::FLAGS_FINALIZED)) {
+        CustomAllocDebug("SweepIsCollectable(%p): already finalized", extraObject);
+        return true;
+    }
+    auto* baseObject = extraObject->GetBaseObject();
+    RuntimeAssert(baseObject->heap(), "SweepIsCollectable on a non-heap object");
+    if (extraObject->getFlag(mm::ExtraObjectData::FLAGS_IN_FINALIZER_QUEUE)) {
+        CustomAllocDebug("SweepIsCollectable(%p): already in finalizer queue, keep base object (%p) alive", extraObject, baseObject);
+        KeepAlive(baseObject);
         return false;
     }
-    return true;
+    if (IsAlive(baseObject)) {
+        return false;
+    }
+    extraObject->ClearWeakReferenceCounter();
+    if (extraObject->HasAssociatedObject()) {
+        extraObject->DetachAssociatedObject();
+        extraObject->setFlag(mm::ExtraObjectData::FLAGS_IN_FINALIZER_QUEUE);
+        finalizerQueue.Push(extraObjectCell);
+        KeepAlive(baseObject);
+        CustomAllocDebug("SweepIsCollectable(%p): add to finalizerQueue", extraObject);
+        return false;
+    } else {
+        if (HasFinalizers(baseObject)) {
+            extraObject->setFlag(mm::ExtraObjectData::FLAGS_IN_FINALIZER_QUEUE);
+            finalizerQueue.Push(extraObjectCell);
+            KeepAlive(baseObject);
+            CustomAllocDebug("SweepIsCollectable(%p): addings to finalizerQueue, keep base object (%p) alive", extraObject, baseObject);
+            return false;
+        }
+        extraObject->Uninstall();
+        CustomAllocDebug("SweepIsCollectable(%p): uninstalled extraObject", extraObject);
+        return true;
+    }
 }
 
 void* SafeAlloc(uint64_t size) noexcept {

kotlin-native/runtime/src/custom_alloc/cpp/GCApi.hpp

@@ -11,10 +11,16 @@
 #include <limits>
 #include <stdlib.h>
 
+#include "AtomicStack.hpp"
+#include "ExtraObjectPage.hpp"
+
 namespace kotlin::alloc {
 
 bool TryResetMark(void* ptr) noexcept;
 
+// Returns true if swept successfully, i.e., if the extraobject can be reclaimed now.
+bool SweepExtraObject(ExtraObjectCell* extraObjectCell, AtomicStack<ExtraObjectCell>& finalizerQueue) noexcept;
+
 void* SafeAlloc(uint64_t size) noexcept;
 
 } // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/Heap.cpp

@@ -12,15 +12,20 @@
 #include <new>
 
 #include "CustomAllocConstants.hpp"
+#include "AtomicStack.hpp"
 #include "CustomLogging.hpp"
-#include "LargePage.hpp"
-#include "MediumPage.hpp"
-#include "SmallPage.hpp"
+#include "ExtraObjectPage.hpp"
 #include "ThreadRegistry.hpp"
 #include "GCImpl.hpp"
 
 namespace kotlin::alloc {
 
+Heap::~Heap() noexcept {
+    ExtraObjectPage* page;
+    while ((page = extraObjectPages_.Pop())) page->Destroy();
+    while ((page = usedExtraObjectPages_.Pop())) page->Destroy();
+}
+
 void Heap::PrepareForGC() noexcept {
     CustomAllocDebug("Heap::PrepareForGC()");
     for (auto& thread : kotlin::mm::ThreadRegistry::Instance().LockForIter()) {
@@ -29,20 +34,35 @@ void Heap::PrepareForGC() noexcept {
 
     mediumPages_.PrepareForGC();
     largePages_.PrepareForGC();
-    for (size_t blockSize = 0; blockSize <= SMALL_PAGE_MAX_BLOCK_SIZE; ++blockSize) {
+    for (int blockSize = 0; blockSize <= SMALL_PAGE_MAX_BLOCK_SIZE; ++blockSize) {
         smallPages_[blockSize].PrepareForGC();
     }
 }
 
 void Heap::Sweep() noexcept {
     CustomAllocDebug("Heap::Sweep()");
-    for (size_t blockSize = 0; blockSize <= SMALL_PAGE_MAX_BLOCK_SIZE; ++blockSize) {
+    for (int blockSize = 0; blockSize <= SMALL_PAGE_MAX_BLOCK_SIZE; ++blockSize) {
         smallPages_[blockSize].Sweep();
     }
     mediumPages_.Sweep();
     largePages_.SweepAndFree();
 }
 
+AtomicStack<ExtraObjectCell> Heap::SweepExtraObjects(gc::GCHandle gcHandle) noexcept {
+    CustomAllocDebug("Heap::SweepExtraObjects()");
+    AtomicStack<ExtraObjectCell> finalizerQueue;
+    ExtraObjectPage* page;
+    while ((page = usedExtraObjectPages_.Pop())) {
+        if (!page->Sweep(finalizerQueue)) {
+            CustomAllocInfo("SweepExtraObjects free(%p)", page);
+            free(page);
+        } else {
+            extraObjectPages_.Push(page);
+        }
+    }
+    return finalizerQueue;
+}
+
 MediumPage* Heap::GetMediumPage(uint32_t cellCount) noexcept {
     CustomAllocDebug("Heap::GetMediumPage()");
     return mediumPages_.GetPage(cellCount);
@@ -58,4 +78,14 @@ LargePage* Heap::GetLargePage(uint64_t cellCount) noexcept {
     return largePages_.NewPage(cellCount);
 }
 
+ExtraObjectPage* Heap::GetExtraObjectPage() noexcept {
+    CustomAllocInfo("CustomAllocator::GetExtraObjectPage()");
+    ExtraObjectPage* page = extraObjectPages_.Pop();
+    if (page == nullptr) {
+        page = ExtraObjectPage::Create();
+    }
+    usedExtraObjectPages_.Push(page);
+    return page;
+}
+
 } // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/Heap.hpp

@@ -9,7 +9,10 @@
 #include <atomic>
 #include <cstring>
 
+#include "AtomicStack.hpp"
 #include "CustomAllocConstants.hpp"
+#include "ExtraObjectPage.hpp"
+#include "GCStatistics.hpp"
 #include "LargePage.hpp"
 #include "MediumPage.hpp"
 #include "PageStore.hpp"
@@ -19,6 +22,8 @@ namespace kotlin::alloc {
 
 class Heap {
 public:
+    ~Heap() noexcept;
+
     // Called once by the GC thread after all mutators have been suspended
     void PrepareForGC() noexcept;
 
@@ -27,14 +32,19 @@ public:
     // seen by one sweeper.
     void Sweep() noexcept;
 
+    AtomicStack<ExtraObjectCell> SweepExtraObjects(gc::GCHandle gcHandle) noexcept;
+
     SmallPage* GetSmallPage(uint32_t cellCount) noexcept;
     MediumPage* GetMediumPage(uint32_t cellCount) noexcept;
     LargePage* GetLargePage(uint64_t cellCount) noexcept;
+    ExtraObjectPage* GetExtraObjectPage() noexcept;
 
 private:
     PageStore<SmallPage> smallPages_[SMALL_PAGE_MAX_BLOCK_SIZE + 1];
     PageStore<MediumPage> mediumPages_;
     PageStore<LargePage> largePages_;
+    AtomicStack<ExtraObjectPage> extraObjectPages_;
+    AtomicStack<ExtraObjectPage> usedExtraObjectPages_;
 };
 
 } // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/HeapTest.cpp

@@ -19,7 +19,7 @@ using SmallPage = typename kotlin::alloc::SmallPage;
 using MediumPage = typename kotlin::alloc::MediumPage;
 using LargePage = typename kotlin::alloc::LargePage;
 
-#define MIN_BLOCK_SIZE 2
+inline constexpr int MIN_BLOCK_SIZE = 2;
 
 void mark(void* obj) {
     reinterpret_cast<uint64_t*>(obj)[0] = 1;

kotlin-native/runtime/src/custom_alloc/cpp/LargePage.hpp

@@ -10,6 +10,7 @@
 #include <cstdint>
 
 #include "AtomicStack.hpp"
+#include "MediumPage.hpp"
 
 namespace kotlin::alloc {
 
@@ -19,10 +20,10 @@ public:
 
     void Destroy() noexcept;
 
-    uint8_t* TryAllocate() noexcept;
-
     uint8_t* Data() noexcept;
 
+    uint8_t* TryAllocate() noexcept;
+
     bool Sweep() noexcept;
 
 private:

kotlin-native/runtime/src/custom_alloc/cpp/MediumPage.hpp

@@ -41,6 +41,7 @@ private:
     Cell* curBlock_;
     Cell cells_[]; // cells_[0] is reserved for an empty block
 };
+
 } // namespace kotlin::alloc
 
 #endif

kotlin-native/runtime/src/custom_alloc/cpp/MediumPageTest.cpp

@@ -19,7 +19,7 @@ using Cell = typename kotlin::alloc::Cell;
 
 TypeInfo fakeType = {.flags_ = 0}; // a type without a finalizer
 
-#define MIN_BLOCK_SIZE (SMALL_PAGE_MAX_BLOCK_SIZE + 1)
+inline constexpr const size_t MIN_BLOCK_SIZE = SMALL_PAGE_MAX_BLOCK_SIZE + 1;
 
 void mark(void* obj) {
     reinterpret_cast<uint64_t*>(obj)[0] = 1;

kotlin-native/runtime/src/custom_alloc/cpp/PageStore.hpp

@@ -16,8 +16,8 @@ template <class T>
 class PageStore {
 public:
     void PrepareForGC() noexcept {
-        unswept_.TransferAllFrom(used_);
-        unswept_.TransferAllFrom(ready_);
+        unswept_.TransferAllFrom(std::move(ready_));
+        unswept_.TransferAllFrom(std::move(used_));
         T* page;
         while ((page = empty_.Pop())) page->Destroy();
     }

kotlin-native/runtime/src/custom_alloc/cpp/SmallPage.cpp

@@ -64,6 +64,7 @@ bool SmallPage::Sweep() noexcept {
             alive = true;
             continue;
         }
+        CustomAllocInfo("SmallPage(%p)::Sweep: reclaim %p", this, cell);
         // Free the current block and insert it into the free list.
         cell->nextFree = *nextFree;
         *nextFree = cell;

kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMarkAndSweep.cpp

@@ -13,8 +13,6 @@
 #include "Logging.hpp"
 #include "MarkAndSweepUtils.hpp"
 #include "Memory.h"
-#include "RootSet.hpp"
-#include "Runtime.h"
 #include "ThreadData.hpp"
 #include "ThreadRegistry.hpp"
 #include "ThreadSuspension.hpp"
@@ -23,6 +21,7 @@
 #include "GCStatistics.hpp"
 
 #ifdef CUSTOM_ALLOCATOR
+#include "CustomFinalizerProcessor.hpp"
 #include "Heap.hpp"
 #endif
 
@@ -100,9 +99,11 @@ gc::ConcurrentMarkAndSweep::ConcurrentMarkAndSweep(
         mm::ObjectFactory<ConcurrentMarkAndSweep>& objectFactory, GCScheduler& gcScheduler) noexcept :
 #ifndef CUSTOM_ALLOCATOR
     objectFactory_(objectFactory),
-#endif
     gcScheduler_(gcScheduler),
     finalizerProcessor_(std_support::make_unique<FinalizerProcessor>([this](int64_t epoch) {
+#else
+    gcScheduler_(gcScheduler), finalizerProcessor_(std_support::make_unique<alloc::CustomFinalizerProcessor>([this](int64_t epoch) {
+#endif
         GCHandle::getByEpoch(epoch).finalizersDone();
         state_.finalized(epoch);
     })) {
@@ -175,24 +176,25 @@ bool gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
     gc::Mark<internal::MarkTraits>(gcHandle, markQueue_);
 
     mm::WaitForThreadsSuspension();
-    mm::ExtraObjectDataFactory& extraObjectDataFactory = mm::GlobalData::Instance().extraObjectDataFactory();
     auto markStats = gcHandle.getMarked();
     scheduler.gcData().UpdateAliveSetBytes(markStats.totalObjectsSize);
 
+#ifndef CUSTOM_ALLOCATOR
+    mm::ExtraObjectDataFactory& extraObjectDataFactory = mm::GlobalData::Instance().extraObjectDataFactory();
     gc::SweepExtraObjects<SweepTraits>(gcHandle, extraObjectDataFactory);
 
-#ifndef CUSTOM_ALLOCATOR
     auto objectFactoryIterable = objectFactory_.LockForIter();
     mm::ResumeThreads();
     gcHandle.threadsAreResumed();
     auto finalizerQueue = gc::Sweep<SweepTraits>(gcHandle, objectFactoryIterable);
+    kotlin::compactObjectPoolInMainThread();
 #else
+    auto finalizerQueue = heap_.SweepExtraObjects(gcHandle);
+
     mm::ResumeThreads();
     gcHandle.threadsAreResumed();
-    SweepTraits::ObjectFactory::FinalizerQueue finalizerQueue;
     heap_.Sweep();
 #endif
-    kotlin::compactObjectPoolInMainThread();
     state_.finish(epoch);
     gcHandle.finalizersScheduled(finalizerQueue.size());
     gcHandle.finished();

kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMarkAndSweep.hpp

@@ -24,12 +24,18 @@
 #ifdef CUSTOM_ALLOCATOR
 #include "CustomAllocator.hpp"
 #include "Heap.hpp"
+
+namespace kotlin::alloc {
+class CustomFinalizerProcessor;
+}
 #endif
 
 namespace kotlin {
 namespace gc {
 
+#ifndef CUSTOM_ALLOCATOR
 class FinalizerProcessor;
+#endif
 
 // Stop-the-world parallel mark + concurrent sweep. The GC runs in a separate thread, finalizers run in another thread of their own.
 // TODO: Also make marking run concurrently with Kotlin threads.
@@ -130,7 +136,11 @@ private:
 
     GCStateHolder state_;
     ScopedThread gcThread_;
+#ifndef CUSTOM_ALLOCATOR
     std_support::unique_ptr<FinalizerProcessor> finalizerProcessor_;
+#else
+    std_support::unique_ptr<alloc::CustomFinalizerProcessor> finalizerProcessor_;
+#endif
 
     MarkQueue markQueue_;
     MarkingBehavior markingBehavior_;

kotlin-native/runtime/src/mm/cpp/ExtraObjectData.cpp

@@ -5,15 +5,18 @@
 
 #include "ExtraObjectData.hpp"
 
-#include "ObjectOps.hpp"
 #include "PointerBits.h"
+#include "ThreadData.hpp"
 #include "Weak.h"
-#include "ExtraObjectDataFactory.hpp"
 
 #ifdef KONAN_OBJC_INTEROP
 #include "ObjCMMAPI.h"
 #endif
 
+#ifdef CUSTOM_ALLOCATOR
+#include "CustomAllocator.hpp"
+#endif
+
 using namespace kotlin;
 
 // static
@@ -30,11 +33,19 @@ mm::ExtraObjectData& mm::ExtraObjectData::Install(ObjHeader* object) noexcept {
     RuntimeCheck(!hasPointerBits(typeInfo, OBJECT_TAG_MASK), "Object must not be tagged");
 
     auto *threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
+#ifdef CUSTOM_ALLOCATOR
+    auto& data = alloc::CustomAllocator::CreateExtraObjectDataForObject(threadData, object, typeInfo);
+
+    if (!compareExchange(object->typeInfoOrMeta_, typeInfo, reinterpret_cast<TypeInfo*>(&data))) {
+        // Somebody else created `mm::ExtraObjectData` for this object.
+        data.setFlag(mm::ExtraObjectData::FLAGS_FINALIZED);
+#else
     auto& data = mm::ExtraObjectDataFactory::Instance().CreateExtraObjectDataForObject(threadData, object, typeInfo);
 
     if (!compareExchange(object->typeInfoOrMeta_, typeInfo, reinterpret_cast<TypeInfo*>(&data))) {
         // Somebody else created `mm::ExtraObjectData` for this object
         mm::ExtraObjectDataFactory::Instance().DestroyExtraObjectData(threadData, data);
+#endif
         return *reinterpret_cast<mm::ExtraObjectData*>(typeInfo);
     }
 

kotlin-native/runtime/src/mm/cpp/ExtraObjectData.hpp

@@ -13,7 +13,6 @@
 #include "Memory.h"
 #include "TypeInfo.h"
 #include "Utils.hpp"
-#include "MultiSourceQueue.hpp"
 #include "Weak.h"
 
 namespace kotlin {
@@ -27,6 +26,7 @@ public:
         FLAGS_FROZEN = 0,
         FLAGS_NEVER_FROZEN = 1,
         FLAGS_IN_FINALIZER_QUEUE = 2,
+        FLAGS_FINALIZED = 3,
     };
 
     static constexpr unsigned WEAK_REF_TAG = 1;
@@ -83,7 +83,6 @@ public:
     }
     ~ExtraObjectData();
 private:
-
     // Must be first to match `TypeInfo` layout.
     const TypeInfo* typeInfo_;
 

kotlin-native/runtime/src/mm/cpp/Memory.cpp

@@ -86,8 +86,10 @@ void ObjHeader::destroyMetaObject(ObjHeader* object) {
     RuntimeAssert(object->has_meta_object(), "Object must have a meta object set");
     auto &extraObject = *mm::ExtraObjectData::Get(object);
     extraObject.Uninstall();
+#ifndef CUSTOM_ALLOCATOR
     auto *threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
     mm::ExtraObjectDataFactory::Instance().DestroyExtraObjectData(threadData, extraObject);
+#endif
 }
 
 ALWAYS_INLINE bool isPermanentOrFrozen(const ObjHeader* obj) {