[K/N] Dispose ObjC objects from main thread on main thread. ^KT-56402

Merge-request: KT-MR-9515
Merged-by: Alexander Shabalin <Alexander.Shabalin@jetbrains.com>

kotlin-native/runtime/src/main/cpp/CompilerConstants.cpp

@@ -30,6 +30,7 @@ RUNTIME_WEAK int32_t Kotlin_printToAndroidLogcat = 1;
 RUNTIME_WEAK int32_t Kotlin_appStateTracking = 0;
 RUNTIME_WEAK int32_t Kotlin_mimallocUseDefaultOptions = 1;
 RUNTIME_WEAK int32_t Kotlin_mimallocUseCompaction = 0;
+RUNTIME_WEAK int32_t Kotlin_objcDisposeOnMain = 0;
 
 ALWAYS_INLINE compiler::DestroyRuntimeMode compiler::destroyRuntimeMode() noexcept {
     return static_cast<compiler::DestroyRuntimeMode>(Kotlin_destroyRuntimeMode);
@@ -73,3 +74,7 @@ ALWAYS_INLINE bool compiler::mimallocUseDefaultOptions() noexcept {
 ALWAYS_INLINE bool compiler::mimallocUseCompaction() noexcept {
     return Kotlin_mimallocUseCompaction != 0;
 }
+
+ALWAYS_INLINE bool compiler::objcDisposeOnMain() noexcept {
+    return Kotlin_objcDisposeOnMain != 0;
+}

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

@@ -115,6 +115,7 @@ AppStateTracking appStateTracking() noexcept;
 int getSourceInfo(void* addr, SourceInfo *result, int result_size) noexcept;
 bool mimallocUseDefaultOptions() noexcept;
 bool mimallocUseCompaction() noexcept;
+bool objcDisposeOnMain() noexcept;
 
 #ifdef KONAN_ANDROID
 bool printToAndroidLogcat() noexcept;

kotlin-native/runtime/src/main/cpp/MainQueueProcessor.cpp

@@ -0,0 +1,46 @@
+/*
+ * 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 "MainQueueProcessor.hpp"
+
+#include <atomic>
+
+#include "KAssert.h"
+
+#if KONAN_SUPPORTS_GRAND_CENTRAL_DISPATCH
+#include <dispatch/dispatch.h>
+#endif
+
+using namespace kotlin;
+
+namespace {
+
+#if KONAN_SUPPORTS_GRAND_CENTRAL_DISPATCH
+std::atomic<bool> isMainQueueProcessed = false;
+#endif
+
+} // namespace
+
+void kotlin::initializeMainQueueProcessor() noexcept {
+#if KONAN_SUPPORTS_GRAND_CENTRAL_DISPATCH
+    dispatch_async_f(
+            dispatch_get_main_queue(), nullptr, [](void*) { isMainQueueProcessed.store(true, std::memory_order_relaxed); });
+#endif
+}
+
+bool kotlin::isMainQueueProcessorAvailable() noexcept {
+#if KONAN_SUPPORTS_GRAND_CENTRAL_DISPATCH
+    return isMainQueueProcessed.load(std::memory_order_relaxed);
+#else
+    return false;
+#endif
+}
+
+void kotlin::runOnMainQueue(void* arg, void (*f)(void*)) noexcept {
+    RuntimeAssert(isMainQueueProcessorAvailable(), "Running on main queue when it's not processed");
+#if KONAN_SUPPORTS_GRAND_CENTRAL_DISPATCH
+    dispatch_async_f(dispatch_get_main_queue(), arg, f);
+#endif
+}

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

@@ -0,0 +1,18 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+namespace kotlin {
+
+void initializeMainQueueProcessor() noexcept;
+
+bool isMainQueueProcessorAvailable() noexcept;
+
+// Run `f(arg)` on main queue without waiting for its completion.
+// Only valid if `isMainQueueProcessorAvailable()` returns true.
+void runOnMainQueue(void* arg, void (*f)(void*)) noexcept;
+
+} // namespace kotlin

kotlin-native/runtime/src/main/cpp/ObjCExport.mm

@@ -108,7 +108,8 @@ ALWAYS_INLINE void send_releaseAsAssociatedObject(void* associatedObject) {
 
 extern "C" ALWAYS_INLINE void Kotlin_ObjCExport_releaseAssociatedObject(void* associatedObject) {
   if (associatedObject != nullptr) {
-    kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
+    // May already be in the native state if was scheduled on the main queue.
+    NativeOrUnregisteredThreadGuard guard(/*reentrant=*/ true);
     send_releaseAsAssociatedObject(associatedObject);
   }
 }

