Merge Log4j2PluginsCacheFileTransformer

Author: Goooler

src/main/java/Log4j2PluginCacheFileTransformer.kt

@@ -4,12 +4,16 @@ import com.github.jengelman.gradle.plugins.shadow.ShadowStats
 import com.github.jengelman.gradle.plugins.shadow.relocation.RelocateClassContext
 import com.github.jengelman.gradle.plugins.shadow.relocation.Relocator
 import com.github.jengelman.gradle.plugins.shadow.transformers.TransformerContext.Companion.getEntryTimestamp
-import java.io.File
 import java.net.URL
+import java.nio.file.Path
 import java.util.Collections
 import java.util.Enumeration
+import kotlin.io.path.createTempFile
+import kotlin.io.path.deleteIfExists
+import kotlin.io.path.outputStream
 import org.apache.commons.io.output.CloseShieldOutputStream
 import org.apache.logging.log4j.core.config.plugins.processor.PluginCache
+import org.apache.logging.log4j.core.config.plugins.processor.PluginEntry
 import org.apache.logging.log4j.core.config.plugins.processor.PluginProcessor
 import org.apache.tools.zip.ZipEntry
 import org.apache.tools.zip.ZipOutputStream
@@ -25,64 +29,101 @@ import org.gradle.api.file.FileTreeElement
  */
 @CacheableTransformer
 public open class Log4j2PluginsCacheFileTransformer : Transformer {
-  private val temporaryFiles = mutableListOf<File>()
-  private val relocators = mutableListOf<Relocator>()
+  /**
+   * Log4j config files to share across the transformation stages.
+   */
+  private val tempFiles = mutableListOf<Path>()
+
+  /**
+   * [Relocator] instances to share across the transformation stages.
+   */
+  private val tempRelocators = mutableListOf<Relocator>()
   private var stats: ShadowStats? = null
 
   override fun canTransformResource(element: FileTreeElement): Boolean {
     return PluginProcessor.PLUGIN_CACHE_FILE == element.name
   }
 
   override fun transform(context: TransformerContext) {
-    val temporaryFile = File.createTempFile("Log4j2Plugins", ".dat")
-    temporaryFile.deleteOnExit()
-    temporaryFiles.add(temporaryFile)
+    val temporaryFile = createTempFile("Log4j2Plugins", ".dat")
+    tempFiles.add(temporaryFile)
     val fos = temporaryFile.outputStream()
     context.inputStream.use {
       it.copyTo(fos)
     }
 
-    relocators.addAll(context.relocators)
+    tempRelocators.addAll(context.relocators)
 
     if (stats == null) {
       stats = context.stats
     }
   }
 
+  /**
+   * @return true if any dat file collected
+   */
   override fun hasTransformedResource(): Boolean {
-    // This functionality matches the original plugin, however, I'm not clear what
-    // the exact logic is. From what I can tell temporaryFiles should be never be empty
-    // if anything has been performed.
-    return temporaryFiles.isNotEmpty() || relocators.isNotEmpty()
+    return tempFiles.isNotEmpty()
   }
 
   override fun modifyOutputStream(os: ZipOutputStream, preserveFileTimestamps: Boolean) {
-    val pluginCache = PluginCache()
-    pluginCache.loadCacheFiles(urlEnumeration)
-    relocatePlugins(pluginCache)
-    val entry = ZipEntry(PluginProcessor.PLUGIN_CACHE_FILE)
-    entry.time = getEntryTimestamp(preserveFileTimestamps, entry.time)
-    os.putNextEntry(entry)
-    pluginCache.writeCache(CloseShieldOutputStream.wrap(os))
-    temporaryFiles.clear()
+    try {
+      val aggregator = PluginCache()
+      aggregator.loadCacheFiles(urlEnumeration)
+      relocatePlugin(tempRelocators, aggregator.allCategories)
+      val entry = ZipEntry(PluginProcessor.PLUGIN_CACHE_FILE)
+      entry.time = getEntryTimestamp(preserveFileTimestamps, entry.time)
+      os.putNextEntry(entry)
+      aggregator.writeCache(CloseShieldOutputStream.wrap(os))
+    } finally {
+      deleteTempFiles()
+    }
   }
 
-  private fun relocatePlugins(pluginCache: PluginCache) {
-    pluginCache.allCategories.values.forEach { currentMap ->
-      currentMap.values.forEach { currentPluginEntry ->
-        val className = currentPluginEntry.className
-        val relocateClassContext = RelocateClassContext(className, requireNotNull(stats))
-        relocators.firstOrNull { it.canRelocateClass(className) }?.let { relocator ->
-          // Then we perform that relocation and update the plugin entry to reflect the new value.
-          currentPluginEntry.className = relocator.relocateClass(relocateClassContext)
+  /**
+   * Applies the given `relocators` to the `aggregator`.
+   *
+   * @param relocators           relocators.
+   * @param aggregatorCategories all categories of the aggregator
+   */
+  private fun relocatePlugin(
+    relocators: List<Relocator>,
+    aggregatorCategories: Map<String, Map<String, PluginEntry>>,
+  ) {
+    for (categoryEntry in aggregatorCategories.entries) {
+      for (pluginMapEntry in categoryEntry.value.entries) {
+        val pluginEntry = pluginMapEntry.value
+        val originalClassName = pluginEntry.className
+        val relocateClassContext = RelocateClassContext(originalClassName, requireNotNull(stats))
+
+        findFirstMatchingRelocator(originalClassName, relocators)?.let {
+          pluginEntry.className = it.relocateClass(relocateClassContext)
         }
       }
     }
   }
 
+  private fun findFirstMatchingRelocator(originalClassName: String, relocators: List<Relocator>): Relocator? {
+    for (relocator in relocators) {
+      if (relocator.canRelocateClass(originalClassName)) {
+        return relocator
+      }
+    }
+    return null
+  }
+
+  private fun deleteTempFiles() {
+    val pathIterator = tempFiles.listIterator()
+    while (pathIterator.hasNext()) {
+      val path = pathIterator.next()
+      path.deleteIfExists()
+      pathIterator.remove()
+    }
+  }
+
   private val urlEnumeration: Enumeration<URL>
     get() {
-      val urls = temporaryFiles.map { it.toURI().toURL() }
+      val urls = tempFiles.map { it.toUri().toURL() }
       return Collections.enumeration(urls)
     }
 }