jp.nephy.kchroner.plugin.KotlinPluginLoader.kt Maven / Gradle / Ivy

package jp.nephy.kchroner.plugin

import jp.nephy.kchroner.KChroner
import jp.nephy.kchroner.api.LoopSubscription
import jp.nephy.kchroner.api.ScheduledSubscription
import jp.nephy.kchroner.api.TwitterSubscription
import org.jetbrains.kotlin.cli.common.environment.setIdeaIoUseFallback
import org.jetbrains.kotlin.cli.common.repl.KotlinJsr223JvmScriptEngineBase
import org.jetbrains.kotlin.cli.common.repl.KotlinJsr223JvmScriptEngineFactoryBase
import org.jetbrains.kotlin.cli.common.repl.ScriptArgsWithTypes
import org.jetbrains.kotlin.script.jsr223.KotlinJsr223JvmLocalScriptEngine
import org.jetbrains.kotlin.script.jsr223.KotlinStandardJsr223ScriptTemplate
import org.jetbrains.kotlin.script.util.scriptCompilationClasspathFromContextOrStlib
import org.jetbrains.kotlin.utils.addToStdlib.measureTimeMillisWithResult
import java.net.URLClassLoader
import java.nio.file.Files
import java.nio.file.Path
import java.nio.file.Paths
import javax.script.*
import kotlin.reflect.full.findAnnotation
import kotlin.reflect.full.memberFunctions
import kotlin.streams.toList

private val logger = KChroner.logger("KChroner.KotlinPluginLoader")

class KotlinPluginLoader {
    companion object {
        val pluginExtensions = arrayOf("kt")
        private val compiledCacheRootDirectory = Paths.get("plugins_cache")
        private val compiledCacheFileRegex = "^Line_\\d+\\.class$".toRegex()
        val additionalLibraryDirectory = Paths.get("lib")!!
    }

    fun load(path: Path): List {
        val compiledCacheDirectory = Paths.get(compiledCacheRootDirectory.toString(), *path.drop(1).map { it.toString() }.toTypedArray())

        if (!Files.isDirectory(compiledCacheDirectory) || enumulateCompiledCaches(compiledCacheDirectory).any { Files.getLastModifiedTime(it) < Files.getLastModifiedTime(path) }) {
            compile(path, compiledCacheDirectory) ?: return emptyList()
        }

        return enumulateCompiledCaches(compiledCacheDirectory).enumulateInstances(compiledCacheDirectory).map { it.buildPlugin() }
    }

    fun load(pluginClass: Class<*>): Plugin? {
        val instance = try {
            if (pluginClass.isAnnotationPresent(jp.nephy.kchroner.api.Plugin::class.java)) {
                pluginClass.newInstance()
            } else {
                null
            }
        } catch (e: Exception) {
            logger.error(e) { "インスタンスの作成に失敗しました. (${pluginClass.name}" }
            null
        }

        return instance?.buildPlugin()
    }

    private fun compile(path: Path, compiledCacheDirectory: Path): KotlinJsr223JvmScriptEngineBase.CompiledKotlinScript? {
        return try {
            val engine = CustomKotlinJsr223JvmLocalScriptEngineFactory().scriptEngine
            setIdeaIoUseFallback()

            val content = path.toFile().reader().use { it.readText() }
            val (compileTimeMs, compiledScript) = measureTimeMillisWithResult {
                (engine as Compilable).compile(content) as KotlinJsr223JvmScriptEngineBase.CompiledKotlinScript
            }
            logger.debug { "コンパイル成功: $path ($compileTimeMs ms)" }

            for (compiledClass in compiledScript.compiledData.classes) {
                val cachePath = Paths.get(compiledCacheDirectory.toString(), compiledClass.path)
                Files.createDirectories(cachePath.parent)
                cachePath.toFile().writeBytes(compiledClass.bytes)
            }
            logger.trace { "$path のコンパイルキャッシュを保存しました." }

            compiledScript
        } catch (e: ScriptException) {
            logger.error(e) { "スクリプトのコンパイルに失敗しました. ($path)" }
            null
        }
    }

    private fun enumulateCompiledCaches(compiledCacheDirectory: Path): List {
        return Files.walk(compiledCacheDirectory).filter { compiledCacheFileRegex.matches(it.fileName.toString()) }.toList()
    }

    private fun List.enumulateInstances(compiledCacheDirectory: Path): List {
        val classLoader = URLClassLoader(
                arrayOf(compiledCacheDirectory.toUri().toURL()) + Files.walk(KotlinPluginLoader.additionalLibraryDirectory).map { it.toUri().toURL() }.toList()
        )
        return mapNotNull {
            val loadedClass = try {
                classLoader.loadClass(it.fileName.toString().removeSuffix(".class"))
            } catch (e: Exception) {
                logger.error(e) { "クラスロードに失敗しました. ($it)" }
                null
            } ?: return@mapNotNull null

            loadedClass.classes.mapNotNull { childClass ->
                try {
                    if (childClass.isAnnotationPresent(jp.nephy.kchroner.api.Plugin::class.java)) {
                        childClass.newInstance()
                    } else {
                        null
                    }
                } catch (e: Exception) {
                    logger.error(e) { "インスタンスの作成に失敗しました. (${childClass.name}" }
                    null
                }
            }
        }.flatten()
    }

    private fun Any.buildPlugin(path: Path? = null): Plugin {
        return Plugin(path, this, this::class.java.getAnnotation(jp.nephy.kchroner.api.Plugin::class.java)).also { plugin ->
            logger.debug { "Kotlinプラグイン: ${plugin.name} のサブスクリプションをロードを開始しました." }

            this::class.memberFunctions.forEach { function ->
                val annotation = function.findAnnotation()
                        ?: function.findAnnotation()
                        ?: function.findAnnotation()
                        ?: return@forEach

                val subscription = Subscription(plugin, function, annotation)
                plugin.register(subscription)
                logger.debug { "サブスクリプション: ${subscription.name} をロードしました." }
            }
        }
    }
}

private class CustomKotlinJsr223JvmLocalScriptEngineFactory: KotlinJsr223JvmScriptEngineFactoryBase() {
    override fun getScriptEngine(): ScriptEngine {
        val classPath = scriptCompilationClasspathFromContextOrStlib("kotlin-script-util.jar", wholeClasspath = true) + Files.walk(KotlinPluginLoader.additionalLibraryDirectory).map { it.toFile() }.toList()

        return KotlinJsr223JvmLocalScriptEngine(
                this, classPath, KotlinStandardJsr223ScriptTemplate::class.qualifiedName!!,
                { ctx, types ->
                    ScriptArgsWithTypes(arrayOf(ctx.getBindings(ScriptContext.ENGINE_SCOPE)), types ?: emptyArray())
                }, arrayOf(Bindings::class)
        )
    }
}