Add structures to support hot restarting

BotCommands-jda-keepalive/build.gradle.kts

@@ -0,0 +1,74 @@
+import org.jetbrains.kotlin.gradle.dsl.JvmTarget
+
+plugins {
+    id("BotCommands-conventions")
+    id("BotCommands-publish-conventions")
+}
+
+dependencies {
+    api(projects.botCommands)
+
+    // Logging
+    implementation(libs.kotlin.logging)
+
+    // -------------------- TEST DEPENDENCIES --------------------
+
+    testImplementation(projects.botCommandsRestarter)
+    testImplementation(libs.mockk)
+    testImplementation(libs.bytebuddy)
+    testImplementation(libs.logback.classic)
+}
+
+fun registerSourceSet(name: String, extendsTestDependencies: Boolean) {
+    sourceSets {
+        register(name) {
+            compileClasspath += sourceSets.main.get().output
+            runtimeClasspath += sourceSets.main.get().output
+        }
+    }
+
+    configurations["${name}Api"].extendsFrom(configurations["api"])
+    configurations["${name}Implementation"].extendsFrom(configurations["implementation"])
+    configurations["${name}CompileOnly"].extendsFrom(configurations["compileOnly"])
+
+    if (extendsTestDependencies) {
+        configurations["${name}Api"].extendsFrom(configurations["testApi"])
+        configurations["${name}Implementation"].extendsFrom(configurations["testImplementation"])
+        configurations["${name}CompileOnly"].extendsFrom(configurations["testCompileOnly"])
+    }
+}
+
+// Register other source sets
+registerSourceSet(name = "testBot", extendsTestDependencies = true)
+
+java {
+    sourceCompatibility = JavaVersion.VERSION_24
+    targetCompatibility = JavaVersion.VERSION_24
+}
+
+kotlin {
+    compilerOptions {
+        jvmTarget = JvmTarget.JVM_24
+
+        freeCompilerArgs.addAll(
+            "-Xcontext-parameters",
+        )
+    }
+}
+
+val jar by tasks.getting(Jar::class) {
+    manifest {
+        attributes(
+            "Premain-Class" to "dev.freya02.botcommands.jda.keepalive.internal.Agent",
+            "Agent-Class" to "dev.freya02.botcommands.jda.keepalive.internal.Agent",
+        )
+    }
+}
+
+tasks.withType<Test> {
+    useJUnitPlatform()
+
+    // Don't use "-javaagent" because [[AgentTest]] requires loading classes before the agent transforms them
+    jvmArgs("-Djdk.attach.allowAttachSelf=true")
+    jvmArgs("-Dbc.jda.keepalive.agentPath=${jar.archiveFile.get().asFile.absolutePath}")
+}

BotCommands-jda-keepalive/src/main/kotlin/dev/freya02/botcommands/jda/keepalive/api/JDAKeepAlive.kt

@@ -0,0 +1,11 @@
+package dev.freya02.botcommands.jda.keepalive.api
+
+import dev.freya02.botcommands.jda.keepalive.internal.Agent
+
+object JDAKeepAlive {
+
+    @JvmStatic
+    fun install() {
+        Agent.load()
+    }
+}

BotCommands-jda-keepalive/src/main/kotlin/dev/freya02/botcommands/jda/keepalive/internal/Agent.kt

