[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

@@ -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();

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;

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)); }

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