Store initializing singletons in a separate thread local small map (#4019)

* Add a test

* Use a specialized storage for shared shadow

* Switch to std::map

* Add another test

* Try using stack

* Remove unneeded addHeapRef

* Fix initializer6 test

* Do not use random in initializers7

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

@@ -921,7 +921,6 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
         }
         val singleton = context.llvmDeclarations.forSingleton(irClass)
         val instanceAddress = singleton.instanceStorage
-        val instanceShadowAddress = singleton.instanceShadowStorage
 
         if (storageKind == ObjectStorageKind.PERMANENT) {
             return loadSlot(instanceAddress.getAddress(this), false)
@@ -938,13 +937,13 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
         val typeInfo = codegen.typeInfoForAllocation(irClass)
         val defaultConstructor = irClass.constructors.single { it.valueParameters.size == 0 }
         val ctor = codegen.llvmFunction(defaultConstructor)
-        val (initFunction, args) =
+        val initFunction =
                 if (storageKind == ObjectStorageKind.SHARED && context.config.threadsAreAllowed) {
-                    val shadowObjectPtr = instanceShadowAddress!!.getAddress(this)
-                    context.llvm.initSharedInstanceFunction to listOf(objectPtr, shadowObjectPtr, typeInfo, ctor)
+                    context.llvm.initSharedInstanceFunction
                 } else {
-                    context.llvm.initInstanceFunction to listOf(objectPtr, typeInfo, ctor)
+                    context.llvm.initInstanceFunction
                 }
+        val args = listOf(objectPtr, typeInfo, ctor)
         val newValue = call(initFunction, args, Lifetime.GLOBAL, exceptionHandler)
         val bbInitResult = currentBlock
         br(bbExit)

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

@@ -67,7 +67,7 @@ internal class ClassLlvmDeclarations(
         val singletonDeclarations: SingletonLlvmDeclarations?,
         val objCDeclarations: KotlinObjCClassLlvmDeclarations?)
 
-internal class SingletonLlvmDeclarations(val instanceStorage: AddressAccess, val instanceShadowStorage: AddressAccess?)
+internal class SingletonLlvmDeclarations(val instanceStorage: AddressAccess)
 
 internal class KotlinObjCClassLlvmDeclarations(
         val classInfoGlobal: StaticData.Global,
@@ -278,12 +278,7 @@ private class DeclarationsGeneratorVisitor(override val context: Context) :
         val symbolName = "kobjref:" + qualifyInternalName(irClass)
         val instanceAddress = addKotlinGlobal(symbolName, getLLVMType(irClass.defaultType), threadLocal = threadLocal)
 
-        val instanceShadowAddress = if (threadLocal || storageKind == ObjectStorageKind.PERMANENT) null else {
-            val shadowSymbolName = "kshadowobjref:" + qualifyInternalName(irClass)
-            addKotlinGlobal(shadowSymbolName, getLLVMType(irClass.defaultType), threadLocal = true)
-        }
-
-        return SingletonLlvmDeclarations(instanceAddress, instanceShadowAddress)
+        return SingletonLlvmDeclarations(instanceAddress)
     }
 
     private fun createKotlinObjCClassDeclarations(irClass: IrClass): KotlinObjCClassLlvmDeclarations {

backend.native/tests/build.gradle

@@ -2672,6 +2672,15 @@ task initializers5(type: KonanLocalTest) {
     source = "runtime/basic/initializers5.kt"
 }
 
+task initializers6(type: KonanLocalTest) {
+    disabled = project.testTarget == 'wasm32' // Needs workers.
+    source = "runtime/basic/initializers6.kt"
+}
+
+task initializers7(type: KonanLocalTest) {
+    source = "runtime/basic/initializers7.kt"
+}
+
 task expression_as_statement(type: KonanLocalTest) {
     expectedFail = (project.testTarget == 'wasm32') // uses exceptions.
     goldValue = "Ok\n"

backend.native/tests/runtime/basic/initializers6.kt

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2010-2020 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.
+ */
+
+package runtime.basic.initializers6
+
+import kotlin.test.*
+
+import kotlin.native.concurrent.*
+
+val aWorkerId = AtomicInt(0)
+val bWorkersCount = 3
+
+val aWorkerUnlocker = AtomicInt(0)
+val bWorkerUnlocker = AtomicInt(0)
+
+object A {
+    init {
+        // Must be called by aWorker only.
+        assertEquals(aWorkerId.value, Worker.current.id)
+        // Only allow b workers to run, when a worker has started initialization.
+        bWorkerUnlocker.increment()
+        // Only proceed with initialization, when all b workers have started executing.
+        while (aWorkerUnlocker.value < bWorkersCount) {}
+        // And now wait a bit, to increase probability of races.
+        Worker.current.park(100 * 1000L)
+    }
+    val a = produceA()
+    val b = produceB()
+}
+
+fun produceA(): String {
+    // Must've been called by aWorker only.
+    assertEquals(aWorkerId.value, Worker.current.id)
+    return "A"
+}
+
+fun produceB(): String {
+    // Must've been called by aWorker only.
+    assertEquals(aWorkerId.value, Worker.current.id)
+    // Also check that it's ok to get A.a while initializing A.b.
+    return "B+${A.a}"
+}
+
+@Test fun runTest() {
+    val aWorker = Worker.start()
+    aWorkerId.value = aWorker.id
+    val bWorkers = Array(bWorkersCount, { _ -> Worker.start() })
+
+    val aFuture = aWorker.execute(TransferMode.SAFE, {}, {
+        A.b
+    })
+    val bFutures = Array(bWorkers.size, {
+        bWorkers[it].execute(TransferMode.SAFE, {}, {
+            // Wait until A has started to initialize.
+            while (bWorkerUnlocker.value < 1) {}
+            // Now allow A initialization to continue.
+            aWorkerUnlocker.increment()
+            // And this should not've tried to init A itself.
+            A.a + A.b
+        })
+    })
+
+    for (future in bFutures) {
+        assertEquals("AB+A", future.result)
+    }
+    assertEquals("B+A", aFuture.result)
+
+    for (worker in bWorkers) {
+        worker.requestTermination().result
+    }
+    aWorker.requestTermination().result
+}

backend.native/tests/runtime/basic/initializers7.kt

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2010-2020 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.
+ */
+
+package runtime.basic.initializers7
+
+import kotlin.test.*
+
+object A {
+    init {
+        assertAUninitialized()
+    }
+    val a1 = 7
+    val a2 = 12
+}
+
+// Check that A is initialized dynamically.
+fun assertAUninitialized() {
+    assertEquals(0, A.a1)
+    assertEquals(0, A.a2)
+}
+
+object B {
+    init {
+        assertBUninitialized()
+    }
+    val b1 = A.a2
+    val b2 = C.c1
+}
+
+// Check that B is initialized dynamically.
+fun assertBUninitialized() {
+    assertEquals(0, B.b1)
+    assertEquals(0, B.b2)
+}
+
+object C {
+    init {
+        assertCUninitialized()
+    }
+    val c1 = 42
+    val c2 = A.a1
+    val c3 = B.b1
+    val c4 = B.b2
+}
+
+// Check that C is initialized dynamically.
+fun assertCUninitialized() {
+    assertEquals(0, C.c1)
+    assertEquals(0, C.c2)
+    assertEquals(0, C.c3)
+    assertEquals(0, C.c4)
+}
+
+@Test fun runTest() {
+    assertEquals(A.a1, C.c2)
+    assertEquals(A.a2, C.c3)
+    assertEquals(C.c1, C.c4)
+}

runtime/src/main/cpp/Memory.cpp

@@ -451,6 +451,9 @@ struct MemoryState {
   uint64_t allocSinceLastGcThreshold;
 #endif // USE_GC
 
+  // A stack of initializing singletons.
+  KStdVector<std::pair<ObjHeader**, ObjHeader*>> initializingSingletons;
+
 #if COLLECT_STATISTIC
   #define CONTAINER_ALLOC_STAT(state, size, container) state->statistic.incAlloc(size, container);
   #define CONTAINER_DESTROY_STAT(state, container) \
@@ -1999,7 +2002,7 @@ OBJ_GETTER(initInstance,
 
 template <bool Strict>
 OBJ_GETTER(initSharedInstance,
-    ObjHeader** location, ObjHeader** localLocation, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
+    ObjHeader** location, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
 #if KONAN_NO_THREADS
   ObjHeader* value = *location;
   if (value != nullptr) {
@@ -2025,25 +2028,31 @@ OBJ_GETTER(initSharedInstance,
   }
 #endif  // KONAN_NO_EXCEPTIONS
 #else  // KONAN_NO_THREADS
-  ObjHeader* value = *localLocation;
-  if (value != nullptr) RETURN_OBJ(value);
+  // Search from the top of the stack.
+  for (auto it = memoryState->initializingSingletons.rbegin(); it != memoryState->initializingSingletons.rend(); ++it) {
+    if (it->first == location) {
+      RETURN_OBJ(it->second);
+    }
+  }
 
   ObjHeader* initializing = reinterpret_cast<ObjHeader*>(1);
 
   // Spin lock.
+  ObjHeader* value = nullptr;
   while ((value = __sync_val_compare_and_swap(location, nullptr, initializing)) == initializing);
   if (value != nullptr) {
     // OK'ish, inited by someone else.
     RETURN_OBJ(value);
   }
   ObjHeader* object = AllocInstance(typeInfo, OBJ_RESULT);
-  UpdateHeapRef(localLocation, object);
+  memoryState->initializingSingletons.push_back(std::make_pair(location, object));
 #if KONAN_NO_EXCEPTIONS
   ctor(object);
   if (Strict)
     FreezeSubgraph(object);
   UpdateHeapRef(location, object);
   synchronize();
+  memoryState->initializingSingletons.pop_back();
   return object;
 #else  // KONAN_NO_EXCEPTIONS
   try {
@@ -2052,11 +2061,12 @@ OBJ_GETTER(initSharedInstance,
       FreezeSubgraph(object);
     UpdateHeapRef(location, object);
     synchronize();
+    memoryState->initializingSingletons.pop_back();
     return object;
   } catch (...) {
     UpdateReturnRef(OBJ_RESULT, nullptr);
     zeroHeapRef(location);
-    zeroHeapRef(localLocation);
+    memoryState->initializingSingletons.pop_back();
     synchronize();
     throw;
   }
@@ -2803,12 +2813,12 @@ OBJ_GETTER(InitInstanceRelaxed,
 }
 
 OBJ_GETTER(InitSharedInstanceStrict,
-    ObjHeader** location, ObjHeader** localLocation, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
-  RETURN_RESULT_OF(initSharedInstance<true>, location, localLocation, typeInfo, ctor);
+    ObjHeader** location, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
+  RETURN_RESULT_OF(initSharedInstance<true>, location, typeInfo, ctor);
 }
 OBJ_GETTER(InitSharedInstanceRelaxed,
-    ObjHeader** location, ObjHeader** localLocation, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
-  RETURN_RESULT_OF(initSharedInstance<false>, location, localLocation, typeInfo, ctor);
+    ObjHeader** location, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
+  RETURN_RESULT_OF(initSharedInstance<false>, location, typeInfo, ctor);
 }
 
 void SetStackRefStrict(ObjHeader** location, const ObjHeader* object) {

runtime/src/main/cpp/Memory.h

@@ -461,11 +461,11 @@ OBJ_GETTER(InitInstance,
     ObjHeader** location, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*));
 
 OBJ_GETTER(InitSharedInstanceStrict,
-    ObjHeader** location, ObjHeader** localLocation, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*));
+    ObjHeader** location, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*));
 OBJ_GETTER(InitSharedInstanceRelaxed,
-    ObjHeader** location, ObjHeader** localLocation, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*));
+    ObjHeader** location, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*));
 OBJ_GETTER(InitSharedInstance,
-    ObjHeader** location, ObjHeader** localLocation, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*));
+    ObjHeader** location, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*));
 
 // Weak reference operations.
 // Atomically clears counter object reference.

runtime/src/relaxed/cpp/MemoryImpl.cpp

@@ -25,8 +25,8 @@ OBJ_GETTER(InitInstance,
 }
 
 OBJ_GETTER(InitSharedInstance,
-    ObjHeader** location, ObjHeader** localLocation, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
-  RETURN_RESULT_OF(InitSharedInstanceRelaxed, location, localLocation, typeInfo, ctor);
+    ObjHeader** location, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
+  RETURN_RESULT_OF(InitSharedInstanceRelaxed, location, typeInfo, ctor);
 }
 
 void ReleaseHeapRef(const ObjHeader* object) {

runtime/src/strict/cpp/MemoryImpl.cpp

@@ -25,8 +25,8 @@ OBJ_GETTER(InitInstance,
 }
 
 OBJ_GETTER(InitSharedInstance,
-    ObjHeader** location, ObjHeader** localLocation, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
-  RETURN_RESULT_OF(InitSharedInstanceStrict, location, localLocation, typeInfo, ctor);
+    ObjHeader** location, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
+  RETURN_RESULT_OF(InitSharedInstanceStrict, location, typeInfo, ctor);
 }
 
 void ReleaseHeapRef(const ObjHeader* object) {