From d88703756b94f6af13fdfa42f602f155df248e8e Mon Sep 17 00:00:00 2001 From: Goooler Date: Sun, 25 Oct 2026 19:53:56 +0800 Subject: [PATCH 1/2] Parallelize source remapping for shadowed sources JAR Apply the same sliding-window coroutine approach used for bytecode remapping in ShadowCopyAction to source file remapping in generateSourcesJar, keeping entry order deterministic. Resources are written directly after draining the pending sources, so they are never buffered. --- .../shadow/internal/ShadowedSourcesJar.kt | 136 ++++++++++++------ .../plugins/shadow/relocation/Relocator.kt | 2 +- .../shadow/internal/ShadowedSourcesJarTest.kt | 38 +++++ 3 files changed, 132 insertions(+), 44 deletions(-) diff --git a/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJar.kt b/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJar.kt index 595cd23eb..85ceab2f7 100644 --- a/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJar.kt +++ b/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJar.kt @@ -4,6 +4,10 @@ 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.Deferred +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.async +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 +54,90 @@ 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 pendingEntries = ArrayDeque() + + 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() } + suspend fun PendingSourceEntry.writeToZip() { + writeToZip(entryName, lastModified, deferredBytes.await()) } - 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() } } + suspend fun flush() { + while (pendingEntries.isNotEmpty()) { + pendingEntries.removeFirst().writeToZip() + } + } + + // Relocates and writes a project or dependency source file, or a resource. + suspend 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. + flush() + writeToZip(relocatedPath, lastModified, readBytes()) + return + } + val rawBytes = readBytes() + val deferred = + async(Dispatchers.Default) { + relocators + .remapSourceFile(rawBytes.decodeToString(), name, relocatedPath) + .toByteArray() } + pendingEntries.addLast( + PendingSourceEntry( + entryName = relocatedPath, + lastModified = lastModified, + deferredBytes = deferred, + ), + ) + if (pendingEntries.size >= MAX_PENDING_ENTRIES) { + pendingEntries.removeFirst().writeToZip() + } } + + 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() } } + } + } + } + + flush() } val entries = zos.entries.map { it.name } @@ -201,6 +239,18 @@ private fun buildSourceToClassesMap( return sourceToClasses } +private class PendingSourceEntry( + val entryName: String, + val lastModified: Long, + val deferredBytes: Deferred, +) + +/** + * Maximum number of in-flight parallel source remapping tasks in the sliding window. Bounds memory + * consumption on large archives while keeping worker threads saturated. + */ +private const val MAX_PENDING_ENTRIES = 128 + // Aligned with the default exclusion of `module-info.class` from the shadowed JAR. private val String.isModuleInfo: Boolean get() = substringAfterLast('/') == "module-info.java" diff --git a/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/relocation/Relocator.kt b/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/relocation/Relocator.kt index cc966785a..dd3b30b15 100644 --- a/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/relocation/Relocator.kt +++ b/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 diff --git a/src/test/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJarTest.kt b/src/test/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJarTest.kt index 5b4cb5011..b1084526e 100644 --- a/src/test/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJarTest.kt +++ b/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( From 7ff45cd1eb02341fc4eab64f913f98e41fddae6c Mon Sep 17 00:00:00 2001 From: Goooler Date: Sun, 25 Oct 2026 19:59:39 +0800 Subject: [PATCH 2/2] Extract OrderedParallelWriter for parallel remapping Share the sliding-window ordered writer between ShadowCopyAction and the shadowed sources JAR generation instead of duplicating the pending entry queue in both. --- .../shadow/internal/OrderedParallelWriter.kt | 42 +++++++++++++++ .../shadow/internal/ShadowedSourcesJar.kt | 49 +++-------------- .../plugins/shadow/tasks/ShadowCopyAction.kt | 53 +++---------------- 3 files changed, 58 insertions(+), 86 deletions(-) create mode 100644 src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/OrderedParallelWriter.kt diff --git a/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/OrderedParallelWriter.kt b/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/OrderedParallelWriter.kt new file mode 100644 index 000000000..b8d1a439d --- /dev/null +++ b/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() + + 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, val write: (ByteArray) -> Unit) diff --git a/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJar.kt b/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJar.kt index 85ceab2f7..be056d953 100644 --- a/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJar.kt +++ b/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/internal/ShadowedSourcesJar.kt @@ -4,9 +4,6 @@ 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.Deferred -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.async import kotlinx.coroutines.runBlocking import org.apache.tools.zip.ZipOutputStream import org.gradle.api.file.FileCollection @@ -55,7 +52,7 @@ internal fun generateSourcesJar( } runBlocking { - val pendingEntries = ArrayDeque() + val parallelWriter = OrderedParallelWriter(this) fun writeToZip(entryName: String, lastModified: Long, bytes: ByteArray) { zos.writeEntry( @@ -68,18 +65,8 @@ internal fun generateSourcesJar( } } - suspend fun PendingSourceEntry.writeToZip() { - writeToZip(entryName, lastModified, deferredBytes.await()) - } - - suspend fun flush() { - while (pendingEntries.isNotEmpty()) { - pendingEntries.removeFirst().writeToZip() - } - } - // Relocates and writes a project or dependency source file, or a resource. - suspend fun addEntry(name: String, lastModified: Long, readBytes: () -> ByteArray) { + fun addEntry(name: String, lastModified: Long, readBytes: () -> ByteArray) { if (name.isModuleInfo) return val isSource = name.isSourceFile() if (isSource && isUnused(name, unusedClasses, sourceToClasses)) return @@ -88,27 +75,19 @@ internal fun generateSourcesJar( if (!visitedFiles.add(relocatedPath)) return if (!isSource) { // Resources gain nothing from parallelism, write them directly to avoid buffering. - flush() + parallelWriter.flush() writeToZip(relocatedPath, lastModified, readBytes()) return } val rawBytes = readBytes() - val deferred = - async(Dispatchers.Default) { + parallelWriter.submit( + compute = { relocators .remapSourceFile(rawBytes.decodeToString(), name, relocatedPath) .toByteArray() - } - pendingEntries.addLast( - PendingSourceEntry( - entryName = relocatedPath, - lastModified = lastModified, - deferredBytes = deferred, - ), + }, + write = { writeToZip(relocatedPath, lastModified, it) }, ) - if (pendingEntries.size >= MAX_PENDING_ENTRIES) { - pendingEntries.removeFirst().writeToZip() - } } filesWithRelPaths @@ -137,7 +116,7 @@ internal fun generateSourcesJar( } } - flush() + parallelWriter.flush() } val entries = zos.entries.map { it.name } @@ -239,18 +218,6 @@ private fun buildSourceToClassesMap( return sourceToClasses } -private class PendingSourceEntry( - val entryName: String, - val lastModified: Long, - val deferredBytes: Deferred, -) - -/** - * Maximum number of in-flight parallel source remapping tasks in the sliding window. Bounds memory - * consumption on large archives while keeping worker threads saturated. - */ -private const val MAX_PENDING_ENTRIES = 128 - // Aligned with the default exclusion of `module-info.class` from the shadowed JAR. private val String.isModuleInfo: Boolean get() = substringAfterLast('/') == "module-info.java" diff --git a/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/tasks/ShadowCopyAction.kt b/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/tasks/ShadowCopyAction.kt index 4fdc22a1d..e9bb57e72 100644 --- a/src/main/kotlin/com/github/jengelman/gradle/plugins/shadow/tasks/ShadowCopyAction.kt +++ b/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() - 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, - ) - 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.",