Refactor jvm evaluation to make module operations cleaner and more robust

in addition: - hide internal configuration keys - add keys for flags controlling evaluation - add initial module implementation for using scripts from a classpath
2019-04-04 17:41:59 +02:00
parent 90444502a6
commit a6f0bd6677
7 changed files with 170 additions and 97 deletions
@@ -50,12 +50,15 @@ val ScriptEvaluationConfigurationKeys.providedProperties by PropertiesCollection
 */
 val ScriptEvaluationConfigurationKeys.constructorArgs by PropertiesCollection.key<List<Any?>>()

+@Deprecated("use scriptsInstancesSharing flag instead", level = DeprecationLevel.ERROR)
+val ScriptEvaluationConfigurationKeys.scriptsInstancesSharingMap by PropertiesCollection.key<MutableMap<KClass<*>, EvaluationResult>>()
+
 /**
- * A map that is used to store evaluated script instances; if provided - the evaluator will try to get imported script from the map and
- * only create/evaluate instances if not found, and evaluator will put newly created instances into the map
+ * If enabled - the evaluator will try to get imported script from a shared container
+ * only create/evaluate instances if not found, and evaluator will put newly created instances into the container
 * This allows to have a single instance of the script if it is imported several times via different import paths.
 */
-val ScriptEvaluationConfigurationKeys.scriptsInstancesSharingMap by PropertiesCollection.key<MutableMap<KClass<*>, EvaluationResult>>()
+val ScriptEvaluationConfigurationKeys.scriptsInstancesSharing by PropertiesCollection.key<Boolean>(false)

 /**
 * The callback that will be called on the script compilation immediately before starting the compilation
@@ -67,7 +70,7 @@ val ScriptEvaluationConfigurationKeys.refineConfigurationBeforeEvaluate by Prope
 */
 fun ScriptEvaluationConfiguration.Builder.enableScriptsInstancesSharing() {
    this {
-        scriptsInstancesSharingMap(HashMap())
+        scriptsInstancesSharing(true)
    }
 }

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package kotlin.script.experimental.jvmhost.impl
+
+import java.io.File
+import java.io.Serializable
+import java.net.URLClassLoader
+
+interface KJvmCompiledModule {
+    fun createClassLoader(baseClassLoader: ClassLoader?): ClassLoader
+}
+
+class KJvmCompiledModuleInMemory(val compilerOutputFiles: Map<String, ByteArray>) : KJvmCompiledModule, Serializable {
+
+    override fun createClassLoader(baseClassLoader: ClassLoader?): ClassLoader =
+        CompiledScriptClassLoader(baseClassLoader, compilerOutputFiles)
+
+    companion object {
+        @JvmStatic
+        private val serialVersionUID = 0L
+    }
+}
+
+class KJvmCompiledModuleFromClassPath(val classpath: Collection<File>) : KJvmCompiledModule {
+
+    override fun createClassLoader(baseClassLoader: ClassLoader?): ClassLoader =
+        URLClassLoader(classpath.map { it.toURI().toURL() }.toTypedArray(), baseClassLoader)
+}
+
+class KJvmCompiledModuleFromLoadedClasses : KJvmCompiledModule {
+
+    override fun createClassLoader(baseClassLoader: ClassLoader?): ClassLoader =
+        baseClassLoader ?: KJvmCompiledModuleFromLoadedClasses::class.java.classLoader
+}
@@ -5,35 +5,19 @@

 package kotlin.script.experimental.jvmhost.impl

-import org.jetbrains.kotlin.codegen.state.GenerationState
 import java.io.ObjectInputStream
 import java.io.ObjectOutputStream
 import java.io.Serializable
-import java.net.URLClassLoader
 import kotlin.reflect.KClass
 import kotlin.script.experimental.api.*
