Parallelize bytecode remapping in ShadowCopyAction

CHANGELOG.md

@@ -28,6 +28,7 @@
   - `XmlAppendingTransformer`
 - Append terminating newline in `ServiceFileTransformer`. ([#2202](https://github.com/GradleUp/shadow/pull/2202))
 - Remove redundant JAR normalization for R8 output. ([#2236](https://github.com/GradleUp/shadow/pull/2236))
+- Parallelize bytecode remapping in `ShadowCopyAction`. ([#2302](https://github.com/GradleUp/shadow/pull/2302))
 
 ### Deprecated
 

build.gradle.kts

@@ -131,6 +131,7 @@ dependencies {
   compileOnly(libs.develocity)
   compileOnly(libs.kotlin.gradlePlugin)
   compileOnly(libs.kotlin.reflect)
+  compileOnly(libs.kotlinx.coroutines)
   api(libs.apache.ant) // Types from Ant are exposed in the public API.
   implementation(libs.apache.log4j)
   implementation(libs.jdependency)
@@ -152,7 +153,10 @@ dependencies {
 
 testing.suites {
   named<JvmTestSuite>("test") {
-    dependencies { implementation(libs.xmlunit) }
+    dependencies {
+      implementation(libs.kotlinx.coroutines)
+      implementation(libs.xmlunit)
+    }
   }
   register<JvmTestSuite>("documentTest") {
     targets.configureEach {

gradle/libs.versions.toml

@@ -12,6 +12,7 @@ jdependency = "org.vafer:jdependency:2.16"
 jdom2 = "org.jdom:jdom2:2.0.6.1"
 kotlin-metadata = { module = "org.jetbrains.kotlin:kotlin-metadata-jvm", version.ref = "kotlin" }
 kotlin-reflect = { module = "org.jetbrains.kotlin:kotlin-reflect", version.ref = "kotlin" }
+kotlinx-coroutines = "org.jetbrains.kotlinx:kotlinx-coroutines-core:1.10.1"
 plexus-utils = "org.codehaus.plexus:plexus-utils:4.1.0"
 plexus-xml = "org.codehaus.plexus:plexus-xml:4.2.0"
 xmlunit = "org.xmlunit:xmlunit-legacy:2.13.0"

src/functionalTest/kotlin/com/github/jengelman/gradle/plugins/shadow/ParallelRelocationTest.kt

@@ -0,0 +1,117 @@
+package com.github.jengelman.gradle.plugins.shadow
+
+import assertk.assertThat
+import assertk.assertions.contains
+import assertk.assertions.isEqualTo
+import com.github.jengelman.gradle.plugins.shadow.testkit.classLoader
+import com.github.jengelman.gradle.plugins.shadow.testkit.containsExactly
+import com.github.jengelman.gradle.plugins.shadow.testkit.loadClass
+import kotlin.io.path.appendText
+import kotlin.io.path.readBytes
+import org.junit.jupiter.api.Test
+
+class ParallelRelocationTest : BasePluginTest() {
+  @Test
+  fun largeNumberOfClassesWithRelocation() {
+    val count = 500
+    val classNames = (1..count).map { "Class%03d".format(it) }
+    val largeJar =
+      buildJar("many-classes.jar") {
+        for (name in classNames) {
+          insert(
+            "com/example/pkg/$name.class",
+            createEmptyClassBytes("com/example/pkg/$name"),
+          )
+        }
+      }
+
+    projectScript.appendText(
+      """
+      |dependencies {
+      |  ${implementationFiles(largeJar)}
+      |}
+      |$shadowJarTask {
+      |  relocate 'com.example.pkg', 'relocated.pkg'
+      |}
+      """
+        .trimMargin()
+    )
+
+    runWithSuccess(shadowJarPath)
+
+    assertThat(outputShadowedJar).useAll {
+      val relocatedEntries = classNames.map { "relocated/pkg/$it.class" }.toTypedArray()
+      containsExactly(
+        "META-INF/MANIFEST.MF",
+        *relocatedEntries,
+        "META-INF/",
+        "relocated/pkg/",
+        "relocated/",
+      )
+      classLoader {
+        loadClass("relocated.pkg.${classNames.first()}")
+        loadClass("relocated.pkg.${classNames.last()}")
+      }
+    }
+  }
+
+  @Test
+  fun deterministicZipEntryOrderAcrossMultipleBuilds() {
+    val count = 150
+    val testJar =
+      buildJar("deterministic-test.jar") {
+        for (i in 1..count) {
+          insert(
+            "com/example/test/TestClass$i.class",
+            createEmptyClassBytes("com/example/test/TestClass$i"),
+          )
+          insert("resources/res_$i.txt", "content $i")
+        }
+      }
+
+    projectScript.appendText(
+      """
+      |dependencies {
+      |  ${implementationFiles(testJar)}
+      |}
+      |$shadowJarTask {
+      |  relocate 'com.example.test', 'shadowed.example.test'
+      |}
+      """
+        .trimMargin()
+    )
+
+    runWithSuccess(shadowJarPath)
+    val firstBytes = path("build/libs/my-1.0-all.jar").readBytes()
+
+    runWithSuccess(shadowJarPath, "--rerun-tasks")
+    val secondBytes = path("build/libs/my-1.0-all.jar").readBytes()
+
+    assertThat(firstBytes).isEqualTo(secondBytes)
+  }
+
+  @Test
+  fun errorPropagationWhenClassIsCorrupted() {
+    val badClassEntry = "corrupt/BadClass.class"
+    val corruptJar =
+      buildJar("corrupt.jar") {
+        insert(badClassEntry, byteArrayOf(0xCA.toByte(), 0xFE.toByte(), 0xBA.toByte()))
+      }
+
+    projectScript.appendText(
+      """
+      |dependencies {
+      |  ${implementationFiles(corruptJar)}
+      |}
+      |$shadowJarTask {
+      |  relocate 'corrupt', 'relocated.corrupt'
+      |}
+      """
+        .trimMargin()
+    )
+
+    val result = runWithFailure(shadowJarPath)
+
+    assertThat(result.output).contains("Error in ASM processing class $badClassEntry")
+  }
+}

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

@@ -74,6 +74,8 @@ internal fun FileTreeElement.inputStream(): InputStream =
     file.inputStream()
   }
 
+internal fun FileTreeElement.readBytes(): ByteArray = inputStream().use(InputStream::readBytes)
+
 internal inline fun <reified V : Any> ObjectFactory.property(
   defaultValue: Any? = null
 ): Property<V> =

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

@@ -13,28 +13,29 @@ import org.vafer.jdeb.shaded.objectweb.asm.commons.Remapper
  * (possibly) remapped class bytes. If no remapping is required, the original bytes are returned.
  */
 internal fun FileCopyDetails.remapClass(relocators: Set<Relocator>): ByteArray =
-  inputStream()
-    .use { it.readBytes() }
-    .let { bytes ->
-      var modified = false
-      val remapper = RelocatorRemapper(relocators) { modified = true }
+  readBytes().remapClass(relocators = relocators, path = path)
 
-      // We don't pass the ClassReader here. This forces the ClassWriter to rebuild the constant
-      // pool. Copying the original constant pool should be avoided because it would keep references
-      // to the original class names. This is not a problem at runtime (because these entries in the
-      // constant pool are never used), but confuses some tools such as Felix's maven-bundle-plugin
-      // that use the constant pool to determine the dependencies of a class.
-      try {
-        val cw = ClassWriter(0)
-        val cr = ClassReader(bytes)
-        val cv = ClassRemapper(cw, remapper)
-        cr.accept(cv, ClassReader.EXPAND_FRAMES)
-        // If we didn't need to change anything, keep the original bytes as-is.
-        if (modified) cw.toByteArray() else bytes
-      } catch (t: Throwable) {
-        gradleError("Error in ASM processing class $path", t)
-      }
+internal fun ByteArray.remapClass(relocators: Set<Relocator>, path: String): ByteArray =
+  let { bytes ->
+    var modified = false
+    val remapper = RelocatorRemapper(relocators) { modified = true }
+
+    // We don't pass the ClassReader here. This forces the ClassWriter to rebuild the constant
+    // pool. Copying the original constant pool should be avoided because it would keep references
+    // to the original class names. This is not a problem at runtime (because these entries in the
+    // constant pool are never used), but confuses some tools such as Felix's maven-bundle-plugin
+    // that use the constant pool to determine the dependencies of a class.
+    try {
+      val cw = ClassWriter(0)
+      val cr = ClassReader(bytes)
+      val cv = ClassRemapper(cw, remapper)
+      cr.accept(cv, ClassReader.EXPAND_FRAMES)
+      // If we didn't need to change anything, keep the original bytes as-is.
+      if (modified) cw.toByteArray() else bytes
+    } catch (t: Throwable) {
+      gradleError("Error in ASM processing class $path", t)
     }
+  }
 
 private class RelocatorRemapper(
   private val relocators: Set<Relocator>,

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

@@ -9,13 +9,22 @@ import com.github.jengelman.gradle.plugins.shadow.internal.entries
 import com.github.jengelman.gradle.plugins.shadow.internal.gradleError
 import com.github.jengelman.gradle.plugins.shadow.internal.inputStream
 import com.github.jengelman.gradle.plugins.shadow.internal.parentDirectoryEntries
+import com.github.jengelman.gradle.plugins.shadow.internal.readBytes
 import com.github.jengelman.gradle.plugins.shadow.internal.remapClass
 import com.github.jengelman.gradle.plugins.shadow.internal.writeEntry
 import com.github.jengelman.gradle.plugins.shadow.relocation.Relocator
 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.CompletableDeferred
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.async
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
 import org.apache.tools.zip.Zip64RequiredException
 import org.apache.tools.zip.ZipOutputStream
 import org.gradle.api.file.FileCopyDetails
@@ -73,7 +82,32 @@ internal constructor(
   override fun execute(stream: CopyActionProcessingStream): WorkResult {
     try {
       zipOutStream.use { zos ->
-        stream.process(StreamAction(zos))
+        runBlocking {
+          val channel = Channel<ProcessItem>(capacity = 128)
+
+          val writer =
+            launch(Dispatchers.Default) {
+              for (item in channel) {
+                zos.writeEntry(
+                  name = item.entryName,
+                  preserveLastModified = isPreserveFileTimestamps,
+                  lastModified = item.lastModified,
+                  unixMode = item.unixMode,
+                ) {
+                  write(item.deferredBytes.await())
+                }
+              }
+            }
+
+          try {
+            stream.process(StreamAction(this, channel))
+          } finally {
+            channel.close()
+          }
+
+          writer.join()
+        }
+
         processTransformers(zos)
         addDirs(zos)
         checkDuplicateEntries(zos)
@@ -150,8 +184,10 @@ internal constructor(
     }
   }
 
-  private inner class StreamAction(private val zipOutStr: ZipOutputStream) :
-    CopyActionProcessingStreamAction {
+  private inner class StreamAction(
+    private val scope: CoroutineScope,
+    private val channel: Channel<ProcessItem>,
+  ) : CopyActionProcessingStreamAction {
     init {
       logger.info("Relocator count: {}.", relocators.size)
     }
@@ -173,27 +209,54 @@ internal constructor(
       when {
         path.endsWith(".class") -> {
           if (isUnused(path)) return
+          val rawBytes = fileDetails.readBytes()
           if (relocators.isEmpty()) {
-            fileDetails.writeToZip(path)
+            fileDetails.sendEntry(
+              entryName = path,
+              deferredBytes = CompletableDeferred(rawBytes),
+            )
           } else {
             // Temporarily remove the multi-release prefix.
             val multiReleasePrefix = multiReleaseRegex.find(path)?.value.orEmpty()
             val pathSuffix = path.removePrefix(multiReleasePrefix)
             val relocatedPath = multiReleasePrefix + relocators.relocatePath(pathSuffix)
-            fileDetails.writeToZip(
+            fileDetails.sendEntry(
               entryName = relocatedPath,
-              bytes = fileDetails.remapClass(relocators = relocators),
+              deferredBytes =
+                scope.async(Dispatchers.Default) {
+                  rawBytes.remapClass(relocators = relocators, path = path)
+                },
             )
           }
         }
         else -> {
           val relocated = relocators.relocatePath(path)
           if (transform(fileDetails, relocated)) return
-          fileDetails.writeToZip(relocated)
+          val rawBytes = fileDetails.readBytes()
+          fileDetails.sendEntry(
+            entryName = relocated,
+            deferredBytes = CompletableDeferred(rawBytes),
+          )
         }
       }
     }
 
+    private fun FileCopyDetails.sendEntry(
+      entryName: String,
+      deferredBytes: Deferred<ByteArray>,
+    ) {
+      runBlocking {
+        channel.send(
+          ProcessItem(
+            entryName = entryName,
+            deferredBytes = deferredBytes,
+            lastModified = lastModified,
+            unixMode = UnixMode.file(permissions.toUnixNumeric()),
+          )
+        )
+      }
+    }
+
     private fun isUnused(classPath: String): Boolean {
       val className = classPath.substringBeforeLast(".").replace('/', '.')
       return unusedClasses.contains(className).also {
@@ -213,23 +276,15 @@ internal constructor(
       }
       return true
     }
-
-    private fun FileCopyDetails.writeToZip(entryName: String, bytes: ByteArray? = null) {
-      zipOutStr.writeEntry(
-        name = entryName,
-        preserveLastModified = isPreserveFileTimestamps,
-        lastModified = lastModified,
-        unixMode = UnixMode.file(permissions.toUnixNumeric()),
-      ) {
-        if (bytes == null) {
-          copyTo(this)
-        } else {
-          write(bytes)
-        }
-      }
-    }
   }
 
+  private class ProcessItem(
+    val entryName: String,
+    val deferredBytes: Deferred<ByteArray>,
+    val lastModified: Long,
+    val unixMode: UnixMode,
+  )
+
   public companion object {
     private val logger = Logging.getLogger(@Suppress("DEPRECATION") ShadowCopyAction::class.java)
     private val multiReleaseRegex = "^META-INF/versions/\\d+/".toRegex()