Parallelize source remapping for shadowed sources JAR

src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/OrderedParallelWriter.kt

@@ -0,0 +1,42 @@
+package com.github.jengelman.gradle.plugins.shadow.internal
+
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.async
+import kotlinx.coroutines.runBlocking
+
+/**
+ * Computes entry bytes in parallel on [Dispatchers.Default] and writes them in submission order,
+ * which keeps the output deterministic.
+ *
+ * Entries written directly bypass the queue, so [flush] must be called before writing them.
+ */
+internal class OrderedParallelWriter(private val scope: CoroutineScope) {
+  private val pendingEntries = ArrayDeque<PendingEntry>()
+
+  fun submit(compute: () -> ByteArray, write: (ByteArray) -> Unit) {
+    pendingEntries.addLast(PendingEntry(scope.async(Dispatchers.Default) { compute() }, write))
+    if (pendingEntries.size >= MAX_PENDING_ENTRIES) {
+      pendingEntries.removeFirst().writeBytes()
+    }
+  }
+
+  fun flush() {
+    while (pendingEntries.isNotEmpty()) {
+      pendingEntries.removeFirst().writeBytes()
+    }
+  }
+
+  private fun PendingEntry.writeBytes() {
+    write(runBlocking { deferredBytes.await() })
+  }
+}
+
+/**
+ * Maximum number of in-flight parallel tasks in the sliding window. Bounds memory consumption on
+ * large archives while keeping worker threads saturated.
+ */
+private const val MAX_PENDING_ENTRIES = 128
+
+private class PendingEntry(val deferredBytes: Deferred<ByteArray>, val write: (ByteArray) -> Unit)

src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJar.kt

@@ -4,6 +4,7 @@ import com.github.jengelman.gradle.plugins.shadow.relocation.Relocator
 import com.github.jengelman.gradle.plugins.shadow.relocation.relocatePath
 import java.io.File
 import java.io.InputStream
+import kotlinx.coroutines.runBlocking
 import org.apache.tools.zip.ZipOutputStream
 import org.gradle.api.file.FileCollection
 import org.vafer.jdeb.shaded.objectweb.asm.ClassReader
@@ -50,56 +51,72 @@ internal fun generateSourcesJar(
         }
       }
 
