diff --git a/build.gradle.kts b/build.gradle.kts deleted file mode 100644 index eddaaea..0000000 --- a/build.gradle.kts +++ /dev/null @@ -1,32 +0,0 @@ -plugins { - kotlin("jvm") version "2.1.20" - `maven-publish` -} - -group = "dev.freya02" -version = "3.0.0-beta.2_DEV" - -repositories { - mavenCentral() -} - -dependencies { - testImplementation(kotlin("test")) -} - -tasks.test { - useJUnitPlatform() -} - -java { - sourceCompatibility = JavaVersion.VERSION_21 - targetCompatibility = JavaVersion.VERSION_21 -} - -publishing { - publications { - create("maven") { - from(components["java"]) - } - } -} \ No newline at end of file diff --git a/buildSrc/build.gradle.kts b/buildSrc/build.gradle.kts new file mode 100644 index 0000000..2e6893b --- /dev/null +++ b/buildSrc/build.gradle.kts @@ -0,0 +1,13 @@ +plugins { + `kotlin-dsl` +} + +repositories { + mavenCentral() + gradlePluginPortal() +} + +dependencies { + // Change in version catalog too + implementation("org.jetbrains.kotlin:kotlin-gradle-plugin:2.2.0-RC") +} \ No newline at end of file diff --git a/buildSrc/src/main/kotlin/BotCommands-Restarter-conventions.gradle.kts b/buildSrc/src/main/kotlin/BotCommands-Restarter-conventions.gradle.kts new file mode 100644 index 0000000..ffa6e82 --- /dev/null +++ b/buildSrc/src/main/kotlin/BotCommands-Restarter-conventions.gradle.kts @@ -0,0 +1,37 @@ +import org.jetbrains.kotlin.gradle.dsl.JvmTarget + +plugins { + kotlin("jvm") +} + +group = "dev.freya02" +version = "3.0.0-beta.2_DEV" + +java { + toolchain { + languageVersion = JavaLanguageVersion.of(24) + } + + sourceCompatibility = JavaVersion.VERSION_21 + targetCompatibility = JavaVersion.VERSION_21 + + withSourcesJar() +} + +repositories { + mavenCentral() + mavenLocal() +} + +dependencies { + testImplementation(kotlin("test")) + testImplementation("org.junit.jupiter:junit-jupiter-params:5.10.0") +} + +kotlin { + compilerOptions { + jvmTarget = JvmTarget.JVM_21 + + freeCompilerArgs.addAll("-Xjsr305=strict") + } +} \ No newline at end of file diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml new file mode 100644 index 0000000..bbf0343 --- /dev/null +++ b/gradle/libs.versions.toml @@ -0,0 +1,16 @@ +[versions] +kotlin = "2.2.0-RC" # Also change in buildSrc +botcommands = "3.0.0-beta.2_DEV" +jda = "5.5.1" +mockk = "1.13.16" +bytebuddy = "1.17.5" +logback-classic = "1.5.18" +kotlin-logging = "7.0.3" + +[libraries] +botcommands = { module = "io.github.freya022:BotCommands", version.ref = "botcommands" } +jda = { module = "net.dv8tion:JDA", version.ref = "jda" } +mockk = { module = "io.mockk:mockk", version.ref = "mockk" } +bytebuddy = { module = "net.bytebuddy:byte-buddy", version.ref = "bytebuddy" } +logback-classic = { module = "ch.qos.logback:logback-classic", version.ref = "logback-classic" } +kotlin-logging = { module = "io.github.oshai:kotlin-logging-jvm", version.ref = "kotlin-logging" } \ No newline at end of file diff --git a/restarter-jda-cache/build.gradle.kts b/restarter-jda-cache/build.gradle.kts new file mode 100644 index 0000000..23aea36 --- /dev/null +++ b/restarter-jda-cache/build.gradle.kts @@ -0,0 +1,55 @@ +import org.jetbrains.kotlin.gradle.dsl.JvmTarget + +plugins { + id("BotCommands-Restarter-conventions") + `maven-publish` +} + +repositories { + maven("https://jitpack.io") +} + +dependencies { + implementation(libs.kotlin.logging) + implementation(libs.botcommands) + + testImplementation(libs.mockk) + testImplementation(libs.bytebuddy) + testImplementation(libs.logback.classic) +} + +java { + sourceCompatibility = JavaVersion.VERSION_24 + targetCompatibility = JavaVersion.VERSION_24 +} + +kotlin { + compilerOptions { + jvmTarget = JvmTarget.JVM_24 + freeCompilerArgs.add("-Xcontext-parameters") + } +} + +val jar by tasks.getting(Jar::class) { + manifest { + attributes( + "Premain-Class" to "dev.freya02.botcommands.restart.jda.cache.Agent", + ) + } +} + +tasks.withType { + useJUnitPlatform() + + jvmArgs("-javaagent:${jar.archiveFile.get().asFile.absolutePath}") +} + +publishing { + publications { + create("maven") { + from(components["java"]) + + artifactId = "BotCommands-Restarter-JDA-Cache" + } + } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/Agent.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/Agent.kt new file mode 100644 index 0000000..9f46ef8 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/Agent.kt @@ -0,0 +1,18 @@ +package dev.freya02.botcommands.restart.jda.cache + +import dev.freya02.botcommands.restart.jda.cache.transformer.BContextImplTransformer +import dev.freya02.botcommands.restart.jda.cache.transformer.JDABuilderTransformer +import dev.freya02.botcommands.restart.jda.cache.transformer.JDAImplTransformer +import dev.freya02.botcommands.restart.jda.cache.transformer.JDAServiceTransformer +import java.lang.instrument.Instrumentation + +object Agent { + + @JvmStatic + fun premain(agentArgs: String?, inst: Instrumentation) { + inst.addTransformer(JDABuilderTransformer) + inst.addTransformer(JDAServiceTransformer) + inst.addTransformer(BContextImplTransformer) + inst.addTransformer(JDAImplTransformer) + } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/BufferingEventManager.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/BufferingEventManager.kt new file mode 100644 index 0000000..32c9581 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/BufferingEventManager.kt @@ -0,0 +1,67 @@ +package dev.freya02.botcommands.restart.jda.cache + +import net.dv8tion.jda.api.events.GenericEvent +import net.dv8tion.jda.api.hooks.IEventManager +import java.util.concurrent.locks.ReentrantLock +import kotlin.concurrent.withLock + +internal class BufferingEventManager @DynamicCall constructor( + delegate: IEventManager, +) : IEventManager { + + private val lock = ReentrantLock() + private val eventBuffer: MutableList = arrayListOf() + + private var delegate: IEventManager? = delegate + + internal fun setDelegate(delegate: IEventManager) { + lock.withLock { + check(delegate !is BufferingEventManager) { + "Tried to delegate to a BufferingEventManager!" + } + + this.delegate = delegate + eventBuffer.forEach(::handle) + } + } + + internal fun detach() { + lock.withLock { + delegate = null + } + } + + override fun register(listener: Any) { + lock.withLock { + val delegate = delegate ?: error("Should not happen, implement a listener queue if necessary") + delegate.register(listener) + } + } + + override fun unregister(listener: Any) { + lock.withLock { + val delegate = delegate ?: error("Should not happen, implement a listener queue if necessary") + delegate.unregister(listener) + } + } + + override fun handle(event: GenericEvent) { + val delegate = lock.withLock { + val delegate = delegate + if (delegate == null) { + eventBuffer += event + return + } + delegate + } + + delegate.handle(event) + } + + override fun getRegisteredListeners(): List { + lock.withLock { + val delegate = delegate ?: error("Should not happen, implement a listener queue if necessary") + return delegate.registeredListeners + } + } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/DynamicCall.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/DynamicCall.kt new file mode 100644 index 0000000..8c393a4 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/DynamicCall.kt @@ -0,0 +1,10 @@ +package dev.freya02.botcommands.restart.jda.cache + +/** + * This member is used by generated code and as such is not directly referenced. + * + * This member must be `public`. + */ +@Retention(AnnotationRetention.SOURCE) +@Target(AnnotationTarget.CONSTRUCTOR, AnnotationTarget.FUNCTION, AnnotationTarget.PROPERTY_GETTER, AnnotationTarget.PROPERTY_SETTER) +internal annotation class DynamicCall diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/JDABuilderConfiguration.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/JDABuilderConfiguration.kt new file mode 100644 index 0000000..9910c7f --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/JDABuilderConfiguration.kt @@ -0,0 +1,130 @@ +package dev.freya02.botcommands.restart.jda.cache + +import io.github.freya022.botcommands.api.core.utils.enumSetOf +import io.github.freya022.botcommands.api.core.utils.enumSetOfAll +import io.github.oshai.kotlinlogging.KotlinLogging +import net.dv8tion.jda.api.OnlineStatus +import net.dv8tion.jda.api.entities.Activity +import net.dv8tion.jda.api.hooks.IEventManager +import net.dv8tion.jda.api.hooks.InterfacedEventManager +import net.dv8tion.jda.api.utils.ChunkingFilter +import net.dv8tion.jda.api.utils.MemberCachePolicy +import net.dv8tion.jda.api.utils.cache.CacheFlag +import java.util.* + +private val logger = KotlinLogging.logger { } + +class JDABuilderConfiguration internal constructor() { + + private val warnedUnsupportedValues: MutableSet = hashSetOf() + + var hasUnsupportedValues = false + private set + + private val builderValues: MutableMap = hashMapOf() + private var _eventManager: IEventManager? = null + val eventManager: IEventManager get() = _eventManager ?: InterfacedEventManager() + + // So we can track the initial token and intents, the constructor will be instrumented and call this method + // The user overriding the values using token/intent setters should not be an issue + @DynamicCall + fun onInit(token: String?, intents: Int) { + builderValues[ValueType.TOKEN] = token + builderValues[ValueType.INTENTS] = intents + builderValues[ValueType.CACHE_FLAGS] = enumSetOfAll() + } + + @DynamicCall + fun markUnsupportedValue(signature: String) { + if (warnedUnsupportedValues.add(signature)) + logger.warn { "Unsupported JDABuilder method '$signature', JDA will not be cached between restarts" } + hasUnsupportedValues = true + } + + @DynamicCall + fun setStatus(status: OnlineStatus) { + builderValues[ValueType.STATUS] = status + } + + @DynamicCall + fun setEventManager(eventManager: IEventManager?) { + _eventManager = eventManager + } + + @DynamicCall + fun setEventPassthrough(enable: Boolean) { + builderValues[ValueType.EVENT_PASSTHROUGH] = enable + } + + @DynamicCall + @Suppress("UNCHECKED_CAST") + fun enableCache(first: CacheFlag, vararg others: CacheFlag) { + (builderValues[ValueType.CACHE_FLAGS] as EnumSet) += enumSetOf(first, *others) + } + + @DynamicCall + @Suppress("UNCHECKED_CAST") + fun enableCache(flags: Collection) { + (builderValues[ValueType.CACHE_FLAGS] as EnumSet) += flags + } + + @DynamicCall + @Suppress("UNCHECKED_CAST") + fun disableCache(first: CacheFlag, vararg others: CacheFlag) { + (builderValues[ValueType.CACHE_FLAGS] as EnumSet) -= enumSetOf(first, *others) + } + + @DynamicCall + @Suppress("UNCHECKED_CAST") + fun disableCache(flags: Collection) { + (builderValues[ValueType.CACHE_FLAGS] as EnumSet) -= flags + } + + @DynamicCall + fun setMemberCachePolicy(memberCachePolicy: MemberCachePolicy?) { + builderValues[ValueType.MEMBER_CACHE_POLICY] = memberCachePolicy + } + + @DynamicCall + fun setChunkingFilter(filter: ChunkingFilter?) { + builderValues[ValueType.CHUNKING_FILTER] = filter + } + + @DynamicCall + fun setLargeThreshold(threshold: Int) { + builderValues[ValueType.LARGE_THRESHOLD] = threshold + } + + @DynamicCall + fun setActivity(activity: Activity?) { + builderValues[ValueType.ACTIVITY] = activity + } + + @DynamicCall + fun setEnableShutdownHook(enable: Boolean) { + builderValues[ValueType.ENABLE_SHUTDOWN_HOOK] = enable + } + + internal infix fun isSameAs(other: JDABuilderConfiguration): Boolean { + if (hasUnsupportedValues) return false + if (other.hasUnsupportedValues) return false + + return builderValues == other.builderValues + } + + private enum class ValueType { + TOKEN, + INTENTS, + STATUS, + EVENT_PASSTHROUGH, + CACHE_FLAGS, + // These two are interfaces, it's fine to compare them by equality, + // their reference will be the same as they are from the app class loader, + // so if two runs uses MemberCachePolicy#VOICE, it'll still be compatible + MEMBER_CACHE_POLICY, + CHUNKING_FILTER, + LARGE_THRESHOLD, + ACTIVITY, + ENABLE_SHUTDOWN_HOOK, + } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/JDABuilderSession.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/JDABuilderSession.kt new file mode 100644 index 0000000..d430112 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/JDABuilderSession.kt @@ -0,0 +1,175 @@ +package dev.freya02.botcommands.restart.jda.cache + +import dev.freya02.botcommands.restart.jda.cache.utils.isJvmShuttingDown +import io.github.freya022.botcommands.api.core.BContext +import io.github.oshai.kotlinlogging.KotlinLogging +import net.dv8tion.jda.api.JDA +import net.dv8tion.jda.api.events.StatusChangeEvent +import net.dv8tion.jda.api.events.guild.GuildReadyEvent +import net.dv8tion.jda.api.events.session.ReadyEvent +import java.util.function.Supplier + +private val logger = KotlinLogging.logger { } + +// TODO there may be an issue with REST requests, +// as the instance will not get shut down, the requester will still run any request currently queued +// so we should find a way to cancel the tasks in the rate limiter + +// TODO a similar feature exists at https://github.com/LorittaBot/DeviousJDA/blob/master/src/examples/java/SessionCheckpointAndGatewayResumeExample.kt +// however as it is a JDA fork, users will not be able to use the latest features, +// there is also a risk that the saved data (checkpoint) could miss fields + +// TODO another way of building this feature is to have the user use an external gateway proxy, such as https://github.com/Gelbpunkt/gateway-proxy +// however such a solution introduces a lot of friction, +// requiring to set up JDA manually, though not complicated, but also docker and that container's config +// An hybrid way would require rewriting that proxy, +// so our module can hook into JDA and set the gateway URL to the proxy's +internal class JDABuilderSession private constructor( + @get:DynamicCall val key: String, +) { + + @get:DynamicCall + val configuration = JDABuilderConfiguration() + var wasBuilt: Boolean = false + private set + + private lateinit var scheduleShutdownSignal: ScheduleShutdownSignalWrapper + + // May also be shutdownNow + @DynamicCall + fun onShutdown(instance: JDA, shutdownFunction: Runnable) { + if (isJvmShuttingDown()) { + // "scheduleShutdownSignal" isn't there yet if this shutdown is triggered by JDA's shutdown hook + return shutdownFunction.run() + } + + if (::scheduleShutdownSignal.isInitialized.not()) { + logger.error { "Expected BContextImpl#scheduleShutdownSignal to be called before shutdown, doing a full shut down" } + return shutdownFunction.run() + } + + // Don't save if this configuration has unsupported values + if (configuration.hasUnsupportedValues) { + scheduleShutdownSignal.runFully() + shutdownFunction.run() + return logger.debug { "Discarding JDA instance as the configuration had unsupported values (key '$key')" } + } + + val eventManager = instance.eventManager as? BufferingEventManager + eventManager?.detach() // If the event manager isn't what we expect, it will be logged when attempting to reuse + + JDACache[key] = JDACache.Data(configuration, instance, shutdownFunction, scheduleShutdownSignal) + } + + /** + * Stores the actual code of BContextImpl#scheduleShutdownSignal and the callback it was passed + * + * [scheduleShutdownSignalFunction] will be called with [afterShutdownSignal] if JDA does get shut down, + * but if JDA is reused, only [afterShutdownSignal] is used. + */ + @DynamicCall + fun onScheduleShutdownSignal(scheduleShutdownSignalFunction: Runnable, afterShutdownSignal: () -> Unit) { + this.scheduleShutdownSignal = ScheduleShutdownSignalWrapper(scheduleShutdownSignalFunction, afterShutdownSignal) + } + + @DynamicCall + fun onBuild(buildFunction: Supplier): JDA { + val jda = buildOrReuse(buildFunction) + wasBuilt = true + return jda + } + + private fun buildOrReuse(buildFunction: Supplier): JDA { + val cachedData = JDACache.remove(key) + + fun createNewInstance(): JDA { + val jda = buildFunction.get() + cachedData?.scheduleShutdownSignal?.runFully() + cachedData?.doShutdown?.run() + return jda + } + + if (configuration.hasUnsupportedValues) { + logger.debug { "Configured JDABuilder has unsupported values, building a new JDA instance (key '$key')" } + return createNewInstance() + } + + if (cachedData == null) { + logger.debug { "Creating a new JDA instance (key '$key')" } + return createNewInstance() + } + + if (cachedData.configuration isSameAs configuration) { + logger.debug { "Reusing JDA instance with compatible configuration (key '$key')" } + val jda = cachedData.jda + val eventManager = jda.eventManager as? BufferingEventManager + ?: run { + logger.warn { "Expected a BufferingEventManager but got a ${jda.eventManager.javaClass.name}, creating a new instance" } + return createNewInstance() + } + + cachedData.scheduleShutdownSignal.runAfterShutdownSignal() + + eventManager.setDelegate(configuration.eventManager) + eventManager.handle(StatusChangeEvent(jda, JDA.Status.LOADING_SUBSYSTEMS, JDA.Status.CONNECTED)) + jda.guildCache.forEachUnordered { eventManager.handle(GuildReadyEvent(jda, -1, it)) } + eventManager.handle(ReadyEvent(jda)) + return jda + } else { + logger.debug { "Creating a new JDA instance as its configuration changed (key '$key')" } + return createNewInstance() + } + } + + internal class ScheduleShutdownSignalWrapper internal constructor( + private val scheduleShutdownSignalFunction: Runnable, + private val afterShutdownSignal: () -> Unit + ) { + + internal fun runFully(): Unit = scheduleShutdownSignalFunction.run() + + internal fun runAfterShutdownSignal(): Unit = afterShutdownSignal() + } + + companion object { + // I would store them in a Map, but JDABuilder has no idea what the key is + private val activeSession: ThreadLocal = + ThreadLocal.withInitial { error("No JDABuilderSession exists for this thread") } + + private val sessions: MutableMap = hashMapOf() + + @JvmStatic + @DynamicCall + fun currentSession(): JDABuilderSession = activeSession.get() + + @JvmStatic + @DynamicCall + fun getSession(key: String): JDABuilderSession { + return sessions[key] ?: error("No JDABuilderSession exists for key '$key'") + } + + @JvmStatic + @DynamicCall + fun getCacheKey(context: BContext): String? = context.config.restartConfig.cacheKey + + @JvmStatic + @DynamicCall + fun withBuilderSession( + key: String, + // Use Java function types to make codegen a bit more reliable + block: Runnable + ) { + val session = JDABuilderSession(key) + sessions[key] = session + activeSession.set(session) + try { + block.run() + if (!session.wasBuilt) { + logger.warn { "Could not save/restore any JDA session as none were built" } + } + } finally { + activeSession.remove() + } + } + } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/JDACache.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/JDACache.kt new file mode 100644 index 0000000..7bc004a --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/JDACache.kt @@ -0,0 +1,21 @@ +package dev.freya02.botcommands.restart.jda.cache + +import net.dv8tion.jda.api.JDA + +internal object JDACache { + + private val cache: MutableMap = hashMapOf() + + internal operator fun set(key: String, data: Data) { + cache[key] = data + } + + internal fun remove(key: String): Data? = cache.remove(key) + + internal class Data internal constructor( + val configuration: JDABuilderConfiguration, + val jda: JDA, + val doShutdown: Runnable, + val scheduleShutdownSignal: JDABuilderSession.ScheduleShutdownSignalWrapper, + ) +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/AbstractClassFileTransformer.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/AbstractClassFileTransformer.kt new file mode 100644 index 0000000..39aa209 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/AbstractClassFileTransformer.kt @@ -0,0 +1,27 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer + +import java.lang.instrument.ClassFileTransformer +import java.security.ProtectionDomain + +internal abstract class AbstractClassFileTransformer protected constructor( + private val target: String +) : ClassFileTransformer { + + override fun transform( + loader: ClassLoader?, + className: String, + classBeingRedefined: Class<*>?, + protectionDomain: ProtectionDomain, + classfileBuffer: ByteArray + ): ByteArray? { + if (className == target) return try { + transform(classfileBuffer) + } catch (e: Throwable) { + e.printStackTrace() + null + } + return null + } + + protected abstract fun transform(classData: ByteArray): ByteArray +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/BContextImplTransformer.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/BContextImplTransformer.kt new file mode 100644 index 0000000..ea3c9c8 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/BContextImplTransformer.kt @@ -0,0 +1,98 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer + +import dev.freya02.botcommands.restart.jda.cache.transformer.utils.* +import io.github.oshai.kotlinlogging.KotlinLogging +import java.lang.classfile.* +import java.lang.constant.ConstantDescs.CD_String +import java.lang.constant.ConstantDescs.CD_void +import java.lang.constant.MethodTypeDesc +import java.lang.reflect.AccessFlag + +private val logger = KotlinLogging.logger { } + +internal object BContextImplTransformer : AbstractClassFileTransformer("io/github/freya022/botcommands/internal/core/BContextImpl") { + + override fun transform(classData: ByteArray): ByteArray { + val classFile = ClassFile.of() + val classModel = classFile.parse(classData) + + return classFile.transformClass( + classModel, + DeferScheduleShutdownSignalTransform(classModel) + ) + } +} + +private class DeferScheduleShutdownSignalTransform(private val classModel: ClassModel) : ContextualClassTransform { + + context(classBuilder: ClassBuilder) + override fun atStartContextual() { + val targetMethod = classModel.findMethod(TARGET_NAME, TARGET_SIGNATURE) + logger.trace { "Moving ${targetMethod.toFullyQualifiedString()} to '$NEW_NAME'" } + targetMethod.transferCodeTo(NEW_NAME) + } + + context(classBuilder: ClassBuilder) + override fun acceptContextual(classElement: ClassElement) { + val methodModel = classElement as? MethodModel ?: return classBuilder.retain(classElement) + if (!methodModel.matches(TARGET_NAME, TARGET_SIGNATURE)) return classBuilder.retain(classElement) + + logger.trace { "Transforming BContextImpl#${TARGET_NAME}${TARGET_SIGNATURE} to defer shutdown signal scheduling" } + classBuilder.transformMethod(methodModel) { methodBuilder, methodElement -> + if (methodElement !is CodeModel) return@transformMethod methodBuilder.retain(methodElement) + + methodBuilder.withFlags(methodModel.flags().flagsMask().withVisibility(AccessFlag.PUBLIC)) + + methodBuilder.withCode { codeBuilder -> + val thisSlot = codeBuilder.receiverSlot() + + val afterShutdownSignalSlot = codeBuilder.parameterSlot(0) + val doScheduleShutdownSignalSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + val sessionKeySlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + val builderSessionSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + + // Runnable doScheduleShutdownSignal = () -> this.doScheduleShutdownSignal(afterShutdownSignal) + codeBuilder.aload(thisSlot) + codeBuilder.aload(afterShutdownSignalSlot) + codeBuilder.invokedynamic( + createLambda( + interfaceMethod = Runnable::run, + targetType = CD_BContextImpl, + targetMethod = NEW_NAME, + targetMethodReturnType = CD_void, + targetMethodArguments = listOf(), + capturedTypes = listOf(CD_Function0), + isStatic = false + ) + ) + codeBuilder.astore(doScheduleShutdownSignalSlot) + + // String sessionKey = JDABuilderSession.getCacheKey(this) + codeBuilder.aload(thisSlot) + codeBuilder.invokestatic(CD_JDABuilderSession, "getCacheKey", MethodTypeDesc.of(CD_String, CD_BContext)) + codeBuilder.astore(sessionKeySlot) + + // JDABuilderSession builderSession = JDABuilderSession.getSession(sessionKey) + codeBuilder.aload(sessionKeySlot) + codeBuilder.invokestatic(CD_JDABuilderSession, "getSession", MethodTypeDesc.of(CD_JDABuilderSession, CD_String)) + codeBuilder.astore(builderSessionSlot) + + // builderSession.onScheduleShutdownSignal(doScheduleShutdownSignal, afterShutdownSignal) + codeBuilder.aload(builderSessionSlot) + codeBuilder.aload(doScheduleShutdownSignalSlot) + codeBuilder.aload(afterShutdownSignalSlot) + codeBuilder.invokevirtual(CD_JDABuilderSession, "onScheduleShutdownSignal", MethodTypeDesc.of(CD_void, CD_Runnable, CD_Function0)) + + // Required + codeBuilder.return_() + } + } + } + + private companion object { + const val TARGET_NAME = "scheduleShutdownSignal" + const val TARGET_SIGNATURE = "(Lkotlin/jvm/functions/Function0;)V" + + const val NEW_NAME = "doScheduleShutdownSignal" + } +} diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/ClassDescriptors.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/ClassDescriptors.kt new file mode 100644 index 0000000..4ddf7b1 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/ClassDescriptors.kt @@ -0,0 +1,32 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer + +import dev.freya02.botcommands.restart.jda.cache.BufferingEventManager +import dev.freya02.botcommands.restart.jda.cache.JDABuilderConfiguration +import dev.freya02.botcommands.restart.jda.cache.JDABuilderSession +import dev.freya02.botcommands.restart.jda.cache.transformer.utils.classDesc +import org.intellij.lang.annotations.Language +import java.lang.constant.ClassDesc + +internal val CD_Function0 = classDescOf("kotlin.jvm.functions.Function0") + +internal val CD_IllegalStateException = classDescOf("java.lang.IllegalStateException") +internal val CD_Runnable = classDescOf("java.lang.Runnable") +internal val CD_Supplier = classDescOf("java.util.function.Supplier") + +internal val CD_BContext = classDescOf("io.github.freya022.botcommands.api.core.BContext") +internal val CD_BContextImpl = classDescOf("io.github.freya022.botcommands.internal.core.BContextImpl") +internal val CD_JDAService = classDescOf("io.github.freya022.botcommands.api.core.JDAService") +internal val CD_BReadyEvent = classDescOf("io.github.freya022.botcommands.api.core.events.BReadyEvent") + +internal val CD_JDA = classDescOf("net.dv8tion.jda.api.JDA") +internal val CD_JDAImpl = classDescOf("net.dv8tion.jda.internal.JDAImpl") +internal val CD_JDABuilder = classDescOf("net.dv8tion.jda.api.JDABuilder") +internal val CD_IEventManager = classDescOf("net.dv8tion.jda.api.hooks.IEventManager") + +internal val CD_BufferingEventManager = classDesc() +internal val CD_JDABuilderSession = classDesc() +internal val CD_JDABuilderConfiguration = classDesc() + +private fun classDescOf(@Language("java", prefix = "import ", suffix = ";") name: String): ClassDesc { + return ClassDesc.of(name) +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/ContextualClassTransform.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/ContextualClassTransform.kt new file mode 100644 index 0000000..e857339 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/ContextualClassTransform.kt @@ -0,0 +1,25 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer + +import java.lang.classfile.ClassBuilder +import java.lang.classfile.ClassElement +import java.lang.classfile.ClassTransform + +internal interface ContextualClassTransform : ClassTransform { + + override fun atStart(builder: ClassBuilder): Unit = context(builder) { atStartContextual() } + + context(classBuilder: ClassBuilder) + fun atStartContextual() { } + + + override fun atEnd(builder: ClassBuilder): Unit = context(builder) { atEndContextual() } + + context(classBuilder: ClassBuilder) + fun atEndContextual() { } + + + override fun accept(builder: ClassBuilder, element: ClassElement): Unit = context(builder) { acceptContextual(element) } + + context(classBuilder: ClassBuilder) + fun acceptContextual(classElement: ClassElement) { } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDABuilderTransformer.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDABuilderTransformer.kt new file mode 100644 index 0000000..0232216 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDABuilderTransformer.kt @@ -0,0 +1,258 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer + +import dev.freya02.botcommands.restart.jda.cache.JDABuilderConfiguration +import dev.freya02.botcommands.restart.jda.cache.transformer.utils.* +import io.github.oshai.kotlinlogging.KotlinLogging +import java.lang.classfile.* +import java.lang.classfile.ClassFile.* +import java.lang.constant.ClassDesc +import java.lang.constant.ConstantDescs.* +import java.lang.constant.MethodTypeDesc +import java.lang.reflect.AccessFlag +import java.util.function.Supplier + +private val logger = KotlinLogging.logger { } + +internal object JDABuilderTransformer : AbstractClassFileTransformer("net/dv8tion/jda/api/JDABuilder") { + + override fun transform(classData: ByteArray): ByteArray { + val classFile = ClassFile.of() + val classModel = classFile.parse(classData) + + return classFile.transformClass( + classModel, + CaptureSetterParametersTransform() + .andThen(CaptureConstructorParametersTransform(classModel)) + .andThen(DeferBuildAndSetBufferingEventManagerTransform(classModel)) + ) + } +} + +private class CaptureConstructorParametersTransform(private val classModel: ClassModel) : ContextualClassTransform { + + context(classBuilder: ClassBuilder) + override fun atStartContextual() { + classModel.findMethod(TARGET_NAME, TARGET_SIGNATURE) + } + + context(classBuilder: ClassBuilder) + override fun acceptContextual(classElement: ClassElement) { + val methodModel = classElement as? MethodModel ?: return classBuilder.retain(classElement) + if (!methodModel.matches(TARGET_NAME, TARGET_SIGNATURE)) return classBuilder.retain(classElement) + + logger.trace { "Transforming ${methodModel.toFullyQualifiedString()} to capture parameters" } + classBuilder.transformMethod(methodModel) { methodBuilder, methodElement -> + val codeModel = methodElement as? CodeModel ?: return@transformMethod methodBuilder.retain(methodElement) + + methodBuilder.withCode { codeBuilder -> + val builderConfigurationSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + val tokenSlot = codeBuilder.parameterSlot(0) + val intentsSlot = codeBuilder.parameterSlot(1) + + // JDABuilderConfiguration configuration = JDABuilderSession.currentSession().getConfiguration(); + codeBuilder.invokestatic(CD_JDABuilderSession, "currentSession", MethodTypeDesc.of(CD_JDABuilderSession)) + codeBuilder.invokevirtual(CD_JDABuilderSession, "getConfiguration", MethodTypeDesc.of(CD_JDABuilderConfiguration)) + codeBuilder.astore(builderConfigurationSlot) + + // configuration.onInit(token, intents); + codeBuilder.aload(builderConfigurationSlot) + codeBuilder.aload(tokenSlot) + codeBuilder.iload(intentsSlot) + codeBuilder.invokevirtual(CD_JDABuilderConfiguration, "onInit", MethodTypeDesc.of(CD_void, CD_String, CD_int)) + + // Add existing instructions + codeModel.forEach { codeBuilder.with(it) } + } + } + } + + private companion object { + const val TARGET_NAME = "" + const val TARGET_SIGNATURE = "(Ljava/lang/String;I)V" + } +} + +private class DeferBuildAndSetBufferingEventManagerTransform(private val classModel: ClassModel) : ContextualClassTransform { + + context(classBuilder: ClassBuilder) + override fun atStartContextual() { + val targetMethod = classModel.findMethod(TARGET_NAME, TARGET_SIGNATURE) + + logger.trace { "Adding JDABuilder#${NEW_NAME}() to set an event manager and build" } + classBuilder.withMethod( + NEW_NAME, + MethodTypeDesc.of(CD_JDA), + ACC_PRIVATE or ACC_SYNTHETIC or ACC_FINAL + ) { methodBuilder -> + val codeModel = targetMethod.code().get() + + methodBuilder.withCode { codeBuilder -> + val thisSlot = codeBuilder.receiverSlot() + + val bufferingEventManagerSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + + // JDABuilder's eventManager is null by default, + // however, the framework mandates setting a framework-provided event manager, + // so let's just throw if it is null. + val nullEventManagerLabel = codeBuilder.newLabel() + codeBuilder.aload(thisSlot) + codeBuilder.getfield(CD_JDABuilder, "eventManager", CD_IEventManager) + codeBuilder.ifnull(nullEventManagerLabel) + + // var bufferingEventManager = new BufferingEventManager + codeBuilder.new_(CD_BufferingEventManager) + codeBuilder.astore(bufferingEventManagerSlot) + + // bufferingEventManager.(eventManager) + codeBuilder.aload(bufferingEventManagerSlot) + codeBuilder.aload(thisSlot) + codeBuilder.getfield(CD_JDABuilder, "eventManager", CD_IEventManager) + codeBuilder.invokespecial(CD_BufferingEventManager, "", MethodTypeDesc.of(CD_void, CD_IEventManager)) + + // this.setEventManager(eventManager) + codeBuilder.aload(thisSlot) + codeBuilder.aload(bufferingEventManagerSlot) + codeBuilder.invokevirtual(CD_JDABuilder, "setEventManager", MethodTypeDesc.of(CD_JDABuilder, CD_IEventManager)) + + // Move the build() code to doBuild() + codeModel.forEach { codeBuilder.with(it) } + + // Branch when "eventManager" is null + codeBuilder.labelBinding(nullEventManagerLabel) + + codeBuilder.new_(CD_IllegalStateException) + codeBuilder.dup() + codeBuilder.ldc("The event manager must be set using the one provided in JDAService#createJDA") + codeBuilder.invokespecial(CD_IllegalStateException, "", MethodTypeDesc.of(CD_void, CD_String)) + codeBuilder.athrow() + } + } + } + + context(classBuilder: ClassBuilder) + override fun acceptContextual(classElement: ClassElement) { + val methodModel = classElement as? MethodModel ?: return classBuilder.retain(classElement) + if (!methodModel.matches(TARGET_NAME, TARGET_SIGNATURE)) return classBuilder.retain(classElement) + + logger.trace { "Transforming ${methodModel.toFullyQualifiedString()} to defer calls" } + classBuilder.transformMethod(methodModel) { methodBuilder, methodElement -> + if (methodElement !is CodeModel) return@transformMethod methodBuilder.retain(methodElement) + + methodBuilder.withCode { codeBuilder -> + val thisSlot = codeBuilder.receiverSlot() + + val builderSessionSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + val doBuildSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + val jdaSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + + // Supplier doBuild = this::doBuild + codeBuilder.aload(thisSlot) + + codeBuilder.invokedynamic( + createLambda( + interfaceMethod = Supplier<*>::get, + targetType = CD_JDABuilder, + targetMethod = NEW_NAME, + targetMethodReturnType = CD_JDA, + targetMethodArguments = listOf(), + capturedTypes = emptyList(), + isStatic = false + ) + ) + codeBuilder.astore(doBuildSlot) + + // JDABuilderSession session = JDABuilderSession.currentSession(); + codeBuilder.invokestatic(CD_JDABuilderSession, "currentSession", MethodTypeDesc.of(CD_JDABuilderSession)) + codeBuilder.astore(builderSessionSlot) + + // var jda = session.onBuild(this::doBuild); + codeBuilder.aload(builderSessionSlot) + codeBuilder.aload(doBuildSlot) + codeBuilder.invokevirtual(CD_JDABuilderSession, "onBuild", MethodTypeDesc.of(CD_JDA, CD_Supplier)) + // Again, prefer using a variable for clarity + codeBuilder.astore(jdaSlot) + + codeBuilder.aload(jdaSlot) + codeBuilder.areturn() + } + } + } + + private companion object { + const val TARGET_NAME = "build" + const val TARGET_SIGNATURE = "()Lnet/dv8tion/jda/api/JDA;" + + const val NEW_NAME = "doBuild" + } +} + +private class CaptureSetterParametersTransform : ContextualClassTransform { + + private val builderSessionMethods: Set = ClassFile.of() + .parse(JDABuilderConfiguration::class.java.getResourceAsStream("JDABuilderConfiguration.class")!!.readAllBytes()) + .methods() + .mapTo(hashSetOf(), ::MethodDesc) + + context(classBuilder: ClassBuilder) + override fun acceptContextual(classElement: ClassElement) { + val methodModel = classElement as? MethodModel ?: return classBuilder.retain(classElement) + if (!methodModel.flags().has(AccessFlag.PUBLIC)) return classBuilder.retain(classElement) + if (methodModel.flags().has(AccessFlag.STATIC)) return classBuilder.retain(classElement) + if (methodModel.methodName().stringValue() == "build") return classBuilder.retain(classElement) + + // Log is done later + classBuilder.transformMethod(methodModel) { methodBuilder, methodElement -> + val codeModel = methodElement as? CodeModel ?: return@transformMethod methodBuilder.retain(methodElement) + + val hasBuilderSessionMethod = methodModel.let(::MethodDesc) in builderSessionMethods + methodBuilder.withCode { codeBuilder -> + val builderConfigurationSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + + // JDABuilderConfiguration configuration = JDABuilderSession.currentSession().getConfiguration(); + codeBuilder.invokestatic(CD_JDABuilderSession, "currentSession", MethodTypeDesc.of(CD_JDABuilderSession)) + codeBuilder.invokevirtual(CD_JDABuilderSession, "getConfiguration", MethodTypeDesc.of(CD_JDABuilderConfiguration)) + codeBuilder.astore(builderConfigurationSlot) + + val methodName = methodModel.methodName().stringValue() + if (hasBuilderSessionMethod) { + logger.trace { "Registering ${methodModel.toFullyQualifiedString()} as a cache-compatible method" } + + // Set return type to "void" because our method won't return JDABuilder, and it doesn't matter anyway + val methodType = methodModel.methodTypeSymbol().changeReturnType(CD_void) + + // configuration.theMethod(parameters); + codeBuilder.aload(builderConfigurationSlot) + methodType.parameterList().forEachIndexed { index, parameter -> + val typeKind = TypeKind.fromDescriptor(parameter.descriptorString()) + val slot = codeBuilder.parameterSlot(index) + codeBuilder.loadLocal(typeKind, slot) + } + codeBuilder.invokevirtual(CD_JDABuilderConfiguration, methodName, methodType) + } else { + logger.trace { "Skipping ${methodModel.toFullyQualifiedString()} as it does not have an equivalent method handler" } + + val signature = methodName + "(${methodModel.methodTypeSymbol().parameterList().joinToString { it.displayName() }})" + + // configuration.markUnsupportedValue() + codeBuilder.aload(builderConfigurationSlot) + codeBuilder.ldc(signature) + codeBuilder.invokevirtual(CD_JDABuilderConfiguration, "markUnsupportedValue", MethodTypeDesc.of(CD_void, CD_String)) + } + + // Add existing instructions + codeModel.forEach { codeBuilder.with(it) } + } + } + } + + // Utility to match methods using their name and parameters, but not return type + private data class MethodDesc( + val name: String, + val paramTypes: List + ) { + constructor(methodModel: MethodModel) : this( + methodModel.methodName().stringValue(), + methodModel.methodTypeSymbol().parameterList(), + ) + } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAImplTransformer.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAImplTransformer.kt new file mode 100644 index 0000000..b144235 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAImplTransformer.kt @@ -0,0 +1,242 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer + +import dev.freya02.botcommands.restart.jda.cache.transformer.utils.* +import io.github.oshai.kotlinlogging.KotlinLogging +import java.lang.classfile.* +import java.lang.classfile.ClassFile.* +import java.lang.classfile.instruction.InvokeInstruction +import java.lang.constant.ConstantDescs.CD_String +import java.lang.constant.ConstantDescs.CD_void +import java.lang.constant.MethodTypeDesc + +private val logger = KotlinLogging.logger { } + +internal object JDAImplTransformer : AbstractClassFileTransformer("net/dv8tion/jda/internal/JDAImpl") { + + override fun transform(classData: ByteArray): ByteArray { + val classFile = ClassFile.of() + val classModel = classFile.parse(classData) + return classFile.transformClass( + classModel, + CaptureSessionKeyTransform() + .andThen(DeferShutdownTransform(classModel)) + .andThen(DeferShutdownNowTransform(classModel)) + .andThen(AwaitShutdownTransform()) + ) + } +} + +private class CaptureSessionKeyTransform : ContextualClassTransform { + + context(classBuilder: ClassBuilder) + override fun atStartContextual() { + logger.trace { "Adding JDAImpl#${CACHE_KEY_NAME}" } + classBuilder.withField(CACHE_KEY_NAME, CD_String, ACC_PRIVATE or ACC_FINAL) + + logger.trace { "Adding JDAImpl#getBuilderSession()" } + classBuilder.withMethod("getBuilderSession", MethodTypeDesc.of(CD_JDABuilderSession), ACC_PUBLIC) { methodBuilder -> + methodBuilder.withCode { codeBuilder -> + codeBuilder.aload(codeBuilder.receiverSlot()) + codeBuilder.getfield(CD_JDAImpl, CACHE_KEY_NAME, CD_String) + codeBuilder.invokestatic(CD_JDABuilderSession, "getSession", MethodTypeDesc.of(CD_JDABuilderSession, CD_String)) + codeBuilder.areturn() + } + } + } + + context(classBuilder: ClassBuilder) + override fun acceptContextual(classElement: ClassElement) { + val methodModel = classElement as? MethodModel ?: return classBuilder.retain(classElement) + // No need to check the signature, we can assign the field in all constructors + if (!methodModel.methodName().equalsString("")) return classBuilder.retain(classElement) + + logger.trace { "Transforming ${methodModel.toFullyQualifiedString()} to store the session key" } + classBuilder.transformMethod(methodModel) { methodBuilder, methodElement -> + val codeModel = methodElement as? CodeModel ?: return@transformMethod methodBuilder.retain(methodElement) + + methodBuilder.withCode { codeBuilder -> + val thisSlot = codeBuilder.receiverSlot() + + // this.cacheKey = JDABuilderSession.currentSession().getKey() + codeBuilder.aload(thisSlot) + codeBuilder.invokestatic(CD_JDABuilderSession, "currentSession", MethodTypeDesc.of(CD_JDABuilderSession)) + codeBuilder.invokevirtual(CD_JDABuilderSession, "getKey", MethodTypeDesc.of(CD_String)) + codeBuilder.putfield(CD_JDAImpl, CACHE_KEY_NAME, CD_String) + + // Add existing instructions + codeModel.forEach { codeBuilder.with(it) } + } + } + } + + private companion object { + const val CACHE_KEY_NAME = "cacheKey" + } +} + +private class DeferShutdownTransform(private val classModel: ClassModel) : ContextualClassTransform { + + context(classBuilder: ClassBuilder) + override fun atStartContextual() { + val targetMethod = classModel.findMethod(TARGET_NAME, TARGET_SIGNATURE) + + logger.trace { "Moving ${targetMethod.toFullyQualifiedString()} to '$NEW_NAME'" } + targetMethod.transferCodeTo(NEW_NAME) + } + + context(classBuilder: ClassBuilder) + override fun acceptContextual(classElement: ClassElement) { + val methodModel = classElement as? MethodModel ?: return classBuilder.retain(classElement) + if (!methodModel.matches(TARGET_NAME, TARGET_SIGNATURE)) return classBuilder.retain(classElement) + + logger.trace { "Transforming ${methodModel.toFullyQualifiedString()} to defer execution" } + classBuilder.transformMethod(methodModel) { methodBuilder, methodElement -> + if (methodElement !is CodeModel) return@transformMethod methodBuilder.retain(methodElement) + + methodBuilder.withCode { codeBuilder -> + val thisSlot = codeBuilder.receiverSlot() + + val doShutdownSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + val builderSessionSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + + // Runnable doShutdown = this::doShutdown + codeBuilder.aload(thisSlot) + codeBuilder.invokedynamic( + createLambda( + interfaceMethod = Runnable::run, + targetType = CD_JDAImpl, + targetMethod = NEW_NAME, + targetMethodReturnType = CD_void, + targetMethodArguments = listOf(), + capturedTypes = listOf(), + isStatic = false + ) + ) + codeBuilder.astore(doShutdownSlot) + + // var builderSession = getBuilderSession() + codeBuilder.aload(thisSlot) + codeBuilder.invokevirtual(CD_JDAImpl, "getBuilderSession", MethodTypeDesc.of(CD_JDABuilderSession)) + codeBuilder.astore(builderSessionSlot) + + // builderSession.onShutdown(this, this::doShutdown); + codeBuilder.aload(builderSessionSlot) + codeBuilder.aload(thisSlot) + codeBuilder.aload(doShutdownSlot) + codeBuilder.invokevirtual(CD_JDABuilderSession, "onShutdown", MethodTypeDesc.of(CD_void, CD_JDA, CD_Runnable)) + + codeBuilder.return_() + } + } + } + + private companion object { + const val TARGET_NAME = "shutdown" + const val TARGET_SIGNATURE = "()V" + + const val NEW_NAME = "doShutdown" + } +} + +private class DeferShutdownNowTransform(private val classModel: ClassModel) : ContextualClassTransform { + + context(classBuilder: ClassBuilder) + override fun atStartContextual() { + val targetMethod = classModel.findMethod(TARGET_NAME, TARGET_SIGNATURE) + + logger.trace { "Moving ${targetMethod.toFullyQualifiedString()} to $NEW_NAME, replacing shutdown() with doShutdown()" } + classBuilder.withMethod( + NEW_NAME, + MethodTypeDesc.of(CD_void), + ACC_PRIVATE or ACC_SYNTHETIC or ACC_FINAL + ) { methodBuilder -> + val codeModel = targetMethod.code().get() + + methodBuilder.withCode { codeBuilder -> + // Move the shutdownNow() code to doShutdownNow() + codeModel.forEach { codeElement -> + // Replace shutdown() with doShutdown() so we don't call [[JDABuilderSession#onShutdown]] more than once + if (codeElement is InvokeInstruction && codeElement.name().equalsString("shutdown")) { + require(codeElement.type().equalsString("()V")) + codeBuilder.invokevirtual(codeElement.owner().asSymbol(), "doShutdown", MethodTypeDesc.of(CD_void)) + return@forEach + } + + codeBuilder.with(codeElement) + } + } + } + } + + context(classBuilder: ClassBuilder) + override fun acceptContextual(classElement: ClassElement) { + val methodModel = classElement as? MethodModel ?: return classBuilder.retain(classElement) + if (!methodModel.matches(TARGET_NAME, TARGET_SIGNATURE)) return classBuilder.retain(classElement) + + logger.trace { "Transforming ${methodModel.toFullyQualifiedString()} to defer execution" } + classBuilder.transformMethod(methodModel) { methodBuilder, methodElement -> + if (methodElement !is CodeModel) return@transformMethod methodBuilder.retain(methodElement) + + methodBuilder.withCode { codeBuilder -> + val thisSlot = codeBuilder.receiverSlot() + + val doShutdownNowSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + val builderSessionSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + + // Runnable doShutdownNow = this::doShutdownNow + codeBuilder.aload(thisSlot) + codeBuilder.invokedynamic( + createLambda( + interfaceMethod = Runnable::run, + targetType = CD_JDAImpl, + targetMethod = NEW_NAME, + targetMethodReturnType = CD_void, + targetMethodArguments = listOf(), + capturedTypes = listOf(), + isStatic = false + ) + ) + codeBuilder.astore(doShutdownNowSlot) + + // var builderSession = getBuilderSession() + codeBuilder.aload(thisSlot) + codeBuilder.invokevirtual(CD_JDAImpl, "getBuilderSession", MethodTypeDesc.of(CD_JDABuilderSession)) + codeBuilder.astore(builderSessionSlot) + + // builderSession.onShutdown(this, this::doShutdownNow); + codeBuilder.aload(builderSessionSlot) + codeBuilder.aload(thisSlot) + codeBuilder.aload(doShutdownNowSlot) + codeBuilder.invokevirtual(CD_JDABuilderSession, "onShutdown", MethodTypeDesc.of(CD_void, CD_JDA, CD_Runnable)) + + codeBuilder.return_() + } + } + } + + private companion object { + const val TARGET_NAME = "shutdownNow" + const val TARGET_SIGNATURE = "()V" + + const val NEW_NAME = "doShutdownNow" + } +} + +private class AwaitShutdownTransform : ContextualClassTransform { + + context(classBuilder: ClassBuilder) + override fun acceptContextual(classElement: ClassElement) { + val methodModel = classElement as? MethodModel ?: return classBuilder.retain(classElement) + if (!methodModel.methodName().equalsString("awaitShutdown")) return classBuilder.retain(classElement) + + logger.trace { "Transforming ${methodModel.toFullyQualifiedString()} to immediately return" } + classBuilder.transformMethod(methodModel) { methodBuilder, methodElement -> + if (methodElement !is CodeModel) return@transformMethod methodBuilder.retain(methodElement) + + methodBuilder.withCode { codeBuilder -> + codeBuilder.iconst_0() + codeBuilder.ireturn() + } + } + } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAServiceTransformer.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAServiceTransformer.kt new file mode 100644 index 0000000..04e032b --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAServiceTransformer.kt @@ -0,0 +1,114 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer + +import dev.freya02.botcommands.restart.jda.cache.transformer.utils.* +import io.github.oshai.kotlinlogging.KotlinLogging +import java.lang.classfile.* +import java.lang.constant.ConstantDescs.CD_String +import java.lang.constant.ConstantDescs.CD_void +import java.lang.constant.MethodTypeDesc + +private val logger = KotlinLogging.logger { } + +internal object JDAServiceTransformer : AbstractClassFileTransformer("io/github/freya022/botcommands/api/core/JDAService") { + + override fun transform(classData: ByteArray): ByteArray { + val classFile = ClassFile.of() + val classModel = classFile.parse(classData) + + return classFile.transformClass( + classModel, + WrapOnReadyEventWithJDABuilderSessionTransform(classModel) + ) + } +} + +private class WrapOnReadyEventWithJDABuilderSessionTransform(private val classModel: ClassModel) : ContextualClassTransform { + + context(classBuilder: ClassBuilder) + override fun atStartContextual() { + // Put the original code of onReadyEvent in the lambda, + // it will be fired by JDABuilderSession.withBuilderSession in onReadyEvent + val targetMethod = classModel.findMethod(TARGET_NAME, TARGET_SIGNATURE) + + logger.trace { "Moving ${targetMethod.toFullyQualifiedString()} to '$NEW_NAME'" } + targetMethod.transferCodeTo(NEW_NAME) + } + + context(classBuilder: ClassBuilder) + override fun acceptContextual(classElement: ClassElement) { + val methodModel = classElement as? MethodModel ?: return classBuilder.retain(classElement) + if (!methodModel.matches(TARGET_NAME, TARGET_SIGNATURE)) return classBuilder.retain(classElement) + + logger.trace { "Transforming ${methodModel.toFullyQualifiedString()} to wrap the code in a build session" } + classBuilder.transformMethod(methodModel) { methodBuilder, methodElement -> + if (methodElement !is CodeModel) return@transformMethod methodBuilder.retain(methodElement) + + methodBuilder.withCode { codeBuilder -> + val thisSlot = codeBuilder.receiverSlot() + + val readyEventSlot = codeBuilder.parameterSlot(0) + val eventManagerSlot = codeBuilder.parameterSlot(1) + + val contextSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + val sessionKeySlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + val sessionRunnableSlot = codeBuilder.allocateLocal(TypeKind.REFERENCE) + + // var context = event.getContext() + // We could inline this to avoid a successive store/load, + // but I think using variables is probably a better practice, let's leave the optimization to the VM + codeBuilder.aload(readyEventSlot) + codeBuilder.invokevirtual(CD_BReadyEvent, "getContext", MethodTypeDesc.of(CD_BContext)) + codeBuilder.astore(contextSlot) + + // var key = JDABuilderSession.getCacheKey(context) + codeBuilder.aload(contextSlot) + codeBuilder.invokestatic(CD_JDABuilderSession, "getCacheKey", MethodTypeDesc.of(CD_String, CD_BContext)) + codeBuilder.astore(sessionKeySlot) + + // THE KEY IS NULLABLE + // If it is, then don't make a session + val nullKeyLabel = codeBuilder.newLabel() + + // if (key == null) -> nullKeyLabel + codeBuilder.aload(sessionKeySlot) + codeBuilder.ifnull(nullKeyLabel) + + // Runnable sessionRunnable = () -> [lambdaName](event, eventManager) + codeBuilder.aload(thisSlot) + codeBuilder.aload(readyEventSlot) + codeBuilder.aload(eventManagerSlot) + codeBuilder.invokedynamic( + createLambda( + interfaceMethod = Runnable::run, + targetType = CD_JDAService, + targetMethod = NEW_NAME, + targetMethodReturnType = CD_void, + targetMethodArguments = listOf(), + capturedTypes = listOf(CD_BReadyEvent, CD_IEventManager), + isStatic = false + ) + ) + codeBuilder.astore(sessionRunnableSlot) + + // JDABuilderSession.withBuilderSession(key, sessionRunnable) + codeBuilder.aload(sessionKeySlot) + codeBuilder.aload(sessionRunnableSlot) + codeBuilder.invokestatic(CD_JDABuilderSession, "withBuilderSession", MethodTypeDesc.of(CD_void, CD_String, CD_Runnable)) + + // Required + codeBuilder.return_() + + // nullKeyLabel code + codeBuilder.labelBinding(nullKeyLabel) + codeBuilder.return_() + } + } + } + + private companion object { + const val TARGET_NAME = $$"onReadyEvent$BotCommands" + const val TARGET_SIGNATURE = "(Lio/github/freya022/botcommands/api/core/events/BReadyEvent;Lnet/dv8tion/jda/api/hooks/IEventManager;)V" + + const val NEW_NAME = $$"lambda$onReadyEvent$BotCommands$withBuilderSession" + } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/utils/CodeBuilderUtils.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/utils/CodeBuilderUtils.kt new file mode 100644 index 0000000..04e4f22 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/utils/CodeBuilderUtils.kt @@ -0,0 +1,87 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer.utils + +import java.lang.classfile.ClassFileBuilder +import java.lang.classfile.ClassFileElement +import java.lang.classfile.CodeBuilder +import java.lang.constant.* +import java.lang.constant.ConstantDescs.CD_String +import java.lang.invoke.* +import kotlin.reflect.KFunction +import kotlin.reflect.jvm.javaMethod + +internal inline fun classDesc(): ClassDesc = ClassDesc.of(T::class.java.name) + +internal fun ClassFileBuilder.retain(element: E) { + with(element) +} + +internal fun CodeBuilder.ldc(string: String) { + ldc(string as java.lang.String) +} + +internal val lambdaMetafactoryDesc = MethodHandleDesc.ofMethod( + DirectMethodHandleDesc.Kind.STATIC, + classDesc(), + "metafactory", + MethodTypeDesc.of( + classDesc(), + classDesc(), + CD_String, + classDesc(), + classDesc(), + classDesc(), + classDesc() + ) +) + +internal fun createLambda( + interfaceMethod: KFunction<*>, + targetType: ClassDesc, + targetMethod: String, + targetMethodReturnType: ClassDesc, + targetMethodArguments: List, + capturedTypes: List, + isStatic: Boolean, +): DynamicCallSiteDesc { + val effectiveCapturedTypes = when { + isStatic -> capturedTypes + else -> listOf(targetType) + capturedTypes + } + + fun Class<*>.toClassDesc(): ClassDesc = + describeConstable().orElseThrow { IllegalArgumentException("$name cannot be transformed to a ClassDesc") } + + val interfaceJavaMethod = interfaceMethod.javaMethod!! + val targetInterface = interfaceJavaMethod.declaringClass.toClassDesc() + val methodReturnType = interfaceJavaMethod.returnType.toClassDesc() + val methodArguments = interfaceJavaMethod.parameterTypes.map { it.toClassDesc() } + + return DynamicCallSiteDesc.of( + lambdaMetafactoryDesc, + // The following parameters are from [[LambdaMetafactory#metafactory]] + // This is the 2nd argument of LambdaMetafactory#metafactory, "interfaceMethodName", + // the method name in Runnable is "run" + interfaceMethod.name, + // This is the 3rd argument of LambdaMetafactory#metafactory, "factoryType", + // the return type is the implemented interface, + // while the parameters are the captured variables + MethodTypeDesc.of(targetInterface, effectiveCapturedTypes), + // Bootstrap arguments (see `javap -c -v ` from a working .java sample) + // This is the 4th argument of LambdaMetafactory#metafactory, "interfaceMethodType", + // which is the signature of the implemented method, in this case, void Runnable.run() + MethodTypeDesc.of(methodReturnType, methodArguments), + // This is the 5th argument of LambdaMetafactory#metafactory, "implementation", + // this is the method to be called when invoking the lambda, + // with the captured variables and parameters + MethodHandleDesc.ofMethod( + if (isStatic) DirectMethodHandleDesc.Kind.STATIC else DirectMethodHandleDesc.Kind.VIRTUAL, + targetType, + targetMethod, + MethodTypeDesc.of(targetMethodReturnType, capturedTypes + targetMethodArguments) + ), + // This is the 6th argument of LambdaMetafactory#metafactory, "dynamicMethodType", + // this is "the signature and return type to be enforced dynamically at invocation type" + // This is usually the same as "interfaceMethodType" + MethodTypeDesc.of(methodReturnType, methodArguments), + ) +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/utils/TransformUtils.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/utils/TransformUtils.kt new file mode 100644 index 0000000..4f59e6f --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/utils/TransformUtils.kt @@ -0,0 +1,42 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer.utils + +import java.lang.classfile.ClassBuilder +import java.lang.classfile.ClassFile.ACC_SYNTHETIC +import java.lang.classfile.ClassModel +import java.lang.classfile.MethodModel +import java.lang.reflect.AccessFlag +import kotlin.jvm.optionals.getOrNull + +internal fun Int.withVisibility(visibility: AccessFlag?): Int { + var flags = this + flags = flags and (AccessFlag.PUBLIC.mask() or AccessFlag.PROTECTED.mask() or AccessFlag.PRIVATE.mask()).inv() + if (visibility != null) // null = package-private + flags = flags or visibility.mask() + return flags +} + +internal fun MethodModel.matches(name: String, signature: String): Boolean { + return methodName().equalsString(name) && methodType().equalsString(signature) +} + +internal fun ClassModel.findMethod(name: String, signature: String): MethodModel { + return this.methods().firstOrNull { it.matches(name, signature) } + ?: error("Could not find ${this.thisClass().name().stringValue()}#$name$signature") +} + +context(classBuilder: ClassBuilder) +internal fun MethodModel.transferCodeTo(targetMethodName: String, visibility: AccessFlag = AccessFlag.PRIVATE) { + classBuilder.withMethodBody( + classBuilder.constantPool().utf8Entry(targetMethodName), + methodType(), + visibility.mask() or ACC_SYNTHETIC // Synthetic so this doesn't require a mock + ) { codeBuilder -> + val codeModel = code().orElseThrow { IllegalArgumentException("Method ${this.toFullyQualifiedString()} does not have code") } + codeModel.forEach { codeBuilder.with(it) } + } +} + +internal fun MethodModel.toFullyQualifiedString(): String { + val className = parent().getOrNull()?.thisClass()?.asSymbol()?.displayName() ?: "" + return "$className#${methodName().stringValue()}${methodTypeSymbol().displayDescriptor()}" +} \ No newline at end of file diff --git a/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/utils/JvmUtils.kt b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/utils/JvmUtils.kt new file mode 100644 index 0000000..bb94188 --- /dev/null +++ b/restarter-jda-cache/src/main/kotlin/dev/freya02/botcommands/restart/jda/cache/utils/JvmUtils.kt @@ -0,0 +1,10 @@ +package dev.freya02.botcommands.restart.jda.cache.utils + +internal fun isJvmShuttingDown() = try { + Runtime.getRuntime().removeShutdownHook(NullShutdownHook) + false +} catch (_: IllegalStateException) { + true +} + +private object NullShutdownHook : Thread() \ No newline at end of file diff --git a/restarter-jda-cache/src/test/java/dev/freya02/botcommands/restart/jda/cache/Test.java b/restarter-jda-cache/src/test/java/dev/freya02/botcommands/restart/jda/cache/Test.java new file mode 100644 index 0000000..ccef89d --- /dev/null +++ b/restarter-jda-cache/src/test/java/dev/freya02/botcommands/restart/jda/cache/Test.java @@ -0,0 +1,34 @@ +package dev.freya02.botcommands.restart.jda.cache; + +import io.github.freya022.botcommands.api.core.JDAService; +import io.github.freya022.botcommands.api.core.events.BReadyEvent; +import net.dv8tion.jda.api.hooks.IEventManager; +import net.dv8tion.jda.api.requests.GatewayIntent; +import net.dv8tion.jda.api.utils.cache.CacheFlag; +import org.jetbrains.annotations.NotNull; + +import java.util.Set; + +public class Test extends JDAService { + + @NotNull + @Override + public Set getIntents() { + return Set.of(); + } + + @NotNull + @Override + public Set getCacheFlags() { + return Set.of(); + } + + @Override + protected void createJDA(@NotNull BReadyEvent bReadyEvent, @NotNull IEventManager iEventManager) { + System.out.println("Test"); + } + + void something(BReadyEvent bReadyEvent, IEventManager iEventManager) { + throw new IllegalStateException("test"); + } +} diff --git a/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/BContextImplTransformerTest.kt b/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/BContextImplTransformerTest.kt new file mode 100644 index 0000000..a03ae2d --- /dev/null +++ b/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/BContextImplTransformerTest.kt @@ -0,0 +1,15 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer + +import org.junit.jupiter.api.assertDoesNotThrow +import kotlin.test.Test + +class BContextImplTransformerTest { + + @Test + fun `BContextImpl is instrumented`() { + assertDoesNotThrow { + Class.forName("io.github.freya022.botcommands.internal.core.BContextImpl") + .getDeclaredMethod("doScheduleShutdownSignal", Function0::class.java) + } + } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDABuilderTransformerTest.kt b/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDABuilderTransformerTest.kt new file mode 100644 index 0000000..f7989ff --- /dev/null +++ b/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDABuilderTransformerTest.kt @@ -0,0 +1,135 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer + +import dev.freya02.botcommands.restart.jda.cache.BufferingEventManager +import dev.freya02.botcommands.restart.jda.cache.JDABuilderConfiguration +import dev.freya02.botcommands.restart.jda.cache.JDABuilderSession +import io.mockk.* +import net.dv8tion.jda.api.JDA +import net.dv8tion.jda.api.JDABuilder +import net.dv8tion.jda.api.OnlineStatus +import net.dv8tion.jda.api.events.GenericEvent +import net.dv8tion.jda.api.hooks.IEventManager +import net.dv8tion.jda.api.requests.GatewayIntent +import okhttp3.OkHttpClient +import org.junit.jupiter.api.assertThrows +import java.util.function.Supplier +import kotlin.test.Test + +class JDABuilderTransformerTest { + + @Test + fun `Constructor is instrumented`() { + val builderConfiguration = mockk(relaxUnitFun = true) + + mockkObject(JDABuilderSession) + every { JDABuilderSession.currentSession().configuration } answers { builderConfiguration } + + JDABuilder.create("MY_TOKEN", setOf(GatewayIntent.GUILD_MEMBERS)) + + verify(exactly = 1) { builderConfiguration.onInit("MY_TOKEN", GatewayIntent.getRaw(GatewayIntent.GUILD_MEMBERS)) } + } + + @Test + fun `Unsupported instance method invalidates cache`() { + // Initial set up, this *may* call "markIncompatible" so we need to do it before really mocking + val builder = createJDABuilder() + + // Actual test + val builderConfiguration = mockk { + every { onInit(any(), any()) } just runs + every { markUnsupportedValue(any()) } just runs + } + + mockkObject(JDABuilderSession) + every { JDABuilderSession.currentSession().configuration } returns builderConfiguration + + builder.setHttpClientBuilder(OkHttpClient.Builder()) + + verify(exactly = 1) { builderConfiguration.markUnsupportedValue(any()) } + } + + @Test + fun `Instance method is instrumented`() { + // Initial set up, this *may* call "markIncompatible" so we need to do it before really mocking + val builder = createJDABuilder() + + // Actual test + val builderConfiguration = mockk { + every { onInit(any(), any()) } just runs + every { setStatus(any()) } just runs + } + + mockkObject(JDABuilderSession) + every { JDABuilderSession.currentSession().configuration } returns builderConfiguration + + builder.setStatus(OnlineStatus.DO_NOT_DISTURB) + + verify(exactly = 1) { builderConfiguration.setStatus(OnlineStatus.DO_NOT_DISTURB) } + } + + @Test + fun `Build method is instrumented`() { + val builderSession = mockk { + every { onBuild(any()) } returns mockk() + every { configuration } returns mockk(relaxUnitFun = true) + } + + mockkObject(JDABuilderSession) + every { JDABuilderSession.currentSession() } returns builderSession + + JDABuilder.createDefault("MY_TOKEN").build() + + verify(exactly = 1) { builderSession.onBuild(any()) } + } + + @Test + fun `Build sets our event manager`() { + val builderConfiguration = mockk(relaxUnitFun = true) + + val builderSession = mockk { + every { onBuild(any()) } answers { arg>(0).get() } + every { configuration } returns builderConfiguration + } + + mockkObject(JDABuilderSession) + every { JDABuilderSession.currentSession() } returns builderSession + + val builder = spyk(JDABuilder.createDefault("MY_TOKEN").setEventManager(DummyEventManager)) + + // The special event manager is set on JDABuilder#build() before any original code is run + // so we'll throw an exception on the first method call of the original code, + // which is checkIntents() + every { builder["checkIntents"]() } throws ExpectedException() + assertThrows { builder.build() } + + verify(exactly = 1) { builder.setEventManager(ofType()) } + } + + /** + * Creates a basic JDABuilder, + * call this on the first line to not record any mocking data before doing the actual test. + */ + private fun createJDABuilder(): JDABuilder { + lateinit var builder: JDABuilder + mockkObject(JDABuilderSession) { + every { JDABuilderSession.currentSession().configuration } returns mockk(relaxUnitFun = true) + + builder = JDABuilder.create("MY_TOKEN", emptySet()) + } + + return builder + } + + private object DummyEventManager : IEventManager { + + override fun register(listener: Any) {} + + override fun unregister(listener: Any) {} + + override fun handle(event: GenericEvent) {} + + override fun getRegisteredListeners(): List = emptyList() + } + + private class ExpectedException : RuntimeException() +} \ No newline at end of file diff --git a/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAImplTransformerTest.kt b/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAImplTransformerTest.kt new file mode 100644 index 0000000..da37b3f --- /dev/null +++ b/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAImplTransformerTest.kt @@ -0,0 +1,26 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer + +import dev.freya02.botcommands.restart.jda.cache.JDABuilderSession +import io.mockk.* +import net.dv8tion.jda.internal.JDAImpl +import kotlin.test.Test + +class JDAImplTransformerTest { + + @Test + fun `Shutdown method is instrumented`() { + val builderSession = mockk { + every { onShutdown(any(), any()) } just runs + } + + val jda = mockk { + // If this getter is missing, then the codegen changed + every { this@mockk["getBuilderSession"]() } returns builderSession + every { shutdown() } answers { callOriginal() } + } + + jda.shutdown() + + verify(exactly = 1) { builderSession.onShutdown(jda, any()) } + } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAServiceTransformerTest.kt b/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAServiceTransformerTest.kt new file mode 100644 index 0000000..9c7badc --- /dev/null +++ b/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/JDAServiceTransformerTest.kt @@ -0,0 +1,87 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer + +import dev.freya02.botcommands.restart.jda.cache.JDABuilderSession +import io.github.freya022.botcommands.api.core.JDAService +import io.github.freya022.botcommands.api.core.events.BReadyEvent +import io.mockk.* +import net.dv8tion.jda.api.hooks.IEventManager +import net.dv8tion.jda.api.requests.GatewayIntent +import net.dv8tion.jda.api.utils.cache.CacheFlag +import kotlin.test.Test + +class JDAServiceTransformerTest { + + class Bot : JDAService() { + + override val intents: Set = emptySet() + override val cacheFlags: Set = emptySet() + + public override fun createJDA(event: BReadyEvent, eventManager: IEventManager) { + println("createJDA") + } + } + + @Test + fun `Event listener is instrumented`() { + mockkObject(JDABuilderSession) + every { JDABuilderSession.withBuilderSession(any(), any()) } answers { callOriginal() } // Will call createJDA + + val onReadyEvent = JDAService::class.java.getDeclaredMethod($$"onReadyEvent$BotCommands", BReadyEvent::class.java, IEventManager::class.java) + val bot = mockk { + every { createJDA(any(), any()) } just runs + every { onReadyEvent.invoke(this@mockk, any(), any()) } answers { callOriginal() } // Will call withBuilderSession + } + + val readyEvent = mockk { + every { context.config.restartConfig.cacheKey } returns "Test" + } + val eventManager = mockk() + + onReadyEvent.invoke(bot, readyEvent, eventManager) + + verify(exactly = 1) { bot.createJDA(readyEvent, eventManager) } + } + + @Test + fun `Cache key enables builder sessions`() { + mockkObject(JDABuilderSession) + every { JDABuilderSession.withBuilderSession(any(), any()) } answers { callOriginal() } + + val onReadyEvent = JDAService::class.java.getDeclaredMethod($$"onReadyEvent$BotCommands", BReadyEvent::class.java, IEventManager::class.java) + val bot = mockk { + every { createJDA(any(), any()) } just runs + every { onReadyEvent.invoke(this@mockk, any(), any()) } answers { callOriginal() } // Will call withBuilderSession + } + + val readyEvent = mockk { + every { context.config.restartConfig.cacheKey } returns "Test" + } + val eventManager = mockk() + + onReadyEvent.invoke(bot, readyEvent, eventManager) + + verify(exactly = 1) { JDABuilderSession.withBuilderSession(any(), any()) } + } + + @Test + fun `Null cache key disables builder sessions`() { + mockkObject(JDABuilderSession) + every { JDABuilderSession.withBuilderSession(any(), any()) } answers { callOriginal() } + every { JDABuilderSession.getCacheKey(any()) } answers { callOriginal() } + + val onReadyEvent = JDAService::class.java.getDeclaredMethod($$"onReadyEvent$BotCommands", BReadyEvent::class.java, IEventManager::class.java) + val bot = mockk { + every { createJDA(any(), any()) } just runs + every { onReadyEvent.invoke(this@mockk, any(), any()) } answers { callOriginal() } // Will call withBuilderSession + } + + val readyEvent = mockk { + every { context.config.restartConfig.cacheKey } returns null + } + val eventManager = mockk() + + onReadyEvent.invoke(bot, readyEvent, eventManager) + + verify(exactly = 0) { JDABuilderSession.withBuilderSession(any(), any()) } + } +} \ No newline at end of file diff --git a/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/utils/CodeBuilderUtilsTest.kt b/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/utils/CodeBuilderUtilsTest.kt new file mode 100644 index 0000000..5931e19 --- /dev/null +++ b/restarter-jda-cache/src/test/kotlin/dev/freya02/botcommands/restart/jda/cache/transformer/utils/CodeBuilderUtilsTest.kt @@ -0,0 +1,121 @@ +package dev.freya02.botcommands.restart.jda.cache.transformer.utils + +import dev.freya02.botcommands.restart.jda.cache.transformer.* +import org.junit.jupiter.params.ParameterizedTest +import org.junit.jupiter.params.provider.Arguments +import org.junit.jupiter.params.provider.MethodSource +import java.lang.constant.* +import java.util.function.Supplier +import kotlin.test.assertEquals + +object CodeBuilderUtilsTest { + + @MethodSource("Test createLambda") + @ParameterizedTest + fun `Test createLambda`(expected: DynamicCallSiteDesc, actual: DynamicCallSiteDesc) { + assertEquals(expected, actual) + } + + @JvmStatic + fun `Test createLambda`(): List = listOf( + Arguments.of( + DynamicCallSiteDesc.of( + lambdaMetafactoryDesc, + "get", + MethodTypeDesc.of(CD_Supplier, CD_JDABuilder), + MethodTypeDesc.of(ConstantDescs.CD_Object), + MethodHandleDesc.ofMethod( + DirectMethodHandleDesc.Kind.VIRTUAL, + CD_JDABuilder, + "doBuild", + MethodTypeDesc.of(CD_JDA) + ), + MethodTypeDesc.of(ConstantDescs.CD_Object), + ), + createLambda( + interfaceMethod = Supplier<*>::get, + targetType = CD_JDABuilder, + targetMethod = "doBuild", + targetMethodReturnType = CD_JDA, + targetMethodArguments = listOf(), + capturedTypes = listOf(), + isStatic = false + ) + ), + + Arguments.of( + DynamicCallSiteDesc.of( + lambdaMetafactoryDesc, + "run", + MethodTypeDesc.of(CD_Runnable, CD_JDAService, CD_BReadyEvent, CD_IEventManager), + MethodTypeDesc.of(ConstantDescs.CD_void), + MethodHandleDesc.ofMethod( + DirectMethodHandleDesc.Kind.VIRTUAL, + CD_JDAService, + $$"lambda$onReadyEvent$BotCommands$withBuilderSession", + MethodTypeDesc.of(ConstantDescs.CD_void, CD_BReadyEvent, CD_IEventManager) + ), + MethodTypeDesc.of(ConstantDescs.CD_void), + ), + createLambda( + interfaceMethod = Runnable::run, + targetType = CD_JDAService, + targetMethod = $$"lambda$onReadyEvent$BotCommands$withBuilderSession", + targetMethodReturnType = ConstantDescs.CD_void, + targetMethodArguments = listOf(), + capturedTypes = listOf(CD_BReadyEvent, CD_IEventManager), + isStatic = false + ) + ), + + Arguments.of( + DynamicCallSiteDesc.of( + lambdaMetafactoryDesc, + "run", + MethodTypeDesc.of(CD_Runnable, CD_JDAImpl), + MethodTypeDesc.of(ConstantDescs.CD_void), + MethodHandleDesc.ofMethod( + DirectMethodHandleDesc.Kind.VIRTUAL, + CD_JDAImpl, + "doShutdown", + MethodTypeDesc.of(ConstantDescs.CD_void) + ), + MethodTypeDesc.of(ConstantDescs.CD_void), + ), + createLambda( + interfaceMethod = Runnable::run, + targetType = CD_JDAImpl, + targetMethod = "doShutdown", + targetMethodReturnType = ConstantDescs.CD_void, + targetMethodArguments = listOf(), + capturedTypes = listOf(), + isStatic = false + ) + ), + + Arguments.of( + DynamicCallSiteDesc.of( + lambdaMetafactoryDesc, + "run", + MethodTypeDesc.of(CD_Runnable, CD_BContextImpl, CD_Function0), + MethodTypeDesc.of(ConstantDescs.CD_void), + MethodHandleDesc.ofMethod( + DirectMethodHandleDesc.Kind.VIRTUAL, + CD_BContextImpl, + "doScheduleShutdownSignal", + MethodTypeDesc.of(ConstantDescs.CD_void, CD_Function0) + ), + MethodTypeDesc.of(ConstantDescs.CD_void), + ), + createLambda( + interfaceMethod = Runnable::run, + targetType = CD_BContextImpl, + targetMethod = "doScheduleShutdownSignal", + targetMethodReturnType = ConstantDescs.CD_void, + targetMethodArguments = listOf(), + capturedTypes = listOf(CD_Function0), + isStatic = false + ) + ), + ) +} \ No newline at end of file diff --git a/restarter-jda-cache/src/test/resources/logback-test.xml b/restarter-jda-cache/src/test/resources/logback-test.xml new file mode 100644 index 0000000..4a6e03c --- /dev/null +++ b/restarter-jda-cache/src/test/resources/logback-test.xml @@ -0,0 +1,14 @@ + + + + + %d{HH:mm:ss.SSS} %boldCyan(%-26.-26thread) %boldYellow(%-20.-20logger{0}) %highlight(%-6level) %msg%n%throwable + + + + + + + + + \ No newline at end of file diff --git a/restarter/build.gradle.kts b/restarter/build.gradle.kts new file mode 100644 index 0000000..6a31a5d --- /dev/null +++ b/restarter/build.gradle.kts @@ -0,0 +1,22 @@ +plugins { + id("BotCommands-Restarter-conventions") + `maven-publish` +} + +repositories { + maven("https://jitpack.io") +} + +dependencies { + implementation(libs.botcommands) +} + +publishing { + publications { + create("maven") { + from(components["java"]) + + artifactId = "BotCommands-Restarter" + } + } +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/ImmediateRestartException.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/ImmediateRestartException.kt new file mode 100644 index 0000000..19edb5b --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/ImmediateRestartException.kt @@ -0,0 +1,29 @@ +package dev.freya02.botcommands.internal.restart + +import java.lang.reflect.InvocationTargetException + +class ImmediateRestartException internal constructor() : RuntimeException("Dummy exception to stop the execution of the first main thread") { + + internal companion object { + internal fun throwAndHandle(): Nothing { + val currentThread = Thread.currentThread() + currentThread.uncaughtExceptionHandler = ExpectedReloadExceptionHandler(currentThread.uncaughtExceptionHandler) + throw ImmediateRestartException() + } + } + + private class ExpectedReloadExceptionHandler(private val delegate: Thread.UncaughtExceptionHandler?) : Thread.UncaughtExceptionHandler { + + override fun uncaughtException(t: Thread, e: Throwable) { + if (e is ImmediateRestartException || (e is InvocationTargetException && e.targetException is ImmediateRestartException)) { + return + } + + if (delegate != null) { + delegate.uncaughtException(t, e) + } else { + e.printStackTrace() + } + } + } +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/LeakSafeExecutor.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/LeakSafeExecutor.kt new file mode 100644 index 0000000..b36bbbe --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/LeakSafeExecutor.kt @@ -0,0 +1,61 @@ +package dev.freya02.botcommands.internal.restart + +import java.util.concurrent.BlockingDeque +import java.util.concurrent.LinkedBlockingDeque +import kotlin.system.exitProcess + +internal class LeakSafeExecutor internal constructor() { + + // As we can only use a Thread once, we put a single LeakSafeThread in a blocking queue, + // then, when a code block runs, a LeakSafeThread is removed from the queue, + // and the LeakSafeThread recreates a new one for the next code block. + // We use a blocking queue to prevent trying to get a LeakSafeThread between the moment it was retrieved and when it'll be added back + private val leakSafeThreads: BlockingDeque = LinkedBlockingDeque() + + init { + leakSafeThreads += LeakSafeThread() + } + + fun callAndWait(callable: () -> V): V = getLeakSafeThread().callAndWait(callable) + + private fun getLeakSafeThread(): LeakSafeThread { + return leakSafeThreads.takeFirst() + } + + /** + * Thread that is created early so not to retain the [RestartClassLoader]. + */ + private inner class LeakSafeThread : Thread() { + + private var callable: (() -> Any?)? = null + + private var result: Any? = null + + init { + isDaemon = false + } + + @Suppress("UNCHECKED_CAST") + fun callAndWait(callable: () -> V): V { + this.callable = callable + start() + try { + join() + return this.result as V + } catch (ex: InterruptedException) { + currentThread().interrupt() + throw IllegalStateException(ex) + } + } + + override fun run() { + try { + this@LeakSafeExecutor.leakSafeThreads.put(LeakSafeThread()) + this.result = this.callable!!.invoke() + } catch (ex: Exception) { + ex.printStackTrace() + exitProcess(1) + } + } + } +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/RestartClassLoader.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/RestartClassLoader.kt new file mode 100644 index 0000000..d803b78 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/RestartClassLoader.kt @@ -0,0 +1,40 @@ +package dev.freya02.botcommands.internal.restart + +import java.net.URL +import java.net.URLClassLoader +import java.util.* + +// STILL SUPER DUPER IMPORTANT TO OVERRIDE SOME STUFF AND DELEGATE +internal class RestartClassLoader internal constructor( + urls: List, + parent: ClassLoader, +) : URLClassLoader(urls.toTypedArray(), parent) { + + override fun getResources(name: String): Enumeration { + return this.parent.getResources(name) + } + + override fun getResource(name: String): URL? { + return findResource(name) ?: super.getResource(name) + } + + override fun findResource(name: String): URL? { + return super.findResource(name) + } + + override fun loadClass(name: String, resolve: Boolean): Class<*> { + return synchronized(getClassLoadingLock(name)) { + val loadedClass = findLoadedClass(name) ?: try { + findClass(name) + } catch (_: ClassNotFoundException) { + Class.forName(name, false, parent) + } + if (resolve) resolveClass(loadedClass) + loadedClass + } + } + + override fun findClass(name: String): Class<*> { + return super.findClass(name) + } +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/RestartClassLoaderFull.kt.disabled b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/RestartClassLoaderFull.kt.disabled new file mode 100644 index 0000000..605a0b4 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/RestartClassLoaderFull.kt.disabled @@ -0,0 +1,116 @@ +package dev.freya02.botcommands.internal.restart + +import dev.freya02.botcommands.internal.restart.sources.DeletedSourceFile +import dev.freya02.botcommands.internal.restart.sources.SourceDirectories +import dev.freya02.botcommands.internal.restart.sources.SourceFile +import java.io.InputStream +import java.net.URL +import java.net.URLClassLoader +import java.net.URLConnection +import java.net.URLStreamHandler +import java.util.* + +internal class RestartClassLoader internal constructor( + urls: List, + parent: ClassLoader, + private val sourceDirectories: SourceDirectories, +) : URLClassLoader(urls.toTypedArray(), parent) { + + override fun getResources(name: String): Enumeration { + val resources = parent.getResources(name) + val updatedFile = sourceDirectories.getFile(name) + + if (updatedFile != null) { + if (resources.hasMoreElements()) { + resources.nextElement() + } + if (updatedFile is SourceFile) { + return MergedEnumeration(createFileUrl(name, updatedFile), resources) + } + } + + return resources + } + + override fun getResource(name: String): URL? { + val updatedFile = sourceDirectories.getFile(name) + if (updatedFile is DeletedSourceFile) { + return null + } + + return findResource(name) ?: super.getResource(name) + } + + override fun findResource(name: String): URL? { + val updatedFile = sourceDirectories.getFile(name) + ?: return super.findResource(name) + return (updatedFile as? SourceFile)?.let { createFileUrl(name, it) } + } + + override fun loadClass(name: String, resolve: Boolean): Class<*> { + val path = "${name.replace('.', '/')}.class" + val updatedFile = sourceDirectories.getFile(path) + if (updatedFile is DeletedSourceFile) + throw ClassNotFoundException(name) + + return synchronized(getClassLoadingLock(name)) { + val loadedClass = findLoadedClass(name) ?: try { + findClass(name) + } catch (_: ClassNotFoundException) { + Class.forName(name, false, parent) + } + if (resolve) resolveClass(loadedClass) + loadedClass + } + } + + override fun findClass(name: String): Class<*> { + val path = "${name.replace('.', '/')}.class" + val updatedFile = sourceDirectories.getFile(path) + ?: return super.findClass(name) + if (updatedFile is DeletedSourceFile) + throw ClassNotFoundException(name) + + updatedFile as SourceFile + return defineClass(name, updatedFile.bytes, 0, updatedFile.bytes.size) + } + + @Suppress("DEPRECATION") // We target Java 17 but JDK 20 deprecates the URL constructors + private fun createFileUrl(name: String, file: SourceFile): URL { + return URL("reloaded", null, -1, "/$name", ClasspathFileURLStreamHandler(file)) + } + + private class ClasspathFileURLStreamHandler( + private val file: SourceFile, + ) : URLStreamHandler() { + + override fun openConnection(u: URL): URLConnection = Connection(u) + + private inner class Connection(url: URL): URLConnection(url) { + + override fun connect() {} + + override fun getInputStream(): InputStream = file.bytes.inputStream() + + override fun getLastModified(): Long = file.lastModified.toEpochMilli() + + override fun getContentLengthLong(): Long = file.bytes.size.toLong() + } + } + + private class MergedEnumeration(private val first: E, private val rest: Enumeration) : Enumeration { + + private var hasConsumedFirst = false + + override fun hasMoreElements(): Boolean = !hasConsumedFirst || rest.hasMoreElements() + + override fun nextElement(): E? { + if (!hasConsumedFirst) { + hasConsumedFirst = true + return first + } else { + return rest.nextElement() + } + } + } +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/RestartListener.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/RestartListener.kt new file mode 100644 index 0000000..4d77390 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/RestartListener.kt @@ -0,0 +1,5 @@ +package dev.freya02.botcommands.internal.restart + +interface RestartListener { + fun beforeStop() +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/Restarter.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/Restarter.kt new file mode 100644 index 0000000..bcdaa5b --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/Restarter.kt @@ -0,0 +1,125 @@ +package dev.freya02.botcommands.internal.restart + +import dev.freya02.botcommands.internal.restart.utils.AppClasspath +import io.github.oshai.kotlinlogging.KotlinLogging +import java.net.URL +import java.util.concurrent.locks.Lock +import java.util.concurrent.locks.ReentrantLock +import kotlin.concurrent.thread +import kotlin.concurrent.withLock + +private val logger = KotlinLogging.logger { } + +class Restarter private constructor( + private val args: List, +) { + + private val appClassLoader: ClassLoader + val appClasspathUrls: List + + private val mainClassName: String + + private val uncaughtExceptionHandler: Thread.UncaughtExceptionHandler + + private val stopLock: Lock = ReentrantLock() + private val listeners: MutableList = arrayListOf() + + private val leakSafeExecutor = LeakSafeExecutor() + + init { + val thread = Thread.currentThread() + + appClassLoader = thread.contextClassLoader + appClasspathUrls = AppClasspath.getPaths().map { it.toUri().toURL() } + + mainClassName = StackWalker.getInstance(StackWalker.Option.RETAIN_CLASS_REFERENCE) + .walk { stream -> stream.filter { it.methodName == "main" }.toList().last() } + .declaringClass.name + + uncaughtExceptionHandler = thread.uncaughtExceptionHandler + } + + fun addListener(listener: RestartListener) { + listeners += listener + } + + private fun initialize(): Nothing { + val throwable = leakSafeExecutor.callAndWait { start() } + if (throwable != null) + throw throwable + ImmediateRestartException.throwAndHandle() + } + + /** + * Runs each [RestartListener.beforeStop] and then starts a new instance of the main class, + * if the new instance fails, the [Throwable] is returned. + */ + fun restart(): Throwable? { + logger.debug { "Restarting application in '$mainClassName'" } + // Do it from the original class loader, so the context is the same as for the initial restart + return leakSafeExecutor.callAndWait { + stop() + start() + } + } + + private fun stop() { + stopLock.withLock { + listeners.forEach { it.beforeStop() } + listeners.clear() + } + // All threads should be stopped at that point + // so the GC should be able to remove all the previous loaded classes + System.gc() + } + + /** + * Starts a new instance of the main class, or returns a [Throwable] if it failed. + */ + private fun start(): Throwable? { + // We use a regular URLClassLoader instead of RestartClassLoaderFull, + // as classpath changes will trigger a restart and thus recreate a new ClassLoader, + // meaning live updating the classes is pointless. + // In contrast, Spring needs their RestartClassLoader because it can override classes remotely, + // but we don't have such a use case. + // However, not using RestartClassLoaderFull, which uses snapshots, has an issue, + // trying to load deleted classes (most likely on shutdown) will fail, + // Spring also has that issue, but it will only happen on classes out of its component scan, + // BC just needs to make sure to at least load the classes on its path too. + val restartClassLoader = RestartClassLoader(appClasspathUrls, appClassLoader) + var error: Throwable? = null + val launchThreads = thread(name = RESTARTED_THREAD_NAME, isDaemon = false, contextClassLoader = restartClassLoader) { + try { + val mainClass = Class.forName(mainClassName, false, restartClassLoader) + val mainMethod = mainClass.getDeclaredMethod("main", Array::class.java) + mainMethod.isAccessible = true + mainMethod.invoke(null, args.toTypedArray()) + } catch (ex: Throwable) { + error = ex + } + } + launchThreads.join() + + return error + } + + companion object { + + const val RESTARTED_THREAD_NAME = "restartedMain" + + private val instanceLock: Lock = ReentrantLock() + lateinit var instance: Restarter + private set + + fun initialize(args: List) { + var newInstance: Restarter? = null + instanceLock.withLock { + if (::instance.isInitialized.not()) { + newInstance = Restarter(args) + instance = newInstance + } + } + newInstance?.initialize() + } + } +} diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/services/RestarterApplicationStartListener.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/services/RestarterApplicationStartListener.kt new file mode 100644 index 0000000..63886c3 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/services/RestarterApplicationStartListener.kt @@ -0,0 +1,16 @@ +package dev.freya02.botcommands.internal.restart.services + +import dev.freya02.botcommands.internal.restart.Restarter +import io.github.freya022.botcommands.api.core.events.ApplicationStartListener +import io.github.freya022.botcommands.api.core.events.BApplicationStartEvent +import io.github.freya022.botcommands.api.core.service.annotations.BService +import io.github.freya022.botcommands.api.core.service.annotations.RequiresDefaultInjection + +@BService +@RequiresDefaultInjection +internal class RestarterApplicationStartListener : ApplicationStartListener { + + override fun onApplicationStart(event: BApplicationStartEvent) { + Restarter.initialize(event.args) + } +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/services/RestarterService.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/services/RestarterService.kt new file mode 100644 index 0000000..61572b2 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/services/RestarterService.kt @@ -0,0 +1,27 @@ +package dev.freya02.botcommands.internal.restart.services + +import dev.freya02.botcommands.internal.restart.RestartListener +import dev.freya02.botcommands.internal.restart.Restarter +import dev.freya02.botcommands.internal.restart.watcher.ClasspathWatcher +import io.github.freya022.botcommands.api.core.BContext +import io.github.freya022.botcommands.api.core.config.BRestartConfig +import io.github.freya022.botcommands.api.core.service.annotations.BService +import io.github.freya022.botcommands.api.core.service.annotations.RequiresDefaultInjection + +@BService +@RequiresDefaultInjection +internal class RestarterService internal constructor ( + context: BContext, + config: BRestartConfig, +) { + + init { + Restarter.instance.addListener(object : RestartListener { + override fun beforeStop() { + context.shutdownNow() + context.awaitShutdown(context.config.shutdownTimeout) + } + }) + ClasspathWatcher.initialize(config.restartDelay) + } +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectories.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectories.kt new file mode 100644 index 0000000..2dc84c9 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectories.kt @@ -0,0 +1,48 @@ +package dev.freya02.botcommands.internal.restart.sources + +import java.nio.file.Path + +internal class SourceDirectories internal constructor() { + private val directories: MutableMap = hashMapOf() + + internal fun getFile(path: String): ISourceFile? { + return directories.firstNotNullOfOrNull { it.value.files[path] } + } + + internal fun setSource(source: SourceDirectory) { + directories[source.directory] = source + } + + internal fun replaceSource(key: Path, directory: SourceDirectory) { + check(key in directories) + + directories[key] = directory + } + + internal fun close() { + directories.values.forEach { it.close() } + } +} + +internal fun SourceDirectories(directories: List, listener: SourceDirectoriesListener): SourceDirectories { + val sourceDirectories = SourceDirectories() + + fun onSourceDirectoryUpdate(directory: Path, sourceFilesFactory: () -> SourceFiles) { + // The command is called when restarting + // so we don't make snapshots before all changes went through + listener.onChange(command = { + val newSourceDirectory = SourceDirectory( + directory, + sourceFilesFactory(), + listener = { onSourceDirectoryUpdate(directory, it) } + ) + sourceDirectories.replaceSource(directory, newSourceDirectory) + }) + } + + directories.forEach { directory -> + sourceDirectories.setSource(SourceDirectory(directory, listener = { onSourceDirectoryUpdate(directory, it) })) + } + + return sourceDirectories +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectoriesListener.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectoriesListener.kt new file mode 100644 index 0000000..75b9f72 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectoriesListener.kt @@ -0,0 +1,7 @@ +package dev.freya02.botcommands.internal.restart.sources + +internal interface SourceDirectoriesListener { + fun onChange(command: () -> Unit) + + fun onCancel() +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectory.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectory.kt new file mode 100644 index 0000000..cc5ab38 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectory.kt @@ -0,0 +1,90 @@ +package dev.freya02.botcommands.internal.restart.sources + +import dev.freya02.botcommands.internal.restart.utils.walkDirectories +import dev.freya02.botcommands.internal.restart.utils.walkFiles +import io.github.oshai.kotlinlogging.KotlinLogging +import java.nio.file.Path +import java.nio.file.StandardWatchEventKinds.* +import kotlin.concurrent.thread +import kotlin.io.path.* + +private val logger = KotlinLogging.logger { } + +@OptIn(ExperimentalPathApi::class) +internal class SourceDirectory internal constructor( + val directory: Path, + val files: SourceFiles, + private val listener: SourceDirectoryListener, +) { + + private val thread: Thread + + init { + require(directory.isDirectory()) + + logger.trace { "Listening to ${directory.absolutePathString()}" } + + val watchService = directory.fileSystem.newWatchService() + directory.walkDirectories { path, attributes -> + path.register(watchService, ENTRY_CREATE, ENTRY_MODIFY, ENTRY_DELETE) + } + + thread = thread(name = "Classpath watcher of '${directory.fileName}'", isDaemon = true) { + try { + watchService.take() // Wait for a change + } catch (_: InterruptedException) { + return@thread logger.trace { "Interrupted watching ${directory.absolutePathString()}" } + } + watchService.close() + + listener.onChange(sourcesFilesFactory = { + val snapshot = directory.takeSnapshot() + + // Exclude deleted files so they don't count as being deleted again + val deletedPaths = files.withoutDeletes().keys - snapshot.keys + if (deletedPaths.isNotEmpty()) { + logger.info { "Deleted files in ${directory.absolutePathString()}: $deletedPaths" } + return@onChange deletedPaths.associateWith { DeletedSourceFile } + snapshot + } + + // Exclude deleted files so they count as being added back + val addedPaths = snapshot.keys - files.withoutDeletes().keys + if (addedPaths.isNotEmpty()) { + logger.info { "Added files in ${directory.absolutePathString()}: $addedPaths" } + return@onChange files + snapshot + } + + val modifiedFiles = snapshot.keys.filter { key -> + val actual = snapshot[key] ?: error("Key from map is missing a value somehow") + val expected = files[key] ?: error("Expected file is missing, should have been detected as deleted") + + // File was deleted (on the 2nd build for example) and got recreated (on the 3rd build for example) + if (expected is DeletedSourceFile) error("Expected file was registered as deleted, should have been detected as added") + expected as SourceFile + + actual as SourceFile // Assertion + + actual.lastModified != expected.lastModified + } + if (modifiedFiles.isNotEmpty()) { + logger.info { "Timestamp changed in ${directory.absolutePathString()}: $modifiedFiles" } + return@onChange files + snapshot + } + + error("Received a file system event but no changes were detected") + }) + } + } + + internal fun close() { + thread.interrupt() + } +} + +internal fun SourceDirectory(directory: Path, listener: SourceDirectoryListener): SourceDirectory { + return SourceDirectory(directory, directory.takeSnapshot(), listener) +} + +private fun Path.takeSnapshot(): SourceFiles = walkFiles().associate { (it, attrs) -> + it.relativeTo(this).pathString to SourceFile(attrs.lastModifiedTime().toInstant()) +}.let(::SourceFiles) diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectoryListener.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectoryListener.kt new file mode 100644 index 0000000..5350211 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceDirectoryListener.kt @@ -0,0 +1,5 @@ +package dev.freya02.botcommands.internal.restart.sources + +internal fun interface SourceDirectoryListener { + fun onChange(sourcesFilesFactory: () -> SourceFiles) +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceFile.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceFile.kt new file mode 100644 index 0000000..648bb57 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceFile.kt @@ -0,0 +1,15 @@ +package dev.freya02.botcommands.internal.restart.sources + +import java.time.Instant + +internal sealed interface ISourceFile + +internal class SourceFile( + val lastModified: Instant, +) : ISourceFile { + + val bytes: ByteArray + get() = throw UnsupportedOperationException("Class data is no longer retained as RestartClassLoader is not used yet") +} + +internal object DeletedSourceFile : ISourceFile \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceFiles.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceFiles.kt new file mode 100644 index 0000000..2ded125 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/sources/SourceFiles.kt @@ -0,0 +1,16 @@ +package dev.freya02.botcommands.internal.restart.sources + +internal class SourceFiles internal constructor( + internal val files: Map, +) { + + val keys: Set get() = files.keys + + internal operator fun get(path: String): ISourceFile? = files[path] + + internal fun withoutDeletes(): SourceFiles = SourceFiles(files.filterValues { it !is DeletedSourceFile }) + + internal operator fun plus(other: SourceFiles): SourceFiles = SourceFiles(files + other.files) +} + +internal operator fun Map.plus(other: SourceFiles): SourceFiles = SourceFiles(this + other.files) \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/utils/AppClasspath.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/utils/AppClasspath.kt new file mode 100644 index 0000000..39eb256 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/utils/AppClasspath.kt @@ -0,0 +1,17 @@ +package dev.freya02.botcommands.internal.restart.utils + +import java.io.File +import java.lang.management.ManagementFactory +import java.nio.file.Path +import kotlin.io.path.Path +import kotlin.io.path.isDirectory + +internal object AppClasspath { + + fun getPaths(): List { + return ManagementFactory.getRuntimeMXBean().classPath + .split(File.pathSeparator) + .map(::Path) + .filter { it.isDirectory() } + } +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/utils/NIO.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/utils/NIO.kt new file mode 100644 index 0000000..cbcbb7d --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/utils/NIO.kt @@ -0,0 +1,66 @@ +package dev.freya02.botcommands.internal.restart.utils + +import java.io.IOException +import java.nio.file.FileVisitResult +import java.nio.file.FileVisitor +import java.nio.file.Files +import java.nio.file.Path +import java.nio.file.attribute.BasicFileAttributes + +// Optimization of Path#walk, cuts CPU usage by 4 +// mostly by eliminating duplicate calls to file attributes +internal fun Path.walkFiles(): List> { + return buildList { + Files.walkFileTree(this@walkFiles, object : FileVisitor { + override fun preVisitDirectory( + dir: Path, + attrs: BasicFileAttributes + ): FileVisitResult = FileVisitResult.CONTINUE + + override fun visitFile( + file: Path, + attrs: BasicFileAttributes + ): FileVisitResult { + add(file to attrs) + return FileVisitResult.CONTINUE + } + + override fun visitFileFailed( + file: Path, + exc: IOException + ): FileVisitResult = FileVisitResult.CONTINUE + + override fun postVisitDirectory( + dir: Path, + exc: IOException? + ): FileVisitResult = FileVisitResult.CONTINUE + }) + } +} + +internal fun Path.walkDirectories(block: (Path, BasicFileAttributes) -> Unit) { + Files.walkFileTree(this@walkDirectories, object : FileVisitor { + override fun preVisitDirectory( + dir: Path, + attrs: BasicFileAttributes + ): FileVisitResult { + block(dir, attrs) + return FileVisitResult.CONTINUE + } + + override fun visitFile( + file: Path, + attrs: BasicFileAttributes + ): FileVisitResult = FileVisitResult.CONTINUE + + override fun visitFileFailed( + file: Path, + exc: IOException + ): FileVisitResult = FileVisitResult.CONTINUE + + override fun postVisitDirectory( + dir: Path, + exc: IOException? + ): FileVisitResult = FileVisitResult.CONTINUE + }) +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/watcher/ClasspathListener.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/watcher/ClasspathListener.kt new file mode 100644 index 0000000..6805bd4 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/watcher/ClasspathListener.kt @@ -0,0 +1,42 @@ +package dev.freya02.botcommands.internal.restart.watcher + +import dev.freya02.botcommands.internal.restart.Restarter +import dev.freya02.botcommands.internal.restart.sources.SourceDirectoriesListener +import io.github.oshai.kotlinlogging.KotlinLogging +import java.util.concurrent.Executors +import java.util.concurrent.ScheduledFuture +import java.util.concurrent.TimeUnit +import kotlin.time.Duration + +private val logger = KotlinLogging.logger { } + +internal class ClasspathListener internal constructor( + private val delay: Duration +) : SourceDirectoriesListener { + + private val scheduler = Executors.newSingleThreadScheduledExecutor() + private lateinit var scheduledRestart: ScheduledFuture<*> + + private val commands: MutableList<() -> Unit> = arrayListOf() + + override fun onChange(command: () -> Unit) { + commands += command + if (::scheduledRestart.isInitialized) scheduledRestart.cancel(false) + + scheduledRestart = scheduler.schedule({ + commands.forEach { it.invoke() } + commands.clear() + + try { + Restarter.instance.restart() + } catch (e: Exception) { + logger.error(e) { "Restart failed, waiting for the next build" } + } + scheduler.shutdown() + }, delay.inWholeMilliseconds, TimeUnit.MILLISECONDS) + } + + override fun onCancel() { + scheduler.shutdownNow() + } +} \ No newline at end of file diff --git a/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/watcher/ClasspathWatcher.kt b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/watcher/ClasspathWatcher.kt new file mode 100644 index 0000000..7184116 --- /dev/null +++ b/restarter/src/main/kotlin/dev/freya02/botcommands/internal/restart/watcher/ClasspathWatcher.kt @@ -0,0 +1,182 @@ +package dev.freya02.botcommands.internal.restart.watcher + +import dev.freya02.botcommands.internal.restart.Restarter +import dev.freya02.botcommands.internal.restart.sources.DeletedSourceFile +import dev.freya02.botcommands.internal.restart.sources.SourceFile +import dev.freya02.botcommands.internal.restart.sources.SourceFiles +import dev.freya02.botcommands.internal.restart.sources.plus +import dev.freya02.botcommands.internal.restart.utils.AppClasspath +import dev.freya02.botcommands.internal.restart.utils.walkDirectories +import dev.freya02.botcommands.internal.restart.utils.walkFiles +import io.github.freya022.botcommands.api.core.utils.joinAsList +import io.github.oshai.kotlinlogging.KotlinLogging +import java.nio.file.FileSystems +import java.nio.file.Path +import java.nio.file.StandardWatchEventKinds.* +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.Executors +import java.util.concurrent.ScheduledFuture +import java.util.concurrent.TimeUnit +import java.util.concurrent.locks.ReentrantLock +import kotlin.concurrent.thread +import kotlin.concurrent.withLock +import kotlin.io.path.absolutePathString +import kotlin.io.path.isDirectory +import kotlin.io.path.pathString +import kotlin.io.path.relativeTo +import kotlin.time.Duration + +private val logger = KotlinLogging.logger { } + +// Lightweight, singleton version of [[SourceDirectories]] + [[ClasspathListener]] +internal class ClasspathWatcher private constructor( + private var settings: Settings?, // null = no instance registered = no restart can be scheduled +) { + + private val scheduler = Executors.newSingleThreadScheduledExecutor() + private lateinit var restartFuture: ScheduledFuture<*> + + private val watchService = FileSystems.getDefault().newWatchService() + private val registeredDirectories: MutableSet = ConcurrentHashMap.newKeySet() + private val snapshots: MutableMap = hashMapOf() + + init { + AppClasspath.getPaths().forEach { classRoot -> + require(classRoot.isDirectory()) + + logger.trace { "Creating snapshot of ${classRoot.absolutePathString()}" } + snapshots[classRoot] = classRoot.takeSnapshot() + + logger.trace { "Listening to ${classRoot.absolutePathString()}" } + registerDirectories(classRoot) + } + + thread(name = "Classpath watcher", isDaemon = true) { + while (true) { + val key = try { + watchService.take() // Wait for a change + } catch (_: InterruptedException) { + return@thread logger.trace { "Interrupted watching classpath" } + } + val pollEvents = key.pollEvents() + if (pollEvents.isNotEmpty()) { + logger.trace { + val affectedList = pollEvents.joinAsList { "${it.kind()}: ${it.context()}" } + "Affected files:\n$affectedList" + } + } else { + // Seems to be empty when a directory gets deleted + // The next watch key *should* be an ENTRY_DELETE of that directory + continue + } + if (!key.reset()) { + logger.warn { "${key.watchable()} is no longer valid" } + continue + } + + scheduleRestart() + } + } + } + + private fun scheduleRestart() { + val settings = settings ?: return // Don't schedule a restart until an instance has registered + if (::restartFuture.isInitialized) restartFuture.cancel(false) + restartFuture = scheduler.schedule(::tryRestart, settings.restartDelay.inWholeMilliseconds, TimeUnit.MILLISECONDS) + } + + private fun tryRestart() { + // Can't set to null after restarting, + // as the restart function only returns after the main method ran + val settings = settings ?: return + try { + logger.debug { "Attempting to restart" } + this.settings = null // Wait until the next instance has given its settings + compareSnapshots() + snapshots.keys.forEach { registerDirectories(it) } + Restarter.instance.restart() + } catch (e: Exception) { + logger.error(e) { "Restart failed, waiting for the next build" } + this.settings = settings // Reuse the old settings to reschedule a new restart + } + } + + private fun compareSnapshots() { + val hasChanges = snapshots.any { (directory, files) -> + val snapshot = directory.takeSnapshot() + + // Exclude deleted files so they don't count as being deleted again + val deletedPaths = files.withoutDeletes().keys - snapshot.keys + if (deletedPaths.isNotEmpty()) { + logger.info { "${deletedPaths.size} files were deleted in ${directory.absolutePathString()}: $deletedPaths" } + snapshots[directory] = deletedPaths.associateWith { DeletedSourceFile } + snapshot + // So we can re-register them in case they are recreated + registeredDirectories.removeAll(deletedPaths.map { directory.resolve(it) }) + return@any true + } + + // Exclude deleted files so they count as being added back + val addedPaths = snapshot.keys - files.withoutDeletes().keys + if (addedPaths.isNotEmpty()) { + logger.info { "${addedPaths.size} files were added in ${directory.absolutePathString()}: $addedPaths" } + snapshots[directory] = files + snapshot + return@any true + } + + val modifiedFiles = snapshot.keys.filter { key -> + val actual = snapshot[key] ?: error("Key from map is missing a value somehow") + val expected = files[key] ?: error("Expected file is missing, should have been detected as deleted") + + // File was deleted (on the 2nd build for example) and got recreated (on the 3rd build for example) + if (expected is DeletedSourceFile) error("Expected file was registered as deleted, should have been detected as added") + expected as SourceFile + + actual as SourceFile // Assertion + + actual.lastModified != expected.lastModified + } + if (modifiedFiles.isNotEmpty()) { + logger.info { "${modifiedFiles.size} files were modified in ${directory.absolutePathString()}: $modifiedFiles" } + snapshots[directory] = files + snapshot + return@any true + } + + false + } + + if (!hasChanges) + error("Received a file system event but no changes were detected") + } + + private fun registerDirectories(directory: Path) { + directory.walkDirectories { path, attributes -> + if (registeredDirectories.add(path)) + path.register(watchService, ENTRY_CREATE, ENTRY_MODIFY, ENTRY_DELETE) + } + } + + private class Settings( + val restartDelay: Duration, + ) + + internal companion object { + private val instanceLock = ReentrantLock() + internal lateinit var instance: ClasspathWatcher + private set + + internal fun initialize(restartDelay: Duration) { + instanceLock.withLock { + val settings = Settings(restartDelay) + if (::instance.isInitialized.not()) { + instance = ClasspathWatcher(settings) + } else { + instance.settings = settings + } + } + } + } +} + +private fun Path.takeSnapshot(): SourceFiles = walkFiles().associate { (it, attrs) -> + it.relativeTo(this).pathString to SourceFile(attrs.lastModifiedTime().toInstant()) +}.let(::SourceFiles) \ No newline at end of file diff --git a/settings.gradle.kts b/settings.gradle.kts index 632066d..f611fa7 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -1 +1,6 @@ -rootProject.name = "BotCommands-Restarter" \ No newline at end of file +rootProject.name = "BotCommands-Restarter" + +enableFeaturePreview("TYPESAFE_PROJECT_ACCESSORS") + +include("restarter") +include("restarter-jda-cache") \ No newline at end of file