kotlin-native/runtime/src/main/cpp/Runtime.cpp

@@ -8,6 +8,7 @@
 #include "CompilerConstants.hpp"
 #include "Exceptions.h"
 #include "KAssert.h"
+#include "MainQueueProcessor.hpp"
 #include "Memory.h"
 #include "ObjCExportInit.h"
 #include "ObjectAlloc.hpp"
@@ -135,6 +136,9 @@ RuntimeState* initRuntime() {
   // Keep global variables in state as well.
   if (firstRuntime) {
     konan::consoleInit();
+    if (compiler::objcDisposeOnMain()) {
+      kotlin::initializeMainQueueProcessor();
+    }
 #if KONAN_OBJC_INTEROP
     Kotlin_ObjCExport_initialize();
 #endif

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

@@ -151,6 +151,7 @@ extern "C" const char* Kotlin_callsCheckerGoodFunctionNames[] = {
         "setenv",
         "unsetenv",
 
+        "dispatch_async_f",
         "dispatch_once",
         "pthread_equal",
         "pthread_main_np",

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

@@ -5,6 +5,7 @@
 
 #include "ExtraObjectData.hpp"
 
+#include "MainQueueProcessor.hpp"
 #include "PointerBits.h"
 #include "ThreadData.hpp"
 
@@ -59,7 +60,13 @@ void mm::ExtraObjectData::Uninstall() noexcept {
             this);
 
 #ifdef KONAN_OBJC_INTEROP
-    Kotlin_ObjCExport_releaseAssociatedObject(associatedObject_);
+    if (getFlag(FLAGS_RELEASE_ON_MAIN_QUEUE) && isMainQueueProcessorAvailable()) {
+        runOnMainQueue(associatedObject_, [](void* obj) {
+            Kotlin_ObjCExport_releaseAssociatedObject(obj);
+        });
+    } else {
+        Kotlin_ObjCExport_releaseAssociatedObject(associatedObject_);
+    }
     associatedObject_ = nullptr;
 #endif
 }

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

@@ -27,6 +27,7 @@ public:
         FLAGS_NEVER_FROZEN = 1,
         FLAGS_IN_FINALIZER_QUEUE = 2,
         FLAGS_FINALIZED = 3,
+        FLAGS_RELEASE_ON_MAIN_QUEUE = 4,
     };
 
     static constexpr unsigned WEAK_REF_TAG = 1;

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

@@ -43,11 +43,27 @@ void* ObjHeader::GetAssociatedObject() const {
 }
 
 void ObjHeader::SetAssociatedObject(void* obj) {
-    return mm::ExtraObjectData::FromMetaObjHeader(meta_object()).AssociatedObject().store(obj, std::memory_order_release);
+    auto& extraObject = mm::ExtraObjectData::FromMetaObjHeader(meta_object());
+    // TODO: Consider additional filtering based on types:
+    //       * have some kind of an allowlist that can be populated by the user
+    //         to specify that objects of these types must be finalized only on
+    //         the main thread.
+    //       * prepopulate it for the system frameworks.
+    //       * if that were to be done at runtime, library authors could register
+    //         their types in a library initialization code.
+    if (pthread_main_np() == 1) {
+        extraObject.setFlag(mm::ExtraObjectData::FLAGS_RELEASE_ON_MAIN_QUEUE);
+    }
+    return extraObject.AssociatedObject().store(obj, std::memory_order_release);
 }
 
 void* ObjHeader::CasAssociatedObject(void* expectedObj, void* obj) {
-    mm::ExtraObjectData::FromMetaObjHeader(meta_object()).AssociatedObject().compare_exchange_strong(expectedObj, obj);
+    auto& extraObject = mm::ExtraObjectData::FromMetaObjHeader(meta_object());
+    bool success = extraObject.AssociatedObject().compare_exchange_strong(expectedObj, obj);
+    // TODO: Consider additional filtering outlined above.
+    if (success && pthread_main_np() == 1) {
+        extraObject.setFlag(mm::ExtraObjectData::FLAGS_RELEASE_ON_MAIN_QUEUE);
+    }
     return expectedObj;
 }