[K/N][runtime][interop] Refactoring to run cinterop in daemon

2021-10-04 11:10:00 +05:00
parent 8907c3506d
commit 47858126da
14 changed files with 302 additions and 87 deletions
@@ -5,13 +5,29 @@

 package org.jetbrains.kotlin.native.interop.indexer

+import kotlinx.cinterop.JvmCInteropCallbacks
+import org.jetbrains.kotlin.konan.util.NativeMemoryAllocator
 import java.io.File
+import kotlin.test.AfterTest
+import kotlin.test.BeforeTest

 open class IndexerTests {
    init {
        System.load(System.getProperty("kotlin.native.llvm.libclang"))
    }

+    @BeforeTest
+    fun init() {
+        NativeMemoryAllocator.init()
+        JvmCInteropCallbacks.init()
+    }
+
+    @AfterTest
+    fun dispose() {
+        JvmCInteropCallbacks.dispose()
+        NativeMemoryAllocator.dispose()
+    }
+
    class TempFiles(name: String) {
        private val tempRootDir = System.getProperty("kotlin.native.interop.indexer.temp") ?: System.getProperty("java.io.tmpdir") ?: "."

@@ -110,6 +110,15 @@ JNIEXPORT jlong JNICALL Java_kotlinx_cinterop_JvmCallbacksKt_ffiTypeStruct0(JNIE
    return (jlong) res;
 }

+/*
+ * Class:     kotlinx_cinterop_JvmCallbacksKt
+ * Method:    ffiFreeTypeStruct0
+ * Signature: (J)V
+ */
+JNIEXPORT void JNICALL Java_kotlinx_cinterop_JvmCallbacksKt_ffiFreeTypeStruct0(JNIEnv *env, jclass cls, jlong ptr) {
+    if (ptr) free((void*)ptr);
+}
+
 /*
 * Class:     kotlinx_cinterop_JvmCallbacksKt
 * Method:    ffiCreateCif0
@@ -120,6 +129,7 @@ JNIEXPORT jlong JNICALL Java_kotlinx_cinterop_JvmCallbacksKt_ffiCreateCif0(JNIEn
    if (res != NULL) {
        ffi_status status = ffi_prep_cif(res, FFI_DEFAULT_ABI, nArgs, (ffi_type*)rType, (ffi_type**)argTypes);
        if (status != FFI_OK) {
+            free(res);
            if (status == FFI_BAD_TYPEDEF) {
                return -(jlong)1;
            } else if (status == FFI_BAD_ABI) {
@@ -132,6 +142,15 @@ JNIEXPORT jlong JNICALL Java_kotlinx_cinterop_JvmCallbacksKt_ffiCreateCif0(JNIEn
    return (jlong) res;
 }

+/*
+ * Class:     kotlinx_cinterop_JvmCallbacksKt
+ * Method:    ffiFreeCif0
+ * Signature: (J)V
+ */
+JNIEXPORT void JNICALL Java_kotlinx_cinterop_JvmCallbacksKt_ffiFreeCif0(JNIEnv *env, jclass cls, jlong ptr) {
+    if (ptr) free((void*)ptr);
+}
+
 static JavaVM *vm = NULL;

 // Returns the JNI env which can be used by the caller.
@@ -192,9 +211,10 @@ static void ffi_fun(ffi_cif *cif, void *ret, void **args, void *user_data) {
 * Method:    ffiCreateClosure0
 * Signature: (JLjava/lang/Object;)J
 */
-JNIEXPORT jlong JNICALL Java_kotlinx_cinterop_JvmCallbacksKt_ffiCreateClosure0(JNIEnv *env, jclass cls, jlong ffiCif, jobject userData) {
+JNIEXPORT jlong JNICALL Java_kotlinx_cinterop_JvmCallbacksKt_ffiCreateClosure0(JNIEnv *env, jclass cls, jlong ffiCif, jlong ffiClosure, jobject userData) {
    jobject userDataGlobalRef = (*env)->NewGlobalRef(env, userData);
    if (userDataGlobalRef == NULL) {
+        *(ffi_closure**)ffiClosure = NULL;
        return (jlong)0;
    }

@@ -204,16 +224,36 @@ JNIEXPORT jlong JNICALL Java_kotlinx_cinterop_JvmCallbacksKt_ffiCreateClosure0(J
    void* res;
    ffi_closure *closure = ffi_closure_alloc(sizeof(ffi_closure), &res);
    if (closure == NULL) {
+        (*env)->DeleteGlobalRef(env, userDataGlobalRef);
+        *(ffi_closure**)ffiClosure = NULL;
        return (jlong)0;
    }
    ffi_status status = ffi_prep_closure_loc(closure, (ffi_cif*)ffiCif, ffi_fun, userDataPtr, res);
    if (status != FFI_OK) {
+        (*env)->DeleteGlobalRef(env, userDataGlobalRef);
+        ffi_closure_free(closure);
+        *(ffi_closure**)ffiClosure = NULL;
        return -(jlong)1;
    }

+    *(ffi_closure**)ffiClosure = closure;
    return (jlong) res;
 }

+/*
+ * Class:     kotlinx_cinterop_JvmCallbacksKt
+ * Method:    ffiFreeClosure0
+ * Signature: (J)V
+ */
+JNIEXPORT void JNICALL Java_kotlinx_cinterop_JvmCallbacksKt_ffiFreeClosure0(JNIEnv *env, jclass cls, jlong ptr) {
+    if (ptr == NULL) return;
+    ffi_closure *closure = (ffi_closure*)ptr;
+    void* userDataPtr = closure->user_data;
+    if (userDataPtr)
+        (*env)->DeleteGlobalRef(env, (jobject)userDataPtr);
+    ffi_closure_free(closure);
+}
+
 /*
 * Class:     kotlinx_cinterop_JvmCallbacksKt
 * Method:    newGlobalRef
@@ -16,6 +16,9 @@

 package kotlinx.cinterop

+import org.jetbrains.kotlin.konan.util.NativeMemoryAllocator
+import org.jetbrains.kotlin.konan.util.ThreadSafeDisposableHelper
+import sun.misc.Unsafe
 import java.util.concurrent.ConcurrentHashMap
 import java.util.function.LongConsumer
 import kotlin.reflect.KClass
@@ -58,9 +61,60 @@ private fun getVariableCType(type: KType): CType<*>? {
    }
 }

-private val structTypeCache = ConcurrentHashMap<Class<*>, CType<*>>()
+internal class Caches {
+    val structTypeCache = ConcurrentHashMap<Class<*>, CType<*>>()
+    val createdStaticFunctions = ConcurrentHashMap<FunctionSpec, CPointer<CFunction<*>>>()

-private fun getStructCType(structClass: KClass<*>): CType<*> = structTypeCache.computeIfAbsent(structClass.java) {
+    // TODO: No concurrent bag or something in Java?
+    private val createdTypeStructs = mutableListOf<NativePtr>()
+    private val createdCifs = mutableListOf<NativePtr>()
+    private val createdClosures = mutableListOf<NativePtr>()
+
+    fun addTypeStruct(ptr: NativePtr) {
+        synchronized(createdTypeStructs) { createdTypeStructs.add(ptr) }
+    }
+
+    fun addCif(ptr: NativePtr) {
+        synchronized(createdCifs) { createdCifs.add(ptr) }
+    }
+
+    fun addClosure(ptr: NativePtr) {
+        synchronized(createdClosures) { createdClosures.add(ptr) }
+    }
+
+    fun disposeFfi() {
+        createdTypeStructs.forEach { ffiFreeTypeStruct0(it) }
+        createdCifs.forEach { ffiFreeCif0(it) }
+        createdClosures.forEach { ffiFreeClosure0(it) }
+    }
+}
+
+@PublishedApi
+internal val jvmCallbacksDisposeHelper = ThreadSafeDisposableHelper(
+        {
+            NativeMemoryAllocator.init()
+            Caches()
+        },
+        {
+            try {
+                it.disposeFfi()
+            } finally {
+                NativeMemoryAllocator.dispose()
+            }
+        }
+)
+
+inline fun <R> usingJvmCInteropCallbacks(block: () -> R) = jvmCallbacksDisposeHelper.usingDisposable(block)
+
+object JvmCInteropCallbacks {
+    fun init() = jvmCallbacksDisposeHelper.create()
+    fun dispose() = jvmCallbacksDisposeHelper.dispose()
+}
+
+private val caches: Caches
+    get() = jvmCallbacksDisposeHelper.holder ?: error("Caches hasn't been created")
+
+private fun getStructCType(structClass: KClass<*>): CType<*> = caches.structTypeCache.computeIfAbsent(structClass.java) {
    // Note that struct classes are not supposed to be user-defined,
    // so they don't require to be checked strictly.

@@ -192,15 +246,13 @@ private fun isStatic(function: Function<*>): Boolean {
    }
 }

-private data class FunctionSpec(val functionClass: Class<*>, val returnType: KType, val parameterTypes: List<KType>)
-
-private val createdStaticFunctions = ConcurrentHashMap<FunctionSpec, CPointer<CFunction<*>>>()
+internal data class FunctionSpec(val functionClass: Class<*>, val returnType: KType, val parameterTypes: List<KType>)

@Suppress("UNCHECKED_CAST")
@PublishedApi
 internal fun <F : Function<*>> staticCFunctionImpl(function: F, returnType: KType, vararg parameterTypes: KType): CPointer<CFunction<F>> {
    val spec = FunctionSpec(function.javaClass, returnType, parameterTypes.asList())
-    return createdStaticFunctions.computeIfAbsent(spec) {
+    return caches.createdStaticFunctions.computeIfAbsent(spec) {
        createStaticCFunction(function, spec)
    } as CPointer<CFunction<F>>
 }
@@ -300,8 +352,8 @@ private inline fun ffiClosureImpl(
 *
 * This description omits the details that are irrelevant for the ABI.
 */
-private abstract class CType<T> internal constructor(val ffiType: ffi_type) {
-    internal constructor(ffiTypePtr: Long) : this(interpretPointed<ffi_type>(ffiTypePtr))
+internal abstract class CType<T> constructor(val ffiType: ffi_type) {
+    constructor(ffiTypePtr: Long) : this(interpretPointed<ffi_type>(ffiTypePtr))
    abstract fun read(location: NativePtr): T
    abstract fun write(location: NativePtr, value: T): Unit
 }
@@ -384,7 +436,7 @@ internal class ffi_type(rawPtr: NativePtr) : COpaque(rawPtr)
 /**
 * Reference to `ffi_cif` struct instance.
 */
-internal class ffi_cif(rawPtr: NativePtr) : COpaque(rawPtr)
+private class ffi_cif(rawPtr: NativePtr) : COpaque(rawPtr)

 private external fun ffiTypeVoid(): Long
 private external fun ffiTypeUInt8(): Long
@@ -398,6 +450,7 @@ private external fun ffiTypeSInt64(): Long
 private external fun ffiTypePointer(): Long

 private external fun ffiTypeStruct0(elements: Long): Long
+private external fun ffiFreeTypeStruct0(ptr: Long)

 /**
 * Allocates and initializes `ffi_type` describing the struct.
@@ -405,15 +458,19 @@ private external fun ffiTypeStruct0(elements: Long): Long
 * @param elements types of the struct elements
 */
 private fun ffiTypeStruct(elementTypes: List<ffi_type>): ffi_type {
-    val elements = persistentHeap.allocArrayOfPointersTo(*elementTypes.toTypedArray(), null)
+    val elements = nativeHeap.allocArrayOfPointersTo(*elementTypes.toTypedArray(), null)
    val res = ffiTypeStruct0(elements.rawValue)
    if (res == 0L) {
        throw OutOfMemoryError()
    }
+
+    caches.addTypeStruct(res)
+
    return interpretPointed(res)
 }

 private external fun ffiCreateCif0(nArgs: Int, rType: Long, argTypes: Long): Long
+private external fun ffiFreeCif0(ptr: Long)

 /**
 * Creates and prepares an `ffi_cif`.
@@ -425,7 +482,7 @@ private external fun ffiCreateCif0(nArgs: Int, rType: Long, argTypes: Long): Lon
 */
 private fun ffiCreateCif(returnType: ffi_type, paramTypes: List<ffi_type>): ffi_cif {
    val nArgs = paramTypes.size
-    val argTypes = persistentHeap.allocArrayOfPointersTo(*paramTypes.toTypedArray(), null)
+    val argTypes = nativeHeap.allocArrayOfPointersTo(*paramTypes.toTypedArray(), null)
    val res = ffiCreateCif0(nArgs, returnType.rawPtr, argTypes.rawValue)

    when (res) {
@@ -435,10 +492,13 @@ private fun ffiCreateCif(returnType: ffi_type, paramTypes: List<ffi_type>): ffi_
        -3L -> throw Error("libffi error occurred")
    }

+    caches.addCif(res)
+
    return interpretPointed(res)
 }

-private external fun ffiCreateClosure0(ffiCif: Long, userData: Any): Long
+private external fun ffiCreateClosure0(ffiCif: Long, ffiClosure: Long, userData: Any): Long
+private external fun ffiFreeClosure0(ptr: Long)

 /**
 * Uses libffi to allocate a native function which will call [impl] when invoked.
@@ -446,34 +506,27 @@ private external fun ffiCreateClosure0(ffiCif: Long, userData: Any): Long
 * @param ffiCif describes the type of the function to create
 */
 private fun ffiCreateClosure(ffiCif: ffi_cif, impl: FfiClosureImpl): NativePtr {
-    val res = ffiCreateClosure0(ffiCif.rawPtr, userData = impl)
+    val ffiClosure = nativeHeap.alloc(Long.SIZE_BYTES, 8)

-    when (res) {
-        0L -> throw OutOfMemoryError()
-        -1L -> throw Error("libffi error occurred")
+    try {
+        val res = ffiCreateClosure0(ffiCif.rawPtr, ffiClosure.rawPtr, userData = impl)
+
+        when (res) {
+            0L -> throw OutOfMemoryError()
+            -1L -> throw Error("libffi error occurred")
+        }
+
+        caches.addClosure(unsafe.getLong(ffiClosure.rawPtr))
+
+        return res
+    } finally {
+        nativeHeap.free(ffiClosure)
    }
-
-    return res
 }

-// Callbacks are globally cached and outlive the memory allocated by nativeHeap,
-// which gets forcibly reclaimed at the end of the compiler invocation.
-// So use an ad hoc allocator that is not affected.
-// Note: this is mostly a workaround, but proper solution would require a significant rework of this machinery.
-private object persistentHeap : NativeFreeablePlacement {
-    override fun alloc(size: Long, align: Int): NativePointed {
-        return interpretOpaquePointed(
-                nativeMemUtils.allocRaw(
-                        if (size == 0L) 1L else size, // It is a hack: `nativeMemUtils.allocRaw` can't allocate zero bytes
-                        align
-                )
-        )
-    }
-
-    override fun free(mem: NativePtr) {
-        nativeMemUtils.freeRaw(mem)
-    }
-
+private val unsafe = with(Unsafe::class.java.getDeclaredField("theUnsafe")) {
+    isAccessible = true
+    return@with this.get(null) as Unsafe
 }

 private external fun newGlobalRef(any: Any): Long
@@ -16,6 +16,7 @@

 package org.jetbrains.kotlin.native.interop.gen.jvm

+import kotlinx.cinterop.usingJvmCInteropCallbacks
 import org.jetbrains.kotlin.konan.TempFiles
 import org.jetbrains.kotlin.konan.exec.Command
 import org.jetbrains.kotlin.konan.util.DefFile
@@ -33,6 +34,7 @@ import org.jetbrains.kotlin.konan.target.CompilerOutputKind
 import org.jetbrains.kotlin.konan.target.Distribution
 import org.jetbrains.kotlin.konan.target.KonanTarget
 import org.jetbrains.kotlin.konan.util.KonanHomeProvider
+import org.jetbrains.kotlin.konan.util.usingNativeMemoryAllocator
 import org.jetbrains.kotlin.library.KotlinLibrary
 import org.jetbrains.kotlin.library.packageFqName
 import org.jetbrains.kotlin.library.resolver.TopologicalLibraryOrder
@@ -62,22 +64,22 @@ fun main(args: Array<String>) {
    val arguments = FullCInteropArguments()
    arguments.argParser.parse(args)
    val flavorName = arguments.flavor
-    processCLib(flavorName, arguments, InternalInteropOptions(arguments.generated, arguments.natives))
+    processCLibSafe(flavorName, arguments, InternalInteropOptions(arguments.generated, arguments.natives))
 }

 fun interop(
        flavor: String, args: Array<String>,
        additionalArgs: InternalInteropOptions
-): Array<String>? = when(flavor) {
-            "jvm", "native" -> {
-                val cinteropArguments = CInteropArguments()
-                cinteropArguments.argParser.parse(args)
-                val platform = KotlinPlatform.values().single { it.name.equals(flavor, ignoreCase = true) }
-                processCLib(platform, cinteropArguments, additionalArgs)
-            }
-            "wasm" -> processIdlLib(args, additionalArgs)
-            else -> error("Unexpected flavor")
-        }
+): Array<String>? = when (flavor) {
+    "jvm", "native" -> {
+        val cinteropArguments = CInteropArguments()
+        cinteropArguments.argParser.parse(args)
+        val platform = KotlinPlatform.values().single { it.name.equals(flavor, ignoreCase = true) }
+        processCLibSafe(platform, cinteropArguments, additionalArgs)
+    }
+    "wasm" -> processIdlLib(args, additionalArgs)
+    else -> error("Unexpected flavor")
+}

 // Options, whose values are space-separated and can be escaped.
 val escapedOptions = setOf("-compilerOpts", "-linkerOpts", "-compiler-options", "-linker-options")
@@ -203,6 +205,14 @@ private fun findFilesByGlobs(roots: List<Path>, globs: List<String>): Map<Path,
    return relativeToRoot
 }

+private fun processCLibSafe(flavor: KotlinPlatform, cinteropArguments: CInteropArguments,
+                            additionalArgs: InternalInteropOptions) =
+        usingNativeMemoryAllocator {
+            usingJvmCInteropCallbacks {
+                processCLib(flavor, cinteropArguments, additionalArgs)
+            }
+        }
+
 private fun processCLib(flavor: KotlinPlatform, cinteropArguments: CInteropArguments,
                        additionalArgs: InternalInteropOptions): Array<String>? {
    val ktGenRoot = additionalArgs.generated
@@ -64,6 +64,19 @@ kotlinNativeInterop {

        pkg 'org.jetbrains.kotlin.backend.konan.files'
    }
+
+
+    env {
+        linker 'clang++'
+        linkOutputs ":kotlin-native:common:${hostName}Env"
+
+        headers fileTree('../common/src/env/headers') {
+            include '**/*.h'
+            include '**/*.hpp'
+        }
+
+        pkg 'org.jetbrains.kotlin.backend.konan.env'
+    }
 }


@@ -128,6 +141,7 @@ dependencies {

    compilerApi kotlinNativeInterop['llvm'].configuration
    compilerApi kotlinNativeInterop['files'].configuration
+    compilerApi kotlinNativeInterop['env'].configuration


    cli_bcApi sourceSets.compiler.output
@@ -142,6 +156,7 @@ jar {
    from sourceSets.cli_bc.output,
            sourceSets.compiler.output,
            sourceSets.filesInteropStubs.output,
+            sourceSets.envInteropStubs.output,
            sourceSets.llvmInteropStubs.output

    dependsOn ':kotlin-native:runtime:hostRuntime', 'external_jars'
@@ -50,12 +50,11 @@ import org.jetbrains.kotlin.ir.declarations.lazy.IrLazyClass
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.IrFieldSymbolImpl
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.konan.library.KonanLibraryLayout
 import org.jetbrains.kotlin.konan.target.Architecture
 import org.jetbrains.kotlin.konan.target.KonanTarget
-import org.jetbrains.kotlin.konan.util.disposeNativeMemoryAllocator
 import org.jetbrains.kotlin.library.SerializedIrModule
 import org.jetbrains.kotlin.resolve.descriptorUtil.isEffectivelyExternal
+import org.jetbrains.kotlin.konan.library.KonanLibraryLayout

 internal class InlineFunctionOriginInfo(val irFile: IrFile, val startOffset: Int, val endOffset: Int)

@@ -383,14 +382,6 @@ internal class Context(config: KonanConfig) : KonanBackendContext(config) {
        llvmDisposed = true
    }

-    private var nativeMemFreed = false
-
-    fun freeNativeMem() {
-        if (nativeMemFreed) return
-        disposeNativeMemoryAllocator()
-        nativeMemFreed = true
-    }
-
    val cStubsManager = CStubsManager(config.target)

    val coverage = CoverageManager(this)
@@ -5,12 +5,14 @@

 package org.jetbrains.kotlin.backend.konan

+import kotlinx.cinterop.usingJvmCInteropCallbacks
 import org.jetbrains.kotlin.analyzer.AnalysisResult
 import org.jetbrains.kotlin.backend.common.phaser.CompilerPhase
 import org.jetbrains.kotlin.backend.common.phaser.invokeToplevel
 import org.jetbrains.kotlin.cli.common.messages.AnalyzerWithCompilerReport
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.config.languageVersionSettings
+import org.jetbrains.kotlin.konan.util.usingNativeMemoryAllocator
 import org.jetbrains.kotlin.utils.addToStdlib.cast

 fun runTopLevelPhases(konanConfig: KonanConfig, environment: KotlinCoreEnvironment) {
@@ -30,13 +32,13 @@ fun runTopLevelPhases(konanConfig: KonanConfig, environment: KotlinCoreEnvironme

    if (!context.frontendPhase()) return

-    try {
-        toplevelPhase.cast<CompilerPhase<Context, Unit, Unit>>().invokeToplevel(context.phaseConfig, context, Unit)
-    } finally {
-        try {
-            context.disposeLlvm()
-        } finally {
-            context.freeNativeMem()
+    usingNativeMemoryAllocator {
+        usingJvmCInteropCallbacks {
+            try {
+                toplevelPhase.cast<CompilerPhase<Context, Unit, Unit>>().invokeToplevel(context.phaseConfig, context, Unit)
+            } finally {
+                context.disposeLlvm()
+            }
        }
    }
 }
@@ -556,8 +556,7 @@ val toplevelPhase: CompilerPhase<*, Unit, Unit> = namedUnitPhase(
                                unitSink()
                ) then
                objectFilesPhase then
-                linkerPhase then
-                freeNativeMemPhase
+                linkerPhase
 )

 internal fun PhaseConfig.disableIf(phase: AnyNamedPhase, condition: Boolean) {
@@ -79,16 +79,6 @@ internal val disposeLLVMPhase = namedUnitPhase(
        }
 )

-internal val freeNativeMemPhase = namedUnitPhase(
-        name = "FreeNativeMem",
-        description = "Free native memory used by interop",
-        lower = object : CompilerPhase<Context, Unit, Unit> {
-            override fun invoke(phaseConfig: PhaseConfig, phaserState: PhaserState<Unit>, context: Context, input: Unit) {
-                context.freeNativeMem()
-            }
-        }
-)
-
 internal val RTTIPhase = makeKonanModuleOpPhase(
        name = "RTTI",
        description = "RTTI generation",
@@ -11,12 +11,16 @@ bitcode {
    create("files"){
        dependsOn(":kotlin-native:dependencies:update")
    }
+    create("env"){
+        dependsOn(":kotlin-native:dependencies:update")
+    }
 }

 val hostName: String by project

 val build by tasks.registering {
    dependsOn("${hostName}Files")
+    dependsOn("${hostName}Env")
 }

 val clean by tasks.registering {
@@ -0,0 +1,23 @@
+/*
+ * Copyright 2010-2021 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 "Env.h"
+
+#ifdef _WIN32
+
+#include <windows.h>
+
+void setEnv(const char* name, const char* value) {
+  SetEnvironmentVariableA(name, value);
+}
+
+#else
+
+#include <cstdlib>
+
+void setEnv(const char* name, const char* value) {
+  setenv(name, value, 1);
+}
+
+#endif
@@ -0,0 +1,18 @@
+/*
+ * Copyright 2010-2021 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.
+ */
+#ifndef COMMON_ENV_H
+#define COMMON_ENV_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+void setEnv(const char* name, const char* value);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // COMMON_ENV_H
@@ -7,15 +7,51 @@ package org.jetbrains.kotlin.konan.util

 import sun.misc.Unsafe

-private val allocatorHolder = ThreadLocal<NativeMemoryAllocator>()
-val nativeMemoryAllocator: NativeMemoryAllocator
-    get() = allocatorHolder.get() ?: NativeMemoryAllocator().also { allocatorHolder.set(it) }
+class ThreadSafeDisposableHelper<T>(create: () -> T, private val dispose: (T) -> Unit) {
+    private val create_ = create

-fun disposeNativeMemoryAllocator() {
-    allocatorHolder.get()?.freeAll()
-    allocatorHolder.remove()
+    var holder: T? = null
+        private set
+
+    private var counter = 0
+    private val lock = Any()
+
+    fun create() {
+        synchronized(lock) {
+            if (counter++ == 0) {
+                check(holder == null)
+                holder = create_()
+            }
+        }
+    }
+
+    fun dispose() {
+        synchronized(lock) {
+            if (--counter == 0) {
+                dispose(holder!!)
+                holder = null
+            }
+        }
+    }
+
+    inline fun <R> usingDisposable(block: () -> R): R {
+        create()
+        return try {
+            block()
+        } finally {
+            dispose()
+        }
+    }
 }

+@PublishedApi
+internal val allocatorDisposeHelper = ThreadSafeDisposableHelper({ NativeMemoryAllocator() }, { it.freeAll() })
+
+inline fun <R> usingNativeMemoryAllocator(block: () -> R) = allocatorDisposeHelper.usingDisposable(block)
+
+val nativeMemoryAllocator: NativeMemoryAllocator
+    get() = allocatorDisposeHelper.holder ?: error("Native memory allocator hasn't been created")
+
 // 256 buckets for sizes <= 2048 padded to 8
 // 256 buckets for sizes <= 64KB padded to 256
 // 256 buckets for sizes <= 1MB padded to 4096
@@ -32,15 +68,23 @@ private const val ChunkHeaderSize = 2 * Int.SIZE_BYTES // chunk size + alignment
 private const val RawChunkSize: Long = 4L * 1024 * 1024

 class NativeMemoryAllocator {
+    companion object {
+        fun init() = allocatorDisposeHelper.create()
+        fun dispose() = allocatorDisposeHelper.dispose()
+    }
+
    private fun alignUp(x: Long, align: Int) = (x + align - 1) and (align - 1).toLong().inv()
    private fun alignUp(x: Int, align: Int) = (x + align - 1) and (align - 1).inv()

    private val smallChunks = LongArray(ChunkBucketSize)
    private val mediumChunks = LongArray(ChunkBucketSize)
    private val bigChunks = LongArray(ChunkBucketSize)
+    private val lock = Any()
+
+    fun alloc(size: Long, align: Int) = synchronized(lock) { allocNoLock(size, align) }

    // Chunk layout: [chunk size,...padding...,diff to start,aligned data start,.....data.....]
-    fun alloc(size: Long, align: Int): Long {
+    private fun allocNoLock(size: Long, align: Int): Long {
        val totalChunkSize = ChunkHeaderSize + size + align
        val ptr = ChunkHeaderSize + when {
            totalChunkSize <= MaxSmallSize -> allocFromFreeList(totalChunkSize.toInt(), SmallChunksSizeAlignment, smallChunks)
@@ -91,7 +135,9 @@ class NativeMemoryAllocator {
        return (rawChunks.last() + rawOffset).also { rawOffset += size }
    }

-    fun free(mem: Long) {
+    fun free(mem: Long) = synchronized(lock) { freeNoLock(mem) }
+
+    private fun freeNoLock(mem: Long) {
        val chunkStart = mem - ChunkHeaderSize - unsafe.getInt(mem - Int.SIZE_BYTES)
        val chunkSize = unsafe.getInt(chunkStart)
        when {
@@ -102,7 +148,7 @@ class NativeMemoryAllocator {
        }
    }

-    fun freeAll() {
+    internal fun freeAll() {
        for (i in 0 until ChunkBucketSize) {
            smallChunks[i] = 0L
            mediumChunks[i] = 0L
@@ -8,6 +8,8 @@ import org.jetbrains.kotlin.native.interop.gen.defFileDependencies
 import org.jetbrains.kotlin.cli.bc.main as konancMain
 import org.jetbrains.kotlin.cli.klib.main as klibMain
 import org.jetbrains.kotlin.cli.bc.mainNoExitWithGradleRenderer as konancMainForGradle
+import org.jetbrains.kotlin.backend.konan.env.setEnv
+import org.jetbrains.kotlin.konan.util.usingNativeMemoryAllocator

 private fun mainImpl(args: Array<String>, konancMain: (Array<String>) -> Unit) {
    val utilityName = args[0]
@@ -59,5 +61,11 @@ private fun mainImpl(args: Array<String>, konancMain: (Array<String>) -> Unit) {

 fun main(args: Array<String>) = mainImpl(args, ::konancMain)

-fun daemonMain(args: Array<String>) = mainImpl(args, ::konancMainForGradle)
+private fun setupClangEnv() {
+    setEnv("LIBCLANG_DISABLE_CRASH_RECOVERY", "1")
+}

+fun daemonMain(args: Array<String>) = usingNativeMemoryAllocator {
+    setupClangEnv() // For in-process invocation have to setup proper environment manually.
+    mainImpl(args, ::konancMainForGradle)
+}