From 8a61eaac9f7d9e0f7c03ae916c928524dd39aed6 Mon Sep 17 00:00:00 2001 From: SettingDust Date: Wed, 27 Nov 2024 14:58:40 +0800 Subject: [PATCH] perf: async remap entries in jar (#82) --- .../kilt/loader/KiltEarlierInitializer.kt | 4 +- .../xyz/bluspring/kilt/loader/KiltLoader.kt | 38 +++-- .../kilt/loader/remap/KiltRemapper.kt | 140 ++++++++++-------- .../xyz/bluspring/kilt/util/CoroutineUtils.kt | 26 ++++ 4 files changed, 133 insertions(+), 75 deletions(-) create mode 100644 src/main/kotlin/xyz/bluspring/kilt/util/CoroutineUtils.kt diff --git a/src/main/kotlin/xyz/bluspring/kilt/loader/KiltEarlierInitializer.kt b/src/main/kotlin/xyz/bluspring/kilt/loader/KiltEarlierInitializer.kt index 55053369..b4c11d69 100644 --- a/src/main/kotlin/xyz/bluspring/kilt/loader/KiltEarlierInitializer.kt +++ b/src/main/kotlin/xyz/bluspring/kilt/loader/KiltEarlierInitializer.kt @@ -1,10 +1,10 @@ package xyz.bluspring.kilt.loader import de.florianmichael.asmfabricloader.api.event.PrePrePreLaunchEntrypoint +import kotlinx.coroutines.runBlocking class KiltEarlierInitializer : PrePrePreLaunchEntrypoint { override fun onLanguageAdapterLaunch() { - KiltLoader.INSTANCE.scanModJob - KiltLoader.INSTANCE.runScanMods() + runBlocking { KiltLoader.INSTANCE.scanMods() } } } \ No newline at end of file diff --git a/src/main/kotlin/xyz/bluspring/kilt/loader/KiltLoader.kt b/src/main/kotlin/xyz/bluspring/kilt/loader/KiltLoader.kt index 9539c08b..f4b09128 100644 --- a/src/main/kotlin/xyz/bluspring/kilt/loader/KiltLoader.kt +++ b/src/main/kotlin/xyz/bluspring/kilt/loader/KiltLoader.kt @@ -3,7 +3,8 @@ package xyz.bluspring.kilt.loader import com.electronwill.nightconfig.core.CommentedConfig import com.electronwill.nightconfig.toml.TomlParser import com.google.gson.JsonParser -import kotlinx.coroutines.* +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.SupervisorJob import net.fabricmc.api.EnvType import net.fabricmc.loader.api.FabricLoader import net.fabricmc.loader.impl.FabricLoaderImpl @@ -14,11 +15,21 @@ import net.minecraft.SharedConstants import net.minecraft.server.Bootstrap import net.minecraftforge.common.ForgeStatesProvider import net.minecraftforge.eventbus.api.Event -import net.minecraftforge.fml.* +import net.minecraftforge.fml.ModList +import net.minecraftforge.fml.ModLoader +import net.minecraftforge.fml.ModLoadingContext +import net.minecraftforge.fml.ModLoadingPhase +import net.minecraftforge.fml.ModLoadingStage import net.minecraftforge.fml.common.Mod import net.minecraftforge.fml.config.ConfigTracker import net.minecraftforge.fml.config.ModConfig -import net.minecraftforge.fml.event.lifecycle.* +import net.minecraftforge.fml.event.lifecycle.FMLClientSetupEvent +import net.minecraftforge.fml.event.lifecycle.FMLCommonSetupEvent +import net.minecraftforge.fml.event.lifecycle.FMLConstructModEvent +import net.minecraftforge.fml.event.lifecycle.FMLDedicatedServerSetupEvent +import net.minecraftforge.fml.event.lifecycle.FMLLoadCompleteEvent +import net.minecraftforge.fml.event.lifecycle.InterModEnqueueEvent +import net.minecraftforge.fml.event.lifecycle.InterModProcessEvent import net.minecraftforge.fml.loading.FMLPaths import net.minecraftforge.fml.loading.moddiscovery.ModAnnotation import net.minecraftforge.fml.loading.moddiscovery.ModClassVisitor @@ -51,7 +62,16 @@ import java.util.function.Consumer import java.util.jar.JarFile import java.util.jar.Manifest import java.util.zip.ZipFile -import kotlin.io.path.* +import kotlin.io.path.createDirectories +import kotlin.io.path.createFile +import kotlin.io.path.div +import kotlin.io.path.exists +import kotlin.io.path.forEachDirectoryEntry +import kotlin.io.path.isDirectory +import kotlin.io.path.name +import kotlin.io.path.nameWithoutExtension +import kotlin.io.path.toPath +import kotlin.io.path.writeBytes import kotlin.system.exitProcess class KiltLoader { @@ -66,15 +86,7 @@ class KiltLoader { val scope = CoroutineScope(SupervisorJob()) - val scanModJob by lazy { - scope.launch(Dispatchers.IO) { scanMods() } - } - - fun runScanMods() { - runBlocking { scanModJob.join() } - } - - private suspend fun scanMods() { + suspend fun scanMods() { Kilt.logger.info("Scanning the mods directory for Forge mods...") DeltaTimeProfiler.push("scanMods") diff --git a/src/main/kotlin/xyz/bluspring/kilt/loader/remap/KiltRemapper.kt b/src/main/kotlin/xyz/bluspring/kilt/loader/remap/KiltRemapper.kt index 5f32bbc2..e548f3d7 100644 --- a/src/main/kotlin/xyz/bluspring/kilt/loader/remap/KiltRemapper.kt +++ b/src/main/kotlin/xyz/bluspring/kilt/loader/remap/KiltRemapper.kt @@ -6,13 +6,19 @@ import com.google.gson.JsonParser import it.unimi.dsi.fastutil.objects.Object2ReferenceFunction import it.unimi.dsi.fastutil.objects.Object2ReferenceMaps import it.unimi.dsi.fastutil.objects.Object2ReferenceOpenHashMap +import kotlinx.atomicfu.locks.synchronized +import kotlinx.coroutines.CoroutineStart import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Job import kotlinx.coroutines.flow.asFlow -import kotlinx.coroutines.flow.filter +import kotlinx.coroutines.flow.collect import kotlinx.coroutines.flow.flowOn +import kotlinx.coroutines.flow.launchIn +import kotlinx.coroutines.flow.toSet import kotlinx.coroutines.joinAll import kotlinx.coroutines.launch +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.stream.consumeAsFlow import kotlinx.coroutines.withContext import net.fabricmc.loader.api.FabricLoader import net.fabricmc.loader.impl.game.GameProviderHelper @@ -42,6 +48,10 @@ import xyz.bluspring.kilt.loader.remap.fixers.WorkaroundFixer import xyz.bluspring.kilt.util.CaseInsensitiveStringHashSet import xyz.bluspring.kilt.util.ClassNameHashSet import xyz.bluspring.kilt.util.KiltHelper +import xyz.bluspring.kilt.util.filterAsync +import xyz.bluspring.kilt.util.flatMapAsync +import xyz.bluspring.kilt.util.mapAsync +import xyz.bluspring.kilt.util.onEachAsync import java.io.ByteArrayOutputStream import java.io.File import java.nio.file.Path @@ -114,32 +124,37 @@ object KiltRemapper { private lateinit var remappedModsDir: Path // SRG name -> (parent class name, intermediary/mapped name) - val srgMappedFields = srgIntermediaryMapping.classes.flatMap { - it.fields.map { f -> - f.original to - if (!forceProductionRemap) - mappingResolver.mapFieldName( - "intermediary", - it.mapped.replace("/", "."), - f.mapped, - f.mappedDescriptor - ) - else - f.mapped - } - }.associateBy { it.first } + val srgMappedFields = runBlocking { + srgIntermediaryMapping.classes.asFlow().flatMapAsync { + it.fields.asFlow().mapAsync { f -> + f.original to + if (!forceProductionRemap) + mappingResolver.mapFieldName( + "intermediary", + it.mapped.replace("/", "."), + f.mapped, + f.mappedDescriptor + ) + else + f.mapped + } + }.toSet().associateBy { it.first } + } // SRG name -> (parent class name, intermediary/mapped name) - val srgMappedMethods = mutableMapOf>() + val srgMappedMethods = + Object2ReferenceMaps.synchronize(Object2ReferenceOpenHashMap>()) init { - srgIntermediaryMapping.classes.forEach { + srgIntermediaryMapping.classes.asFlow().flowOn(Dispatchers.IO).onEachAsync { it.methods.forEach m@{ f -> // otherwise FunctionalInterface methods don't get remapped properly??? if (!f.mapped.startsWith("method_") && !FabricLoader.getInstance().isDevelopmentEnvironment) return@m - val map = srgMappedMethods.computeIfAbsent(f.original) { mutableMapOf() } + val map = srgMappedMethods.getOrPut(f.original) { + Object2ReferenceMaps.synchronize(Object2ReferenceOpenHashMap()) + } val mapped = if (!forceProductionRemap) (mappingResolver.mapMethodName( "intermediary", @@ -152,7 +167,7 @@ object KiltRemapper { map[f.parent.original] = mapped } - } + }.launchIn(Kilt.loader.scope) } suspend fun remapMods(modLoadingQueue: ConcurrentLinkedQueue, remappedModsDir: Path): List { @@ -233,7 +248,7 @@ object KiltRemapper { }.build() val remapper = KiltEnhancedRemapper(classProvider, srgIntermediaryMapping, logConsumer) - val entryToClassNodes = Object2ReferenceOpenHashMap() + val entryToClassNodes = Object2ReferenceMaps.synchronize(Object2ReferenceOpenHashMap()) val mixinClasses = ClassNameHashSet() val refmaps = CaseInsensitiveStringHashSet() @@ -250,9 +265,11 @@ object KiltRemapper { manifest.entries.keys.removeIf { it == "SHA-256-Digest" || it == "SHA-1-Digest" } - jarOutput.putNextEntry(manifestEntry) - jarOutput.write(ByteArrayOutputStream().also { manifest.write(it) }.toByteArray()) - jarOutput.closeEntry() + synchronized(jarOutput) { + jarOutput.putNextEntry(manifestEntry) + jarOutput.write(ByteArrayOutputStream().also { manifest.write(it) }.toByteArray()) + jarOutput.closeEntry() + } return@withContext manifest } @@ -492,17 +509,19 @@ object KiltRemapper { this.add("named:intermediary", newMappings) }) - jarOutput.putNextEntry(entry) - jarOutput.write(Kilt.gson.toJson(refmapData).toByteArray()) - jarOutput.closeEntry() + synchronized(jarOutput) { + jarOutput.putNextEntry(entry) + jarOutput.write(Kilt.gson.toJson(refmapData).toByteArray()) + jarOutput.closeEntry() + } return } - jar.entries().iterator().asFlow() + jar.stream().consumeAsFlow() .flowOn(Dispatchers.IO) - .filter { !it.name.equals("META-INF/MANIFEST.MF", true) } - .filter { + .filterAsync { !it.name.equals("META-INF/MANIFEST.MF", true) } + .filterAsync { val isHash = it.name.endsWith(".rsa", true) || it.name.endsWith(".sf", true) if (isHash) { // ignore JAR signatures. @@ -512,19 +531,21 @@ object KiltRemapper { } !isHash } - .collect { entry -> + .onEachAsync { entry -> when { entry.name in refmaps -> remapRefmap(jar, entry, remapper, jarOutput) // Keep the other resources - !entry.name.endsWith(".class") -> withContext(Dispatchers.IO) { - jarOutput.putNextEntry(entry) - jarOutput.write(jar.getInputStream(entry).readAllBytes()) - jarOutput.closeEntry() + !entry.name.endsWith(".class") -> { + synchronized(jarOutput) { + jarOutput.putNextEntry(entry) + jarOutput.write(jar.getInputStream(entry).readAllBytes()) + jarOutput.closeEntry() + } } else -> { - val classReader = ClassReader(withContext(Dispatchers.IO) { jar.getInputStream(entry) }) + val classReader = ClassReader(jar.getInputStream(entry)) // we need the info for this for the class writer val classNode = ClassNode(Opcodes.ASM9) @@ -534,6 +555,7 @@ object KiltRemapper { } } } + .collect() val classesToProcess = entryToClassNodes @@ -541,7 +563,7 @@ object KiltRemapper { .intersect(KiltHelper.getForgeClassNodes().toSet()) .toList() - suspend fun remapClass( + fun remapClass( remapper: KiltEnhancedRemapper, originalNode: ClassNode, mixinClasses: ClassNameHashSet, @@ -571,7 +593,7 @@ object KiltRemapper { val classWriter = ClassWriter(0) remappedNode.accept(classWriter) - withContext(Dispatchers.IO) { + synchronized(jarOutput) { jarOutput.putNextEntry(entry) jarOutput.write(classWriter.toByteArray()) jarOutput.closeEntry() @@ -582,9 +604,10 @@ object KiltRemapper { } } - for ((entry, originalNode) in entryToClassNodes) { - remapClass(remapper, originalNode, mixinClasses, classesToProcess, jarOutput, entry, exceptions) - } + entryToClassNodes.asSequence().asFlow().flowOn(Dispatchers.IO) + .onEachAsync { (entry, originalNode) -> + remapClass(remapper, originalNode, mixinClasses, classesToProcess, jarOutput, entry, exceptions) + }.collect() mod.remappedModFile = modifiedJarFile.toFile() withContext(Dispatchers.IO) { jarOutput.close() } @@ -595,7 +618,7 @@ object KiltRemapper { fun remapModAsync( mod: ForgeMod, mods: Collection - ): Job = Kilt.loader.scope.launch(Dispatchers.IO) { + ): Job = Kilt.loader.scope.launch(Dispatchers.IO, start = CoroutineStart.LAZY) { runCatching { mod.dependencies.mapNotNull { if (modToJob[it.modId] == null) @@ -760,38 +783,35 @@ object KiltRemapper { runCatching { srgFile.createFile() } withContext(Dispatchers.IO) { JarOutputStream(srgFile.outputStream()) }.use { outputJar -> - for (entry in gameJar.entries()) { - if (entry.name.endsWith(".class")) { - val classReader = ClassReader(withContext(Dispatchers.IO) { - gameJar.getInputStream(entry) - }) + gameJar.stream().consumeAsFlow().flowOn(Dispatchers.IO) + .onEachAsync { entry -> + if (entry.name.endsWith(".class")) { + val classReader = ClassReader(gameJar.getInputStream(entry)) - val classNode = ClassNode(Opcodes.ASM9) - classReader.accept(classNode, 0) + val classNode = ClassNode(Opcodes.ASM9) + classReader.accept(classNode, 0) - val classWriter = ClassWriter(0) + val classWriter = ClassWriter(0) - val visitor = - EnhancedClassRemapper(classWriter, srgRemapper, RenamingTransformer(srgRemapper, false)) - classNode.accept(visitor) - ConflictingStaticMethodFixer.fixClass(classNode) + val visitor = + EnhancedClassRemapper(classWriter, srgRemapper, RenamingTransformer(srgRemapper, false)) + classNode.accept(visitor) + ConflictingStaticMethodFixer.fixClass(classNode) - // We need to remap to the SRG name, otherwise the remapper completely fails in production environments. - val srgName = intermediarySrgMapping.remapClass(entry.name.removePrefix("/").removeSuffix(".class")) + // We need to remap to the SRG name, otherwise the remapper completely fails in production environments. + val srgName = + intermediarySrgMapping.remapClass(entry.name.removePrefix("/").removeSuffix(".class")) - withContext(Dispatchers.IO) { outputJar.putNextEntry(JarEntry("$srgName.class")) outputJar.write(classWriter.toByteArray()) outputJar.closeEntry() - } - } else { - withContext(Dispatchers.IO) { + } else { outputJar.putNextEntry(entry) outputJar.write(gameJar.getInputStream(entry).readAllBytes()) outputJar.closeEntry() } } - } + .collect() } logger.info("Remapped Minecraft from Intermediary to SRG. (took ${System.currentTimeMillis() - startTime} ms)") diff --git a/src/main/kotlin/xyz/bluspring/kilt/util/CoroutineUtils.kt b/src/main/kotlin/xyz/bluspring/kilt/util/CoroutineUtils.kt new file mode 100644 index 00000000..244a7fdf --- /dev/null +++ b/src/main/kotlin/xyz/bluspring/kilt/util/CoroutineUtils.kt @@ -0,0 +1,26 @@ +package xyz.bluspring.kilt.util + +import kotlinx.coroutines.async +import kotlinx.coroutines.flow.Flow +import kotlinx.coroutines.flow.channelFlow +import kotlinx.coroutines.flow.map +import kotlinx.coroutines.flow.transform + +// https://github.com/Kotlin/kotlinx.coroutines/issues/1147 + +fun Flow.mapAsync(transform: suspend (value: T) -> R): Flow = + channelFlow { + collect { e -> send(async { transform(e) }) } + }.map { it.await() } + +fun Flow.filterAsync(predicate: suspend (value: T) -> Boolean): Flow = + mapAsync { it to predicate(it) }.transform { if (it.second) emit(it.first) } + +fun Flow.onEachAsync(action: suspend (value: T) -> Unit): Flow = + mapAsync { + action(it) + it + } + +fun Flow.flatMapAsync(transform: suspend (value: T) -> Flow): Flow = + mapAsync { transform(it) }.transform { it.collect { emit(it) } } \ No newline at end of file