-      // Relocates and writes a project or dependency source file, or a resource.
-      fun addEntry(name: String, lastModified: Long, readBytes: () -> ByteArray) {
-        if (name.isModuleInfo) return
-        val isSource = name.isSourceFile()
-        if (isSource && isUnused(name, unusedClasses, sourceToClasses)) return
-        val relocatedPath =
-          if (isSource) relocators.relocateSourcePath(name) else relocators.relocatePath(name)
-        if (!visitedFiles.add(relocatedPath)) return
-        val bytes =
-          if (isSource) {
-            relocators
-              .remapSourceFile(readBytes().decodeToString(), name, relocatedPath)
-              .toByteArray()
-          } else {
-            readBytes()
+      runBlocking {
+        val parallelWriter = OrderedParallelWriter(this)
+
+        fun writeToZip(entryName: String, lastModified: Long, bytes: ByteArray) {
+          zos.writeEntry(
+            name = entryName,
+            preserveLastModified = preserveFileTimestamps,
+            lastModified = lastModified,
+            unixMode = UnixMode.file(),
+          ) {
+            write(bytes)
           }
-        zos.writeEntry(
-          name = relocatedPath,
-          preserveLastModified = preserveFileTimestamps,
-          lastModified = lastModified,
-          unixMode = UnixMode.file(),
-        ) {
-          write(bytes)
         }
-      }
 
-      filesWithRelPaths
-        .sortedBy { it.second }
-        .forEach { (file, relPath) ->
-          addEntry(relPath, file.lastModified()) { file.readBytes() }
+        // Relocates and writes a project or dependency source file, or a resource.
+        fun addEntry(name: String, lastModified: Long, readBytes: () -> ByteArray) {
+          if (name.isModuleInfo) return
+          val isSource = name.isSourceFile()
+          if (isSource && isUnused(name, unusedClasses, sourceToClasses)) return
+          val relocatedPath =
+            if (isSource) relocators.relocateSourcePath(name) else relocators.relocatePath(name)
+          if (!visitedFiles.add(relocatedPath)) return
+          if (!isSource) {
+            // Resources gain nothing from parallelism, write them directly to avoid buffering.
+            parallelWriter.flush()
+            writeToZip(relocatedPath, lastModified, readBytes())
+            return
+          }
+          val rawBytes = readBytes()
+          parallelWriter.submit(
+            compute = {
+              relocators
+                .remapSourceFile(rawBytes.decodeToString(), name, relocatedPath)
+                .toByteArray()
+            },
+            write = { writeToZip(relocatedPath, lastModified, it) },
+          )
         }
 
-      sourcesJars.forEach { jarFile ->
-        jarFile.useZip {
-          entries()
-            .toList()
-            .filterNot { entry ->
-              val name = entry.name
-              entry.isDirectory ||
-                name == "META-INF/MANIFEST.MF" ||
-                name.endsWith(".class") ||
-                name.startsWith("META-INF/INDEX.LIST") ||
-                (name.startsWith("META-INF/") &&
-                  (name.endsWith(".SF") || name.endsWith(".DSA") || name.endsWith(".RSA")))
-            }
-            .sortedBy { it.name }
-            .forEach { entry ->
-              addEntry(entry.name, entry.time) { getInputStream(entry).use { it.readBytes() } }
-            }
+        filesWithRelPaths
+          .sortedBy { it.second }
+          .forEach { (file, relPath) ->
+            addEntry(relPath, file.lastModified()) { file.readBytes() }
+          }
+
+        sourcesJars.forEach { jarFile ->
+          jarFile.useZip {
+            entries()
+              .toList()
+              .filterNot { entry ->
+                val name = entry.name
+                entry.isDirectory ||
+                  name == "META-INF/MANIFEST.MF" ||
+                  name.endsWith(".class") ||
+                  name.startsWith("META-INF/INDEX.LIST") ||
+                  (name.startsWith("META-INF/") &&
+                    (name.endsWith(".SF") || name.endsWith(".DSA") || name.endsWith(".RSA")))
+              }
+              .sortedBy { it.name }
+              .forEach { entry ->
+                addEntry(entry.name, entry.time) { getInputStream(entry).use { it.readBytes() } }
+              }
+          }
         }
+
+        parallelWriter.flush()
       }
 
       val entries = zos.entries.map { it.name }

src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/relocation/Relocator.kt

@@ -9,7 +9,7 @@ import org.gradle.api.tasks.Input
  * [org.apache.maven.plugins.shade.relocation.Relocator.java](https://github.com/apache/maven-shade-plugin/blob/master/src/main/java/org/apache/maven/plugins/shade/relocation/Relocator.java).
  *
  * Implementations of [Relocator] must be thread-safe, as their methods may be invoked concurrently
- * across multiple worker threads during parallel bytecode remapping.
+ * across multiple worker threads during parallel bytecode and source remapping.
  *
  * @author Jason van Zyl
  * @author John Engelman

src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/tasks/ShadowCopyAction.kt

@@ -4,6 +4,7 @@
 
 package com.github.jengelman.gradle.plugins.shadow.tasks
 
+import com.github.jengelman.gradle.plugins.shadow.internal.OrderedParallelWriter
 import com.github.jengelman.gradle.plugins.shadow.internal.UnixMode
 import com.github.jengelman.gradle.plugins.shadow.internal.entries
 import com.github.jengelman.gradle.plugins.shadow.internal.gradleError
@@ -17,10 +18,6 @@ import com.github.jengelman.gradle.plugins.shadow.relocation.relocatePath
 import com.github.jengelman.gradle.plugins.shadow.transformers.ResourceTransformer
 import com.github.jengelman.gradle.plugins.shadow.transformers.TransformerContext
 import java.io.File
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Deferred
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.async
 import kotlinx.coroutines.runBlocking
 import org.apache.tools.zip.Zip64RequiredException
 import org.apache.tools.zip.ZipOutputStream
@@ -80,7 +77,7 @@ internal constructor(
     try {
       zipOutStream.use { zos ->
         runBlocking {
-          val action = StreamAction(zos, this)
+          val action = StreamAction(zos, OrderedParallelWriter(this))
           stream.process(action)
           action.flush()
         }
@@ -163,10 +160,8 @@ internal constructor(
 
   private inner class StreamAction(
     private val zipOutStr: ZipOutputStream,
-    private val scope: CoroutineScope,
+    private val parallelWriter: OrderedParallelWriter,
   ) : CopyActionProcessingStreamAction {
-    private val pendingEntries = ArrayDeque<PendingEntry>()
-
     init {
       logger.info("Relocator count: {}.", relocators.size)
     }
@@ -196,38 +191,22 @@ internal constructor(
             val pathSuffix = path.removePrefix(multiReleasePrefix)
             val relocatedPath = multiReleasePrefix + relocators.relocatePath(pathSuffix)
             val rawBytes = fileDetails.readBytes()
-            val deferred =
-              scope.async(Dispatchers.Default) {
-                rawBytes.remapClass(relocators = relocators, path = path)
-              }
-            pendingEntries.addLast(
-              PendingEntry(
-                entryName = relocatedPath,
-                fileDetails = fileDetails,
-                deferredBytes = deferred,
-              ),
+            parallelWriter.submit(
+              compute = { rawBytes.remapClass(relocators = relocators, path = path) },
+              write = { fileDetails.writeToZip(relocatedPath, it) },
             )
-            if (pendingEntries.size >= MAX_PENDING_ENTRIES) {
-              runBlocking { pendingEntries.removeFirst().writeToZip() }
-            }
           }
         }
         else -> {
           val relocated = relocators.relocatePath(path)
           if (transform(fileDetails, relocated)) return
-          if (pendingEntries.isNotEmpty()) {
-            runBlocking { flush() }
-          }
+          parallelWriter.flush()
           fileDetails.writeToZip(relocated)
         }
       }
     }
 
-    suspend fun flush() {
-      while (pendingEntries.isNotEmpty()) {
-        pendingEntries.removeFirst().writeToZip()
-      }
-    }
+    fun flush() = parallelWriter.flush()
 
     private fun isUnused(classPath: String): Boolean {
       val className = classPath.substringBeforeLast(".").replace('/', '.')
@@ -249,10 +228,6 @@ internal constructor(
       return true
     }
 
-    private suspend fun PendingEntry.writeToZip() {
-      fileDetails.writeToZip(entryName, deferredBytes.await())
-    }
-
     private fun FileCopyDetails.writeToZip(entryName: String, bytes: ByteArray? = null) {
       zipOutStr.writeEntry(
         name = entryName,
@@ -269,22 +244,10 @@ internal constructor(
     }
   }
 
-  private class PendingEntry(
-    val entryName: String,
-    val fileDetails: FileCopyDetails,
-    val deferredBytes: Deferred<ByteArray>,
-  )
-
   public companion object {
     private val logger = Logging.getLogger(@Suppress("DEPRECATION") ShadowCopyAction::class.java)
     private val multiReleaseRegex = "^META-INF/versions/\\d+/".toRegex()
 
-    /**
-     * Maximum number of in-flight parallel class remapping tasks in the sliding window. Bounds
-     * memory consumption on large archives while keeping worker threads saturated.
-     */
-    private const val MAX_PENDING_ENTRIES = 128
-
     @Deprecated(
       message =
         "Use `ShadowJar.CONSTANT_TIME_FOR_ZIP_ENTRIES` constant instead. This will be removed in Shadow 10.",

src/test/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJarTest.kt

@@ -278,6 +278,44 @@ class ShadowedSourcesJarTest {
       getContent("shaded/dep/resource.txt").isEqualTo("dep resource")
     }
   }
+
+  @Test
+  fun remapsManySourcesInParallel() {
+    val srcDir = tempDir.resolve("src").createDirectories()
+    val names = (1..500).map { "Class%03d".format(it) }
+    names.forEach { name ->
+      srcDir.resolve("com/example/$name.java").apply {
+        createParentDirectories()
+        writeText("package com.example;\npublic class $name {}")
+      }
+      srcDir.resolve("com/example/$name.txt").writeText("resource $name")
+    }
+
+    val outputJar = tempDir.resolve("output-sources.jar")
+    generateSourcesJar(
+      sourcesJarFile = outputJar.toFile(),
+      sourceSetsSourceDirs = testObjectFactory.fileCollection().from(srcDir),
+      relocators = listOf(SimpleRelocator("com.example", "shaded.example")),
+    )
+
+    val expectedEntries = names.flatMap {
+      listOf("shaded/example/$it.java", "shaded/example/$it.txt")
+    }
+    assertThat(JarPath(outputJar)).useAll {
+      containsExactly(
+        "META-INF/MANIFEST.MF",
+        *expectedEntries.toTypedArray(),
+        "META-INF/",
+        "shaded/",
+        "shaded/example/",
+      )
+      names.forEach { name ->
+        getContent("shaded/example/$name.java")
+          .isEqualTo("package shaded.example;\npublic class $name {}")
+        getContent("shaded/example/$name.txt").isEqualTo("resource $name")
+      }
+    }
+  }
 }
 
 private fun generateSourcesJar(