Add cog#unload, pass ObjectStorage to classes for persistent storage (needs testing)

Author: devoxin

src/main/kotlin/me/devoxin/flight/api/entities/Cog.kt

@@ -26,4 +26,11 @@ interface Cog {
      */
     fun localCheck(ctx: Context, command: CommandFunction): Boolean = true
 
+
+    /**
+     * Invoked when this Cog gets unloaded, usually through [CommandRegistry.unload].
+     * This can be used as a last-ditch attempt to clean up, or shut down any resources.
+     */
+    fun unload(): Unit = Unit
+
 }

src/main/kotlin/me/devoxin/flight/api/entities/CommandRegistry.kt

@@ -9,8 +9,11 @@ import net.dv8tion.jda.api.interactions.commands.build.CommandData
 import net.dv8tion.jda.api.interactions.commands.build.Commands
 import net.dv8tion.jda.api.interactions.commands.build.OptionData
 import net.dv8tion.jda.api.interactions.commands.build.SubcommandData
+import org.slf4j.LoggerFactory
 
 class CommandRegistry : HashMap<String, CommandFunction>() {
+    val objectStorage = ObjectStorage()
+
     fun toDiscordCommands(): List<CommandData> {
         val commands = mutableListOf<CommandData>()
 
@@ -77,7 +80,7 @@ class CommandRegistry : HashMap<String, CommandFunction>() {
     }
 
     fun findCogByName(name: String): Cog? {
-        return this.values.firstOrNull { it.cog::class.simpleName == name }?.cog
+        return this.values.firstOrNull { it.cog.name() == name || it.cog::class.simpleName == name }?.cog
     }
 
     fun findCommandsByCog(cog: Cog): List<CommandFunction> {
@@ -86,12 +89,15 @@ class CommandRegistry : HashMap<String, CommandFunction>() {
 
     fun unload(commandFunction: CommandFunction) {
         this.values.remove(commandFunction)
+        doUnload(commandFunction.cog)
     }
 
     fun unload(cog: Cog) {
         val commands = this.values.filter { it.cog == cog }
         this.values.removeAll(commands)
 
+        commands.map(CommandFunction::cog).let(::doUnload)
+
         val jar = commands.firstOrNull { it.jar != null }?.jar
             ?: return // No commands loaded from jar, thus no classloader to close.
 
@@ -107,21 +113,23 @@ class CommandRegistry : HashMap<String, CommandFunction>() {
         val commands = this.values.filter { it.jar == jar }
         this.values.removeAll(commands)
 
+        commands.map(CommandFunction::cog).let(::doUnload)
+
         jar.close()
     }
 
     fun register(packageName: String) {
         val indexer = Indexer(packageName)
 
-        for (cog in indexer.getCogs()) {
+        for (cog in indexer.getCogs(objectStorage)) {
             register(cog, indexer)
         }
     }
 
     fun register(jarPath: String, packageName: String) {
         val indexer = Indexer(packageName, jarPath)
 
-        for (cog in indexer.getCogs()) {
+        for (cog in indexer.getCogs(objectStorage)) {
             register(cog, indexer)
         }
     }
@@ -140,4 +148,24 @@ class CommandRegistry : HashMap<String, CommandFunction>() {
             this[cmd.name] = cmd
         }
     }
+
+    private fun doUnload(cogs: Iterable<Cog>) {
+        val uniqueCogs = cogs.distinctBy(Cog::name)
+
+        for (cog in uniqueCogs) {
+            doUnload(cog)
+        }
+    }
+
+    private fun doUnload(cog: Cog) {
+        try {
+            cog.unload()
+        } catch (t: Throwable) {
+            log.error("An error occurred whilst unloading cog \"{}\"", cog.name() ?: cog::class.java.simpleName, t)
+        }
+    }
+
+    companion object {
+        private val log = LoggerFactory.getLogger(CommandRegistry::class.java)
+    }
 }

src/main/kotlin/me/devoxin/flight/api/entities/ObjectStorage.kt

@@ -0,0 +1,37 @@
+package me.devoxin.flight.api.entities
+
+import java.util.concurrent.ConcurrentHashMap
+
+class ObjectStorage {
+    private val map = ConcurrentHashMap<String, Any>()
+
+    val size: Int
+        get() = map.size
+
+    operator fun set(key: String, value: Any) {
+        map[key] = value
+    }
+
+    operator fun get(key: String): Any? {
+        return map[key]
+    }
+
+    fun <T> get(key: String, cls: Class<T>): T? {
+        val value = map[key]
+
+        if (cls.isInstance(value)) {
+            return cls.cast(value)
+        }
+
+        return null
+    }
+
+    @Suppress("UNCHECKED_CAST")
+    fun <T : Any> computeIfAbsent(key: String, initializer: (String) -> T): T {
+        return map.computeIfAbsent(key, initializer) as T
+    }
+
+    fun clear() {
+        map.clear()
+    }
+}

src/main/kotlin/me/devoxin/flight/internal/utils/Indexer.kt

@@ -7,11 +7,13 @@ import me.devoxin.flight.api.annotations.*
 import me.devoxin.flight.internal.arguments.Argument
 import me.devoxin.flight.internal.entities.Jar
 import me.devoxin.flight.api.entities.Cog
+import me.devoxin.flight.api.entities.ObjectStorage
 import org.reflections.Reflections
 import org.reflections.scanners.MethodParameterNamesScanner
 import org.reflections.scanners.Scanners
 import org.slf4j.LoggerFactory
 import java.io.File
+import java.lang.reflect.Constructor
 import java.lang.reflect.Modifier
 import java.net.URL
 import java.net.URLClassLoader
@@ -47,13 +49,13 @@ class Indexer {
         reflections = Reflections(packageName, this.classLoader, MethodParameterNamesScanner(), Scanners.SubTypes)
     }
 
-    fun getCogs(): List<Cog> {
+    fun getCogs(objectStorage: ObjectStorage): List<Cog> {
         val cogs = reflections.getSubTypesOf(Cog::class.java)
         log.debug("Discovered ${cogs.size} cogs in $packageName")
 
         return cogs
             .filter { !Modifier.isAbstract(it.modifiers) && !it.isInterface && Cog::class.java.isAssignableFrom(it) }
-            .map { it.getDeclaredConstructor().newInstance() }
+            .map { construct(it, objectStorage) }
     }
 
     fun getCommands(cog: Cog): List<KFunction<*>> {
@@ -147,6 +149,14 @@ class Indexer {
         return arguments
     }
 
+    private fun construct(cls: Class<out Cog>, objectStorage: ObjectStorage): Cog {
+        return try {
+            cls.getDeclaredConstructor(ObjectStorage::class.java).newInstance(objectStorage)
+        } catch (t: NoSuchMethodException) {
+            cls.getDeclaredConstructor().newInstance()
+        }
+    }
+
     companion object {
         private val log = LoggerFactory.getLogger(Indexer::class.java)
     }