[K/N] concurrent weak sweep ^KT-57772

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/BinaryOptions.kt

@@ -39,6 +39,8 @@ object BinaryOptions : BinaryOptionRegistry() {
 
     val gcMarkSingleThreaded by booleanOption()
 
+    val concurrentWeakSweep by booleanOption()
+
     val linkRuntime by option<RuntimeLinkageStrategyBinaryOption>()
 
     val bundleId by stringOption()

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/KonanConfig.kt

@@ -143,7 +143,10 @@ class KonanConfig(val project: Project, val configuration: CompilerConfiguration
     }
 
     val gcMarkSingleThreaded: Boolean
-        get() = configuration.get(BinaryOptions.gcMarkSingleThreaded) == true
+        get() = configuration.get(BinaryOptions.gcMarkSingleThreaded) ?: false
+
+    val concurrentWeakSweep: Boolean
+        get() = configuration.get(BinaryOptions.concurrentWeakSweep) ?: false
 
     val irVerificationMode: IrVerificationMode
         get() = configuration.getNotNull(KonanConfigKeys.VERIFY_IR)

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/IrToBitcode.kt

@@ -2998,6 +2998,7 @@ internal fun NativeGenerationState.generateRuntimeConstantsModule() : LLVMModule
     setRuntimeConstGlobal("Kotlin_runtimeLogs", runtimeLogs)
     setRuntimeConstGlobal("Kotlin_freezingEnabled", llvm.constInt32(if (config.freezing.enableFreezeAtRuntime) 1 else 0))
     setRuntimeConstGlobal("Kotlin_freezingChecksEnabled", llvm.constInt32(if (config.freezing.enableFreezeChecks) 1 else 0))
+    setRuntimeConstGlobal("Kotlin_concurrentWeakSweep", llvm.constInt32(if (context.config.concurrentWeakSweep) 1 else 0))
 
     return llvmModule
 }

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

@@ -0,0 +1,97 @@
+/*
+ * Copyright 2010-2023 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 "Barriers.hpp"
+
+#include <algorithm>
+#include <atomic>
+
+#include "GCImpl.hpp"
+#include "SafePoint.hpp"
+#include "ThreadData.hpp"
+#include "ThreadRegistry.hpp"
+
+using namespace kotlin;
+
+namespace {
+
+std::atomic<ObjHeader* (*)(ObjHeader*)> weakRefBarrier = nullptr;
+
+ObjHeader* weakRefBarrierImpl(ObjHeader* weakReferee) noexcept {
+    if (!weakReferee) return nullptr;
+    // When weak ref barriers are enabled, marked state cannot change and the
+    // object cannot be deleted.
+    if (!gc::isMarked(weakReferee)) {
+        return nullptr;
+    }
+    return weakReferee;
+}
+
+NO_INLINE ObjHeader* weakRefReadSlowPath(std::atomic<ObjHeader*>& weakReferee) noexcept {
+    // reread an action to avoid register pollution outside the function
+    auto barrier = weakRefBarrier.load(std::memory_order_seq_cst);
+    auto* weak = weakReferee.load(std::memory_order_relaxed);
+    return barrier ? barrier(weak) : weak;
+}
+
+void waitForThreadsToReachCheckpoint() {
+    // Reset checkpoint on all threads.
+    for (auto& thr : mm::ThreadRegistry::Instance().LockForIter()) {
+        thr.gc().impl().gc().barriers().resetCheckpoint();
+    }
+
+    mm::SafePointActivator safePointActivator;
+
+    // Disable new threads coming and going.
+    auto threads = mm::ThreadRegistry::Instance().LockForIter();
+    // And wait for all threads to either have passed safepoint or to be in the native state.
+    // Either of these mean that none of them are inside a weak reference accessing code.
+    while (!std::all_of(threads.begin(), threads.end(), [](mm::ThreadData& thread) noexcept {
+        return thread.gc().impl().gc().barriers().visitedCheckpoint() || thread.suspensionData().suspendedOrNative();
+    })) {
+        std::this_thread::yield();
+    }
+}
+
+} // namespace
+
+void gc::BarriersThreadData::onCheckpoint() noexcept {
+    visitedCheckpoint_.store(true, std::memory_order_release);
+}
+
+void gc::BarriersThreadData::resetCheckpoint() noexcept {
+    visitedCheckpoint_.store(false, std::memory_order_release);
+}
+
+bool gc::BarriersThreadData::visitedCheckpoint() const noexcept {
+    return visitedCheckpoint_.load(std::memory_order_acquire);
+}
+
+void gc::EnableWeakRefBarriers() noexcept {
+    weakRefBarrier.store(weakRefBarrierImpl, std::memory_order_seq_cst);
+}
+
+void gc::DisableWeakRefBarriers() noexcept {
+    weakRefBarrier.store(nullptr, std::memory_order_seq_cst);
+    waitForThreadsToReachCheckpoint();
+}
+
+OBJ_GETTER(kotlin::gc::WeakRefRead, std::atomic<ObjHeader*>& weakReferee) noexcept {
+    // TODO: Make this work with GCs that can stop thread at any point.
+
+    if (!compiler::concurrentWeakSweep()) {
+        RETURN_OBJ(weakReferee.load(std::memory_order_relaxed));
+    }
+
+    // Copying the scheme from SafePoint.cpp: branch + indirect call.
+    auto barrier = weakRefBarrier.load(std::memory_order_relaxed);
+    ObjHeader* result;
+    if (__builtin_expect(barrier != nullptr, false)) {
+        result = weakRefReadSlowPath(weakReferee);
+    } else {
+        result = weakReferee.load(std::memory_order_relaxed);
+    }
+    RETURN_OBJ(result);
+}

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

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2010-2023 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.
+ */
+
+#pragma once
+
+#include <atomic>
+#include <optional>
+
+#include "Memory.h"
+#include "Utils.hpp"
+
+namespace kotlin::gc {
+
+class BarriersThreadData : private Pinned {
+public:
+    void onCheckpoint() noexcept;
+    void resetCheckpoint() noexcept;
+    bool visitedCheckpoint() const noexcept;
+
+private:
+    std::atomic<bool> visitedCheckpoint_ = false;
+};
+
+// Must be called during STW.
+void EnableWeakRefBarriers() noexcept;
+
+// Must be called outside STW.
+void DisableWeakRefBarriers() noexcept;
+
+OBJ_GETTER(WeakRefRead, std::atomic<ObjHeader*>& weakReferee) noexcept;
+
+} // namespace kotlin::gc

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

