[K/N] Remember StableRefs, released during RC colleciton (KT-70159)

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

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

@@ -126,8 +126,14 @@ void gc::barriers::disableBarriers() noexcept {
 namespace {
 
 // TODO decide whether it's really beneficial to NO_INLINE the slow path
-NO_INLINE void beforeHeapRefUpdateSlowPath(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {
-    auto prev = ref.load();
+NO_INLINE void beforeHeapRefUpdateSlowPath(mm::DirectRefAccessor ref, ObjHeader* value, bool loadAtomic) noexcept {
+    ObjHeader* prev;
+    if (loadAtomic) {
+        prev = ref.loadAtomic(std::memory_order_relaxed);
+    } else {
+        prev = ref.load();
+    }
+
     if (prev != nullptr && prev->heap()) {
         // TODO Redundant if the destination object is black.
         //      Yet at the moment there is now efficient way to distinguish black and gray objects.
@@ -143,11 +149,11 @@ NO_INLINE void beforeHeapRefUpdateSlowPath(mm::DirectRefAccessor ref, ObjHeader*
 
 } // namespace
 
-ALWAYS_INLINE void gc::barriers::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {
+ALWAYS_INLINE void gc::barriers::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value, bool loadAtomic) noexcept {
     auto phase = currentPhase();
     BarriersLogDebug(phase, "Write *%p <- %p (%p overwritten)", ref.location(), value, ref.load());
     if (__builtin_expect(phase == BarriersPhase::kMarkClosure, false)) {
-        beforeHeapRefUpdateSlowPath(ref, value);
+        beforeHeapRefUpdateSlowPath(ref, value, loadAtomic);
     }
 }
 
@@ -175,7 +181,7 @@ NO_INLINE ObjHeader* weakRefReadInWeakSweepSlowPath(ObjHeader* weakReferee) noex
 
 } // namespace
 
-ALWAYS_INLINE ObjHeader* gc::barriers::weakRefReadBarrier(std::atomic<ObjHeader*>& weakReferee) noexcept {
+ALWAYS_INLINE ObjHeader* gc::barriers::weakRefReadBarrier(std_support::atomic_ref<ObjHeader*> weakReferee) noexcept {
     if (__builtin_expect(currentPhase() != BarriersPhase::kDisabled, false)) {
         // Mark dispatcher requires weak reads be protected by the following:
         auto weakReadProtector = markDispatcher().weakReadProtector();

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

@@ -34,8 +34,8 @@ void enableBarriers(int64_t epoch) noexcept;
 void switchToWeakProcessingBarriers() noexcept;
 void disableBarriers() noexcept;
 
-void beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept;
+void beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value, bool loadAtomic) noexcept;
 
-ObjHeader* weakRefReadBarrier(std::atomic<ObjHeader*>& weakReferee) noexcept;
+ObjHeader* weakRefReadBarrier(std_support::atomic_ref<ObjHeader*> weakReferee) noexcept;
 
 } // namespace kotlin::gc::barriers

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

@@ -6,6 +6,7 @@
 #include "ConcurrentMarkAndSweep.hpp"
 
 #include <mutex>
+#include <atomic>
 
 #include "gtest/gtest.h"
 
@@ -110,5 +111,41 @@ TYPED_TEST_P(TracingGCTest, WeakResurrectionAtMarkTermination) {
     }
 }
 
-REGISTER_TYPED_TEST_SUITE_WITH_LISTS(TracingGCTest, TRACING_GC_TEST_LIST, WeakResurrectionAtMarkTermination);
+TYPED_TEST_P(TracingGCTest, ReleaseStableRefDuringRSCollection) {
+    std::vector<Mutator> mutators(kDefaultThreadCount);
+
+    std::atomic<size_t> readyThreads = 0;
+    std::atomic<bool> gcDone = false;
+
+    std::vector<std::future<void>> mutatorFutures;
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        mutatorFutures.push_back(mutators[i].Execute([&](mm::ThreadData& threadData, Mutator& mutator) {
+            auto& obj = AllocateObject(threadData);
+            auto stableRef = mm::StableRef::create(obj.header());
+
+            ++readyThreads;
+            while (!mm::IsThreadSuspensionRequested()) {}
+
+            mm::safePoint();
+
+            mutator.AddStackRoot(obj.header());
+            std::move(stableRef).dispose();
+
+            while (!gcDone) { mm::safePoint(); }
+
+            EXPECT_THAT(mutator.Alive(), testing::Contains(obj.header()));
+        }));
+    }
+
+    while (readyThreads < kDefaultThreadCount) {}
+
+    mm::GlobalData::Instance().gcScheduler().scheduleAndWaitFinalized();
+    gcDone = true;
+
+    for (auto& future : mutatorFutures) {
+        future.wait();
+    }
+}
+
+REGISTER_TYPED_TEST_SUITE_WITH_LISTS(TracingGCTest, TRACING_GC_TEST_LIST, WeakResurrectionAtMarkTermination, ReleaseStableRefDuringRSCollection);
 INSTANTIATE_TYPED_TEST_SUITE_P(CMS, TracingGCTest, ConcurrentMarkAndSweepTest);

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

@@ -87,11 +87,11 @@ bool gc::GC::mainThreadFinalizerProcessorAvailable() noexcept {
     return impl_->gc().mainThreadFinalizerProcessor().available();
 }
 
-ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {
-    barriers::beforeHeapRefUpdate(ref, value);
+ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value, bool loadAtomic) noexcept {
+    barriers::beforeHeapRefUpdate(ref, value, loadAtomic);
 }
 
-ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept {
+ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std_support::atomic_ref<ObjHeader*> weakReferee) noexcept {
     RETURN_OBJ(gc::barriers::weakRefReadBarrier(weakReferee));
 }
 

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

@@ -87,8 +87,8 @@ class GC : private Pinned {
     std::unique_ptr<Impl> impl_;
 };
 
-void beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept;
-OBJ_GETTER(weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept;
+void beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value, bool loadAtomic) noexcept;
+OBJ_GETTER(weakRefReadBarrier, std_support::atomic_ref<ObjHeader*> weakReferee) noexcept;
 
 bool isMarked(ObjHeader* object) noexcept;
 

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

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

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

@@ -62,9 +62,9 @@ bool gc::GC::mainThreadFinalizerProcessorAvailable() noexcept {
     return false;
 }
 
-ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {}
+ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value, bool loadAtomic) noexcept {}
 
-ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept {
+ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std_support::atomic_ref<ObjHeader*> weakReferee) noexcept {
     RETURN_OBJ(weakReferee.load(std::memory_order_relaxed));
 }
 

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

@@ -30,7 +30,7 @@ ObjHeader* weakRefBarrierImpl(ObjHeader* weakReferee) noexcept {
     return weakReferee;
 }
 
-NO_INLINE ObjHeader* weakRefReadSlowPath(std::atomic<ObjHeader*>& weakReferee) noexcept {
+NO_INLINE ObjHeader* weakRefReadSlowPath(std_support::atomic_ref<ObjHeader*> weakReferee) noexcept {
     // reread an action to avoid register pollution outside the function
     auto barrier = weakRefBarrier.load(std::memory_order_seq_cst);
 
@@ -107,7 +107,7 @@ void gc::DisableWeakRefBarriers() noexcept {
     }
 }
 
-OBJ_GETTER(gc::WeakRefRead, std::atomic<ObjHeader*>& weakReferee) noexcept {
+OBJ_GETTER(gc::WeakRefRead, std_support::atomic_ref<ObjHeader*> weakReferee) noexcept {
     if (!compiler::concurrentWeakSweep()) {
         RETURN_OBJ(weakReferee.load(std::memory_order_relaxed));
     }

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

@@ -32,6 +32,6 @@ class BarriersThreadData : private Pinned {
 void EnableWeakRefBarriers(int64_t epoch) noexcept;
 void DisableWeakRefBarriers() noexcept;
 
-OBJ_GETTER(WeakRefRead, std::atomic<ObjHeader*>& weakReferee) noexcept;
+OBJ_GETTER(WeakRefRead, std_support::atomic_ref<ObjHeader*> weakReferee) noexcept;
 
 } // namespace kotlin::gc

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

@@ -87,9 +87,9 @@ bool gc::GC::mainThreadFinalizerProcessorAvailable() noexcept {
     return impl_->gc().mainThreadFinalizerProcessor().available();
 }
 
-ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {}
+ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value, bool loadAtomic) noexcept {}
 
-ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept {
+ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std_support::atomic_ref<ObjHeader*> weakReferee) noexcept {
     RETURN_RESULT_OF(gc::WeakRefRead, weakReferee);
 }
 

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

@@ -80,9 +80,9 @@ bool gc::GC::mainThreadFinalizerProcessorAvailable() noexcept {
     return impl_->gc().mainThreadFinalizerProcessor().available();
 }
 
-ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value) noexcept {}
+ALWAYS_INLINE void gc::beforeHeapRefUpdate(mm::DirectRefAccessor ref, ObjHeader* value, bool loadAtomic) noexcept {}
 
-ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept {
+ALWAYS_INLINE OBJ_GETTER(gc::weakRefReadBarrier, std_support::atomic_ref<ObjHeader*> weakReferee) noexcept {
     RETURN_OBJ(weakReferee.load(std::memory_order_relaxed));
 }
 

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

@@ -210,6 +210,6 @@ class RefField : private Pinned {
     ObjHeader* value_ = nullptr;
 };
 
-OBJ_GETTER(weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept;
+OBJ_GETTER(weakRefReadBarrier, std_support::atomic_ref<ObjHeader*> weakReferee) noexcept;
 
 }

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

@@ -17,12 +17,12 @@ template<> ALWAYS_INLINE void mm::RefAccessor<true>::afterLoad() noexcept {}
 
 // on heap
 template<> ALWAYS_INLINE void mm::RefAccessor<false>::beforeStore(ObjHeader* value) noexcept {
-    gc::beforeHeapRefUpdate(direct(), value);
+    gc::beforeHeapRefUpdate(direct(), value, false);
 }
 template<> ALWAYS_INLINE void mm::RefAccessor<false>::afterStore(ObjHeader*) noexcept {}
 template<> ALWAYS_INLINE void mm::RefAccessor<false>::beforeLoad() noexcept {}
 template<> ALWAYS_INLINE void mm::RefAccessor<false>::afterLoad() noexcept {}
 
-ALWAYS_INLINE OBJ_GETTER(mm::weakRefReadBarrier, std::atomic<ObjHeader*>& weakReferee) noexcept {
+ALWAYS_INLINE OBJ_GETTER(mm::weakRefReadBarrier, std_support::atomic_ref<ObjHeader*> weakReferee) noexcept {
     RETURN_RESULT_OF(gc::weakRefReadBarrier, weakReferee);
 }

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

@@ -56,7 +56,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_.load(std::memory_order_relaxed) != nullptr, "candidate cannot have a null obj_");
+            RuntimeAssert(candidate->objAtomic().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

@@ -88,7 +88,7 @@ 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);
+                    auto* obj = objAtomic().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
@@ -107,19 +107,19 @@ 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_.load(std::memory_order_relaxed);
+            return objAtomic().load(std::memory_order_relaxed);
         }
 
         OBJ_GETTER0(tryRef) noexcept {
             AssertThreadState(ThreadState::kRunnable);
-            RETURN_RESULT_OF(mm::weakRefReadBarrier, obj_);
+            RETURN_RESULT_OF(mm::weakRefReadBarrier, objAtomic());
         }
 
         void retainRef() noexcept {
             auto rc = rc_.fetch_add(1, std::memory_order_relaxed);
             RuntimeAssert(rc >= 0, "Retaining StableRef@%p with rc %d", this, rc);
             if (rc == 0) {
-                if (!obj_.load(std::memory_order_relaxed)) {
+                if (!objAtomic().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
@@ -129,27 +129,37 @@ class SpecialRefRegistry : private Pinned {
                     return;
                 }
 
-                // 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`.
-                //       If 0->1 happened during construction:
-                //       * First of all, currently it's impossible because the `Node` is created with rc=1 and not inserted
-                //         into the roots list until publishing.
-                //       * Even if the above changes, for the construction, the object must have been passed in from somewhere,
-                //         so it must be reachable anyway.
-                //       If 0->1 happened because an object is passing through the interop border for the second time (or more)
-                //       (e.g. accessing a non-permanent global a couple of times). Follows the construction case above:
-                //       "the object must have been passed in from somewhere, so it must be reachable anyway".
+                // With the current CMS implementation no barrier is required here.
+                // The CMS builds Snapshot-at-the-beginning mark closure,
+                // which means, it has to remember only the concurrent deletion of references, not creation.
+                // TODO: A write-into-root-set barrier might be required here for other concurrent mark strategies.
+
+                // In case of non-concurrent root set scanning, it is only required for the object to already be in roots.
+                // If 0->1 happened from `[ObjCClass _tryRetain]`, it would first hold the object
+                // on the stack via `tryRef`.
+                // If 0->1 happened during construction:
+                // * First of all, currently it's impossible because the `Node` is created with rc=1 and not inserted
+                //   into the roots list until publishing.
+                // * Even if the above changes, for the construction, the object must have been passed in from somewhere,
+                //   so it must be reachable anyway.
+                // If 0->1 happened because an object is passing through the interop border for the second time (or more)
+                // (e.g. accessing a non-permanent global a couple of times). Follows the construction case above:
+                // "the object must have been passed in from somewhere, so it must be reachable anyway".
 
                 // 0->1 changes require putting this node into the root set.
                 SpecialRefRegistry::instance().insertIntoRootsHead(*this);
             }
         }
 
         void releaseRef() noexcept {
-            auto rc = rc_.fetch_sub(1, std::memory_order_relaxed);
-            RuntimeAssert(rc > 0, "Releasing StableRef@%p with rc %d", this, rc);
+            auto rcBefore = rc_.fetch_sub(1, std::memory_order_relaxed);
+            RuntimeAssert(rcBefore > 0, "Releasing StableRef@%p with rc %d", this, rcBefore);
+            if (rcBefore == 1) {
+                // It's potentially a removal from global root set.
+                // The CMS GC scans global root set concurrently.
+                // Notify GC about the removal.
+                gc::beforeHeapRefUpdate(mm::DirectRefAccessor(obj_), nullptr, true);
+            }
         }
 
         RawSpecialRef* asRaw() noexcept { return reinterpret_cast<RawSpecialRef*>(this); }
@@ -169,7 +179,10 @@ class SpecialRefRegistry : private Pinned {
         //   Synchronization between GC and mutators happens via enabling/disabling
         //   the barriers.
         // TODO: Try to handle it atomically only when the GC is in progress.
-        std::atomic<ObjHeader*> obj_;
+        std_support::atomic_ref<ObjHeader*> objAtomic() noexcept  { return std_support::atomic_ref{obj_}; }
+        std_support::atomic_ref<ObjHeader* const> objAtomic() const noexcept { return std_support::atomic_ref{obj_}; }
+        ObjHeader* obj_;
+
         // Only ever updated using relaxed memory ordering. Any synchronization
         // with nextRoot_ is achieved via acquire-release of nextRoot_.
         std::atomic<Rc> rc_; // After dispose() will be disposedMarker.
@@ -221,7 +234,7 @@ class SpecialRefRegistry : private Pinned {
         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_.load(std::memory_order_relaxed);
+            return node_->objAtomic().load(std::memory_order_relaxed);
         }
 
         RootsIterator& operator++() noexcept {
@@ -258,7 +271,7 @@ class SpecialRefRegistry : private Pinned {
 
     class Iterator {
     public:
-        std::atomic<ObjHeader*>& operator*() noexcept { return iterator_->obj_; }
+        std_support::atomic_ref<ObjHeader*> operator*() noexcept { return iterator_->objAtomic(); }
 
         Iterator& operator++() noexcept {
             iterator_ = owner_->findAliveNode(std::next(iterator_));

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

@@ -58,7 +58,7 @@ class SpecialRefRegistryTest : public testing::Test {
     std::vector<ObjHeader*> all(Invalidated&&... invalidated) noexcept {
         std::set<ObjHeader*> invalidatedSet({std::forward<Invalidated>(invalidated)...});
         std::vector<ObjHeader*> result;
-        for (auto& obj : mm::SpecialRefRegistry::instance().lockForIter()) {
+        for (auto obj : mm::SpecialRefRegistry::instance().lockForIter()) {
             if (invalidatedSet.find(obj) != invalidatedSet.end()) {
                 obj = nullptr;
             }