From f2501ad424e5705d1e65da96acebd49b0a3ff36c Mon Sep 17 00:00:00 2001
From: "Aleksei.Glushko" <aleksei.glushko@jetbrains.com>
Date: Wed, 13 Mar 2024 14:43:59 +0000
Subject: [PATCH] [K/N] Use already dereferenced weakReferee (KT-66371)

    And prevent weak barriers from execution during STW

Merge-request: KT-MR-14797
Merged-by: Alexey Glushko <aleksei.glushko@jetbrains.com>
---
 kotlin-native/runtime/src/gc/cms/cpp/Barriers.cpp       | 8 +++++++-
 kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMark.cpp | 3 +++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/kotlin-native/runtime/src/gc/cms/cpp/Barriers.cpp b/kotlin-native/runtime/src/gc/cms/cpp/Barriers.cpp
index 29105fe5b2f..df747e4c143 100644
--- a/kotlin-native/runtime/src/gc/cms/cpp/Barriers.cpp
+++ b/kotlin-native/runtime/src/gc/cms/cpp/Barriers.cpp
@@ -187,13 +187,19 @@ ALWAYS_INLINE ObjHeader* gc::barriers::weakRefReadBarrier(std::atomic<ObjHeader*
         auto phase = currentPhase();
         BarriersLogDebug(phase, "Weak read %p", weak);
 
+        // The weak read protector switches the thread state to native, thus making this code able to execute during STW.
+        // However, this is only possible with the disabled barriers.
+        // Be extra cautious not to access or modify heap structure here. e.g. do not allocate objects
+        AssertThreadState(ThreadState::kNative);
+        RuntimeAssert(!mm::IsThreadSuspensionRequested() || phase == BarriersPhase::kDisabled, "Unexpected barriers phase during STW: %s", toString(phase));
+
         if (__builtin_expect(phase == BarriersPhase::kMarkClosure, false)) {
             weakRefReadInMarkSlowPath(weak);
         } else {
             if (__builtin_expect(phase == BarriersPhase::kWeakProcessing, false)) {
                 // TODO reread the referee here under the barrier guard
                 //      if `disableBarriers` would be possible outside of STW
-                return weakRefReadInWeakSweepSlowPath(weakReferee);
+                return weakRefReadInWeakSweepSlowPath(weak);
             }
         }
         return weak;
diff --git a/kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMark.cpp b/kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMark.cpp
index 26e8f09a72f..7754e34fc87 100644
--- a/kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMark.cpp
+++ b/kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMark.cpp
@@ -87,6 +87,9 @@ void gc::mark::ConcurrentMark::runMainInSTW() {
     gc::processWeaks<DefaultProcessWeaksTraits>(gcHandle(), mm::SpecialRefRegistry::instance());
 
     if (!terminateInSTW) {
+        // Mutator threads execute weak barrier in "native" state. This hack maes them stop with the rest of the world.
+        std::unique_lock markTerminationGuard(markTerminationMutex_);
+
         stopTheWorld(gcHandle(), "GC stop the world #2: prepare to sweep");
     }