@@ -184,25 +184,38 @@ void gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
     auto markStats = gcHandle.getMarked();
     scheduler.gcData().UpdateAliveSetBytes(markStats.markedSizeBytes);
 
-    gc::processWeaks<ProcessWeaksTraits>(gcHandle, mm::SpecialRefRegistry::instance());
-
 #ifndef CUSTOM_ALLOCATOR
     // Taking the locks before the pause is completed. So that any destroying thread
     // would not publish into the global state at an unexpected time.
     std::optional extraObjectFactoryIterable = mm::GlobalData::Instance().extraObjectDataFactory().LockForIter();
     std::optional objectFactoryIterable = objectFactory_.LockForIter();
-    mm::ResumeThreads();
-    gcHandle.threadsAreResumed();
+#endif
 
+    if (compiler::concurrentWeakSweep()) {
+        // Expected to happen inside STW.
+        gc::EnableWeakRefBarriers();
+
+        mm::ResumeThreads();
+        gcHandle.threadsAreResumed();
+    }
+
+    gc::processWeaks<ProcessWeaksTraits>(gcHandle, mm::SpecialRefRegistry::instance());
+
+    if (compiler::concurrentWeakSweep()) {
+        // Expected to happen outside STW.
+        gc::DisableWeakRefBarriers();
+    } else {
+        mm::ResumeThreads();
+        gcHandle.threadsAreResumed();
+    }
+
+#ifndef CUSTOM_ALLOCATOR
     gc::SweepExtraObjects<SweepTraits>(gcHandle, *extraObjectFactoryIterable);
     extraObjectFactoryIterable = std::nullopt;
     auto finalizerQueue = gc::Sweep<SweepTraits>(gcHandle, *objectFactoryIterable);
     objectFactoryIterable = std::nullopt;
     kotlin::compactObjectPoolInMainThread();
 #else
-    mm::ResumeThreads();
-    gcHandle.threadsAreResumed();
-
     // also sweeps extraObjects
     auto finalizerQueue = heap_.Sweep(gcHandle);
 #endif

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

@@ -9,6 +9,7 @@
 #include <cstddef>
 
 #include "Allocator.hpp"
+#include "Barriers.hpp"
 #include "FinalizerProcessor.hpp"
 #include "GCScheduler.hpp"
 #include "GCState.hpp"
@@ -91,14 +92,19 @@ public:
 
         void OnSuspendForGC() noexcept;
 
+        void safePoint() noexcept { barriers_.onCheckpoint(); }
+
         Allocator CreateAllocator() noexcept { return Allocator(gc::Allocator(), *this); }
 
+        BarriersThreadData& barriers() noexcept { return barriers_; }
+
     private:
         friend ConcurrentMarkAndSweep;
         ConcurrentMarkAndSweep& gc_;
         mm::ThreadData& threadData_;
         gcScheduler::GCSchedulerThreadData& gcScheduler_;
         std::atomic<bool> marking_;
+        BarriersThreadData barriers_;
     };
 
     using Allocator = ThreadData::Allocator;
@@ -130,7 +136,8 @@ public:
     alloc::Heap& heap() noexcept { return heap_; }
 #endif
 
-    void Schedule() noexcept { state_.schedule(); }
+    int64_t Schedule() noexcept { return state_.schedule(); }
+    void WaitFinalized(int64_t epoch) noexcept { state_.waitEpochFinalized(epoch); }
 
 private:
     void PerformFullGC(int64_t epoch) noexcept;

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