@@ -0,0 +1,99 @@
+package dev.freya02.botcommands.jda.keepalive.internal
+
+import com.sun.tools.attach.VirtualMachine
+import dev.freya02.botcommands.jda.keepalive.internal.exceptions.AlreadyLoadedClassesException
+import dev.freya02.botcommands.jda.keepalive.internal.exceptions.AttachSelfDeniedException
+import dev.freya02.botcommands.jda.keepalive.internal.exceptions.IllegalAgentContainerException
+import dev.freya02.botcommands.jda.keepalive.internal.transformer.*
+import io.github.freya022.botcommands.api.core.utils.joinAsList
+import java.lang.instrument.Instrumentation
+import java.lang.management.ManagementFactory
+import java.nio.file.Path
+import java.util.concurrent.locks.ReentrantLock
+import kotlin.concurrent.withLock
+import kotlin.io.path.Path
+import kotlin.io.path.absolutePathString
+import kotlin.io.path.extension
+import kotlin.io.path.toPath
+
+internal object Agent {
+
+    internal val transformers = mapOf(
+        CD_JDABuilder to JDABuilderTransformer,
+        CD_JDAImpl to JDAImplTransformer,
+        CD_JDAService to JDAServiceTransformer,
+        CD_BContextImpl to BContextImplTransformer,
+    )
+
+    private val lock = ReentrantLock()
+    internal var isLoaded = false
+        private set
+
+    @JvmStatic
+    fun premain(agentArgs: String?, inst: Instrumentation) {
+        lock.withLock {
+            isLoaded = true
+        }
+
+        transformers.values.forEach(inst::addTransformer)
+    }
+
+    @JvmStatic
+    fun agentmain(agentArgs: String?, inst: Instrumentation) {
+        lock.withLock {
+            if (isLoaded) return
+            isLoaded = true
+        }
+
+        checkNoLoadedClassesAreToBeTransformed(inst.allLoadedClasses)
+
+        transformers.values.forEach(inst::addTransformer)
+    }
+
+    internal fun checkNoLoadedClassesAreToBeTransformed(allLoadedClasses: Array<Class<*>>) {
+        val transformedClasses = transformers.keys.mapTo(hashSetOf()) { it.packageName() + "." + it.displayName() }
+        val earlyLoadedClasses = allLoadedClasses.filter { it.name in transformedClasses }
+        if (earlyLoadedClasses.isNotEmpty()) {
+            // TODO wiki link (of the whole loading mechanism probably)
+            throw AlreadyLoadedClassesException(
+                "Dynamically loaded agents must be loaded before the classes it transforms, although it is recommended to add the agent via the command line, offending classes:\n" +
+                        earlyLoadedClasses.joinAsList { it.name }
+            )
+        }
+    }
+
+    internal fun load() {
+        lock.withLock {
+            if (isLoaded) return
+        }
+
+        // Check self-attaching agents are allowed
+        if (System.getProperty("jdk.attach.allowAttachSelf") != "true") {
+            // TODO wiki link (show how to do in IJ)
+            throw AttachSelfDeniedException("Can only dynamically load an agent with the '-Djdk.attach.allowAttachSelf=true' VM argument")
+        }
+
+        // Get the agent JAR
+        // In a user's dev environment, this should be the dependency's JAR
+        // but in *our* test environment, we need a property to the JAR produced by Gradle,
+        // as it still uses the directory in the classpath
+        val agentSourcePath: Path = run {
+            // Property to test instrumentation is applied
+            System.getProperty("bc.jda.keepalive.agentPath")?.let { return@run Path(it) }
+
+            javaClass.protectionDomain.codeSource.location.toURI().toPath()
+        }
+        if (agentSourcePath.extension != "jar") {
+            // TODO add wiki link about including the dependency only in dev envs
+            throw IllegalAgentContainerException(
+                "Can only dynamically load an agent using a JAR, this agent should only be loaded in development, please see the wiki\n" +
+                        "Agent source @ ${agentSourcePath.absolutePathString()}"
+            )
+        }
+
+        // Load agent on ourselves
+        val jvm = VirtualMachine.attach(ManagementFactory.getRuntimeMXBean().pid.toString())
+        jvm.loadAgent(agentSourcePath.absolutePathString())
+        jvm.detach()
+    }
+}

BotCommands-jda-keepalive/src/main/kotlin/dev/freya02/botcommands/jda/keepalive/internal/BufferingEventManager.kt

@@ -0,0 +1,67 @@
+package dev.freya02.botcommands.jda.keepalive.internal
+
+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<GenericEvent> = 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<Any?> {
+        lock.withLock {
+            val delegate = delegate ?: error("Should not happen, implement a listener queue if necessary")
+            return delegate.registeredListeners
+        }
+    }
+}

BotCommands-jda-keepalive/src/main/kotlin/dev/freya02/botcommands/jda/keepalive/internal/DynamicCall.kt

@@ -0,0 +1,10 @@
+package dev.freya02.botcommands.jda.keepalive.internal
+
+/**
+ * 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

BotCommands-jda-keepalive/src/main/kotlin/dev/freya02/botcommands/jda/keepalive/internal/JDABuilderConfiguration.kt

@@ -0,0 +1,130 @@
+package dev.freya02.botcommands.jda.keepalive.internal
+
+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<String> = hashSetOf()
+
+    var hasUnsupportedValues = false
+        private set
+
+    private val builderValues: MutableMap<ValueType, Any?> = 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<CacheFlag>()
+    }
+
+    @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<CacheFlag>) += enumSetOf(first, *others)
+    }
+
+    @DynamicCall
+    @Suppress("UNCHECKED_CAST")
+    fun enableCache(flags: Collection<CacheFlag>) {
+        (builderValues[ValueType.CACHE_FLAGS] as EnumSet<CacheFlag>) += flags
+    }
+
+    @DynamicCall
+    @Suppress("UNCHECKED_CAST")
+    fun disableCache(first: CacheFlag, vararg others: CacheFlag) {
+        (builderValues[ValueType.CACHE_FLAGS] as EnumSet<CacheFlag>) -= enumSetOf(first, *others)
+    }
+
+    @DynamicCall
+    @Suppress("UNCHECKED_CAST")
+    fun disableCache(flags: Collection<CacheFlag>) {
+        (builderValues[ValueType.CACHE_FLAGS] as EnumSet<CacheFlag>) -= 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,
+    }
+}