Implement compiled script caching with tests

2018-09-05 13:16:49 +02:00
parent 836ea28fe7
commit ae107339fd
8 changed files with 357 additions and 65 deletions
@@ -31,6 +31,8 @@ open class PropertiesCollection(private val properties: Map<Key<*>, Any> = empty
    fun <T> getNoDefault(key: PropertiesCollection.Key<T>): T? =
        properties[key]?.let { it as T }

+    fun entries(): Set<Map.Entry<Key<*>, Any>> = properties.entries
+
    companion object {
        fun <T> key(defaultValue: T? = null) = PropertyKeyDelegate(defaultValue)
        fun <T> keyCopy(source: Key<T>) = PropertyKeyCopyDelegate(source)
@@ -0,0 +1,35 @@
+/*
+ * Copyright 2010-2018 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 org.jetbrains.kotlin.codegen.BytesUrlUtils
+import java.io.ByteArrayInputStream
+import java.io.InputStream
+import java.net.URL
+import java.util.*
+
+internal class CompiledScriptClassLoader(parent: ClassLoader, private val entries: Map<String, ByteArray>) : ClassLoader(parent) {
+
+    override fun findClass(name: String): Class<*>? {
+        val classPathName = name.replace('.', '/') + ".class"
+        val classBytes = entries[classPathName] ?: return null
+        return defineClass(name, classBytes, 0, classBytes.size)
+    }
+
+    override fun getResourceAsStream(name: String): InputStream? =
+        entries[name]?.let(::ByteArrayInputStream) ?: super.getResourceAsStream(name)
+
+    override fun findResources(name: String?): Enumeration<URL>? {
+        val fromParent = super.findResources(name)
+
+        val url = entries[name]?.let { BytesUrlUtils.createBytesUrl(it) } ?: return fromParent
+
+        return Collections.enumeration(listOf(url) + fromParent.asSequence())
+    }
+
+    override fun findResource(name: String?): URL? =
+        entries[name]?.let { BytesUrlUtils.createBytesUrl(it) } ?: super.findResource(name)
+}
@@ -0,0 +1,81 @@
+/*
+ * Copyright 2010-2018 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 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.JvmScriptEvaluationConfiguration
+import kotlin.script.experimental.jvmhost.baseClassLoader
+
+class KJvmCompiledScript<out ScriptBase : Any>(
+    compilationConfiguration: ScriptCompilationConfiguration,
+    generationState: GenerationState,
+    private var scriptClassFQName: String
+) : CompiledScript<ScriptBase>, Serializable {
+
+    private var _compilationConfiguration: ScriptCompilationConfiguration? = compilationConfiguration
+    private var compilerOutputFiles: Map<String, ByteArray> = run {
+        val res = sortedMapOf<String, ByteArray>()
+        for (it in generationState.factory.asList()) {
+            res[it.relativePath] = it.asByteArray()
+        }
+        res
+    }
+
+    override val compilationConfiguration: ScriptCompilationConfiguration
+        get() = _compilationConfiguration!!
+
+    override suspend fun getClass(scriptEvaluationConfiguration: ScriptEvaluationConfiguration?): ResultWithDiagnostics<KClass<*>> = try {
+        val baseClassLoader = scriptEvaluationConfiguration?.get(JvmScriptEvaluationConfiguration.baseClassLoader)
+            ?: Thread.currentThread().contextClassLoader
+        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)
+        val classLoader = CompiledScriptClassLoader(classLoaderWithDeps, compilerOutputFiles)
+
+        val clazz = classLoader.loadClass(scriptClassFQName).kotlin
+        clazz.asSuccess()
+    } catch (e: Throwable) {
+        ResultWithDiagnostics.Failure(
+            ScriptDiagnostic(
+                "Unable to instantiate class $scriptClassFQName",
+                exception = e
+            )
+        )
+    }
+
+    // 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(compilerOutputFiles)
+        outputStream.writeObject(scriptClassFQName)
+    }
+
+    private fun readObject(inputStream: ObjectInputStream) {
+        _compilationConfiguration = null
+        compilerOutputFiles = inputStream.readObject() as Map<String, ByteArray>
+        scriptClassFQName = inputStream.readObject() as String
+    }
+
+    companion object {
+        @JvmStatic
+        private val serialVersionUID = 0L
+    }
+}
@@ -25,7 +25,6 @@ import org.jetbrains.kotlin.cli.jvm.config.JvmModulePathRoot
 import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
 import org.jetbrains.kotlin.cli.jvm.modules.CoreJrtFileSystem
 import org.jetbrains.kotlin.codegen.ClassBuilderFactories
-import org.jetbrains.kotlin.codegen.GeneratedClassLoader
 import org.jetbrains.kotlin.codegen.KotlinCodegenFacade
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.*
@@ -35,8 +34,6 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.script.KotlinScriptDefinition
 import org.jetbrains.kotlin.script.util.KotlinJars
 import java.io.File
-import java.net.URLClassLoader
-import kotlin.reflect.KClass
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.dependencies.DependenciesResolver
 import kotlin.script.experimental.host.ScriptingHostConfiguration
@@ -46,35 +43,9 @@ import kotlin.script.experimental.jvm.JvmDependency
 import kotlin.script.experimental.jvm.impl.BridgeDependenciesResolver
 import kotlin.script.experimental.jvm.javaHome
 import kotlin.script.experimental.jvm.jvm
-import kotlin.script.experimental.jvmhost.JvmScriptEvaluationConfiguration
 import kotlin.script.experimental.jvmhost.KJvmCompilerProxy
-import kotlin.script.experimental.jvmhost.baseClassLoader
 import kotlin.script.experimental.util.getOrError

-class KJvmCompiledScript<out ScriptBase : Any>(
-    override val compilationConfiguration: ScriptCompilationConfiguration,
-    private val generationState: GenerationState,
-    private val scriptClassFQName: String
-) : CompiledScript<ScriptBase> {
-
-    override suspend fun getClass(scriptEvaluationConfiguration: ScriptEvaluationConfiguration?): ResultWithDiagnostics<KClass<*>> = try {
-        val baseClassLoader = scriptEvaluationConfiguration?.get(JvmScriptEvaluationConfiguration.baseClassLoader)
-            ?: Thread.currentThread().contextClassLoader
-        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)
-        val classLoader = GeneratedClassLoader(generationState.factory, classLoaderWithDeps)
-
-        val clazz = classLoader.loadClass(scriptClassFQName).kotlin
-        clazz.asSuccess()
-    } catch (e: Throwable) {
-        ResultWithDiagnostics.Failure(ScriptDiagnostic("Unable to instantiate class $scriptClassFQName", exception = e))
-    }
-}
-
 class KJvmCompilerImpl(val hostConfiguration: ScriptingHostConfiguration) : KJvmCompilerProxy {

    override fun compile(
@@ -204,7 +175,7 @@ class KJvmCompilerImpl(val hostConfiguration: ScriptingHostConfiguration) : KJvm
    }
 }

-class ScriptDiagnosticsMessageCollector : MessageCollector {
+internal class ScriptDiagnosticsMessageCollector : MessageCollector {

    private val _diagnostics = arrayListOf<ScriptDiagnostic>()

@@ -44,7 +44,7 @@ open class BasicJvmScriptEvaluator : ScriptEvaluator {
                    scriptEvaluationConfiguration?.get(ScriptEvaluationConfiguration.constructorArgs)?.let {
                        args.addAll(it)
                    }
-                    val ctor = scriptClass.java.constructors.first()
+                    val ctor = scriptClass.java.constructors.single()
                    val instance = ctor.newInstance(*args.toArray())

                    // TODO: fix result value
@@ -5,14 +5,21 @@

 package kotlin.script.experimental.jvmhost.test

+import kotlinx.coroutines.experimental.runBlocking
+import org.jetbrains.kotlin.daemon.common.toHexString
 import org.junit.Assert
 import org.junit.Test
-import java.io.ByteArrayOutputStream
-import java.io.PrintStream
-import kotlin.script.experimental.api.ResultWithDiagnostics
+import java.io.*
+import java.nio.file.Files
+import java.security.MessageDigest
+import kotlin.reflect.KClass
+import kotlin.script.experimental.api.*
+import kotlin.script.experimental.host.BasicScriptingHost
 import kotlin.script.experimental.host.toScriptSource
-import kotlin.script.experimental.jvmhost.BasicJvmScriptingHost
-import kotlin.script.experimental.jvmhost.createCompilationConfigurationFromTemplate
+import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
+import kotlin.script.experimental.jvmhost.*
+import kotlin.script.experimental.jvmhost.impl.CompiledScriptClassLoader
+import kotlin.script.experimental.jvmhost.impl.KJvmCompiledScript
 import kotlin.script.templates.standard.SimpleScriptTemplate

 class ScriptingHostTest {
@@ -22,20 +29,188 @@ class ScriptingHostTest {
        val greeting = "Hello from script!"
        val output = captureOut {
            evalScript("println(\"$greeting\")")
-        }.trim()
+        }
        Assert.assertEquals(greeting, output)
    }
-}

-internal fun evalScript(script: String) {
-    val compilationConfiguration = createCompilationConfigurationFromTemplate<SimpleScriptTemplate>()
-    val res = BasicJvmScriptingHost().eval(script.toScriptSource(), compilationConfiguration, null)
-    if (res is ResultWithDiagnostics.Failure) {
-        throw Exception("Compilation/evaluation failed:\n  ${res.reports.joinToString("\n  ") { it.exception?.toString() ?: it.message }}")
+    @Test
+    fun testMemoryCache() {
+        val script = "val x = 1\nprintln(\"x = \$x\")"
+        val cache = SimpleMemoryScriptsCache()
+        val compiler = JvmScriptCompiler(defaultJvmScriptingHostConfiguration, cache = cache)
+        val evaluator = BasicJvmScriptEvaluator()
+        val host = BasicJvmScriptingHost(compiler = compiler, evaluator = evaluator)
+        Assert.assertTrue(cache.data.isEmpty())
+
+        val output = captureOut { evalScript(script, host) }
+        Assert.assertEquals("x = 1", output)
+
+        Assert.assertEquals(1, cache.data.size)
+        val compiled = cache.data.values.first()
+
+        val output2 = captureOut { runBlocking { evaluator(compiled, null) } }
+        Assert.assertEquals(output, output2)
+
+        // TODO: check if cached script is actually used
+        val output3 = captureOut { evalScript(script, host) }.trim()
+        Assert.assertEquals(output, output3)
+    }
+
+    @Test
+    fun testFileCache() {
+        val script = "val x = 1\nprintln(\"x = \$x\")"
+        val cacheDir = Files.createTempDirectory("scriptingTestCache").toFile()
+        try {
+            val cache = FileBasedScriptCache(cacheDir)
+            val compiler = JvmScriptCompiler(defaultJvmScriptingHostConfiguration, cache = cache)
+            val evaluator = BasicJvmScriptEvaluator()
+            val host = BasicJvmScriptingHost(compiler = compiler, evaluator = evaluator)
+            Assert.assertTrue(cache.baseDir.listFiles().isEmpty())
+
+            val scriptCompilationConfiguration = createCompilationConfigurationFromTemplate<SimpleScriptTemplate>()
+
+            var compiledScript: CompiledScript<*>? = null
+            val output = captureOut {
+                runBlocking {
+                    compiler(script.toScriptSource(), scriptCompilationConfiguration).onSuccess {
+                        compiledScript = it
+                        evaluator(it, null)
+                    }.throwOnFailure()
+                }
+            }
+            Assert.assertEquals("x = 1", output)
+
+            val cachedCompiledScript = cache.baseDir.listFiles().let { files ->
+                Assert.assertEquals(1, files.size)
+                files.first().readCompiledScript(scriptCompilationConfiguration)
+            }
+
+            runBlocking {
+                Assert.assertEquals(
+                    compiledScript!!.getClass(null).resultOrNull()?.qualifiedName,
+                    cachedCompiledScript.getClass(null).resultOrNull()?.qualifiedName
+                )
+            }
+
+            val output2 = captureOut {
+                runBlocking {
+                    evaluator(cachedCompiledScript, null).throwOnFailure()
+                }
+            }
+            Assert.assertEquals(output, output2)
+
+            // TODO: check if cached script is actually used
+            val output3 = captureOut { evalScript(script, host) }.trim()
+            Assert.assertEquals(output, output3)
+        } finally {
+            cacheDir.deleteRecursively()
+        }
+    }
+
+    @Test
+    fun testCompiledScriptClassLoader() {
+        val script = "val x = 1"
+        val scriptCompilationConfiguration = createCompilationConfigurationFromTemplate<SimpleScriptTemplate>()
+        val compiler = JvmScriptCompiler(defaultJvmScriptingHostConfiguration)
+        val compiledScript = runBlocking {
+            val res = compiler(script.toScriptSource(), scriptCompilationConfiguration).throwOnFailure()
+            (res as ResultWithDiagnostics.Success<CompiledScript<*>>).value
+        }
+        val compiledScriptClass = runBlocking { compiledScript.getClass(null).throwOnFailure().resultOrNull()!! as KClass<*> }
+        val classLoader = compiledScriptClass.java.classLoader
+
+        Assert.assertTrue(classLoader is CompiledScriptClassLoader)
+        val anotherClass = classLoader.loadClass(compiledScriptClass.qualifiedName)
+
+        Assert.assertEquals(compiledScriptClass.java, anotherClass)
+
+        val classResourceName = compiledScriptClass.qualifiedName!!.replace('.', '/') + ".class"
+        val classAsResourceUrl = classLoader.getResource(classResourceName)
+        val classAssResourceStream = classLoader.getResourceAsStream(classResourceName)
+
+        Assert.assertNotNull(classAsResourceUrl)
+        Assert.assertNotNull(classAssResourceStream)
+
+        val classAsResourceData = classAsResourceUrl.openConnection().getInputStream().readBytes()
+        val classAsResourceStreamData = classAssResourceStream.readBytes()
+
+        Assert.assertArrayEquals(classAsResourceData, classAsResourceStreamData)
+
+        // TODO: consider testing getResources as well
    }
 }

-internal fun captureOut(body: () -> Unit): String {
+fun ResultWithDiagnostics<*>.throwOnFailure(): ResultWithDiagnostics<*> = apply {
+    if (this is ResultWithDiagnostics.Failure) {
+        throw Exception("Compilation/evaluation failed:\n  ${reports.joinToString("\n  ") { it.exception?.toString() ?: it.message }}")
+    }
+}
+
+private fun evalScript(script: String, host: BasicScriptingHost = BasicJvmScriptingHost()) {
+    val compilationConfiguration = createCompilationConfigurationFromTemplate<SimpleScriptTemplate>()
+    host.eval(script.toScriptSource(), compilationConfiguration, null).throwOnFailure()
+}
+
+
+private class SimpleMemoryScriptsCache : CompiledJvmScriptsCache {
+
+    internal val data = hashMapOf<Pair<SourceCode, ScriptCompilationConfiguration>, CompiledScript<*>>()
+
+    override fun get(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript<*>? =
+        data[script to scriptCompilationConfiguration]
+
+    override fun store(
+        compiledScript: CompiledScript<*>,
+        script: SourceCode,
+        scriptCompilationConfiguration: ScriptCompilationConfiguration
+    ) {
+        data[script to scriptCompilationConfiguration] = compiledScript
+    }
+}
+
+private fun File.readCompiledScript(scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript<*> {
+    return inputStream().use { fs ->
+        ObjectInputStream(fs).use { os ->
+            (os.readObject() as KJvmCompiledScript<*>).apply {
+                setCompilationConfiguration(scriptCompilationConfiguration)
+            }
+        }
+    }
+}
+
+
+private class FileBasedScriptCache(val baseDir: File) : CompiledJvmScriptsCache {
+
+    internal fun uniqueHash(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): String {
+        val digestWrapper = MessageDigest.getInstance("MD5")
+        digestWrapper.update(script.text.toByteArray())
+        scriptCompilationConfiguration.entries().sortedBy { it.key.name }.forEach {
+            digestWrapper.update(it.key.name.toByteArray())
+            digestWrapper.update(it.value.toString().toByteArray())
+        }
+        return digestWrapper.digest().toHexString()
+    }
+
+    override fun get(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript<*>? {
+        val file = File(baseDir, uniqueHash(script, scriptCompilationConfiguration))
+        return if (!file.exists()) null else file.readCompiledScript(scriptCompilationConfiguration)
+    }
+
+    override fun store(
+        compiledScript: CompiledScript<*>,
+        script: SourceCode,
+        scriptCompilationConfiguration: ScriptCompilationConfiguration
+    ) {
+        val file = File(baseDir, uniqueHash(script, scriptCompilationConfiguration))
+        file.outputStream().use { fs ->
+            ObjectOutputStream(fs).use { os ->
+                os.writeObject(compiledScript)
+            }
+        }
+    }
+}
+
+private fun captureOut(body: () -> Unit): String {
    val outStream = ByteArrayOutputStream()
    val prevOut = System.out
    System.setOut(PrintStream(outStream))
@@ -45,5 +220,5 @@ internal fun captureOut(body: () -> Unit): String {
        System.out.flush()
        System.setOut(prevOut)
    }
-    return outStream.toString()
+    return outStream.toString().trim()
 }