-import kotlin.script.experimental.jvm.JvmDependency
-import kotlin.script.experimental.jvmhost.*
-
-class KJvmCompiledModule(
-    generationState: GenerationState
-) : Serializable {
-    val compilerOutputFiles: Map<String, ByteArray> =
-        generationState.factory.asList()
-            .associateTo(sortedMapOf<String, ByteArray>()) { it.relativePath to it.asByteArray() }
-
-    companion object {
-        @JvmStatic
-        private val serialVersionUID = 0L
-    }
-}
+import kotlin.script.experimental.jvmhost.getOrCreateActualClassloader

 class KJvmCompiledScript<out ScriptBase : Any>(
    sourceLocationId: String?,
    compilationConfiguration: ScriptCompilationConfiguration,
    private var scriptClassFQName: String,
    otherScripts: List<CompiledScript<*>> = emptyList(),
-    internal var compiledModule: KJvmCompiledModule? = null
+    internal var compiledModule: KJvmCompiledModule
 ) : CompiledScript<ScriptBase>, Serializable {

    private var _sourceLocationId: String? = sourceLocationId
@@ -54,21 +38,7 @@ class KJvmCompiledScript<out ScriptBase : Any>(
    override suspend fun getClass(scriptEvaluationConfiguration: ScriptEvaluationConfiguration?): ResultWithDiagnostics<KClass<*>> = try {
        // ensuring proper defaults are used
        val actualEvaluationConfiguration = scriptEvaluationConfiguration ?: ScriptEvaluationConfiguration()
-        val classLoader = actualEvaluationConfiguration[ScriptEvaluationConfiguration.jvm.actualClassLoader]
-            ?: run {
-                if (compiledModule == null)
-                    return ResultWithDiagnostics.Failure(
-                        "Unable to load class $scriptClassFQName: no compiled module is provided".asErrorDiagnostics(path = sourceLocationId)
-                    )
-                val baseClassLoader = actualEvaluationConfiguration[ScriptEvaluationConfiguration.jvm.baseClassLoader]
-                val dependencies = compilationConfiguration[ScriptCompilationConfiguration.dependencies]
-                    ?.flatMap { (it as? JvmDependency)?.classpath?.map { it.toURI().toURL() } ?: emptyList() }
-                // TODO: previous dependencies and classloaders should be taken into account here
-                val classLoaderWithDeps =
-                    if (dependencies == null) baseClassLoader
-                    else URLClassLoader(dependencies.toTypedArray(), baseClassLoader)
-                CompiledScriptClassLoader(classLoaderWithDeps, compiledModule!!.compilerOutputFiles)
-            }
+        val classLoader = getOrCreateActualClassloader(actualEvaluationConfiguration)

        val clazz = classLoader.loadClass(scriptClassFQName).kotlin
        clazz.asSuccess()
@@ -82,14 +52,8 @@ class KJvmCompiledScript<out ScriptBase : Any>(
        )
    }

-    // This method is exposed because the compilation configuration is not generally serializable (yet), but since it is supposed to
-    // be deserialized only from the cache, the configuration could be assigned from the cache.load method
-    fun setCompilationConfiguration(configuration: ScriptCompilationConfiguration) {
-        if (_compilationConfiguration != null) throw IllegalStateException("This method is applicable only in deserialization context")
-        _compilationConfiguration = configuration
-    }
-
    private fun writeObject(outputStream: ObjectOutputStream) {
+        outputStream.writeObject(compilationConfiguration)
        outputStream.writeObject(sourceLocationId)
        outputStream.writeObject(otherScripts)
        outputStream.writeObject(compiledModule)
@@ -98,15 +62,15 @@ class KJvmCompiledScript<out ScriptBase : Any>(

    @Suppress("UNCHECKED_CAST")
    private fun readObject(inputStream: ObjectInputStream) {
-        _compilationConfiguration = null
+        _compilationConfiguration = inputStream.readObject() as ScriptCompilationConfiguration
        _sourceLocationId = inputStream.readObject() as String?
        _otherScripts = inputStream.readObject() as List<CompiledScript<*>>
-        compiledModule = inputStream.readObject() as KJvmCompiledModule?
+        compiledModule = inputStream.readObject() as KJvmCompiledModuleInMemory
        scriptClassFQName = inputStream.readObject() as String
    }

    companion object {
        @JvmStatic
-        private val serialVersionUID = 1L
+        private val serialVersionUID = 2L
    }
 }
@@ -351,6 +351,7 @@ class KJvmCompilerImpl(val hostConfiguration: ScriptingHostConfiguration) : KJvm
            sourceDependencies: List<ScriptsCompilationDependencies.SourceDependencies>,
            getScriptConfiguration: (KtFile) -> ScriptCompilationConfiguration
        ): KJvmCompiledScript<Any> {
+            val module = makeCompiledModule(generationState)
            val scriptDependenciesStack = ArrayDeque<KtScript>()

            fun makeOtherScripts(script: KtScript): List<KJvmCompiledScript<*>> {
@@ -368,7 +369,8 @@ class KJvmCompilerImpl(val hostConfiguration: ScriptingHostConfiguration) : KJvm
                                containingKtFile.virtualFile?.path,
                                getScriptConfiguration(sourceFile),
                                it.fqName.asString(),
-                                makeOtherScripts(it)
+                                makeOtherScripts(it),
+                                module
                            )
                        }
                    } ?: emptyList()
@@ -382,7 +384,7 @@ class KJvmCompilerImpl(val hostConfiguration: ScriptingHostConfiguration) : KJvm
                getScriptConfiguration(ktScript.containingKtFile),
                ktScript.fqName.asString(),
                makeOtherScripts(ktScript),
-                KJvmCompiledModule(generationState)
+                module
            )
        }

