From fffc921d42c2bde31767e8066bb9b56f955b7e36 Mon Sep 17 00:00:00 2001
From: Troels Bjerre Lund <troels@google.com>
Date: Fri, 9 Jun 2023 13:01:00 +0000
Subject: [PATCH] [K/N] custom_alloc: fix race on extraobject ^KT-55364

This avoids that the thread sweeping the base object can get stale
locals, in the case that the finalizer thread unlinks the extraobject at
the same time. This is done by moving the responsibility of unlinking
the extraobject to the sweeping of the base object.


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

Merge-request: KT-MR-10518
Merged-by: Alexander Shabalin <Alexander.Shabalin@jetbrains.com>
---
 kotlin-native/runtime/src/custom_alloc/cpp/GCApi.cpp | 7 +++++--
 kotlin-native/runtime/src/mm/cpp/ExtraObjectData.cpp | 9 ++++++++-
 kotlin-native/runtime/src/mm/cpp/ExtraObjectData.hpp | 3 +++
 kotlin-native/runtime/src/mm/cpp/Memory.cpp          | 5 +++--
 4 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/kotlin-native/runtime/src/custom_alloc/cpp/GCApi.cpp b/kotlin-native/runtime/src/custom_alloc/cpp/GCApi.cpp
index b4526779687..767c7f3bd62 100644
--- a/kotlin-native/runtime/src/custom_alloc/cpp/GCApi.cpp
+++ b/kotlin-native/runtime/src/custom_alloc/cpp/GCApi.cpp
@@ -38,7 +38,8 @@ bool SweepObject(uint8_t* object, FinalizerQueue& finalizerQueue, gc::GCHandle::
         gcHandle.addKeptObject();
         return true;
     }
-    auto* extraObject = mm::ExtraObjectData::Get(&objHeader->object);
+    auto* baseObject = &objHeader->object;
+    auto* extraObject = mm::ExtraObjectData::Get(baseObject);
     if (extraObject) {
         if (!extraObject->getFlag(mm::ExtraObjectData::FLAGS_IN_FINALIZER_QUEUE)) {
             CustomAllocDebug("SweepObject(%p): needs to be finalized, extraObject at %p", object, extraObject);
@@ -49,11 +50,13 @@ bool SweepObject(uint8_t* object, FinalizerQueue& finalizerQueue, gc::GCHandle::
             gcHandle.addMarkedObject();
             return true;
         }
-        if (!extraObject->getFlag(mm::ExtraObjectData::FLAGS_SWEEPABLE)) {
+        if (!extraObject->getFlag(mm::ExtraObjectData::FLAGS_FINALIZED)) {
             CustomAllocDebug("SweepObject(%p): already waiting to be finalized", object);
             gcHandle.addMarkedObject();
             return true;
         }
+        extraObject->UnlinkFromBaseObject();
+        extraObject->setFlag(mm::ExtraObjectData::FLAGS_SWEEPABLE);
     }
     CustomAllocDebug("SweepObject(%p): can be reclaimed", object);
     gcHandle.addSweptObject();
diff --git a/kotlin-native/runtime/src/mm/cpp/ExtraObjectData.cpp b/kotlin-native/runtime/src/mm/cpp/ExtraObjectData.cpp
index b2f70845d36..5e4d2332978 100644
--- a/kotlin-native/runtime/src/mm/cpp/ExtraObjectData.cpp
+++ b/kotlin-native/runtime/src/mm/cpp/ExtraObjectData.cpp
@@ -52,13 +52,15 @@ mm::ExtraObjectData& mm::ExtraObjectData::Install(ObjHeader* object) noexcept {
     return data;
 }
 
-void mm::ExtraObjectData::Uninstall() noexcept {
+void mm::ExtraObjectData::UnlinkFromBaseObject() noexcept {
     auto *object = GetBaseObject();
     atomicSetRelease(const_cast<const TypeInfo**>(&object->typeInfoOrMeta_), typeInfo_);
     RuntimeAssert(
             !object->has_meta_object(), "Object %p has metaobject %p after removing metaobject %p", object, object->meta_object_or_null(),
             this);
+}
 
+void mm::ExtraObjectData::ReleaseAssociatedObject() noexcept {
 #ifdef KONAN_OBJC_INTEROP
     if (void* associatedObject = associatedObject_) {
         if (getFlag(FLAGS_RELEASE_ON_MAIN_QUEUE) && isMainQueueProcessorAvailable()) {
@@ -73,6 +75,11 @@ void mm::ExtraObjectData::Uninstall() noexcept {
 #endif
 }
 
+void mm::ExtraObjectData::Uninstall() noexcept {
+    UnlinkFromBaseObject();
+    ReleaseAssociatedObject();
+}
+
 bool mm::ExtraObjectData::HasAssociatedObject() noexcept {
 #ifdef KONAN_OBJC_INTEROP
     return associatedObject_ != nullptr;
diff --git a/kotlin-native/runtime/src/mm/cpp/ExtraObjectData.hpp b/kotlin-native/runtime/src/mm/cpp/ExtraObjectData.hpp
index b5a36d01832..5118b14f17f 100644
--- a/kotlin-native/runtime/src/mm/cpp/ExtraObjectData.hpp
+++ b/kotlin-native/runtime/src/mm/cpp/ExtraObjectData.hpp
@@ -28,6 +28,7 @@ public:
         FLAGS_IN_FINALIZER_QUEUE = 2,
         FLAGS_SWEEPABLE = 3,
         FLAGS_RELEASE_ON_MAIN_QUEUE = 4,
+        FLAGS_FINALIZED = 5,
     };
 
     static constexpr unsigned WEAK_REF_TAG = 1;
@@ -43,11 +44,13 @@ public:
 
     static ExtraObjectData& Install(ObjHeader* object) noexcept;
     void Uninstall() noexcept;
+    void UnlinkFromBaseObject() noexcept;
 
 #ifdef KONAN_OBJC_INTEROP
     std::atomic<void*>& AssociatedObject() noexcept { return associatedObject_; }
 #endif
     bool HasAssociatedObject() noexcept;
+    void ReleaseAssociatedObject() noexcept;
 
     bool getFlag(Flags value) noexcept { return (flags_.load() & (1u << static_cast<uint32_t>(value))) != 0; }
     void setFlag(Flags value) noexcept { flags_.fetch_or(1u << static_cast<uint32_t>(value)); }
diff --git a/kotlin-native/runtime/src/mm/cpp/Memory.cpp b/kotlin-native/runtime/src/mm/cpp/Memory.cpp
index fb68843403e..388c8c8a9f0 100644
--- a/kotlin-native/runtime/src/mm/cpp/Memory.cpp
+++ b/kotlin-native/runtime/src/mm/cpp/Memory.cpp
@@ -78,10 +78,11 @@ MetaObjHeader* ObjHeader::createMetaObject(ObjHeader* object) {
 void ObjHeader::destroyMetaObject(ObjHeader* object) {
     RuntimeAssert(object->has_meta_object(), "Object must have a meta object set");
     auto &extraObject = *mm::ExtraObjectData::Get(object);
-    extraObject.Uninstall();
 #ifdef CUSTOM_ALLOCATOR
-    extraObject.setFlag(mm::ExtraObjectData::FLAGS_SWEEPABLE);
+    extraObject.ReleaseAssociatedObject();
+    extraObject.setFlag(mm::ExtraObjectData::FLAGS_FINALIZED);
 #else
+    extraObject.Uninstall();
     auto *threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
     mm::ExtraObjectDataFactory::Instance().DestroyExtraObjectData(threadData, extraObject);
 #endif