@@ -728,17 +728,19 @@ TEST_P(ConcurrentMarkAndSweepTest, MultipleMutatorsCollect) {
                 .wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
+    std_support::vector<std::future<void>> gcFutures;
 
-    gcFutures[0] = mutators[0].Execute([](mm::ThreadData& threadData, Mutator& mutator) { threadData.gc().ScheduleAndWaitFullGCWithFinalizers(); });
-
-    // Spin until thread suspension is requested.
-    while (!mm::IsThreadSuspensionRequested()) {
+    std::atomic<bool> gcDone = false;
+    for (auto& mutator : mutators) {
+        gcFutures.emplace_back(mutator.Execute([&](mm::ThreadData& threadData, Mutator& mutator) noexcept {
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+            }
+        }));
     }
 
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([](mm::ThreadData& threadData, Mutator& mutator) { mm::safePoint(threadData); });
-    }
+    mm::GlobalData::Instance().gc().ScheduleAndWaitFullGCWithFinalizers();
+    gcDone.store(true, std::memory_order_relaxed);
 
     for (auto& future : gcFutures) {
         future.wait();
@@ -783,15 +785,15 @@ TEST_P(ConcurrentMarkAndSweepTest, MultipleMutatorsAllCollect) {
                 .wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
+    std_support::vector<std::future<void>> gcFutures;
 
-    for (int i = 0; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([](mm::ThreadData& threadData, Mutator& mutator) {
+    for (auto& mutator : mutators) {
+        gcFutures.emplace_back(mutator.Execute([](mm::ThreadData& threadData, Mutator& mutator) {
             threadData.gc().ScheduleAndWaitFullGCWithFinalizers();
             // If GC starts before all thread executed line above, two gc will be run
-            // So we are temporary switch threads to native state and then return them back after all GC runs are done
+            // So we temporary switch threads to native state and then return them back after all GC runs are done
             SwitchThreadState(mm::GetMemoryState(), kotlin::ThreadState::kNative);
-        });
+        }));
     }
 
     for (auto& future : gcFutures) {
@@ -846,34 +848,33 @@ TEST_P(ConcurrentMarkAndSweepTest, MultipleMutatorsAddToRootSetAfterCollectionRe
         mutator.AddStackRoot(locals[i]);
     };
 
-    mutators[0]
-            .Execute([expandRootSet, allocateInHeap](mm::ThreadData& threadData, Mutator& mutator) {
-                allocateInHeap(threadData, mutator, 0);
-                expandRootSet(threadData, mutator, 0);
-            })
-            .wait();
-
     // Allocate everything in heap before scheduling the GC.
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
         mutators[i]
                 .Execute([allocateInHeap, i](mm::ThreadData& threadData, Mutator& mutator) { allocateInHeap(threadData, mutator, i); })
                 .wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
-    gcFutures[0] = mutators[0].Execute([](mm::ThreadData& threadData, Mutator& mutator) { threadData.gc().ScheduleAndWaitFullGCWithFinalizers(); });
+    std_support::vector<std::future<void>> gcFutures;
+    auto epoch = mm::GlobalData::Instance().gc().Schedule();
+    std::atomic<bool> gcDone = false;
 
     // Spin until thread suspension is requested.
     while (!mm::IsThreadSuspensionRequested()) {
     }
 
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([i, expandRootSet](mm::ThreadData& threadData, Mutator& mutator) {
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        gcFutures.emplace_back(mutators[i].Execute([&gcDone, i, expandRootSet](mm::ThreadData& threadData, Mutator& mutator) {
             expandRootSet(threadData, mutator, i);
-            mm::safePoint(threadData);
-        });
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+            }
+        }));
     }
 
+    mm::GlobalData::Instance().gc().WaitFinalizers(epoch);
+    gcDone.store(true, std::memory_order_relaxed);
+
     for (auto& future : gcFutures) {
         future.wait();
     }
@@ -924,17 +925,19 @@ TEST_P(ConcurrentMarkAndSweepTest, CrossThreadReference) {
                 .wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
+    std_support::vector<std::future<void>> gcFutures;
 
-    gcFutures[0] = mutators[0].Execute([](mm::ThreadData& threadData, Mutator& mutator) { threadData.gc().ScheduleAndWaitFullGCWithFinalizers(); });
-
-    // Spin until thread suspension is requested.
-    while (!mm::IsThreadSuspensionRequested()) {
+    std::atomic<bool> gcDone = false;
+    for (auto& mutator : mutators) {
+        gcFutures.emplace_back(mutator.Execute([&](mm::ThreadData& threadData, Mutator& mutator) noexcept {
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+            }
+        }));
     }
 
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([](mm::ThreadData& threadData, Mutator& mutator) { mm::safePoint(threadData); });
-    }
+    mm::GlobalData::Instance().gc().ScheduleAndWaitFullGCWithFinalizers();
+    gcDone.store(true, std::memory_order_relaxed);
 
     for (auto& future : gcFutures) {
         future.wait();
@@ -984,24 +987,26 @@ TEST_P(ConcurrentMarkAndSweepTest, MultipleMutatorsWeaks) {
         mutators[i].Execute([](mm::ThreadData& threadData, Mutator& mutator) {}).wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
-
-    gcFutures[0] = mutators[0].Execute([weak](mm::ThreadData& threadData, Mutator& mutator) {
-        threadData.gc().ScheduleAndWaitFullGCWithFinalizers();
-        EXPECT_THAT(weak->get(), nullptr);
-    });
+    std_support::vector<std::future<void>> gcFutures;
+    auto epoch = mm::GlobalData::Instance().gc().Schedule();
+    std::atomic<bool> gcDone = false;
 
     // Spin until thread suspension is requested.
     while (!mm::IsThreadSuspensionRequested()) {
     }
 
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([weak](mm::ThreadData& threadData, Mutator& mutator) {
-            mm::safePoint(threadData);
-            EXPECT_THAT(weak->get(), nullptr);
-        });
+    for (auto& mutator : mutators) {
+        gcFutures.emplace_back(mutator.Execute([&](mm::ThreadData& threadData, Mutator& mutator) noexcept {
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+                EXPECT_THAT(weak->get(), nullptr);
+            }
+        }));
     }
 
+    mm::GlobalData::Instance().gc().WaitFinalizers(epoch);
+    gcDone.store(true, std::memory_order_relaxed);
+
     for (auto& future : gcFutures) {
         future.wait();
     }
@@ -1036,9 +1041,9 @@ TEST_P(ConcurrentMarkAndSweepTest, NewThreadsWhileRequestingCollection) {
                 .wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
-
-    gcFutures[0] = mutators[0].Execute([](mm::ThreadData& threadData, Mutator& mutator) { threadData.gc().ScheduleAndWaitFullGCWithFinalizers(); });
+    std_support::vector<std::future<void>> gcFutures;
+    auto epoch = mm::GlobalData::Instance().gc().Schedule();
+    std::atomic<bool> gcDone = false;
 
     // Spin until thread suspension is requested.
     while (!mm::IsThreadSuspensionRequested()) {
@@ -1049,14 +1054,27 @@ TEST_P(ConcurrentMarkAndSweepTest, NewThreadsWhileRequestingCollection) {
     std_support::vector<std::future<void>> attachFutures(kDefaultThreadCount);
 
     for (int i = 0; i < kDefaultThreadCount; ++i) {
-        attachFutures[i] = newMutators[i].Execute([i, expandRootSet](mm::ThreadData& threadData, Mutator& mutator) { expandRootSet(threadData, mutator, i + kDefaultThreadCount); });
+        attachFutures[i] = newMutators[i].Execute([&gcDone, i, expandRootSet](mm::ThreadData& threadData, Mutator& mutator) {
+            expandRootSet(threadData, mutator, i + kDefaultThreadCount);
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+            }
+        });
     }
 
     // All the other threads are stopping at safe points.
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([](mm::ThreadData& threadData, Mutator& mutator) { mm::safePoint(threadData); });
+    for (auto& mutator : mutators) {
+        gcFutures.emplace_back(mutator.Execute([&gcDone](mm::ThreadData& threadData, Mutator& mutator) {
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+            }
+        }));
     }
 
+    // Wait for the GC to be done.
+    mm::GlobalData::Instance().gc().WaitFinalizers(epoch);
+    gcDone.store(true, std::memory_order_relaxed);
+
     // GC will be completed first
     for (auto& future : gcFutures) {
         future.wait();

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

@@ -8,6 +8,7 @@
 #include "GC.hpp"
 #include "GCStatistics.hpp"
 #include "MarkAndSweepUtils.hpp"
+#include "ObjectOps.hpp"
 #include "ThreadSuspension.hpp"
 #include "std_support/Memory.hpp"
 
@@ -62,6 +63,10 @@ void gc::GC::ThreadData::OnSuspendForGC() noexcept {
     impl_->gc().OnSuspendForGC();
 }
 
+void gc::GC::ThreadData::safePoint() noexcept {
+    impl_->gc().safePoint();
+}
+
 gc::GC::GC(gcScheduler::GCScheduler& gcScheduler) noexcept : impl_(std_support::make_unique<Impl>(gcScheduler)) {}
 
 gc::GC::~GC() = default;
@@ -114,6 +119,19 @@ ALWAYS_INLINE void gc::GC::processFieldInMark(void* state, ObjHeader* field) noe
     gc::internal::processFieldInMark<gc::internal::MarkTraits>(state, field);
 }
 
-void gc::GC::Schedule() noexcept {
-    impl_->gc().Schedule();
+int64_t gc::GC::Schedule() noexcept {
+    return impl_->gc().Schedule();
+}
+
+void gc::GC::WaitFinalizers(int64_t epoch) noexcept {
+    impl_->gc().WaitFinalized(epoch);
+}
+
+bool gc::isMarked(ObjHeader* object) noexcept {
+    auto& objectData = mm::ObjectFactory<gc::ConcurrentMarkAndSweep>::NodeRef::From(object).ObjectData();
+    return objectData.marked();
+}
+
+ALWAYS_INLINE OBJ_GETTER(gc::tryRef, std::atomic<ObjHeader*>& object) noexcept {
+    RETURN_RESULT_OF(gc::WeakRefRead, object);
 }

kotlin-native/runtime/src/gc/common/cpp/GC.hpp

@@ -5,6 +5,8 @@
 
 #pragma once
 
+#include <atomic>
+
 #include "GCScheduler.hpp"
 #include "Memory.h"
 #include "Types.h"
@@ -44,6 +46,8 @@ public:
 
         void OnSuspendForGC() noexcept;
 
+        void safePoint() noexcept;
+
     private:
         std_support::unique_ptr<Impl> impl_;
     };
@@ -67,13 +71,18 @@ public:
     static void processArrayInMark(void* state, ArrayHeader* array) noexcept;
     static void processFieldInMark(void* state, ObjHeader* field) noexcept;
 
-    // TODO: This should be moved into the scheduler.
-    void Schedule() noexcept;
+    // TODO: These should be moved into the scheduler.
+    int64_t Schedule() noexcept;
+    void WaitFinalizers(int64_t epoch) noexcept;
+    void ScheduleAndWaitFullGCWithFinalizers() noexcept { WaitFinalizers(Schedule()); }
 
 private:
     std_support::unique_ptr<Impl> impl_;
 };
 
+bool isMarked(ObjHeader* object) noexcept;
+OBJ_GETTER(tryRef, std::atomic<ObjHeader*>& object) noexcept;
+
 inline constexpr bool kSupportsMultipleMutators = true;
 
 } // namespace gc

kotlin-native/runtime/src/gc/common/cpp/MarkAndSweepUtils.hpp

@@ -220,7 +220,7 @@ template <typename Traits>
 void processWeaks(GCHandle gcHandle, mm::SpecialRefRegistry& registry) noexcept {
     auto handle = gcHandle.processWeaks();
     for (auto& object : registry.lockForIter()) {
-        auto* obj = object;
+        auto* obj = object.load(std::memory_order_relaxed);
         if (!obj) {
             // We already processed it at some point.
             handle.addUndisposed();
@@ -233,7 +233,7 @@ void processWeaks(GCHandle gcHandle, mm::SpecialRefRegistry& registry) noexcept
             continue;
         }
         // Object is not alive. Clear it out.
-        object = nullptr;
+        object.store(nullptr, std::memory_order_relaxed);
         handle.addNulled();
     }
 }

kotlin-native/runtime/src/gc/noop/cpp/GCImpl.cpp

@@ -9,6 +9,7 @@
 #include "std_support/Memory.hpp"
 #include "GlobalData.hpp"
 #include "GCStatistics.hpp"
+#include "ObjectOps.hpp"
 
 using namespace kotlin;
 
@@ -47,6 +48,8 @@ ALWAYS_INLINE ArrayHeader* gc::GC::ThreadData::CreateArray(const TypeInfo* typeI
 
 void gc::GC::ThreadData::OnSuspendForGC() noexcept { }
 
+void gc::GC::ThreadData::safePoint() noexcept {}
+
 gc::GC::GC(gcScheduler::GCScheduler&) noexcept : impl_(std_support::make_unique<Impl>()) {}
 
 gc::GC::~GC() = default;
@@ -82,4 +85,16 @@ ALWAYS_INLINE void gc::GC::processArrayInMark(void* state, ArrayHeader* array) n
 // static
 ALWAYS_INLINE void gc::GC::processFieldInMark(void* state, ObjHeader* field) noexcept {}
 
-void gc::GC::Schedule() noexcept {}
+int64_t gc::GC::Schedule() noexcept {
+    return 0;
+}
+void gc::GC::WaitFinalizers(int64_t epoch) noexcept {}
+
+bool gc::isMarked(ObjHeader* object) noexcept {
+    RuntimeAssert(false, "Should not reach here");
+    return true;
+}
+
+ALWAYS_INLINE OBJ_GETTER(gc::tryRef, std::atomic<ObjHeader*>& object) noexcept {
+    RETURN_OBJ(object.load(std::memory_order_relaxed));
+}

kotlin-native/runtime/src/gc/stms/cpp/GCImpl.cpp

@@ -10,6 +10,7 @@
 #include "std_support/Memory.hpp"
 #include "GlobalData.hpp"
 #include "GCStatistics.hpp"
+#include "ObjectOps.hpp"
 
 using namespace kotlin;
 
@@ -48,6 +49,8 @@ ALWAYS_INLINE ArrayHeader* gc::GC::ThreadData::CreateArray(const TypeInfo* typeI
 
 void gc::GC::ThreadData::OnSuspendForGC() noexcept { }
 
+void gc::GC::ThreadData::safePoint() noexcept {}
+
 gc::GC::GC(gcScheduler::GCScheduler& gcScheduler) noexcept : impl_(std_support::make_unique<Impl>(gcScheduler)) {}
 
 gc::GC::~GC() = default;
@@ -94,6 +97,19 @@ ALWAYS_INLINE void gc::GC::processFieldInMark(void* state, ObjHeader* field) noe
     gc::internal::processFieldInMark<gc::internal::MarkTraits>(state, field);
 }
 
-void gc::GC::Schedule() noexcept {
-    impl_->gc().Schedule();
+int64_t gc::GC::Schedule() noexcept {
+    return impl_->gc().Schedule();
+}
+
+void gc::GC::WaitFinalizers(int64_t epoch) noexcept {
+    impl_->gc().WaitFinalized(epoch);
+}
+
+bool gc::isMarked(ObjHeader* object) noexcept {
+    auto& objectData = mm::ObjectFactory<gc::SameThreadMarkAndSweep>::NodeRef::From(object).ObjectData();
+    return objectData.marked();
+}
+
+ALWAYS_INLINE OBJ_GETTER(gc::tryRef, std::atomic<ObjHeader*>& object) noexcept {
+    RETURN_OBJ(object.load(std::memory_order_relaxed));
 }

kotlin-native/runtime/src/gc/stms/cpp/SameThreadMarkAndSweep.hpp

@@ -102,7 +102,8 @@ public:
     void StopFinalizerThreadIfRunning() noexcept;
     bool FinalizersThreadIsRunning() noexcept;
 
-    void Schedule() noexcept { state_.schedule(); }
+    int64_t Schedule() noexcept { return state_.schedule(); }
+    void WaitFinalized(int64_t epoch) noexcept { state_.waitEpochFinalized(epoch); }
 
 private:
     void PerformFullGC(int64_t epoch) noexcept;

kotlin-native/runtime/src/gc/stms/cpp/SameThreadMarkAndSweepTest.cpp

@@ -724,18 +724,19 @@ TEST_F(SameThreadMarkAndSweepTest, MultipleMutatorsCollect) {
                 .wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
+    std_support::vector<std::future<void>> gcFutures;
 
-    gcFutures[0] = mutators[0].Execute(
-            [](mm::ThreadData& threadData, Mutator& mutator) { threadData.gc().ScheduleAndWaitFullGCWithFinalizers(); });
-
-    // Spin until thread suspension is requested.
-    while (!mm::IsThreadSuspensionRequested()) {
+    std::atomic<bool> gcDone = false;
+    for (auto& mutator : mutators) {
+        gcFutures.emplace_back(mutator.Execute([&](mm::ThreadData& threadData, Mutator& mutator) noexcept {
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+            }
+        }));
     }
 
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([](mm::ThreadData& threadData, Mutator& mutator) { mm::safePoint(threadData); });
-    }
+    mm::GlobalData::Instance().gc().ScheduleAndWaitFullGCWithFinalizers();
+    gcDone.store(true, std::memory_order_relaxed);
 
     for (auto& future : gcFutures) {
         future.wait();
@@ -780,15 +781,15 @@ TEST_F(SameThreadMarkAndSweepTest, MultipleMutatorsAllCollect) {
                 .wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
+    std_support::vector<std::future<void>> gcFutures;
 
-    for (int i = 0; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([](mm::ThreadData& threadData, Mutator& mutator) {
+    for (auto& mutator : mutators) {
+        gcFutures.emplace_back(mutator.Execute([](mm::ThreadData& threadData, Mutator& mutator) {
             threadData.gc().ScheduleAndWaitFullGCWithFinalizers();
             // If GC starts before all thread executed line above, two gc will be run
-            // So we are temporary switch threads to native state and then return them back after all GC runs are done
+            // So we temporary switch threads to native state and then return them back after all GC runs are done
             SwitchThreadState(mm::GetMemoryState(), kotlin::ThreadState::kNative);
-        });
+        }));
     }
 
     for (auto& future : gcFutures) {
@@ -843,35 +844,33 @@ TEST_F(SameThreadMarkAndSweepTest, MultipleMutatorsAddToRootSetAfterCollectionRe
         mutator.AddStackRoot(locals[i]);
     };
 
-    mutators[0]
-            .Execute([expandRootSet, allocateInHeap](mm::ThreadData& threadData, Mutator& mutator) {
-                allocateInHeap(threadData, mutator, 0);
-                expandRootSet(threadData, mutator, 0);
-            })
-            .wait();
-
     // Allocate everything in heap before scheduling the GC.
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
         mutators[i]
                 .Execute([allocateInHeap, i](mm::ThreadData& threadData, Mutator& mutator) { allocateInHeap(threadData, mutator, i); })
                 .wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
-    gcFutures[0] = mutators[0].Execute(
-            [](mm::ThreadData& threadData, Mutator& mutator) { threadData.gc().ScheduleAndWaitFullGCWithFinalizers(); });
+    std_support::vector<std::future<void>> gcFutures;
+    auto epoch = mm::GlobalData::Instance().gc().Schedule();
+    std::atomic<bool> gcDone = false;
 
     // Spin until thread suspension is requested.
     while (!mm::IsThreadSuspensionRequested()) {
     }
 
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([i, expandRootSet](mm::ThreadData& threadData, Mutator& mutator) {
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        gcFutures.emplace_back(mutators[i].Execute([&gcDone, i, expandRootSet](mm::ThreadData& threadData, Mutator& mutator) {
             expandRootSet(threadData, mutator, i);
-            mm::safePoint(threadData);
-        });
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+            }
+        }));
     }
 
+    mm::GlobalData::Instance().gc().WaitFinalizers(epoch);
+    gcDone.store(true, std::memory_order_relaxed);
+
     for (auto& future : gcFutures) {
         future.wait();
     }
@@ -922,18 +921,19 @@ TEST_F(SameThreadMarkAndSweepTest, CrossThreadReference) {
                 .wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
+    std_support::vector<std::future<void>> gcFutures;
 
-    gcFutures[0] = mutators[0].Execute(
-            [](mm::ThreadData& threadData, Mutator& mutator) { threadData.gc().ScheduleAndWaitFullGCWithFinalizers(); });
-
-    // Spin until thread suspension is requested.
-    while (!mm::IsThreadSuspensionRequested()) {
+    std::atomic<bool> gcDone = false;
+    for (auto& mutator : mutators) {
+        gcFutures.emplace_back(mutator.Execute([&](mm::ThreadData& threadData, Mutator& mutator) noexcept {
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+            }
+        }));
     }
 
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([](mm::ThreadData& threadData, Mutator& mutator) { mm::safePoint(threadData); });
-    }
+    mm::GlobalData::Instance().gc().ScheduleAndWaitFullGCWithFinalizers();
+    gcDone.store(true, std::memory_order_relaxed);
 
     for (auto& future : gcFutures) {
         future.wait();
@@ -983,24 +983,26 @@ TEST_F(SameThreadMarkAndSweepTest, MultipleMutatorsWeaks) {
         mutators[i].Execute([](mm::ThreadData& threadData, Mutator& mutator) {}).wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
-
-    gcFutures[0] = mutators[0].Execute([weak](mm::ThreadData& threadData, Mutator& mutator) {
-        threadData.gc().ScheduleAndWaitFullGCWithFinalizers();
-        EXPECT_THAT(weak->get(), nullptr);
-    });
+    std_support::vector<std::future<void>> gcFutures;
+    auto epoch = mm::GlobalData::Instance().gc().Schedule();
+    std::atomic<bool> gcDone = false;
 
     // Spin until thread suspension is requested.
     while (!mm::IsThreadSuspensionRequested()) {
     }
 
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([weak](mm::ThreadData& threadData, Mutator& mutator) {
-            mm::safePoint(threadData);
-            EXPECT_THAT(weak->get(), nullptr);
-        });
+    for (auto& mutator : mutators) {
+        gcFutures.emplace_back(mutator.Execute([&](mm::ThreadData& threadData, Mutator& mutator) noexcept {
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+                EXPECT_THAT(weak->get(), nullptr);
+            }
+        }));
     }
 
+    mm::GlobalData::Instance().gc().WaitFinalizers(epoch);
+    gcDone.store(true, std::memory_order_relaxed);
+
     for (auto& future : gcFutures) {
         future.wait();
     }
@@ -1035,10 +1037,9 @@ TEST_F(SameThreadMarkAndSweepTest, NewThreadsWhileRequestingCollection) {
                 .wait();
     }
 
-    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
-
-    gcFutures[0] = mutators[0].Execute(
-            [](mm::ThreadData& threadData, Mutator& mutator) { threadData.gc().ScheduleAndWaitFullGCWithFinalizers(); });
+    std_support::vector<std::future<void>> gcFutures;
+    auto epoch = mm::GlobalData::Instance().gc().Schedule();
+    std::atomic<bool> gcDone = false;
 
     // Spin until thread suspension is requested.
     while (!mm::IsThreadSuspensionRequested()) {
@@ -1049,14 +1050,27 @@ TEST_F(SameThreadMarkAndSweepTest, NewThreadsWhileRequestingCollection) {
     std_support::vector<std::future<void>> attachFutures(kDefaultThreadCount);
 
     for (int i = 0; i < kDefaultThreadCount; ++i) {
-        attachFutures[i] = newMutators[i].Execute([i, expandRootSet](mm::ThreadData& threadData, Mutator& mutator) { expandRootSet(threadData, mutator, i + kDefaultThreadCount); });
+        attachFutures[i] = newMutators[i].Execute([&gcDone, i, expandRootSet](mm::ThreadData& threadData, Mutator& mutator) {
+            expandRootSet(threadData, mutator, i + kDefaultThreadCount);
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+            }
+        });
     }
 
     // All the other threads are stopping at safe points.
-    for (int i = 1; i < kDefaultThreadCount; ++i) {
-        gcFutures[i] = mutators[i].Execute([](mm::ThreadData& threadData, Mutator& mutator) { mm::safePoint(threadData); });
+    for (auto& mutator : mutators) {
+        gcFutures.emplace_back(mutator.Execute([&gcDone](mm::ThreadData& threadData, Mutator& mutator) {
+            while (!gcDone.load(std::memory_order_relaxed)) {
+                mm::safePoint(threadData);
+            }
+        }));
     }
 
+    // Wait for the GC to be done.
+    mm::GlobalData::Instance().gc().WaitFinalizers(epoch);
+    gcDone.store(true, std::memory_order_relaxed);
+
     // GC will be completed first
     for (auto& future : gcFutures) {
         future.wait();

kotlin-native/runtime/src/main/cpp/CompilerConstants.hpp

@@ -37,6 +37,7 @@ extern "C" const int32_t Kotlin_needDebugInfo;
 extern "C" const int32_t Kotlin_runtimeAssertsMode;
 extern "C" const int32_t Kotlin_disableMmap;
 extern "C" const char* const Kotlin_runtimeLogs;
+extern "C" const int32_t Kotlin_concurrentWeakSweep;
 extern "C" const int32_t Kotlin_freezingEnabled;
 extern "C" const int32_t Kotlin_freezingChecksEnabled;
 
@@ -98,6 +99,10 @@ ALWAYS_INLINE inline bool freezingChecksEnabled() noexcept {
     return Kotlin_freezingChecksEnabled != 0;
 }
 
+ALWAYS_INLINE inline bool concurrentWeakSweep() noexcept {
+    return Kotlin_concurrentWeakSweep != 0;
+}
+
 WorkerExceptionHandling workerExceptionHandling() noexcept;
 DestroyRuntimeMode destroyRuntimeMode() noexcept;
 bool gcMarkSingleThreaded() noexcept;

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

@@ -26,6 +26,7 @@ void safePointActionImpl(mm::ThreadData& threadData) noexcept {
     AutoReset guard(&recursion, true);
 
     mm::GlobalData::Instance().gcScheduler().safePoint();
+    threadData.gc().safePoint();
     threadData.suspensionData().suspendIfRequested();
 }
 

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

@@ -67,7 +67,7 @@ mm::SpecialRefRegistry::Node* mm::SpecialRefRegistry::nextRoot(Node* current) no
         auto [candidatePrev, candidateNext] = eraseFromRoots(current, candidate);
         // We removed candidate. But should we have?
         if (candidate->rc_.load(std::memory_order_relaxed) > 0) {
-            RuntimeAssert(candidate->obj_ != nullptr, "candidate cannot have a null obj_");
+            RuntimeAssert(candidate->obj_.load(std::memory_order_relaxed) != nullptr, "candidate cannot have a null obj_");
             // Ooops. Let's put it back. Okay to put into the head.
             insertIntoRootsHead(*candidate);
         }

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

@@ -7,6 +7,7 @@
 
 #include <atomic>
 
+#include "GC.hpp"
 #include "Memory.h"
 #include "RawPtr.hpp"
 #include "ThreadRegistry.hpp"
@@ -77,13 +78,14 @@ class SpecialRefRegistry : private Pinned {
             auto rc = rc_.exchange(disposedMarker, std::memory_order_release);
             if (compiler::runtimeAssertsEnabled()) {
                 if (rc > 0) {
+                    auto* obj = obj_.load(std::memory_order_relaxed);
                     // In objc export if ObjCClass extends from KtClass
                     // doing retain+autorelease inside [ObjCClass dealloc] will cause
                     // this->dispose() be called after this->retain() but before
                     // subsequent this->release().
                     // However, since this happens in dealloc, the stored object must
                     // have been cleared already.
-                    RuntimeAssert(obj_ == nullptr, "Disposing StableRef@%p with rc %d and uncleaned object %p", this, rc, obj_);
+                    RuntimeAssert(obj == nullptr, "Disposing StableRef@%p with rc %d and uncleaned object %p", this, rc, obj);
                 }
                 RuntimeAssert(rc >= 0, "Disposing StableRef@%p with rc %d", this, rc);
             }
@@ -95,13 +97,12 @@ class SpecialRefRegistry : private Pinned {
                 auto rc = rc_.load(std::memory_order_relaxed);
                 RuntimeAssert(rc >= 0, "Dereferencing StableRef@%p with rc %d", this, rc);
             }
-            return obj_;
+            return obj_.load(std::memory_order_relaxed);
         }
 
         OBJ_GETTER0(tryRef) noexcept {
             AssertThreadState(ThreadState::kRunnable);
-            // TODO: Weak read barrier with CMS.
-            RETURN_OBJ(obj_);
+            RETURN_RESULT_OF(kotlin::gc::tryRef, obj_);
         }
 
         void retainRef() noexcept {
@@ -112,7 +113,7 @@ class SpecialRefRegistry : private Pinned {
                         position_ == std_support::list<Node>::iterator{},
                         "Retaining StableRef@%p with fast deletion optimization is disallowed", this);
 
-                if (!obj_) {
+                if (!obj_.load(std::memory_order_relaxed)) {
                     // In objc export if ObjCClass extends from KtClass
                     // calling retain inside [ObjCClass dealloc] will cause
                     // node.retainRef() be called after node.obj_ was cleared but
@@ -161,7 +162,8 @@ class SpecialRefRegistry : private Pinned {
         //   be nulled, and disable the barriers when the phase is completed.
         //   Synchronization between GC and mutators happens via enabling/disabling
         //   the barriers.
-        ObjHeader* obj_ = nullptr;
+        // TODO: Try to handle it atomically only when the GC is in progress.
+        std::atomic<ObjHeader*> obj_ = nullptr;
         // Only ever updated using relaxed memory ordering. Any synchronization
         // with nextRoot_ is achieved via acquire-release of nextRoot_.
         std::atomic<Rc> rc_ = 0; // After dispose() will be disposedMarker.
@@ -244,7 +246,7 @@ public:
         ObjHeader* operator*() const noexcept {
             // Ignoring rc here. If someone nulls out rc during root
             // scanning, it's okay to be conservative and still make it a root.
-            return node_->obj_;
+            return node_->obj_.load(std::memory_order_relaxed);
         }
 
         RootsIterator& operator++() noexcept {
@@ -281,7 +283,7 @@ public:
 
     class Iterator {
     public:
-        ObjHeader*& operator*() noexcept { return iterator_->obj_; }
+        std::atomic<ObjHeader*>& operator*() noexcept { return iterator_->obj_; }
 
         Iterator& operator++() noexcept {
             iterator_ = owner_->findAliveNode(std::next(iterator_));

kotlin-native/runtime/src/test_support/cpp/CompilerGenerated.cpp

@@ -74,6 +74,7 @@ extern const int32_t Kotlin_disableMmap = 1;
 extern const int32_t Kotlin_disableMmap = 0;
 #endif
 extern const char* const Kotlin_runtimeLogs = nullptr;
+extern const int32_t Kotlin_concurrentWeakSweep = 1;
 extern const int32_t Kotlin_freezingChecksEnabled = 1;
 extern const int32_t Kotlin_freezingEnabled = 1;