@@ -566,4 +568,9 @@ internal class ScriptLightVirtualFile(name: String, private val _path: String?,

    override fun getPath(): String = _path ?: super.getPath()
    override fun getCanonicalPath(): String? = path
-}
+}
+
+private fun makeCompiledModule(generationState: GenerationState) = KJvmCompiledModuleInMemory(
+    generationState.factory.asList()
+        .associateTo(sortedMapOf<String, ByteArray>()) { it.relativePath to it.asByteArray() }
+)
@@ -7,8 +7,11 @@

 package kotlin.script.experimental.jvmhost

+import java.net.URLClassLoader
 import kotlin.reflect.KClass
 import kotlin.script.experimental.api.*
+import kotlin.script.experimental.jvm.JvmDependency
+import kotlin.script.experimental.jvmhost.impl.KJvmCompiledScript
 import kotlin.script.experimental.util.PropertiesCollection

 interface JvmScriptEvaluationConfigurationKeys
@@ -18,58 +21,84 @@ open class JvmScriptEvaluationConfigurationBuilder : PropertiesCollection.Builde
    companion object : JvmScriptEvaluationConfigurationBuilder()
 }

+/**
+ * The base classloader to use for script classes loading
+ */
 val JvmScriptEvaluationConfigurationKeys.baseClassLoader by PropertiesCollection.key<ClassLoader?>(Thread.currentThread().contextClassLoader)

-val JvmScriptEvaluationConfigurationKeys.actualClassLoader by PropertiesCollection.key<ClassLoader?>()
+/**
+ * Load script dependencies before evaluation, true by default
+ * If false, it is assumed that the all dependencies will be provided via baseClassLoader
+ */
+val JvmScriptEvaluationConfigurationKeys.loadDependencies by PropertiesCollection.key<Boolean>(true)
+
+internal val JvmScriptEvaluationConfigurationKeys.actualClassLoader by PropertiesCollection.key<ClassLoader?>()
+
+internal val JvmScriptEvaluationConfigurationKeys.scriptsInstancesSharingMap by PropertiesCollection.key<MutableMap<KClass<*>, EvaluationResult>>()

 val ScriptEvaluationConfigurationKeys.jvm get() = JvmScriptEvaluationConfigurationBuilder()

