Implement thread-safe tracking of Obj-C references to Kotlin objects

Fix #2052, #3204

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

@@ -442,6 +442,11 @@ internal class Llvm(val context: Context, val llvmModule: LLVMModuleRef) {
  val freezeSubgraph = importRtFunction("FreezeSubgraph")
  val checkMainThread = importRtFunction("CheckIsMainThread")
 
+ val kRefSharedHolderInitLocal = importRtFunction("KRefSharedHolder_initLocal")
+ val kRefSharedHolderInit = importRtFunction("KRefSharedHolder_init")
+ val kRefSharedHolderDispose = importRtFunction("KRefSharedHolder_dispose")
+ val kRefSharedHolderRef = importRtFunction("KRefSharedHolder_ref")
+
  val createKotlinObjCClass by lazy { importRtFunction("CreateKotlinObjCClass") }
  val getObjCKotlinTypeInfo by lazy { importRtFunction("GetObjCKotlinTypeInfo") }
  val missingInitImp by lazy { importRtFunction("MissingInitImp") }

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/objcexport/BlockPointerSupport.kt

@@ -136,9 +136,11 @@ private val BlockPointerBridge.nameSuffix: String
 internal class BlockAdapterToFunctionGenerator(val objCExportCodeGenerator: ObjCExportCodeGenerator) {
  private val codegen get() = objCExportCodeGenerator.codegen
 
+ private val kRefSharedHolderType = LLVMGetTypeByName(codegen.runtime.llvmModule, "class.KRefSharedHolder")!!
+
  private val blockLiteralType = structType(
  codegen.runtime.getStructType("Block_literal_1"),
- codegen.kObjHeaderPtr
+ kRefSharedHolderType
  )
 
  private val blockDescriptorType = codegen.runtime.getStructType("Block_descriptor_1")
@@ -149,8 +151,8 @@ internal class BlockAdapterToFunctionGenerator(val objCExportCodeGenerator: ObjC
  "blockDisposeHelper"
  ) {
  val blockPtr = bitcast(pointerType(blockLiteralType), param(0))
- val slot = structGep(blockPtr, 1)
- storeHeapRef(kNullObjHeaderPtr, slot) // TODO: can dispose_helper write to the block?
+ val refHolder = structGep(blockPtr, 1)
+ call(context.llvm.kRefSharedHolderDispose, listOf(refHolder))
 
  ret(null)
  }.also {
@@ -163,15 +165,22 @@ internal class BlockAdapterToFunctionGenerator(val objCExportCodeGenerator: ObjC
  "blockCopyHelper"
  ) {
  val dstBlockPtr = bitcast(pointerType(blockLiteralType), param(0))
- val dstSlot = structGep(dstBlockPtr, 1)
+ val dstRefHolder = structGep(dstBlockPtr, 1)
 
  val srcBlockPtr = bitcast(pointerType(blockLiteralType), param(1))
- val srcSlot = structGep(srcBlockPtr, 1)
+ val srcRefHolder = structGep(srcBlockPtr, 1)
+
+ // Note: in current implementation copy helper is invoked only for stack-allocated blocks from the same thread,
+ // so it is technically not necessary to check owner.
+ // However this is not guaranteed by Objective-C runtime, so keep it suboptimal but reliable:
+ val ref = call(
+ context.llvm.kRefSharedHolderRef,
+ listOf(srcRefHolder),
+ exceptionHandler = ExceptionHandler.Caller,
+ verbatim = true
+ )
 
- // Kotlin reference was `memcpy`ed from src to dst, "revert" this:
- storeRefUnsafe(kNullObjHeaderPtr, dstSlot)
- // and copy properly:
- storeHeapRef(loadSlot(srcSlot, isVar = false), dstSlot)
+ call(context.llvm.kRefSharedHolderInit, listOf(dstRefHolder, ref))
 
  ret(null)
  }.also {
@@ -211,22 +220,17 @@ internal class BlockAdapterToFunctionGenerator(val objCExportCodeGenerator: ObjC
 
 
 
- private fun FunctionGenerationContext.storeRefUnsafe(value: LLVMValueRef, slot: LLVMValueRef) {
- assert(value.type == kObjHeaderPtr)
- assert(slot.type == kObjHeaderPtrPtr)
-
- store(
- bitcast(int8TypePtr, value),
- bitcast(pointerType(int8TypePtr), slot)
- )
- }
-
  private fun ObjCExportCodeGenerator.generateInvoke(bridge: BlockPointerBridge): ConstPointer {
  val numberOfParameters = bridge.numberOfParameters
 
  val result = generateFunction(codegen, bridge.blockInvokeLlvmType, "invokeBlock${bridge.nameSuffix}") {
  val blockPtr = bitcast(pointerType(blockLiteralType), param(0))
- val kotlinFunction = loadSlot(structGep(blockPtr, 1), isVar = false)
+ val kotlinFunction = call(
+ context.llvm.kRefSharedHolderRef,
+ listOf(structGep(blockPtr, 1)),
+ exceptionHandler = ExceptionHandler.Caller,
+ verbatim = true
+ )
 
  val args = (1 .. numberOfParameters).map { index ->
  objCReferenceToKotlin(param(index), Lifetime.ARGUMENT)
@@ -282,16 +286,15 @@ internal class BlockAdapterToFunctionGenerator(val objCExportCodeGenerator: ObjC
 
  val blockOnStack = alloca(blockLiteralType)
  val blockOnStackBase = structGep(blockOnStack, 0)
- val slot = structGep(blockOnStack, 1)
+ val refHolder = structGep(blockOnStack, 1)
 
  listOf(bitcast(int8TypePtr, isa), flags, reserved, invoke, descriptor).forEachIndexed { index, value ->
  // Although value is actually on the stack, it's not in normal slot area, so we cannot handle it
  // as if it was on the stack.
  store(value, structGep(blockOnStackBase, index))
  }
 
- // Note: it is the slot in the block located on stack, so no need to manage it properly:
- storeRefUnsafe(kotlinRef, slot)
+ call(context.llvm.kRefSharedHolderInitLocal, listOf(refHolder, kotlinRef))
 
  val copiedBlock = callFromBridge(retainBlock, listOf(bitcast(int8TypePtr, blockOnStack)))
 

backend.native/tests/framework/values/expectedLazy.h

@@ -960,6 +960,29 @@ __attribute__((swift_name("TestWeakRefs")))
 - (NSArray<id> *)createCycle __attribute__((swift_name("createCycle()")));
 @end;
 
+__attribute__((objc_subclassing_restricted))
+__attribute__((swift_name("SharedRefs")))
+@interface ValuesSharedRefs : KotlinBase
+- (instancetype)init __attribute__((swift_name("init()"))) __attribute__((objc_designated_initializer));
++ (instancetype)new __attribute__((availability(swift, unavailable, message="use object initializers instead")));
+- (ValuesSharedRefsMutableData *)createRegularObject __attribute__((swift_name("createRegularObject()")));
+- (void (^)(void))createLambda __attribute__((swift_name("createLambda()")));
+- (NSMutableArray<id> *)createCollection __attribute__((swift_name("createCollection()")));
+- (ValuesSharedRefsMutableData *)createFrozenRegularObject __attribute__((swift_name("createFrozenRegularObject()")));
+- (void (^)(void))createFrozenLambda __attribute__((swift_name("createFrozenLambda()")));
+- (NSMutableArray<id> *)createFrozenCollection __attribute__((swift_name("createFrozenCollection()")));
+- (BOOL)hasAliveObjects __attribute__((swift_name("hasAliveObjects()")));
+@end;
+
+__attribute__((objc_subclassing_restricted))
+__attribute__((swift_name("SharedRefs.MutableData")))
+@interface ValuesSharedRefsMutableData : KotlinBase
+- (instancetype)init __attribute__((swift_name("init()"))) __attribute__((objc_designated_initializer));
++ (instancetype)new __attribute__((availability(swift, unavailable, message="use object initializers instead")));
+- (void)update __attribute__((swift_name("update()")));
+@property int32_t x __attribute__((swift_name("x")));
+@end;
+
 @interface ValuesEnumeration (ValuesKt)
 - (ValuesEnumeration *)getAnswer __attribute__((swift_name("getAnswer()")));
 @end;

backend.native/tests/framework/values/values.kt

@@ -10,6 +10,7 @@ package conversions
 
 import kotlin.native.concurrent.freeze
 import kotlin.native.concurrent.isFrozen
+import kotlin.native.ref.WeakReference
 import kotlin.properties.ReadWriteProperty
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
@@ -738,3 +739,38 @@ class TestWeakRefs(private val frozen: Boolean) {
 
  private class Node(var next: Node?)
 }
+
+class SharedRefs {
+ class MutableData {
+ var x = 0
+
+ fun update() { x += 1 }
+ }
+
+ fun createRegularObject(): MutableData = create { MutableData() }
+
+ fun createLambda(): () -> Unit = create {
+ var mutableData = 0
+ {
+ println(mutableData++)
+ }
+ }
+
+ fun createCollection(): MutableList<Any> = create {
+ mutableListOf()
+ }
+
+ fun createFrozenRegularObject() = createRegularObject().freeze()
+ fun createFrozenLambda() = createLambda().freeze()
+ fun createFrozenCollection() = createCollection().freeze()
+
+ fun hasAliveObjects(): Boolean {
+ kotlin.native.internal.GC.collect()
+ return mustBeRemoved.any { it.get() != null }
+ }
+
+ private fun <T : Any> create(block: () -> T) = block()
+ .also { mustBeRemoved += WeakReference(it) }
+
+ private val mustBeRemoved = mutableListOf<WeakReference<*>>()
+}

backend.native/tests/framework/values/values.swift

@@ -808,6 +808,175 @@ func testWeakRefs0(frozen: Bool) throws {
  // try test4()
 }
 
+var falseFlag = false
+
+class TestSharedRefs {
+ private func testLambdaSimple() throws {
+ func getClosure() -> (() -> Void) {
+ let lambda = autoreleasepool {
+ SharedRefs().createLambda()
+ }
+ return { if falseFlag { lambda() } }
+ }
+
+ DispatchQueue.global().async(execute: getClosure())
+ }
+
+ private static func runInNewThread(initializeKotlinRuntime: Bool, block: @escaping () -> Void) {
+ class Closure {
+ static var currentBlock: (() -> Void)? = nil
+ static var initializeKotlinRuntime: Bool = false
+ }
+
+ Closure.currentBlock = block
+ Closure.initializeKotlinRuntime = initializeKotlinRuntime
+
+ var thread: pthread_t? = nil
+ let createCode = pthread_create(&thread, nil, { _ in
+ if Closure.initializeKotlinRuntime {
+ let ignore = SharedRefs() // Ensures that Kotlin runtime gets initialized.
+ }
+
+ Closure.currentBlock!()
+ Closure.currentBlock = nil
+
+ return nil
+ }, nil)
+ try! assertEquals(actual: createCode, expected: 0)
+
+ let joinCode = pthread_join(thread!, nil)
+ try! assertEquals(actual: joinCode, expected: 0)
+ }
+
+ private func runInNewThread(initializeKotlinRuntime: Bool, block: @escaping () -> Void) {
+ return TestSharedRefs.runInNewThread(initializeKotlinRuntime: initializeKotlinRuntime, block: block)
+ }
+
+ private func testObjectPartialRelease() {
+ let object = autoreleasepool { SharedRefs().createRegularObject() }
+ var objectVar: AnyObject? = object
+
+ runInNewThread(initializeKotlinRuntime: true) {
+ objectVar = nil
+ }
+ }
+
+ private func testRunRefCount<T>(
+ run: (@escaping () -> Void) -> Void,
+ createObject: @escaping (SharedRefs) -> T
+ ) throws {
+ let refs = SharedRefs()
+
+ var objectVar1: T? = autoreleasepool { createObject(refs) }
+ var objectVar2: T? = nil
+
+ try assertTrue(refs.hasAliveObjects())
+
+ run {
+ objectVar2 = objectVar1!
+ objectVar1 = nil
+ }
+
+ try assertTrue(refs.hasAliveObjects())
+
+ run {
+ objectVar2 = nil
+ }
+
+ try assertFalse(refs.hasAliveObjects())
+ }
+
+ private func testBackgroundRefCount<T>(createObject: @escaping (SharedRefs) -> T) throws {
+ try testRunRefCount(
+ run: { runInNewThread(initializeKotlinRuntime: false, block: $0) },
+ createObject: createObject
+ )
+
+ try testRunRefCount(
+ run: { runInNewThread(initializeKotlinRuntime: true, block: $0) },
+ createObject: createObject
+ )
+ }
+
+ private func testReferenceOutlivesThread(releaseWithKotlinRuntime: Bool) throws {
+ var objectVar: AnyObject? = nil
+ weak var objectWeakVar: AnyObject? = nil
+ var collection: AnyObject? = nil
+
+ runInNewThread(initializeKotlinRuntime: false) {
+ autoreleasepool {
+ let refs = SharedRefs()
+ collection = refs.createCollection()
+
+ let object = refs.createRegularObject()
+ objectVar = object
+ objectWeakVar = object
+
+ try! assertTrue(objectWeakVar === object)
+ }
+ }
+
+ runInNewThread(initializeKotlinRuntime: releaseWithKotlinRuntime) {
+ objectVar = nil
+ collection = nil
+ ValuesKt.gc()
+ try! assertTrue(objectWeakVar === nil)
+ }
+
+ }
+
+ private func testMoreWorkBeforeThreadExit() throws {
+ class Deinit {
+ static var object1: AnyObject? = nil
+ static var object2: AnyObject? = nil
+ static weak var weakVar2: AnyObject? = nil
+
+ deinit {
+ TestSharedRefs.runInNewThread(initializeKotlinRuntime: false) {
+ Deinit.object2 = nil
+ }
+ }
+ }
+
+ runInNewThread(initializeKotlinRuntime: false) {
+ autoreleasepool {
+ let object1 = SharedRefs.MutableData()
+ Deinit.object1 = object1
+ setAssociatedObject(object: object1, value: Deinit())
+
+ let object2 = SharedRefs.MutableData()
+ Deinit.object2 = object2
+ Deinit.weakVar2 = object2
+ }
+
+ TestSharedRefs.runInNewThread(initializeKotlinRuntime: false) {
+ Deinit.object1 = nil
+ }
+ }
+
+ try assertTrue(Deinit.weakVar2 === nil)
+ }
+
+ func test() throws {
+ try testLambdaSimple()
+ try testObjectPartialRelease()
+
+ try testBackgroundRefCount(createObject: { $0.createLambda() })
+ try testBackgroundRefCount(createObject: { $0.createRegularObject() })
+ try testBackgroundRefCount(createObject: { $0.createCollection() })
+
+ try testBackgroundRefCount(createObject: { $0.createFrozenLambda() })
+ try testBackgroundRefCount(createObject: { $0.createFrozenRegularObject() })
+ try testBackgroundRefCount(createObject: { $0.createFrozenCollection() })
+
+ try testReferenceOutlivesThread(releaseWithKotlinRuntime: false)
+ try testReferenceOutlivesThread(releaseWithKotlinRuntime: true)
+ try testMoreWorkBeforeThreadExit()
+
+ usleep(300 * 1000)
+ }
+}
+
 // See https:/JetBrains/kotlin-native/issues/2931
 func testGH2931() throws {
  for i in 0..<50000 {
@@ -877,6 +1046,7 @@ class ValuesTests : TestProvider {
  TestCase(name: "TestInvalidIdentifiers", method: withAutorelease(testInvalidIdentifiers)),
  TestCase(name: "TestDeprecation", method: withAutorelease(testDeprecation)),
  TestCase(name: "TestWeakRefs", method: withAutorelease(testWeakRefs)),
+ TestCase(name: "TestSharedRefs", method: withAutorelease(TestSharedRefs().test)),
  TestCase(name: "TestGH2931", method: withAutorelease(testGH2931)),
  ]
  }