Cleanup plugin loading code (#1009)

* cleanup code * apply feedback * don't iterate jars twice * Update LavalinkServer/src/main/java/lavalink/server/bootstrap/PluginManager.kt Co-authored-by: Toπ <[email protected]> * remove stray `")` --------- Co-authored-by: Toπ <[email protected]>
lavalink-devs · Feb 6, 2024 · b9eae8a · b9eae8a
1 parent 59b7ac8
commit b9eae8a
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 82 deletions.
diff --git a/.gitignore b/.gitignore
@@ -21,4 +21,5 @@ gradle.properties
 application.yml
 LavalinkServer/plugins
 .cache/
-site/
+site/
+.DS_Store
diff --git a/LavalinkServer/src/main/java/lavalink/server/Launcher.kt b/LavalinkServer/src/main/java/lavalink/server/Launcher.kt
@@ -143,7 +143,7 @@ object Launcher {
             .properties(properties)
             .web(WebApplicationType.SERVLET)
             .bannerMode(Banner.Mode.OFF)
-            .resourceLoader(DefaultResourceLoader(pluginManager.classLoader))
+            .resourceLoader(DefaultResourceLoader(pluginManager::class.java.classLoader))
             .listeners(
                 ApplicationListener { event: Any ->
                     when (event) {

diff --git a/LavalinkServer/src/main/java/lavalink/server/bootstrap/PluginManager.kt b/LavalinkServer/src/main/java/lavalink/server/bootstrap/PluginManager.kt
@@ -15,13 +15,11 @@ import java.util.jar.JarFile
 
 @SpringBootApplication
 class PluginManager(val config: PluginsConfig) {
-
     companion object {
         private val log: Logger = LoggerFactory.getLogger(PluginManager::class.java)
     }
 
     final val pluginManifests: MutableList<PluginManifest> = mutableListOf()
-    var classLoader: ClassLoader = PluginManager::class.java.classLoader
 
     init {
         manageDownloads()
@@ -33,57 +31,59 @@ class PluginManager(val config: PluginsConfig) {
 
     private fun manageDownloads() {
         if (config.plugins.isEmpty()) return
+
         val directory = File(config.pluginsDir)
         directory.mkdir()
 
-        data class PluginJar(val manifest: PluginManifest, val file: File)
-
-        val pluginJars = directory.listFiles()!!.filter { it.extension == "jar" }.map {
-            JarFile(it).use { jar ->
-                loadPluginManifests(jar).map { manifest ->
-                    log.info("Found plugin '${manifest.name}' version ${manifest.version}")
-                    PluginJar(manifest, it)
+        val pluginJars = directory.listFiles()?.filter { it.extension == "jar" }
+            ?.flatMap { file ->
+                JarFile(file).use { jar ->
+                    loadPluginManifests(jar).map { manifest -> PluginJar(manifest, file) }
                 }
             }
-        }.flatten()
-
-        data class Declaration(val group: String, val name: String, val version: String, val repository: String)
+            ?.onEach { log.info("Found plugin '${it.manifest.name}' version ${it.manifest.version}") }
+            ?: return
 
         val declarations = config.plugins.map { declaration ->
             if (declaration.dependency == null) throw RuntimeException("Illegal dependency declaration: null")
             val fragments = declaration.dependency!!.split(":")
             if (fragments.size != 3) throw RuntimeException("Invalid dependency \"${declaration.dependency}\"")
 
-            var repository = declaration.repository
-                ?: if (declaration.snapshot) config.defaultPluginSnapshotRepository else config.defaultPluginRepository
-            repository = if (repository.endsWith("/")) repository else "$repository/"
-            Declaration(fragments[0], fragments[1], fragments[2], repository)
+            val repository = declaration.repository
+                ?: config.defaultPluginSnapshotRepository.takeIf { declaration.snapshot }
+                ?: config.defaultPluginRepository
+
+            Declaration(fragments[0], fragments[1], fragments[2], "${repository.removeSuffix("/")}/")
         }.distinctBy { "${it.group}:${it.name}" }
 
-        declarations.forEach declarationLoop@{ declaration ->
-            var hasVersion = false
-            pluginJars.filter { jar -> declaration.name == jar.manifest.name }
-                .forEach pluginLoop@{ jar ->
-                    if (declaration.version == jar.manifest.version && !hasVersion) {
-                        hasVersion = true
-                        // We already have this jar so don't redownload it
-                        return@pluginLoop
-                    }
-
-                    // Delete jar of different versions
-                    if (!jar.file.delete()) throw RuntimeException("Failed to delete ${jar.file.path}")
-                    log.info("Deleted ${jar.file.path}")
+        for (declaration in declarations) {
+            val jars = pluginJars.filter { it.manifest.name == declaration.name }
+            var hasCurrentVersion = false
+
+            for (jar in jars) {
+                if (jar.manifest.version == declaration.version) {
+                    hasCurrentVersion = true
+                    // Don't clean up the jar if it's a current version.
+                    continue
                 }
-            if (hasVersion) return@declarationLoop
 
-            val url = declaration.run { "$repository${group.replace(".", "/")}/$name/$version/$name-$version.jar" }
-            val file = File(directory, declaration.run { "$name-$version.jar" })
-            downloadJar(file, url)
+                // Delete versions of the plugin that aren't the same as declared version.
+                if (!jar.file.delete()) throw RuntimeException("Failed to delete ${jar.file.path}")
+                log.info("Deleted ${jar.file.path} (new version: ${declaration.version})")
+
+            }
+
+            if (!hasCurrentVersion) {
+                val url = declaration.url
+                val file = File(directory, declaration.canonicalJarName)
+                downloadJar(file, url)
+            }
         }
     }
 
     private fun downloadJar(output: File, url: String) {
         log.info("Downloading $url")
+
         Channels.newChannel(URL(url).openStream()).use {
             FileOutputStream(output).channel.transferFrom(it, 0, Long.MAX_VALUE)
         }
@@ -92,61 +92,43 @@ class PluginManager(val config: PluginsConfig) {
     private fun readClasspathManifests(): List<PluginManifest> {
         return PathMatchingResourcePatternResolver()
             .getResources("classpath*:lavalink-plugins/*.properties")
-            .map map@{ r ->
-                val manifest = parsePluginManifest(r.inputStream)
-                log.info("Found plugin '${manifest.name}' version ${manifest.version}")
-                return@map manifest
-            }
+            .map { parsePluginManifest(it.inputStream) }
+            .onEach { log.info("Found plugin '${it.name}' version ${it.version}") }
     }
 
     private fun loadJars(): List<PluginManifest> {
-        val directory = File(config.pluginsDir)
-        if (!directory.isDirectory) return emptyList()
-        val jarsToLoad = mutableListOf<File>()
-
-        directory.listFiles()?.forEach { file ->
-            if (!file.isFile) return@forEach
-            if (file.extension != "jar") return@forEach
-            jarsToLoad.add(file)
-        }
+        val directory = File(config.pluginsDir).takeIf { it.isDirectory }
+            ?: return emptyList()
 
-        if (jarsToLoad.isEmpty()) return emptyList()
+        val jarsToLoad = directory.listFiles()?.filter { it.isFile && it.extension == "jar" }
+            ?.takeIf { it.isNotEmpty() }
+            ?: return emptyList()
 
-        val cl = URLClassLoader.newInstance(
+        val classLoader = URLClassLoader.newInstance(
             jarsToLoad.map { URL("jar:file:${it.absolutePath}!/") }.toTypedArray(),
             javaClass.classLoader
         )
-        classLoader = cl
-
-        val manifests = mutableListOf<PluginManifest>()
-        jarsToLoad.forEach { file ->
-            try {
-                manifests.addAll(loadJar(file, cl))
-            } catch (e: Exception) {
-                throw RuntimeException("Error loading $file", e)
-            }
-        }
 
-
-        return manifests
+        return jarsToLoad.flatMap { loadJar(it, classLoader) }
     }
 
     private fun loadJar(file: File, cl: URLClassLoader): List<PluginManifest> {
-        var classCount = 0
         val jar = JarFile(file)
-        var manifests: List<PluginManifest>
+        val manifests = loadPluginManifests(jar)
+        var classCount = 0
 
         jar.use {
-            manifests = loadPluginManifests(jar)
             if (manifests.isEmpty()) {
                 throw RuntimeException("No plugin manifest found in ${file.path}")
             }
-            val allowedPaths = manifests.map { it.path.replace(".", "/") }
 
-            jar.entries().asIterator().forEach { entry ->
-                if (entry.isDirectory) return@forEach
-                if (!entry.name.endsWith(".class")) return@forEach
-                if (!allowedPaths.any { entry.name.startsWith(it) }) return@forEach
+            val allowedPaths = manifests.map { manifest -> manifest.path.replace(".", "/") }
+
+            for (entry in it.entries()) {
+                if (entry.isDirectory ||
+                    !entry.name.endsWith(".class") ||
+                    allowedPaths.none(entry.name::startsWith)) continue
+
                 cl.loadClass(entry.name.dropLast(6).replace("/", "."))
                 classCount++
             }
@@ -157,16 +139,10 @@ class PluginManager(val config: PluginsConfig) {
     }
 
     private fun loadPluginManifests(jar: JarFile): List<PluginManifest> {
-        val manifests = mutableListOf<PluginManifest>()
-
-        jar.entries().asIterator().forEach { entry ->
-            if (entry.isDirectory) return@forEach
-            if (!entry.name.startsWith("lavalink-plugins/")) return@forEach
-            if (!entry.name.endsWith(".properties")) return@forEach
-
-            manifests.add(parsePluginManifest(jar.getInputStream(entry)))
-        }
-        return manifests
+        return jar.entries().asSequence()
+            .filter { !it.isDirectory && it.name.startsWith("lavalink-plugins/") && it.name.endsWith(".properties") }
+            .map { parsePluginManifest(jar.getInputStream(it)) }
+            .toList()
     }
 
     private fun parsePluginManifest(stream: InputStream): PluginManifest {
@@ -179,4 +155,10 @@ class PluginManager(val config: PluginsConfig) {
         val version = props.getProperty("version") ?: throw RuntimeException("Manifest is missing 'version'")
         return PluginManifest(name, path, version)
     }
-}
+
+    private data class PluginJar(val manifest: PluginManifest, val file: File)
+    private data class Declaration(val group: String, val name: String, val version: String, val repository: String) {
+        val canonicalJarName = "$name-$version.jar"
+        val url = "$repository${group.replace(".", "/")}/$name/$version/$name-$version.jar"
+    }
+}