[K/N] Mark barriers implementation

Merge-request: KT-MR-13409
Merged-by: Alexey Glushko <aleksei.glushko@jetbrains.com>

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

@@ -28,7 +28,9 @@ enum class LoggingTag(val ord: Int) {
     TLS(4),
     Pause(5),
     Alloc(6),
-    Balancing(7);
+    Balancing(7),
+    Barriers(8),
+    ;
 
     companion object {
         fun parse(str: String) = entries.firstOrNull {

kotlin-native/runtime/src/alloc/custom/cpp/GCApi.hpp

@@ -27,7 +27,10 @@ struct HeapObjHeader {
 
     static HeapObjHeader& from(gc::GC::ObjectData& objectData) noexcept { return *descriptor().fromField<0>(&objectData); }
 
-    static HeapObjHeader& from(ObjHeader* object) noexcept { return *descriptor().fromField<1>(object); }
+    static HeapObjHeader& from(ObjHeader* object) noexcept {
+        RuntimeAssert(object->heap(), "Object %p does not reside in the heap", object);
+        return *descriptor().fromField<1>(object);
+    }
 
     gc::GC::ObjectData& objectData() noexcept { return *descriptor().field<0>(this).second; }
 

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

@@ -9,100 +9,96 @@
 #include <atomic>
 
 #include "GCImpl.hpp"
-#include "SafePoint.hpp"
 #include "ThreadData.hpp"
 #include "ThreadRegistry.hpp"
 
 using namespace kotlin;
 
+inline constexpr auto kTagBarriers = logging::Tag::kBarriers;
+#define BarriersLogDebug(active, format, ...) RuntimeLogDebug({kTagBarriers}, "%s" format, active ? "[active] " : "", ##__VA_ARGS__)
+
 namespace {
 
-std::atomic<ObjHeader* (*)(ObjHeader*)> weakRefBarrier = nullptr;
+std::atomic<bool> markBarriersEnabled = false;
-std::atomic<int64_t> weakProcessingEpoch = 0;
+std::atomic<int64_t> markingEpoch = 0;
-
-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;
-}
 
 } // namespace
 
-void gc::BarriersThreadData::onThreadRegistration() noexcept {
+void gc::barriers::BarriersThreadData::onThreadRegistration() noexcept {
-    if (weakRefBarrier.load(std::memory_order_acquire) != nullptr) {
+    if (markBarriersEnabled.load(std::memory_order_acquire)) {
-        startMarkingNewObjects(GCHandle::getByEpoch(weakProcessingEpoch.load(std::memory_order_relaxed)));
+        startMarkingNewObjects(GCHandle::getByEpoch(markingEpoch.load(std::memory_order_relaxed)));
     }
 }
 
-ALWAYS_INLINE void gc::BarriersThreadData::onSafePoint() noexcept {}
+ALWAYS_INLINE void gc::barriers::BarriersThreadData::onSafePoint() noexcept {}
 
-void gc::BarriersThreadData::startMarkingNewObjects(gc::GCHandle gcHandle) noexcept {
+void gc::barriers::BarriersThreadData::startMarkingNewObjects(gc::GCHandle gcHandle) noexcept {
-    RuntimeAssert(weakRefBarrier.load(std::memory_order_relaxed) != nullptr, "New allocations marking may only be requested by weak ref barriers");
+    RuntimeAssert(markBarriersEnabled.load(std::memory_order_relaxed), "New allocations marking may only be requested by mark barriers");
     markHandle_ = gcHandle.mark();
 }
 
-void gc::BarriersThreadData::stopMarkingNewObjects() noexcept {
+void gc::barriers::BarriersThreadData::stopMarkingNewObjects() noexcept {
-    RuntimeAssert(weakRefBarrier.load(std::memory_order_relaxed) == nullptr, "New allocations marking could only been requested by weak ref barriers");
+    RuntimeAssert(!markBarriersEnabled.load(std::memory_order_relaxed), "New allocations marking could only been requested by mark barriers");
     markHandle_ = std::nullopt;
 }
 
-bool gc::BarriersThreadData::shouldMarkNewObjects() const noexcept {
+bool gc::barriers::BarriersThreadData::shouldMarkNewObjects() const noexcept {
     return markHandle_.has_value();
 }
 
-ALWAYS_INLINE void gc::BarriersThreadData::onAllocation(ObjHeader* allocated) {
+ALWAYS_INLINE void gc::barriers::BarriersThreadData::onAllocation(ObjHeader* allocated) {
-    if (compiler::concurrentWeakSweep()) {
+    bool shouldMark = shouldMarkNewObjects();
-        bool shouldMark = shouldMarkNewObjects();
+    RuntimeAssert(shouldMark == markBarriersEnabled.load(std::memory_order_relaxed), "New allocations marking must happen with and only with mark barriers");
-        bool barriersEnabled = weakRefBarrier.load(std::memory_order_relaxed) != nullptr;
+    BarriersLogDebug(shouldMark, "Allocation %p", allocated);
-        RuntimeAssert(shouldMark == barriersEnabled, "New allocations marking must happen with and only with weak ref barriers");
+    if (shouldMark) {
-        if (shouldMark) {
+        auto& objectData = alloc::objectDataForObject(allocated);
-            auto& objectData = alloc::objectDataForObject(allocated);
+        objectData.markUncontended();
-            objectData.markUncontended();
+        markHandle_->addObject();
-            markHandle_->addObject();
-        }
     }
 }
 
-void gc::EnableWeakRefBarriers(int64_t epoch) noexcept {
+void gc::barriers::enableMarkBarriers(int64_t epoch) noexcept {
     auto mutators = mm::ThreadRegistry::Instance().LockForIter();
-    weakProcessingEpoch.store(epoch, std::memory_order_relaxed);
+    markingEpoch.store(epoch, std::memory_order_relaxed);
-    weakRefBarrier.store(weakRefBarrierImpl, std::memory_order_seq_cst);
+    markBarriersEnabled.store(true, std::memory_order_release);
     for (auto& mutator: mutators) {
         mutator.gc().impl().gc().barriers().startMarkingNewObjects(GCHandle::getByEpoch(epoch));
     }
 }
 
-void gc::DisableWeakRefBarriers() noexcept {
+void gc::barriers::disableMarkBarriers() noexcept {
     auto mutators = mm::ThreadRegistry::Instance().LockForIter();
-    weakRefBarrier.store(nullptr, std::memory_order_seq_cst);
+    markBarriersEnabled.store(false, std::memory_order_release);
     for (auto& mutator: mutators) {
         mutator.gc().impl().gc().barriers().stopMarkingNewObjects();
     }
 }
 
-OBJ_GETTER(gc::WeakRefRead, std::atomic<ObjHeader*>& weakReferee) noexcept {
+namespace {
-    if (!compiler::concurrentWeakSweep()) {
-        RETURN_OBJ(weakReferee.load(std::memory_order_relaxed));
-    }
 
-    // Copying the scheme from SafePoint.cpp: branch + indirect call.
+// TODO decide whether it's really beneficial to NO_INLINE the slow path
-    auto barrier = weakRefBarrier.load(std::memory_order_relaxed);
+NO_INLINE void beforeHeapRefUpdateSlowPath(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {
-    ObjHeader* result;
+    auto prev = ref.load();
-    if (__builtin_expect(barrier != nullptr, false)) {
+    BarriersLogDebug(true, "Write *%p <- %p (%p overwritten)", ref.location(), value, prev);
-        result = weakRefReadSlowPath(weakReferee);
+    if (prev != nullptr && prev->heap()) {
-    } else {
+        // TODO Redundant if the destination object is black.
-        result = weakReferee.load(std::memory_order_relaxed);
+        //      Yet at the moment there is now efficient way to distinguish black and gray objects.
+
+        // TODO perhaps it would be better to path the thread data from outside
+        auto& threadData = *mm::ThreadRegistry::Instance().CurrentThreadData();
+        auto& markQueue = *threadData.gc().impl().gc().mark().markQueue();
+        gc::mark::ParallelMark::MarkTraits::tryEnqueue(markQueue, prev);
+        // No need to add the marked object in statistics here.
+        // Objects will be counted on dequeue.
+    }
+}
+
+} // namespace
+
+ALWAYS_INLINE void gc::barriers::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {
+    if (__builtin_expect(markBarriersEnabled.load(std::memory_order_acquire), false)) {
+        beforeHeapRefUpdateSlowPath(ref, value);
+    } else {
+        BarriersLogDebug(false, "Write *%p <- %p (%p overwritten)", ref.location(), value, ref.load());
     }
-    RETURN_OBJ(result);
 }

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

@@ -8,11 +8,11 @@
 #include <atomic>
 #include <optional>
 
-#include "Memory.h"
 #include "Utils.hpp"
 #include "GCStatistics.hpp"
+#include "ReferenceOps.hpp"
 
-namespace kotlin::gc {
+namespace kotlin::gc::barriers {
 
 class BarriersThreadData : private Pinned {
 public:
@@ -29,9 +29,11 @@ private:
 };
 
 // Must be called during STW.
-void EnableWeakRefBarriers(int64_t epoch) noexcept;
+void enableMarkBarriers(int64_t epoch) noexcept;
-void DisableWeakRefBarriers() noexcept;
+void disableMarkBarriers() noexcept;
 
-OBJ_GETTER(WeakRefRead, std::atomic<ObjHeader*>& weakReferee) noexcept;
+void beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept;
+
+// TODO re-introduce weak barriers again
 
 } // namespace kotlin::gc

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

@@ -0,0 +1,155 @@
+/*
+ * 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 "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include "Barriers.hpp"
+#include "ParallelMark.hpp"
+#include "GCImpl.hpp"
+#include "ManuallyScoped.hpp"
+#include "ObjectTestSupport.hpp"
+#include "ObjectOps.hpp"
+#include "ReferenceOps.hpp"
+#include "TestSupport.hpp"
+#include "ObjectData.hpp"
+
+
+using namespace kotlin;
+
+namespace {
+
+auto gcHandle = gc::GCHandle::createFakeForTests();
+
+struct Payload {
+    mm::RefField field1;
+    mm::RefField field2;
+    mm::RefField field3;
+
+    static constexpr std::array kFields = {
+            &Payload::field1,
+            &Payload::field2,
+            &Payload::field3,
+    };
+};
+
+test_support::TypeInfoHolder typeHolder{test_support::TypeInfoHolder::ObjectBuilder<Payload>()};
+
+test_support::Object<Payload>& AllocateObject(mm::ThreadData& threadData) {
+    ObjHolder holder;
+    mm::AllocateObject(&threadData, typeHolder.typeInfo(), holder.slot());
+    return test_support::Object<Payload>::FromObjHeader(holder.obj());
+}
+
+class BarriersTest : public testing::Test {
+public:
+
+    ~BarriersTest() override {
+        mm::SpecialRefRegistry::instance().clearForTests();
+        mm::GlobalData::Instance().allocator().clearForTests();
+    }
+
+    void initMutatorMarkQueue(mm::ThreadData& thread) {
+        auto& markData = thread.gc().impl().gc().mark();
+        markData.markQueue().construct(parProc_);
+    }
+
+private:
+    gc::mark::ParallelMark::ParallelProcessor parProc_;
+};
+
+} // namespace
+
+TEST_F(BarriersTest, Deletion) {
+    RunInNewThread([this](mm::ThreadData& threadData) {
+        initMutatorMarkQueue(threadData);
+        auto& prevObj = AllocateObject(threadData);
+        auto& newObj = AllocateObject(threadData);
+
+        ObjHeader* ref = prevObj.header();
+
+        EXPECT_THAT(gc::isMarked(prevObj.header()), false);
+        EXPECT_THAT(gc::isMarked(newObj.header()), false);
+
+        {
+            ThreadStateGuard guard(ThreadState::kNative); // pretend to be the GC thread
+            gc::barriers::enableMarkBarriers(gcHandle.getEpoch());
+        }
+
+        UpdateHeapRef(&ref, newObj.header());
+
+        EXPECT_THAT(gc::isMarked(prevObj.header()), true);
+        EXPECT_THAT(gc::isMarked(newObj.header()), false);
+
+        {
+            ThreadStateGuard guard(ThreadState::kNative); // pretend to be the GC thread
+            gc::barriers::disableMarkBarriers();
+        }
+
+    });
+}
+
+TEST_F(BarriersTest, AllocationDuringMarkBarreirs) {
+    gc::barriers::enableMarkBarriers(gcHandle.getEpoch());
+
+    RunInNewThread([this](mm::ThreadData& threadData) {
+        initMutatorMarkQueue(threadData);
+        auto& obj = AllocateObject(threadData);
+        EXPECT_THAT(gc::isMarked(obj.header()), true);
+    });
+
+    gc::barriers::disableMarkBarriers();
+}
+
+TEST_F(BarriersTest, ConcurrentDeletion) {
+    constexpr auto kObjsPerThread = 100;
+
+    ObjHeader* ref = nullptr;
+
+    RunInNewThread([&](mm::ThreadData& threadData) {
+        auto& obj = AllocateObject(threadData);
+        UpdateHeapRef(&ref, obj.header());
+        threadData.allocator().prepareForGC();
+    });
+
+    EXPECT_THAT(gc::isMarked(ref), false);
+
+    std::atomic<bool> canStart = false;
+    std::atomic<std::size_t> finished = 0;
+
+    gc::barriers::enableMarkBarriers(gcHandle.getEpoch());
+
+    std::vector<ScopedThread> threads;
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        threads.emplace_back([&]() noexcept {
+            ScopedMemoryInit memory;
+            mm::ThreadData& threadData = *memory.memoryState()->GetThreadData();
+            initMutatorMarkQueue(threadData);
+
+            while (!canStart.load()) std::this_thread::yield();
+
+            for (int j = 0; j < kObjsPerThread; ++j) {
+                auto& obj = AllocateObject(threadData);
+                // auto&& accessor = mm::RefFieldAccessor(&ref);
+                // accessor.storeAtomic(obj.header(), std::memory_order_release);
+                UpdateHeapRef(&ref, obj.header());
+            }
+
+            finished += 1;
+
+            threadData.allocator().prepareForGC();
+        });
+    }
+
+    canStart = true;
+
+    while (finished.load() < threads.size()) {
+        std::this_thread::yield();
+    }
+
+    gc::barriers::disableMarkBarriers();
+
+    EXPECT_THAT(gc::isMarked(ref), true);
+}

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

@@ -156,22 +156,14 @@ void gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
 
     markDispatcher_.endMarkingEpoch();
 
-    if (compiler::concurrentWeakSweep()) {
+    // TODO re-enable concurrent weak sweep again
-        // Expected to happen inside STW.
-        gc::EnableWeakRefBarriers(epoch);
-        resumeTheWorld(gcHandle);
-    }
 
     gc::processWeaks<DefaultProcessWeaksTraits>(gcHandle, mm::SpecialRefRegistry::instance());
 
-    if (compiler::concurrentWeakSweep()) {
-        stopTheWorld(gcHandle);
-        gc::DisableWeakRefBarriers();
-    }
-
     // TODO outline as mark_.isolateMarkedHeapAndFinishMark()
     // By this point all the alive heap must be marked.
     // All the mutations (incl. allocations) after this method will be subject for the next GC.
+
     // This should really be done by each individual thread while waiting
     int threadCount = 0;
     for (auto& thread : kotlin::mm::ThreadRegistry::Instance().LockForIter()) {

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

@@ -50,15 +50,14 @@ public:
 
         auto& commonThreadData() const noexcept { return threadData_; }
         auto& barriers() noexcept { return barriers_; }
-        // TODO use in concurrent mark
+        auto& mark() noexcept { return mark_; }
-        [[maybe_unused]] auto& markQueue() noexcept { return markQueue_; }
 
     private:
         friend ConcurrentMarkAndSweep;
         ConcurrentMarkAndSweep& gc_;
         mm::ThreadData& threadData_;
-        BarriersThreadData barriers_;
+        barriers::BarriersThreadData barriers_;
-        ManuallyScoped<mark::ParallelMark::MutatorQueue> markQueue_;
+        mark::ParallelMark::ThreadData mark_;
 
         std::atomic<bool> rootSetLocked_ = false;
         std::atomic<bool> published_ = false;

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

@@ -46,7 +46,7 @@ TYPED_TEST_P(TracingGCTest, CMSMultipleMutatorsWeak) {
                 auto& object = AllocateObject(threadData);
                 auto& objectWeak = ([&threadData, &object]() -> test_support::RegularWeakReferenceImpl& {
                     ObjHolder holder;
-                    return InstallWeakReference(threadData, object.header(), holder.slot());
+                    return test_support::InstallWeakReference(threadData, object.header(), holder.slot());
                 })();
                 global->field1 = objectWeak.header();
                 weak = &objectWeak;

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

@@ -19,10 +19,6 @@ gc::GC::ThreadData::ThreadData(GC& gc, mm::ThreadData& threadData) noexcept : im
 
 gc::GC::ThreadData::~ThreadData() = default;
 
-ALWAYS_INLINE void gc::GC::ThreadData::onAllocation(ObjHeader* object) noexcept {
-    impl().gc().barriers().onAllocation(object);
-}
-
 void gc::GC::ThreadData::OnSuspendForGC() noexcept {
     impl_->gc().OnSuspendForGC();
 }
@@ -35,6 +31,10 @@ void gc::GC::ThreadData::onThreadRegistration() noexcept {
     impl_->gc().onThreadRegistration();
 }
 
+ALWAYS_INLINE void gc::GC::ThreadData::onAllocation(ObjHeader* object) noexcept {
+    impl().gc().barriers().onAllocation(object);
+}
+
 gc::GC::GC(alloc::Allocator& allocator, gcScheduler::GCScheduler& gcScheduler) noexcept :
     impl_(std::make_unique<Impl>(allocator, gcScheduler)) {}
 
@@ -84,12 +84,16 @@ void gc::GC::WaitFinalizers(int64_t epoch) noexcept {
     impl_->gc().state().waitEpochFinalized(epoch);
 }
 
-bool gc::isMarked(ObjHeader* object) noexcept {
+ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {
-    return alloc::objectDataForObject(object).marked();
+    barriers::beforeHeapRefUpdate(ref, value);
 }
 
-ALWAYS_INLINE OBJ_GETTER(gc::tryRef, std::atomic<ObjHeader*>& object) noexcept {
+ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept {
-    RETURN_RESULT_OF(gc::WeakRefRead, object);
+    RETURN_OBJ(weakReferee.load(std::memory_order_relaxed));
+}
+
+bool gc::isMarked(ObjHeader* object) noexcept {
+    return alloc::objectDataForObject(object).marked();
 }
 
 ALWAYS_INLINE bool gc::tryResetMark(GC::ObjectData& objectData) noexcept {

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

@@ -76,16 +76,17 @@ private:
 class ParallelMark : private Pinned {
     using MarkStackImpl = intrusive_forward_list<GC::ObjectData>;
     // work balancing parameters were chosen pretty arbitrary
-    using ParallelProcessor = ParallelProcessor<MarkStackImpl, 512, 4096>;
-public:
 
+public:
+    using ParallelProcessor = ParallelProcessor<MarkStackImpl, 512, 4096>;
     using MutatorQueue = ParallelProcessor::WorkSource;
 
     class MarkTraits {
     public:
         using MarkQueue = ParallelProcessor::Worker;
+        using AnyQueue = ParallelProcessor::WorkSource;
 
-        static void clear(MarkQueue& queue) noexcept {
+        static void clear(AnyQueue& queue) noexcept {
             RuntimeAssert(queue.localEmpty(), "Mark queue must be empty");
         }
 
@@ -97,7 +98,7 @@ public:
             return nullptr;
         }
 
-        static ALWAYS_INLINE bool tryEnqueue(MarkQueue& queue, ObjHeader* object) noexcept {
+        static ALWAYS_INLINE bool tryEnqueue(AnyQueue& queue, ObjHeader* object) noexcept {
             auto& objectData = alloc::objectDataForObject(object);
             return compiler::gcMarkSingleThreaded() ? queue.tryPushLocal(objectData) : queue.tryPush(objectData);
         }
@@ -114,7 +115,14 @@ public:
         }
     };
 
-    ParallelMark(bool mutatorsCooperate);
+    class ThreadData : private Pinned {
+    public:
+        auto& markQueue() noexcept { return markQueue_; }
+    private:
+        ManuallyScoped<MutatorQueue> markQueue_{};
+    };
+
+    explicit ParallelMark(bool mutatorsCooperate);
 
     void beginMarkingEpoch(gc::GCHandle gcHandle);
     void endMarkingEpoch();

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

@@ -11,7 +11,7 @@
 
 #include "ExtraObjectData.hpp"
 #include "GCScheduler.hpp"
-#include "Memory.h"
+#include "ReferenceOps.hpp"
 #include "Utils.hpp"
 
 namespace kotlin {
@@ -83,8 +83,10 @@ private:
     std::unique_ptr<Impl> impl_;
 };
 
+void beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept;
+OBJ_GETTER(weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept;
+
 bool isMarked(ObjHeader* object) noexcept;
-OBJ_GETTER(tryRef, std::atomic<ObjHeader*>& object) noexcept;
 
 // This will drop the mark bit if it was set and return `true`.
 // If the mark bit was unset, this will return `false`.

kotlin-native/runtime/src/gc/common/cpp/GCStatistics.cpp

@@ -43,6 +43,8 @@ namespace {
 constexpr KNativePtr heapPoolKey = const_cast<KNativePtr>(static_cast<const void*>("heap"));
 constexpr KNativePtr extraPoolKey = const_cast<KNativePtr>(static_cast<const void*>("extra"));
 
+constexpr auto kInvalidEpoch = std::numeric_limits<uint64_t>::max();
+
 struct MemoryUsage {
     uint64_t sizeBytes;
 };
@@ -178,7 +180,7 @@ GCHandle GCHandle::create(uint64_t epoch) {
     current.memoryUsageBefore.heap = currentHeapUsage();
     return getByEpoch(epoch);
 }
-GCHandle GCHandle::createFakeForTests() { return getByEpoch(invalid().getEpoch() - 1); }
+GCHandle GCHandle::createFakeForTests() { return getByEpoch(kInvalidEpoch - 1); }
 GCHandle GCHandle::getByEpoch(uint64_t epoch) {
     GCHandle handle{epoch};
     RuntimeAssert(handle.isValid(), "Must be valid");
@@ -195,7 +197,7 @@ std::optional<gc::GCHandle> gc::GCHandle::currentEpoch() noexcept {
 }
 
 GCHandle GCHandle::invalid() {
-    return GCHandle{std::numeric_limits<uint64_t>::max()};
+    return GCHandle{kInvalidEpoch};
 }
 void GCHandle::ClearForTests() {
     std::lock_guard guard(lock);
@@ -203,7 +205,7 @@ void GCHandle::ClearForTests() {
     last = {};
 }
 bool GCHandle::isValid() const {
-    return epoch_ != GCHandle::invalid().epoch_;
+    return epoch_ != kInvalidEpoch;
 }
 void GCHandle::finished() {
     std::lock_guard guard(lock);

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

@@ -52,7 +52,10 @@ public:
     static GCHandle invalid();
     static void ClearForTests();
 
-    uint64_t getEpoch() { return epoch_; }
+    uint64_t getEpoch() const {
+        RuntimeAssert(isValid(), "Invalid GC handle");
+        return epoch_;
+    }
     bool isValid() const;
     void finished();
     void finalizersDone();
@@ -84,6 +87,11 @@ private:
 
 class GCHandle::GCStageScopeBase : private MoveOnly {
 public:
+    ALWAYS_INLINE void requireValid() const {
+        RuntimeAssert(handle_.isValid(), "Invalid GC handle accessed");
+    }
+
+protected:
     explicit GCStageScopeBase(GCHandle gcHandle) : handle_(gcHandle) {}
 
     friend void swap(GCStageScopeBase& first, GCStageScopeBase& second) noexcept {
@@ -100,62 +108,113 @@ public:
         return *this;
     }
 
-protected:
     uint64_t getStageTime() const { return (konan::getTimeMicros() - startTime_); }
 
     GCHandle handle_;
     uint64_t startTime_ = konan::getTimeMicros();
 };
 
-class GCHandle::GCSweepScope : GCStageScopeBase {
+class GCHandle::GCSweepScope : public GCStageScopeBase {
     SweepStats stats_;
     uint64_t markedCount_ = 0;
 
 public:
     explicit GCSweepScope(GCHandle handle);
-    GCSweepScope(GCSweepScope&& that) = default;
+    GCSweepScope(GCSweepScope&& that) noexcept : GCSweepScope(GCHandle::invalid()) {
-    GCSweepScope& operator=(GCSweepScope&& that) = default;
+        swap(*this, that);
+    }
+    GCSweepScope& operator=(GCSweepScope&& that) noexcept {
+        auto tmp = std::move(that);
+        swap(*this, tmp);
+        return *this;
+    }
+    friend void swap(GCSweepScope& first, GCSweepScope& second) noexcept {
+        using std::swap;
+        swap(reinterpret_cast<GCStageScopeBase&>(first), reinterpret_cast<GCStageScopeBase&>(second));
+        swap(first.stats_, second.stats_);
+        swap(first.markedCount_, second.markedCount_);
+    }
     ~GCSweepScope();
 
-    void addSweptObject() noexcept { stats_.sweptCount += 1; }
+    void addSweptObject() noexcept {
+        requireValid();
+        stats_.sweptCount += 1;
+    }
     void addKeptObject(size_t sizeBytes) noexcept {
+        requireValid();
         stats_.keptCount += 1;
         stats_.keptSizeBytes += sizeBytes;
     }
     // Custom allocator only. To be finalized objects are kept alive.
-    void addMarkedObject() noexcept { markedCount_ += 1; }
+    void addMarkedObject() noexcept {
+        requireValid();
+        markedCount_ += 1;
+    }
 };
 
-class GCHandle::GCSweepExtraObjectsScope : private GCStageScopeBase {
+class GCHandle::GCSweepExtraObjectsScope : public GCStageScopeBase {
     SweepStats stats_;
 
 public:
     explicit GCSweepExtraObjectsScope(GCHandle handle);
-    GCSweepExtraObjectsScope(GCSweepExtraObjectsScope&& that) = default;
+    GCSweepExtraObjectsScope(GCSweepExtraObjectsScope&& that) noexcept : GCSweepExtraObjectsScope(GCHandle::invalid()) {
-    GCSweepExtraObjectsScope& operator=(GCSweepExtraObjectsScope&& that) = default;
+        swap(*this, that);
+    }
+    GCSweepExtraObjectsScope& operator=(GCSweepExtraObjectsScope&& that) noexcept {
+        auto tmp = std::move(that);
+        swap(*this, tmp);
+        return *this;
+    }
+    friend void swap(GCSweepExtraObjectsScope& first, GCSweepExtraObjectsScope& second) noexcept {
+        using std::swap;
+        swap(reinterpret_cast<GCStageScopeBase&>(first), reinterpret_cast<GCStageScopeBase&>(second));
+        swap(first.stats_, second.stats_);
+    }
     ~GCSweepExtraObjectsScope();
 
-    void addSweptObject() noexcept { stats_.sweptCount += 1; }
+    void addSweptObject() noexcept {
+        requireValid();
+        stats_.sweptCount += 1;
+    }
     void addKeptObject(size_t sizeBytes) noexcept {
+        requireValid();
         stats_.keptCount += 1;
         stats_.keptSizeBytes += sizeBytes;
     }
 };
 
-class GCHandle::GCGlobalRootSetScope : private GCStageScopeBase {
+class GCHandle::GCGlobalRootSetScope : public GCStageScopeBase {
     uint64_t globalRoots_ = 0;
     uint64_t stableRoots_ = 0;
 
 public:
     explicit GCGlobalRootSetScope(GCHandle handle);
-    GCGlobalRootSetScope(GCGlobalRootSetScope&& that) = default;
+    GCGlobalRootSetScope(GCGlobalRootSetScope&& that) noexcept : GCGlobalRootSetScope(GCHandle::invalid()) {
-    GCGlobalRootSetScope& operator=(GCGlobalRootSetScope&& that) = default;
+        swap(*this, that);
+    }
+    GCGlobalRootSetScope& operator=(GCGlobalRootSetScope&& that) noexcept {
+        auto tmp = std::move(that);
+        swap(*this, tmp);
+        return *this;
+    }
+    friend void swap(GCGlobalRootSetScope& first, GCGlobalRootSetScope& second) noexcept {
+        using std::swap;
+        swap(reinterpret_cast<GCStageScopeBase&>(first), reinterpret_cast<GCStageScopeBase&>(second));
+        swap(first.globalRoots_, second.globalRoots_);
+        swap(first.stableRoots_, second.stableRoots_);
+    }
     ~GCGlobalRootSetScope();
-    void addGlobalRoot() { globalRoots_++; }
+    void addGlobalRoot() {
-    void addStableRoot() { stableRoots_++; }
+        requireValid();
+        globalRoots_++;
+    }
+    void addStableRoot() {
+        requireValid();
+        stableRoots_++;
+    }
 };
 
-class GCHandle::GCThreadRootSetScope : private GCStageScopeBase {
+class GCHandle::GCThreadRootSetScope : public GCStageScopeBase {
     mm::ThreadData& threadData_;
     uint64_t stackRoots_ = 0;
     uint64_t threadLocalRoots_ = 0;
@@ -165,36 +224,78 @@ public:
     GCThreadRootSetScope(GCThreadRootSetScope&& that) = default;
     GCThreadRootSetScope& operator=(GCThreadRootSetScope&& that) = delete;
     ~GCThreadRootSetScope();
-    void addStackRoot() { stackRoots_++; }
+    void addStackRoot() {
-    void addThreadLocalRoot() { threadLocalRoots_++; }
+        requireValid();
+        stackRoots_++;
+    }
+    void addThreadLocalRoot() {
+        requireValid();
+        threadLocalRoots_++;
+    }
 };
 
-class GCHandle::GCMarkScope : private GCStageScopeBase {
+class GCHandle::GCMarkScope : public GCStageScopeBase {
     MarkStats stats_;
 
 public:
     explicit GCMarkScope(GCHandle handle);
-    GCMarkScope(GCMarkScope&& that) = default;
+    GCMarkScope(GCMarkScope&& that) noexcept : GCMarkScope(GCHandle::invalid()) {
-    GCMarkScope& operator=(GCMarkScope&& that) = default;
+        swap(*this, that);
+    }
+    GCMarkScope& operator=(GCMarkScope&& that) noexcept {
+        auto tmp = std::move(that);
+        swap(*this, tmp);
+        return *this;
+    }
+    friend void swap(GCMarkScope& first, GCMarkScope& second) noexcept {
+        using std::swap;
+        swap(reinterpret_cast<GCStageScopeBase&>(first), reinterpret_cast<GCStageScopeBase&>(second));
+        swap(first.stats_, second.stats_);
+    }
     ~GCMarkScope();
 
-    void addObject() noexcept { ++stats_.markedCount; }
+    void addObject() noexcept {
+        requireValid();
+        ++stats_.markedCount;
+    }
 };
 
-class GCHandle::GCProcessWeaksScope : private GCStageScopeBase {
+class GCHandle::GCProcessWeaksScope : public GCStageScopeBase {
     uint64_t undisposedCount_ = 0;
     uint64_t aliveCount_ = 0;
     uint64_t nulledCount_ = 0;
 
 public:
     explicit GCProcessWeaksScope(GCHandle handle) noexcept;
-    GCProcessWeaksScope(GCProcessWeaksScope&& that) = default;
+    GCProcessWeaksScope(GCProcessWeaksScope&& that) noexcept : GCProcessWeaksScope(GCHandle::invalid()) {
-    GCProcessWeaksScope& operator=(GCProcessWeaksScope&& that) = default;
+        swap(*this, that);
+    }
+    GCProcessWeaksScope& operator=(GCProcessWeaksScope&& that) noexcept {
+        auto tmp = std::move(that);
+        swap(*this, tmp);
+        return *this;
+    }
+    friend void swap(GCProcessWeaksScope& first, GCProcessWeaksScope& second) noexcept {
+        using std::swap;
+        swap(reinterpret_cast<GCStageScopeBase&>(first), reinterpret_cast<GCStageScopeBase&>(second));
+        swap(first.undisposedCount_, second.undisposedCount_);
+        swap(first.aliveCount_, second.aliveCount_);
+        swap(first.nulledCount_, second.nulledCount_);
+    }
     ~GCProcessWeaksScope();
 
-    void addUndisposed() noexcept { ++undisposedCount_; }
+    void addUndisposed() noexcept {
-    void addAlive() noexcept { ++aliveCount_; }
+        requireValid();
-    void addNulled() noexcept { ++nulledCount_; }
+        ++undisposedCount_;
+    }
+    void addAlive() noexcept {
+        requireValid();
+        ++aliveCount_;
+    }
+    void addNulled() noexcept {
+        requireValid();
+        ++nulledCount_;
+    }
 };
 
 }

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

@@ -179,17 +179,6 @@ std::vector<ObjHeader*> Alive(mm::ThreadData& threadData) {
     return alloc::test_support::allocatedObjects(threadData);
 }
 
-test_support::RegularWeakReferenceImpl& InstallWeakReference(mm::ThreadData& threadData, ObjHeader* objHeader, ObjHeader** location) {
-    mm::AllocateObject(&threadData, theRegularWeakReferenceImplTypeInfo, location);
-    auto& weakReference = test_support::RegularWeakReferenceImpl::FromObjHeader(*location);
-    auto& extraObjectData = mm::ExtraObjectData::GetOrInstall(objHeader);
-    weakReference->weakRef = static_cast<mm::RawSpecialRef*>(mm::WeakRef::create(objHeader));
-    weakReference->referred = objHeader;
-    auto* setWeakRef = extraObjectData.GetOrSetRegularWeakReferenceImpl(objHeader, weakReference.header());
-    EXPECT_EQ(setWeakRef, weakReference.header());
-    return weakReference;
-}
-
 
 template <typename FixtureImpl>
 class TracingGCTest : public testing::Test {
@@ -324,7 +313,7 @@ TYPED_TEST_P(TracingGCTest, FreeObjectWithFreeWeak) {
         auto& object1 = AllocateObject(threadData);
         auto& weak1 = ([&threadData, &object1]() -> test_support::RegularWeakReferenceImpl& {
             ObjHolder holder;
-            return InstallWeakReference(threadData, object1.header(), holder.slot());
+            return test_support::InstallWeakReference(threadData, object1.header(), holder.slot());
         })();
 
         ASSERT_THAT(Alive(threadData), testing::UnorderedElementsAre(object1.header(), weak1.header()));
@@ -343,7 +332,7 @@ TYPED_TEST_P(TracingGCTest, FreeObjectWithHoldedWeak) {
     RunInNewThread([](mm::ThreadData& threadData) {
         auto& object1 = AllocateObject(threadData);
         StackObjectHolder stack{threadData};
-        auto& weak1 = InstallWeakReference(threadData, object1.header(), stack->field1.ptr());
+        auto& weak1 = test_support::InstallWeakReference(threadData, object1.header(), stack->field1.ptr());
 
         ASSERT_THAT(Alive(threadData), testing::UnorderedElementsAre(object1.header(), weak1.header(), stack.header()));
         ASSERT_THAT(gc::isMarked(object1.header()), false);
@@ -1074,7 +1063,7 @@ TYPED_TEST_P(TracingGCTest, FreeObjectWithFreeWeakReversedOrder) {
         while (object1.load() == nullptr) {
         }
         ObjHolder holder;
-        auto& weak_local = InstallWeakReference(threadData, object1.load()->header(), holder.slot());
+        auto& weak_local = test_support::InstallWeakReference(threadData, object1.load()->header(), holder.slot());
         weak = &weak_local;
         *holder.slot() = nullptr;
         while (!done) mm::safePoint(threadData);
@@ -1122,7 +1111,7 @@ TYPED_TEST_P(STWMarkGCTest, MultipleMutatorsWeaks) {
                 auto& object = AllocateObject(threadData);
                 auto& objectWeak = ([&threadData, &object]() -> test_support::RegularWeakReferenceImpl& {
                     ObjHolder holder;
-                    return InstallWeakReference(threadData, object.header(), holder.slot());
+                    return test_support::InstallWeakReference(threadData, object.header(), holder.slot());
                 })();
                 global->field1 = objectWeak.header();
                 weak = &objectWeak;
@@ -1187,7 +1176,7 @@ TYPED_TEST_P(STWMarkGCTest, MultipleMutatorsWeakNewObj) {
             auto& object = AllocateObject(threadData);
             auto& objectWeak = ([&threadData, &object]() -> test_support::RegularWeakReferenceImpl& {
                 ObjHolder holder;
-                return InstallWeakReference(threadData, object.header(), holder.slot());
+                return test_support::InstallWeakReference(threadData, object.header(), holder.slot());
             })();
             EXPECT_NE(objectWeak.get(), nullptr);
 

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

@@ -16,14 +16,14 @@ gc::GC::ThreadData::ThreadData(GC& gc, mm::ThreadData& threadData) noexcept {}
 
 gc::GC::ThreadData::~ThreadData() = default;
 
-ALWAYS_INLINE void gc::GC::ThreadData::onAllocation(ObjHeader* object) noexcept {}
-
 void gc::GC::ThreadData::OnSuspendForGC() noexcept { }
 
 void gc::GC::ThreadData::safePoint() noexcept {}
 
 void gc::GC::ThreadData::onThreadRegistration() noexcept {}
 
+ALWAYS_INLINE void gc::GC::ThreadData::onAllocation(ObjHeader* object) noexcept {}
+
 gc::GC::GC(alloc::Allocator&, gcScheduler::GCScheduler&) noexcept {
     RuntimeLogInfo({kTagGC}, "No-op GC initialized");
 }
@@ -59,15 +59,17 @@ void gc::GC::WaitFinished(int64_t epoch) noexcept {}
 
 void gc::GC::WaitFinalizers(int64_t epoch) noexcept {}
 
+ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {}
+
+ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept {
+    RETURN_OBJ(weakReferee.load(std::memory_order_relaxed));
+}
+
 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));
-}
-
 ALWAYS_INLINE bool gc::tryResetMark(GC::ObjectData& objectData) noexcept {
     RuntimeAssert(false, "Should not reach here");
     return true;

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

@@ -19,10 +19,6 @@ gc::GC::ThreadData::ThreadData(GC& gc, mm::ThreadData& threadData) noexcept : im
 
 gc::GC::ThreadData::~ThreadData() = default;
 
-ALWAYS_INLINE void gc::GC::ThreadData::onAllocation(ObjHeader* object) noexcept {
-    impl().gc().barriers().onAllocation(object);
-}
-
 void gc::GC::ThreadData::OnSuspendForGC() noexcept {
     impl_->gc().OnSuspendForGC();
 }
@@ -35,6 +31,10 @@ void gc::GC::ThreadData::onThreadRegistration() noexcept {
     impl_->gc().onThreadRegistration();
 }
 
+ALWAYS_INLINE void gc::GC::ThreadData::onAllocation(ObjHeader* object) noexcept {
+    impl().gc().barriers().onAllocation(object);
+}
+
 gc::GC::GC(alloc::Allocator& allocator, gcScheduler::GCScheduler& gcScheduler) noexcept :
     impl_(std::make_unique<Impl>(allocator, gcScheduler)) {}
 
@@ -84,12 +84,14 @@ void gc::GC::WaitFinalizers(int64_t epoch) noexcept {
     impl_->gc().state().waitEpochFinalized(epoch);
 }
 
-bool gc::isMarked(ObjHeader* object) noexcept {
+ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {}
-    return alloc::objectDataForObject(object).marked();
+
+ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept {
+    RETURN_RESULT_OF(gc::WeakRefRead, weakReferee);
 }
 
-ALWAYS_INLINE OBJ_GETTER(gc::tryRef, std::atomic<ObjHeader*>& object) noexcept {
+bool gc::isMarked(ObjHeader* object) noexcept {
-    RETURN_RESULT_OF(gc::WeakRefRead, object);
+    return alloc::objectDataForObject(object).marked();
 }
 
 ALWAYS_INLINE bool gc::tryResetMark(GC::ObjectData& objectData) noexcept {

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

@@ -20,14 +20,14 @@ gc::GC::ThreadData::ThreadData(GC& gc, mm::ThreadData& threadData) noexcept : im
 
 gc::GC::ThreadData::~ThreadData() = default;
 
-ALWAYS_INLINE void gc::GC::ThreadData::onAllocation(ObjHeader* object) noexcept {}
-
 void gc::GC::ThreadData::OnSuspendForGC() noexcept { }
 
 void gc::GC::ThreadData::safePoint() noexcept {}
 
 void gc::GC::ThreadData::onThreadRegistration() noexcept {}
 
+ALWAYS_INLINE void gc::GC::ThreadData::onAllocation(ObjHeader* object) noexcept {}
+
 gc::GC::GC(alloc::Allocator& allocator, gcScheduler::GCScheduler& gcScheduler) noexcept :
     impl_(std::make_unique<Impl>(allocator, gcScheduler)) {}
 
@@ -77,12 +77,14 @@ void gc::GC::WaitFinalizers(int64_t epoch) noexcept {
     impl_->gc().state().waitEpochFinalized(epoch);
 }
 
-bool gc::isMarked(ObjHeader* object) noexcept {
+ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {}
-    return alloc::objectDataForObject(object).marked();
+
+ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept {
+    RETURN_OBJ(weakReferee.load(std::memory_order_relaxed));
 }
 
-ALWAYS_INLINE OBJ_GETTER(gc::tryRef, std::atomic<ObjHeader*>& object) noexcept {
+bool gc::isMarked(ObjHeader* object) noexcept {
-    RETURN_OBJ(object.load(std::memory_order_relaxed));
+    return alloc::objectDataForObject(object).marked();
 }
 
 ALWAYS_INLINE bool gc::tryResetMark(GC::ObjectData& objectData) noexcept {

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

@@ -37,8 +37,9 @@ enum class Tag : int32_t {
     kPause = 5,
     kAlloc = 6,
     kBalancing = 7,
+    kBarriers = 8,
 
-    kEnumSize = 8
+    kEnumSize = 9
 };
 
 namespace internal {
@@ -63,6 +64,7 @@ inline const char* name(Tag tag) {
         case Tag::kPause: return "pause";
         case Tag::kAlloc: return "alloc";
         case Tag::kBalancing: return "balancing";
+        case Tag::kBarriers: return "barriers";
 
         case Tag::kEnumSize: break;
     }

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

@@ -41,10 +41,8 @@ private:
 
         ALWAYS_INLINE bool tryPush(typename ListImpl::reference value) noexcept {
             RuntimeAssert(!full(), "Batch overflow");
-            bool pushed = elems_.try_push_front(value);
+            const bool pushed = elems_.try_push_front(value);
-            if (pushed) {
+            elemsCount_ += static_cast<int>(pushed);
-                ++elemsCount_;
-            }
             return pushed;
         }
 

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

@@ -8,17 +8,20 @@
 #include "Memory.h"
 #include "std_support/AtomicRef.hpp"
 
-// Concurrent GC may cause conflicting unordered accesses to references in heap.
+#if __has_feature(thread_sanitizer)
-// C++ memory model declares that accesses have data race unless they are atomic.
+#include <sanitizer/tsan_interface.h>
-// Thus TSAN would report such non-atomic accesses.
+#endif
-// TODO find out if compiler may take advantage of the theoretical UB here.
+
+// C++ memory model is in some sence stricter than the memmory model of real target CPUs.
+// For example all the ptr-sized memory accesses on intel x86 and arm CPUs are atomic.
+// Another case is the release-consume memory ordering, which can be achieved without additional memory fences on consume.
 //
-// However, in practice all the ptr-sized loads and stores are atomic on CPU-level
+// However, LLVM often fails to properly optimize atomic operations.
-// even if they are not std::atomic.
+// So we have to allow some imeplementation-defined UB here.
-// And as far as we aware,
+//
-// std::atomic operations are not optimized by LLVM even if they have relaxed memory order.
+// Under this flag all tha operations with references in the kotlin heap
-// So we don't want to compile every heap reference access into std::atomic access.
+// are implemented in complete complience with C++ memory model.
-#define ALWAYS_ATOMIC_REFS __has_feature(thread_sanitizer)
+#define STRICT_ATOMICS_IN_HEAP __has_feature(thread_sanitizer)
 
 namespace kotlin::mm {
 
@@ -46,15 +49,22 @@ public:
     ALWAYS_INLINE ObjHeader* operator=(ObjHeader* desired) noexcept { store(desired); return desired; }
 
     ALWAYS_INLINE ObjHeader* load() const noexcept {
-#if ALWAYS_ATOMIC_REFS
+#if STRICT_ATOMICS_IN_HEAP
-        return loadAtomic(std::memory_order_relaxed);
+        // Consume stores in the object, that were released on the object's allocation
+        // See `ObjectOps.cpp`
+        auto loaded = loadAtomic(std::memory_order_consume);
+#if __has_feature(thread_sanitizer)
+        // The stores were released by an atomic_thread_fence, TSAN doesn't support fences.
+        __tsan_acquire(loaded);
+#endif
+        return loaded;
 #else
         return ref_;
 #endif
     }
 
     ALWAYS_INLINE void store(ObjHeader* desired) noexcept {
-#if ALWAYS_ATOMIC_REFS
+#if STRICT_ATOMICS_IN_HEAP
         storeAtomic(desired, std::memory_order_relaxed);
 #else
         ref_ = desired;
@@ -200,6 +210,6 @@ private:
     ObjHeader* value_ = nullptr;
 };
 
-OBJ_GETTER(weakRefReadBarrier, std::atomic<ObjHeader*>& referee) noexcept;
+OBJ_GETTER(weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept;
 
 }

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

@@ -9,15 +9,24 @@
 #include "ThreadData.hpp"
 #include "ThreadState.hpp"
 
+#if __has_feature(thread_sanitizer)
+#include <sanitizer/tsan_interface.h>
+#endif
+
 using namespace kotlin;
 
 OBJ_GETTER(mm::AllocateObject, ThreadData* threadData, const TypeInfo* typeInfo) noexcept {
     AssertThreadState(threadData, ThreadState::kRunnable);
     // TODO: Make this work with GCs that can stop thread at any point.
     auto* object = threadData->allocator().allocateObject(typeInfo);
-    // Prevents unsafe class publication (see KT-58995).
-    std::atomic_thread_fence(std::memory_order_release);
     threadData->gc().onAllocation(object);
+    // Prevents unsafe class publication (see KT-58995).
+    // Also important in case of the concurrent GC mark phase.
+    std::atomic_thread_fence(std::memory_order_release);
+#if __has_feature(thread_sanitizer)
+    // TSAN doesn't support fences.
+    __tsan_release(object);
+#endif
     RETURN_OBJ(object);
 }
 
@@ -25,8 +34,13 @@ OBJ_GETTER(mm::AllocateArray, ThreadData* threadData, const TypeInfo* typeInfo,
     AssertThreadState(threadData, ThreadState::kRunnable);
     // TODO: Make this work with GCs that can stop thread at any point.
     auto* array = threadData->allocator().allocateArray(typeInfo, static_cast<uint32_t>(elements));
-    // Prevents unsafe class publication (see KT-58995).
-    std::atomic_thread_fence(std::memory_order_release);
     threadData->gc().onAllocation(array->obj());
+    // Prevents unsafe class publication (see KT-58995).
+    // Also important in case of the concurrent GC mark phase.
+    std::atomic_thread_fence(std::memory_order_release);
+#if __has_feature(thread_sanitizer)
+    // TSAN doesn't support fences.
+    __tsan_release(array);
+#endif
     RETURN_OBJ(array->obj());
 }

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

@@ -10,17 +10,19 @@
 using namespace kotlin;
 
 // on stack
-template<> void mm::RefAccessor<true>::beforeStore(ObjHeader*) noexcept {}
+template<> ALWAYS_INLINE void mm::RefAccessor<true>::beforeStore(ObjHeader*) noexcept {}
-template<> void mm::RefAccessor<true>::afterStore(ObjHeader*) noexcept {}
+template<> ALWAYS_INLINE void mm::RefAccessor<true>::afterStore(ObjHeader*) noexcept {}
-template<> void mm::RefAccessor<true>::beforeLoad() noexcept {}
+template<> ALWAYS_INLINE void mm::RefAccessor<true>::beforeLoad() noexcept {}
-template<> void mm::RefAccessor<true>::afterLoad() noexcept {}
+template<> ALWAYS_INLINE void mm::RefAccessor<true>::afterLoad() noexcept {}
 
 // on heap
-template<> void mm::RefAccessor<false>::beforeStore(ObjHeader*) noexcept {}
+template<> ALWAYS_INLINE void mm::RefAccessor<false>::beforeStore(ObjHeader* value) noexcept {
-template<> void mm::RefAccessor<false>::afterStore(ObjHeader*) noexcept {}
+    gc::beforeHeapRefUpdate(direct(), value);
-template<> void mm::RefAccessor<false>::beforeLoad() noexcept {}
+}
-template<> void mm::RefAccessor<false>::afterLoad() noexcept {}
+template<> ALWAYS_INLINE void mm::RefAccessor<false>::afterStore(ObjHeader*) noexcept {}
-
+template<> ALWAYS_INLINE void mm::RefAccessor<false>::beforeLoad() noexcept {}
-ALWAYS_INLINE OBJ_GETTER(mm::weakRefReadBarrier, std::atomic<ObjHeader*>& referee) noexcept {
+template<> ALWAYS_INLINE void mm::RefAccessor<false>::afterLoad() noexcept {}
-    RETURN_RESULT_OF(kotlin::gc::tryRef, referee);
+
+ALWAYS_INLINE OBJ_GETTER(mm::weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept {
+    RETURN_RESULT_OF(gc::weakRefReadBarrier, weakReferee);
 }

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

@@ -13,6 +13,7 @@
 
 #include "ShadowStack.hpp"
 #include "StableRef.hpp"
+#include "TestSupport.hpp"
 
 using namespace kotlin;
 
@@ -92,54 +93,58 @@ TEST(ThreadRootSetTest, Empty) {
 }
 
 TEST(GlobalRootSetTest, Basic) {
-    mm::GlobalsRegistry globals;
+    RunInNewThread([](mm::ThreadData& threadData) {
-    mm::GlobalsRegistry::ThreadQueue globalsProducer(globals);
+        mm::GlobalsRegistry globals;
-    ObjHeader* global1 = reinterpret_cast<ObjHeader*>(1);
+        mm::GlobalsRegistry::ThreadQueue globalsProducer(globals);
-    ObjHeader* global2 = reinterpret_cast<ObjHeader*>(2);
+        ObjHeader* global1 = reinterpret_cast<ObjHeader*>(1);
-    globalsProducer.Insert(&global1);
+        ObjHeader* global2 = reinterpret_cast<ObjHeader*>(2);
-    globalsProducer.Insert(&global2);
+        globalsProducer.Insert(&global1);
+        globalsProducer.Insert(&global2);
 
-    mm::SpecialRefRegistry specialRefsRegistry;
+        mm::SpecialRefRegistry specialRefsRegistry;
-    mm::SpecialRefRegistry::ThreadQueue stableRefsProducer(specialRefsRegistry);
+        mm::SpecialRefRegistry::ThreadQueue stableRefsProducer(specialRefsRegistry);
-    ObjHeader* stableRef1 = reinterpret_cast<ObjHeader*>(3);
+        ObjHeader* stableRef1 = reinterpret_cast<ObjHeader*>(3);
-    ObjHeader* stableRef2 = reinterpret_cast<ObjHeader*>(4);
+        ObjHeader* stableRef2 = reinterpret_cast<ObjHeader*>(4);
-    ObjHeader* stableRef3 = reinterpret_cast<ObjHeader*>(5);
+        ObjHeader* stableRef3 = reinterpret_cast<ObjHeader*>(5);
-    auto stableRefHandle1 = stableRefsProducer.createStableRef(stableRef1);
+        auto stableRefHandle1 = stableRefsProducer.createStableRef(stableRef1);
-    auto stableRefHandle2 = stableRefsProducer.createStableRef(stableRef2);
+        auto stableRefHandle2 = stableRefsProducer.createStableRef(stableRef2);
-    auto stableRefHandle3 = stableRefsProducer.createStableRef(stableRef3);
+        auto stableRefHandle3 = stableRefsProducer.createStableRef(stableRef3);
 
-    globalsProducer.Publish();
+        globalsProducer.Publish();
-    stableRefsProducer.publish();
+        stableRefsProducer.publish();
 
-    mm::GlobalRootSet iter(globals, specialRefsRegistry);
+        mm::GlobalRootSet iter(globals, specialRefsRegistry);
 
-    std::vector<mm::GlobalRootSet::Value> actual;
+        std::vector<mm::GlobalRootSet::Value> actual;
-    for (auto object : iter) {
+        for (auto object : iter) {
-        actual.push_back(object);
+            actual.push_back(object);
-    }
+        }
 
-    auto asGlobal = [](ObjHeader*& object) -> mm::GlobalRootSet::Value { return {object, mm::GlobalRootSet::Source::kGlobal}; };
+        auto asGlobal = [](ObjHeader*& object) -> mm::GlobalRootSet::Value { return {object, mm::GlobalRootSet::Source::kGlobal}; };
-    auto asStableRef = [](ObjHeader*& object) -> mm::GlobalRootSet::Value { return {object, mm::GlobalRootSet::Source::kStableRef}; };
+        auto asStableRef = [](ObjHeader*& object) -> mm::GlobalRootSet::Value { return {object, mm::GlobalRootSet::Source::kStableRef}; };
-    EXPECT_THAT(
+        EXPECT_THAT(
-            actual,
+                actual,
-            testing::UnorderedElementsAre(
+                testing::UnorderedElementsAre(
-                    asGlobal(global1), asGlobal(global2), asStableRef(stableRef1), asStableRef(stableRef2), asStableRef(stableRef3)));
+                        asGlobal(global1), asGlobal(global2), asStableRef(stableRef1), asStableRef(stableRef2), asStableRef(stableRef3)));
 
-    std::move(stableRefHandle1).dispose();
+        std::move(stableRefHandle1).dispose();
-    std::move(stableRefHandle2).dispose();
+        std::move(stableRefHandle2).dispose();
-    std::move(stableRefHandle3).dispose();
+        std::move(stableRefHandle3).dispose();
+    });
 }
 
 TEST(GlobalRootSetTest, Empty) {
-    mm::GlobalsRegistry globals;
+    RunInNewThread([](mm::ThreadData& threadData) {
-    mm::SpecialRefRegistry specialRefsRegistry;
+        mm::GlobalsRegistry globals;
+        mm::SpecialRefRegistry specialRefsRegistry;
 
-    mm::GlobalRootSet iter(globals, specialRefsRegistry);
+        mm::GlobalRootSet iter(globals, specialRefsRegistry);
 
-    std::vector<mm::GlobalRootSet::Value> actual;
+        std::vector<mm::GlobalRootSet::Value> actual;
-    for (auto object : iter) {
+        for (auto object : iter) {
-        actual.push_back(object);
+            actual.push_back(object);
-    }
+        }
 
-    EXPECT_THAT(actual, testing::IsEmpty());
+        EXPECT_THAT(actual, testing::IsEmpty());
+    });
 }

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

@@ -10,6 +10,7 @@
 
 #include "GC.hpp"
 #include "Memory.h"
+#include "ReferenceOps.hpp"
 #include "RawPtr.hpp"
 #include "ReferenceOps.hpp"
 #include "ThreadRegistry.hpp"
@@ -124,7 +125,7 @@ class SpecialRefRegistry : private Pinned {
                     return;
                 }
 
-                // TODO: With CMS barrier for marking `obj_` should be here.
+                // TODO: With CMS root-set-write barrier for marking `obj_` should be here.
                 //       Until we have the barrier, the object must already be in the roots.
                 //       If 0->1 happened from `[ObjCClass _tryRetain]`, it would first hold the object
                 //       on the stack via `tryRef`.
@@ -206,7 +207,9 @@ public:
                     // roots. So, it'll either publish the dying thread itself, or
                     // if the dying thread has already deregistered, it means it published
                     // itself. In any case, global root scanning happens afterwards.
-                    // TODO: With CMS barrier for marking `node.obj_` should be here.
+                    // TODO: With CMS root-set-write barrier for marking `node.obj_` should be here.
+                    // TODO: No barrier that uses mm::ThreadData can be placed here.
+
                     owner_.insertIntoRootsHead(node);
                 }
             }

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

@@ -10,9 +10,11 @@
 #include "GC.hpp"
 #include "GlobalData.hpp"
 #include "GlobalsRegistry.hpp"
+#include "ObjectOps.hpp"
 #include "TestSupport.hpp"
 #include "ThreadData.hpp"
 #include "ThreadState.hpp"
+#include "ObjectTestSupport.hpp"
 
 using namespace kotlin;
 
@@ -62,3 +64,16 @@ void kotlin::DeinitMemoryForTests(MemoryState* memoryState) {
 std::ostream& kotlin::operator<<(std::ostream& stream, ThreadState state) {
     return stream << ThreadStateName(state);
 }
+
+test_support::RegularWeakReferenceImpl& test_support::InstallWeakReference(
+        mm::ThreadData& threadData, ObjHeader* objHeader, ObjHeader** location)
+{
+    mm::AllocateObject(&threadData, theRegularWeakReferenceImplTypeInfo, location);
+    auto& weakReference = test_support::RegularWeakReferenceImpl::FromObjHeader(*location);
+    auto& extraObjectData = mm::ExtraObjectData::GetOrInstall(objHeader);
+    weakReference->weakRef = static_cast<mm::RawSpecialRef*>(mm::WeakRef::create(objHeader));
+    weakReference->referred = objHeader;
+    auto* setWeakRef = extraObjectData.GetOrSetRegularWeakReferenceImpl(objHeader, weakReference.header());
+    EXPECT_EQ(setWeakRef, weakReference.header());
+    return weakReference;
+}

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

@@ -10,7 +10,9 @@
 #include <ostream>
 
 #include "MemoryPrivate.hpp"
+#include "ObjectTestSupport.hpp"
 #include "ThreadData.hpp"
+#include "WeakRef.hpp"
 
 namespace kotlin {
 
@@ -24,4 +26,10 @@ inline void RunInNewThread(std::function<void(mm::ThreadData&)> f) {
 // to pretty print ThreadState constants.
 std::ostream& operator<<(std::ostream& stream, ThreadState state);
 
+namespace test_support {
+
+RegularWeakReferenceImpl& InstallWeakReference(mm::ThreadData& threadData, ObjHeader* objHeader, ObjHeader** location);
+
+} // namespace test_support
+
 } // namespace kotlin