[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/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/BinaryOptions.kt

@@ -52,6 +52,8 @@ object BinaryOptions : BinaryOptionRegistry() {
     val mimallocUseCompaction by booleanOption()
 
     val compileBitcodeWithXcodeLlvm by booleanOption()
+
+    val objcDisposeOnMain by booleanOption()
 }
 
 open class BinaryOption<T : Any>(

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/KonanConfig.kt

@@ -194,6 +194,10 @@ class KonanConfig(val project: Project, val configuration: CompilerConfiguration
         configuration.get(BinaryOptions.mimallocUseCompaction) ?: false
     }
 
+    val objcDisposeOnMain: Boolean by lazy {
+        configuration.get(BinaryOptions.objcDisposeOnMain) ?: true
+    }
+
     init {
         if (!platformManager.isEnabled(target)) {
             error("Target ${target.visibleName} is not available on the ${HostManager.hostName} host")

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/IrToBitcode.kt

@@ -2743,6 +2743,7 @@ internal class CodeGeneratorVisitor(
         overrideRuntimeGlobal("Kotlin_appStateTracking", llvm.constInt32(context.config.appStateTracking.value))
         overrideRuntimeGlobal("Kotlin_mimallocUseDefaultOptions", llvm.constInt32(if (context.config.mimallocUseDefaultOptions) 1 else 0))
         overrideRuntimeGlobal("Kotlin_mimallocUseCompaction", llvm.constInt32(if (context.config.mimallocUseCompaction) 1 else 0))
+        overrideRuntimeGlobal("Kotlin_objcDisposeOnMain", llvm.constInt32(if (context.config.objcDisposeOnMain) 1 else 0))
     }
 
     //-------------------------------------------------------------------------//

kotlin-native/backend.native/tests/build.gradle

@@ -4442,6 +4442,10 @@ if (PlatformInfo.isAppleTarget(project)) {
     createInterop("frameworkForwardDeclarations") {
         it.defFile 'framework/forwardDeclarations/clib.def'
     }
+    createInterop("objc_kt56402") {
+        it.defFile 'interop/objc/kt56402/objclib.def'
+        it.headers "$projectDir/interop/objc/kt56402/objclib.h"
+    }
 }
 
 createInterop("withSpaces") {
@@ -5284,6 +5288,27 @@ if (PlatformInfo.isAppleTarget(project)) {
         useGoldenData = true
         UtilsKt.dependsOnPlatformLibs(it)
     }
+
+    interopTest("interop_objc_kt56402") {
+        // Test depends on macOS-specific AppKit
+        enabled = (project.testTarget == 'macos_x64' || project.testTarget == 'macos_arm64' || project.testTarget == null) && isExperimentalMM
+        source = 'interop/objc/kt56402/main.kt'
+        interop = "objc_kt56402"
+        flags = ["-tr", "-e", "runAllTests"] // Generate test suites, but use custom entrypoint.
+
+        doBeforeBuild {
+            mkdir(buildDir)
+            project.extensions.execClang.execClangForCompilerTests(project.target) {
+                args "$projectDir/interop/objc/kt56402/objclib.m"
+                args "-lobjc", '-fobjc-arc'
+                args '-fPIC', '-shared', '-o', "$buildDir/libobjc_kt56402.dylib"
+                args '-framework', 'AppKit'
+            }
+            if (UtilsKt.isSimulatorTarget(project, project.target)) {
+                UtilsKt.codesign(project, "$buildDir/libobjc_kt56402.dylib")
+            }
+        }
+    }
 }
 
 tasks.register("KT-50983", KonanDriverTest) {

kotlin-native/backend.native/tests/interop/objc/kt56402/main.kt

@@ -0,0 +1,186 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+import objclib.*
+
+import kotlin.native.concurrent.*
+import kotlin.native.internal.test.testLauncherEntryPoint
+import kotlin.system.exitProcess
+import kotlin.test.*
+import kotlinx.cinterop.*
+
+class LockedSet<T> {
+    private val lock = AtomicInt(0)
+    private val impl = mutableSetOf<T>()
+
+    private inline fun <R> locked(f: () -> R): R {
+        while (!lock.compareAndSet(0, 1)) {}
+        try {
+            return f()
+        } finally {
+            lock.value = 0
+        }
+    }
+
+    fun add(id: T) = locked {
+        assertFalse(id in impl)
+        impl.add(id)
+    }
+
+    fun remove(id: T) = locked {
+        assertTrue(id in impl)
+        impl.remove(id)
+    }
+
+    operator fun contains(id: T) = locked {
+        id in impl
+    }
+}
+
+class OnDestroyHookSub(onDestroy: (ULong) -> Unit) : OnDestroyHook(onDestroy)
+
+val aliveObjectIds = LockedSet<ULong>()
+
+fun alloc(ctor: ((ULong) -> Unit) -> ULong): ULong = autoreleasepool {
+    val id = ctor {
+        aliveObjectIds.remove(it)
+    }
+    aliveObjectIds.add(id)
+    id
+}
+
+fun waitDestruction(id: ULong) {
+    assertTrue(isMainThread())
+    kotlin.native.internal.GC.collect()
+    while (true) {
+        spin()
+        if (!(id in aliveObjectIds)) {
+            break
+        }
+    }
+}
+
+@Test
+fun testOnMainThread() {
+    assertTrue(isMainThread())
+    val id = alloc { onDestroy ->
+        OnDestroyHook {
+            assertTrue(isMainThread())
+            onDestroy(it)
+        }.identity()
+    }
+    waitDestruction(id)
+}
+
+@Test
+fun testOnSecondaryThread() {
+    val id = withWorker {
+        execute(TransferMode.SAFE, {}) {
+            assertFalse(isMainThread())
+            alloc { onDestroy ->
+                OnDestroyHook {
+                    assertFalse(isMainThread())
+                    onDestroy(it)
+                }.identity()
+            }
+        }.result
+    }
+    waitDestruction(id)
+}
+
+@Test
+fun testSubOnMainThread() {
+    assertTrue(isMainThread())
+    val id = alloc { onDestroy ->
+        OnDestroyHookSub {
+            assertTrue(isMainThread())
+            onDestroy(it)
+        }.identity()
+    }
+    waitDestruction(id)
+}
+
+@Test
+fun testSubOnSecondaryThread() {
+    val id = withWorker {
+        execute(TransferMode.SAFE, {}) {
+            assertFalse(isMainThread())
+            alloc { onDestroy ->
+                OnDestroyHookSub {
+                    assertFalse(isMainThread())
+                    onDestroy(it)
+                }.identity()
+            }
+        }.result
+    }
+    waitDestruction(id)
+}
+
+@Test
+fun testGlobalOnMainThread() {
+    assertTrue(isMainThread())
+    val id = alloc { onDestroy ->
+        val obj = newGlobal {
+            assertTrue(isMainThread())
+            onDestroy(it)
+        }!!
+        clearGlobal()
+        obj.identity()
+    }
+    waitDestruction(id)
+}
+
+@Test
+fun testGlobalOnSecondaryThread() {
+    val id = withWorker {
+        execute(TransferMode.SAFE, {}) {
+            assertFalse(isMainThread())
+            alloc { onDestroy ->
+                val obj = newGlobal {
+                    assertFalse(isMainThread())
+                    onDestroy(it)
+                }!!
+                clearGlobal()
+                obj.identity()
+            }
+        }.result
+    }
+    waitDestruction(id)
+}
+
+@Test
+fun testProtocolOnMainThread() {
+    assertTrue(isMainThread())
+    val id = alloc { onDestroy ->
+        newOnDestroyHook {
+            assertTrue(isMainThread())
+            onDestroy(it)
+        }!!.identity()
+    }
+    waitDestruction(id)
+}
+
+@Test
+fun testProtocolOnSecondaryThread() {
+    val id = withWorker {
+        execute(TransferMode.SAFE, {}) {
+            assertFalse(isMainThread())
+            alloc { onDestroy ->
+                newOnDestroyHook {
+                    assertFalse(isMainThread())
+                    onDestroy(it)
+                }!!.identity()
+            }
+        }.result
+    }
+    waitDestruction(id)
+}
+
+fun runAllTests(args: Array<String>) = startApp {
+    val exitCode = testLauncherEntryPoint(args)
+    if (exitCode != 0) {
+        exitProcess(exitCode)
+    }
+}

kotlin-native/backend.native/tests/interop/objc/kt56402/objclib.def

@@ -0,0 +1,3 @@
+language = Objective-C
+headerFilter = **/objclib.h
+linkerOpts = -lobjc_kt56402

kotlin-native/backend.native/tests/interop/objc/kt56402/objclib.h

@@ -0,0 +1,21 @@
+#include <inttypes.h>
+#include <objc/NSObject.h>
+
+@interface OnDestroyHook: NSObject
+-(instancetype)init:(void(^)(uintptr_t))onDestroy;
+-(uintptr_t)identity;
+@end
+
+@protocol OnDestroyHook
+-(uintptr_t)identity;
+@end
+
+void startApp(void(^task)());
+uint64_t currentThreadId();
+BOOL isMainThread();
+void spin();
+
+OnDestroyHook* newGlobal(void(^onDestroy)(uintptr_t));
+void clearGlobal();
+
+id<OnDestroyHook> newOnDestroyHook(void(^onDestroy)(uintptr_t));

kotlin-native/backend.native/tests/interop/objc/kt56402/objclib.m

@@ -0,0 +1,101 @@
+#include "objclib.h"
+
+#include <assert.h>
+#include <dispatch/dispatch.h>
+#include <pthread.h>
+#import <AppKit/NSApplication.h>
+#import <Foundation/NSRunLoop.h>
+#import <Foundation/NSThread.h>
+
+@implementation OnDestroyHook {
+    void (^onDestroy_)(uintptr_t);
+}
+
+- (uintptr_t)identity {
+    return (uintptr_t)self;
+}
+
+- (instancetype)init:(void (^)(uintptr_t))onDestroy {
+    if (self = [super init]) {
+        onDestroy_ = onDestroy;
+    }
+    return self;
+}
+
+- (void)dealloc {
+    onDestroy_([self identity]);
+}
+
+@end
+
+@interface OnDestroyHookPrivate : NSObject<OnDestroyHook>
+@end
+
+@implementation OnDestroyHookPrivate {
+    void (^onDestroy_)(uintptr_t);
+}
+
+- (uintptr_t)identity {
+    return (uintptr_t)self;
+}
+
+- (instancetype)init:(void (^)(uintptr_t))onDestroy {
+    if (self = [super init]) {
+        onDestroy_ = onDestroy;
+    }
+    return self;
+}
+
+- (void)dealloc {
+    onDestroy_([self identity]);
+}
+
+@end
+
+void startApp(void (^task)()) {
+    dispatch_async(dispatch_get_main_queue(), ^{
+        // At this point all other scheduled main queue tasks were already executed.
+        // Executing via performBlock to allow a recursive run loop in `spin()`.
+        [[NSRunLoop currentRunLoop] performBlock:^{
+            task();
+            [NSApp terminate:NULL];
+        }];
+    });
+    [[NSApplication sharedApplication] run];
+}
+
+uint64_t currentThreadId() {
+    uint64_t result;
+    int ret = pthread_threadid_np(NULL, &result);
+    assert(ret == 0);
+    return result;
+}
+
+BOOL isMainThread() {
+    return [NSThread isMainThread];
+}
+
+void spin() {
+    if ([NSRunLoop currentRunLoop] != [NSRunLoop mainRunLoop]) {
+        fprintf(stderr, "Must spin main run loop\n");
+        exit(1);
+    }
+    [[NSRunLoop currentRunLoop]
+           runMode:NSDefaultRunLoopMode
+        beforeDate:[NSDate dateWithTimeIntervalSinceNow:0.1]];
+}
+
+OnDestroyHook* global = NULL;
+
+OnDestroyHook* newGlobal(void(^onDestroy)(uintptr_t)) {
+    global = [[OnDestroyHook alloc] init:onDestroy];
+    return global;
+}
+
+void clearGlobal() {
+    global = NULL;
+}
+
+id<OnDestroyHook> newOnDestroyHook(void(^onDestroy)(uintptr_t)) {
+    return [[OnDestroyHookPrivate alloc] init:onDestroy];
+}

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