+class JvmScriptEvaluationContext(
+    val classLoader: ClassLoader,
+    val scriptsInstancesSharingMap: MutableMap<KClass<*>, EvaluationResult>?
+)
+
 open class BasicJvmScriptEvaluator : ScriptEvaluator {

    override suspend operator fun invoke(
        compiledScript: CompiledScript<*>,
        scriptEvaluationConfiguration: ScriptEvaluationConfiguration
-    ): ResultWithDiagnostics<EvaluationResult> =
-        try {
-            compiledScript.getClass(scriptEvaluationConfiguration).onSuccess { scriptClass ->
-                // in the future, when (if) we'll stop to compile everything into constructor
-                // run as SAM
-                // return res
+    ): ResultWithDiagnostics<EvaluationResult> = try {
+        val configuration = getConfigurationWithClassloader(compiledScript, scriptEvaluationConfiguration)

-                // for other scripts we need evaluation configuration with actualClassloader set,
-                // so they are loaded in the same classloader as the "main" script
-                val updatedEvalConfiguration =
-                    if (scriptEvaluationConfiguration.containsKey(ScriptEvaluationConfiguration.jvm.actualClassLoader))
-                        scriptEvaluationConfiguration
-                    else
-                        ScriptEvaluationConfiguration(scriptEvaluationConfiguration) {
-                            ScriptEvaluationConfiguration.jvm.actualClassLoader(scriptClass.java.classLoader)
-                        }
+        compiledScript.getClass(configuration).onSuccess { scriptClass ->
+            // in the future, when (if) we'll stop to compile everything into constructor
+            // run as SAM
+            // return res

-                val sharedScripts = scriptEvaluationConfiguration[ScriptEvaluationConfiguration.scriptsInstancesSharingMap]
+            val sharedScripts = configuration[ScriptEvaluationConfiguration.jvm.scriptsInstancesSharingMap]

-                sharedScripts?.get(scriptClass)?.asSuccess()
-                    ?: compiledScript.otherScripts.mapSuccess {
-                        invoke(it, updatedEvalConfiguration)
-                    }.onSuccess { importedScriptsEvalResults ->
+            sharedScripts?.get(scriptClass)?.asSuccess()
+                ?: compiledScript.otherScripts.mapSuccess {
+                    invoke(it, configuration)
+                }.onSuccess { importedScriptsEvalResults ->

-                        val refinedEvalConfiguration =
-                            updatedEvalConfiguration[ScriptEvaluationConfiguration.refineConfigurationBeforeEvaluate]
-                                ?.handler?.invoke(ScriptEvaluationConfigurationRefinementContext(compiledScript, updatedEvalConfiguration))
-                                ?.onFailure {
-                                    return@invoke ResultWithDiagnostics.Failure(it.reports)
-                                }
-                                ?.resultOrNull()
-                                ?: updatedEvalConfiguration
+                    val refinedEvalConfiguration =
+                        configuration[ScriptEvaluationConfiguration.refineConfigurationBeforeEvaluate]
+                            ?.handler?.invoke(ScriptEvaluationConfigurationRefinementContext(compiledScript, configuration))
+                            ?.onFailure {
+                                return@invoke ResultWithDiagnostics.Failure(it.reports)
+                            }
+                            ?.resultOrNull()
+                            ?: configuration

-                        scriptClass.evalWithConfigAndOtherScriptsResults(refinedEvalConfiguration, importedScriptsEvalResults).let {
-                            sharedScripts?.put(scriptClass, it)
-                            ResultWithDiagnostics.Success(it)
-                        }
+                    scriptClass.evalWithConfigAndOtherScriptsResults(refinedEvalConfiguration, importedScriptsEvalResults).let {
+                        sharedScripts?.put(scriptClass, it)
+                        ResultWithDiagnostics.Success(it)
                    }
+                }
+        }
+    } catch (e: Throwable) {
+        ResultWithDiagnostics.Failure(e.asDiagnostics("Error evaluating script", path = compiledScript.sourceLocationId))
+    }
+
+    private fun getConfigurationWithClassloader(
+        script: CompiledScript<*>, baseConfiguration: ScriptEvaluationConfiguration
+    ): ScriptEvaluationConfiguration =
+        if (baseConfiguration.containsKey(ScriptEvaluationConfiguration.jvm.actualClassLoader))
+            baseConfiguration
+        else {
+            val jvmScript = (script as? KJvmCompiledScript<*>)
+                ?: throw IllegalArgumentException("Unexpected compiled script type: $script")
+
+            val classloader = jvmScript.getOrCreateActualClassloader(baseConfiguration)
+
+            ScriptEvaluationConfiguration(baseConfiguration) {
+                ScriptEvaluationConfiguration.jvm.actualClassLoader(classloader)
+                if (baseConfiguration[ScriptEvaluationConfiguration.scriptsInstancesSharing] == true) {
+                    ScriptEvaluationConfiguration.jvm.scriptsInstancesSharingMap(mutableMapOf())
+                }
            }
-        } catch (e: Throwable) {
-            ResultWithDiagnostics.Failure(e.asDiagnostics("Error evaluating script", path = compiledScript.sourceLocationId))
        }

    private fun KClass<*>.evalWithConfigAndOtherScriptsResults(
@@ -99,3 +128,32 @@ open class BasicJvmScriptEvaluator : ScriptEvaluator {
    }
 }

+internal fun KJvmCompiledScript<*>.getOrCreateActualClassloader(evaluationConfiguration: ScriptEvaluationConfiguration): ClassLoader =
+    evaluationConfiguration[ScriptEvaluationConfiguration.jvm.actualClassLoader] ?: run {
+        val baseClassLoader = evaluationConfiguration[ScriptEvaluationConfiguration.jvm.baseClassLoader]
+        val classLoaderWithDeps =
+            if (evaluationConfiguration[ScriptEvaluationConfiguration.jvm.loadDependencies] == false) baseClassLoader
+            else makeClassLoaderFromDependencies(baseClassLoader)
+        compiledModule.createClassLoader(classLoaderWithDeps)
+    }
+
+private fun CompiledScript<*>.makeClassLoaderFromDependencies(baseClassLoader: ClassLoader?): ClassLoader? {
+    val processedScripts = mutableSetOf<CompiledScript<*>>()
+    fun seq(res: Sequence<CompiledScript<*>>, script: CompiledScript<*>): Sequence<CompiledScript<*>> {
+        if (processedScripts.contains(script)) return res
+        processedScripts.add(script)
+        return script.otherScripts.asSequence().fold(res + script, ::seq)
+    }
+
+    val dependencies = seq(emptySequence(), this).flatMap { script ->
+        script.compilationConfiguration[ScriptCompilationConfiguration.dependencies]
+            ?.asSequence()
+            ?.flatMap { dep ->
+                (dep as? JvmDependency)?.classpath?.asSequence()?.map { it.toURI().toURL() } ?: emptySequence()
+            }
+            ?: emptySequence()
+    }.distinct()
+    // TODO: previous dependencies and classloaders should be taken into account here
+    return if (dependencies.none()) baseClassLoader
+    else URLClassLoader(dependencies.toList().toTypedArray(), baseClassLoader)
+}
@@ -12,6 +12,7 @@ import java.util.jar.JarOutputStream
 import java.util.jar.Manifest
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.jvm.JvmDependency
+import kotlin.script.experimental.jvmhost.impl.KJvmCompiledModuleInMemory
 import kotlin.script.experimental.jvmhost.impl.KJvmCompiledScript

 // TODO: generate execution code (main)
@@ -24,8 +25,9 @@ open class BasicJvmScriptClassFilesGenerator(val outputDir: File) : ScriptEvalua
    ): ResultWithDiagnostics<EvaluationResult> {
        try {
            if (compiledScript !is KJvmCompiledScript<*>)
-                return failure("Unsupported compiled script type ${compiledScript::class.qualifiedName}")
-            val module = compiledScript.compiledModule ?: return failure("Cannot generate classes: module is null")
+                return failure("Cannot generate classes: unsupported compiled script type $compiledScript")
+            val module = (compiledScript.compiledModule as? KJvmCompiledModuleInMemory)
+                ?: return failure("Cannot generate classes: unsupported module type ${compiledScript.compiledModule}")
            for ((path, bytes) in module.compilerOutputFiles) {
                File(outputDir, path).apply {
                    if (!parentFile.isDirectory) {
@@ -37,7 +39,7 @@ open class BasicJvmScriptClassFilesGenerator(val outputDir: File) : ScriptEvalua
            return ResultWithDiagnostics.Success(EvaluationResult(ResultValue.Unit, scriptEvaluationConfiguration))
        } catch (e: Throwable) {
            return ResultWithDiagnostics.Failure(
-                e.asDiagnostics("Error saving compiled script classes", path = compiledScript.sourceLocationId)
+                e.asDiagnostics("Cannot generate script classes: ${e.message}", path = compiledScript.sourceLocationId)
            )
        }
    }
@@ -52,8 +54,9 @@ open class BasicJvmScriptJarGenerator(val outputJar: File) : ScriptEvaluator {
    ): ResultWithDiagnostics<EvaluationResult> {
        try {
            if (compiledScript !is KJvmCompiledScript<*>)
-                return failure("Unsupported compiled script type ${compiledScript::class.qualifiedName}")
-            val module = compiledScript.compiledModule ?: return failure("Cannot generate classes: module is null")
+                return failure("Cannot generate jar: unsupported compiled script type $compiledScript")
+            val module = (compiledScript.compiledModule as? KJvmCompiledModuleInMemory)
+                ?: return failure("Cannot generate jar: unsupported module type ${compiledScript.compiledModule}")
            val dependencies = compiledScript.compilationConfiguration[ScriptCompilationConfiguration.dependencies]
                ?.filterIsInstance<JvmDependency>()
                ?.flatMap { it.classpath }
@@ -78,7 +81,7 @@ open class BasicJvmScriptJarGenerator(val outputJar: File) : ScriptEvaluator {
            return ResultWithDiagnostics.Success(EvaluationResult(ResultValue.Unit, scriptEvaluationConfiguration))
        } catch (e: Throwable) {
            return ResultWithDiagnostics.Failure(
-                e.asDiagnostics("Error saving compiled script jar", path = compiledScript.sourceLocationId)
+                e.asDiagnostics("Cannot generate script jar: ${e.message}", path = compiledScript.sourceLocationId)
            )
        }
    }
@@ -293,8 +293,11 @@ class ScriptingHostTest : TestCase() {
        Assert.assertNotNull(cachedScript)
        Assert.assertEquals(1, cache.retrievedScripts)

-        val compiledScriptClass = runBlocking { compiledScript!!.getClass(null).resultOrNull() }
-        val cachedScriptClass = runBlocking { cachedScript!!.getClass(null).resultOrNull() }
+        val compiledScriptClassRes = runBlocking { compiledScript!!.getClass(null) }
+        val cachedScriptClassRes = runBlocking { cachedScript!!.getClass(null) }
+
+        val compiledScriptClass = compiledScriptClassRes.resultOrNull()
+        val cachedScriptClass = cachedScriptClassRes.resultOrNull()

        Assert.assertEquals(compiledScriptClass!!.qualifiedName, cachedScriptClass!!.qualifiedName)
        Assert.assertEquals(compiledScriptClass!!.supertypes, cachedScriptClass!!.supertypes)
@@ -416,10 +419,8 @@ private fun ScriptCompilationConfiguration.Builder.makeSimpleConfigurationWithTe

 private fun File.readCompiledScript(scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript<*> {
    return inputStream().use { fs ->
-        ObjectInputStream(fs).use { os ->
-            (os.readObject() as KJvmCompiledScript<*>).apply {
-                setCompilationConfiguration(scriptCompilationConfiguration)
-            }
+        ObjectInputStream(fs).use {
+            it.readObject() as KJvmCompiledScript<*>
        }
    }
 }