Make SourceRemapper threaded (#446)

Should be a little bit faster, but nothing magicial.

src/main/java/net/fabricmc/loom/configuration/LoomDependencyManager.java

@@ -25,6 +25,7 @@
 package net.fabricmc.loom.configuration;
 
 import java.io.File;
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -169,7 +170,15 @@ public class LoomDependencyManager {
 
 		ModCompileRemapper.remapDependencies(project, mappingsKey, extension, sourceRemapper);
 
-		sourceRemapper.remapAll();
+		long start = System.currentTimeMillis();
+
+		try {
+			sourceRemapper.remapAll();
+		} catch (IOException exception) {
+			throw new RuntimeException("Failed to remap mod sources", exception);
+		}
+
+		project.getLogger().info("Source remapping took: %dms".formatted(System.currentTimeMillis() - start));
 
 		for (Runnable runnable : afterTasks) {
 			runnable.run();

src/main/java/net/fabricmc/loom/configuration/RemapConfiguration.java

@@ -108,7 +108,6 @@ public class RemapConfiguration {
 
 				rootProject.getTasks().register(remapAllSourcesTaskName, RemapAllSourcesTask.class, task -> {
 					task.sourceRemapper = sourceRemapper;
-					task.doLast(t -> sourceRemapper.remapAll());
 				});
 
 				parentTask = rootProject.getTasks().getByName(remapAllSourcesTaskName);

src/main/java/net/fabricmc/loom/task/RemapAllSourcesTask.java

@@ -24,7 +24,10 @@
 
 package net.fabricmc.loom.task;
 
+import java.io.IOException;
+
 import org.gradle.api.tasks.Internal;
+import org.gradle.api.tasks.TaskAction;
 
 import net.fabricmc.loom.util.SourceRemapper;
 
@@ -35,4 +38,13 @@ public class RemapAllSourcesTask extends AbstractLoomTask {
 	public SourceRemapper getSourceRemapper() {
 		return sourceRemapper;
 	}
+
+	@TaskAction
+	public void remap() {
+		try {
+			sourceRemapper.remapAll();
+		} catch (IOException exception) {
+			throw new RuntimeException("Failed to remap mod", exception);
+		}
+	}
 }

src/main/java/net/fabricmc/loom/util/SourceRemapper.java

@@ -29,20 +29,27 @@ import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
-import java.util.function.Consumer;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentLinkedDeque;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 
 import org.cadixdev.lorenz.MappingSet;
 import org.cadixdev.mercury.Mercury;
+import org.cadixdev.mercury.SourceProcessor;
 import org.cadixdev.mercury.remapper.MercuryRemapper;
 import org.gradle.api.Project;
+import org.gradle.api.plugins.JavaPlugin;
 import org.zeroturnaround.zip.ZipUtil;
 
 import net.fabricmc.loom.LoomGradleExtension;
 import net.fabricmc.loom.configuration.RemappedConfigurationEntry;
 import net.fabricmc.loom.configuration.providers.mappings.MappingsProviderImpl;
-import net.fabricmc.loom.util.gradle.ProgressLogger;
 import net.fabricmc.lorenztiny.TinyMappingsReader;
 import net.fabricmc.mapping.tree.TinyTree;
 import net.fabricmc.stitch.util.StitchUtil;
@@ -50,115 +57,171 @@ import net.fabricmc.stitch.util.StitchUtil;
 public class SourceRemapper {
 	private final Project project;
 	private final boolean toNamed;
-	private final List<Consumer<ProgressLogger>> remapTasks = new ArrayList<>();
+	private final List<RemapTask> remapTasks = new ArrayList<>();
-
+	private boolean hasStartedRemap = false;
-	private Mercury mercury;
 
 	public SourceRemapper(Project project, boolean toNamed) {
 		this.project = project;
 		this.toNamed = toNamed;
 	}
 
-	public static void remapSources(Project project, File input, File output, boolean named, boolean reproducibleFileOrder, boolean preserveFileTimestamps) {
+	public static void remapSources(Project project, File input, File output, boolean named, boolean reproducibleFileOrder, boolean preserveFileTimestamps) throws IOException {
 		SourceRemapper sourceRemapper = new SourceRemapper(project, named);
 		sourceRemapper.scheduleRemapSources(input, output, reproducibleFileOrder, preserveFileTimestamps);
 		sourceRemapper.remapAll();
 	}
 
 	public void scheduleRemapSources(File source, File destination, boolean reproducibleFileOrder, boolean preserveFileTimestamps) {
-		remapTasks.add((logger) -> {
+		this.scheduleRemapSources(new RemapTask(source, destination, reproducibleFileOrder, preserveFileTimestamps));
-			try {
-				logger.progress("remapping sources - " + source.getName());
-				remapSourcesInner(source, destination);
-				ZipReprocessorUtil.reprocessZip(destination, reproducibleFileOrder, preserveFileTimestamps);
-
-				// Set the remapped sources creation date to match the sources if we're likely succeeded in making it
-				destination.setLastModified(source.lastModified());
-			} catch (Exception e) {
-				// Failed to remap, lets clean up to ensure we try again next time
-				destination.delete();
-				throw new RuntimeException("Failed to remap sources for " + source, e);
-			}
-		});
 	}
 
-	public void remapAll() {
+	public synchronized void scheduleRemapSources(RemapTask data) {
+		if (hasStartedRemap) {
+			throw new UnsupportedOperationException("Cannot add data after remapping has started");
+		}
+
+		this.remapTasks.add(data);
+	}
+
+	public void remapAll() throws IOException {
+		hasStartedRemap = true;
+
 		if (remapTasks.isEmpty()) {
 			return;
 		}
 
-		project.getLogger().lifecycle(":remapping sources");
+		int threads = Math.min(remapTasks.size(), Math.max(2, Runtime.getRuntime().availableProcessors() / 2));
+		ExecutorService ioExecutor = Executors.newFixedThreadPool(2);
+		ExecutorService remapExecutor = Executors.newFixedThreadPool(threads);
 
-		ProgressLogger progressLogger = ProgressLogger.getProgressFactory(project, SourceRemapper.class.getName());
+		List<CompletableFuture<Void>> completableFutures = new ArrayList<>();
-		progressLogger.start("Remapping dependency sources", "sources");
 
-		remapTasks.forEach(consumer -> consumer.accept(progressLogger));
+		{
+			// We have to build the Mercury instances on the main thread as createMercuryRemapper resolves gradle stuff
+			// TODO refactor this a bit to not do that.
+			var mercuryQueue = new ConcurrentLinkedDeque<Mercury>();
 
-		progressLogger.completed();
+			final var remapper = createMercuryRemapper();
+			final var inputClasspath = getInputClasspath(project);
+
+			for (int i = 0; i < threads; i++) {
+				Mercury mercury = createMercuryWithClassPath(project, toNamed);
+				mercury.getProcessors().add(remapper);
+				mercury.getClassPath().addAll(inputClasspath);
+
+				mercuryQueue.add(mercury);
+			}
+
+			ThreadLocal<Mercury> mercuryThreadLocal = ThreadLocal.withInitial(() -> Objects.requireNonNull(mercuryQueue.pop(), "No Mercury instances left"));
+
+			for (RemapTask remapTask : this.remapTasks) {
+				completableFutures.add(
+						CompletableFuture.supplyAsync(() ->
+							prepareForRemap(remapTask), ioExecutor)
+							.thenApplyAsync(remapData -> doRemap(mercuryThreadLocal.get(), remapData), remapExecutor)
+							.thenAcceptAsync(remapData -> completeRemap(remapData, remapTask), ioExecutor)
+				);
+			}
+		}
+
+		for (CompletableFuture<Void> future : completableFutures) {
+			try {
+				future.join();
+			} catch (CompletionException e) {
+				try {
+					throw e.getCause();
+				} catch (IOException ioe) {
+					throw ioe;
+				} catch (Throwable unknown) {
+					throw new RuntimeException("An unknown exception occured when remapping", unknown);
+				}
+			}
+		}
+
+		ioExecutor.shutdown();
+		remapExecutor.shutdown();
 
 		// TODO: FIXME - WORKAROUND https://github.com/FabricMC/fabric-loom/issues/45
 		System.gc();
 	}
 
-	private void remapSourcesInner(File source, File destination) throws Exception {
+	private RemapData prepareForRemap(RemapTask remapTask) {
-		project.getLogger().info(":remapping source jar");
-		Mercury mercury = getMercuryInstance();
-
-		if (source.equals(destination)) {
-			if (source.isDirectory()) {
-				throw new RuntimeException("Directories must differ!");
-			}
-
-			source = new File(destination.getAbsolutePath().substring(0, destination.getAbsolutePath().lastIndexOf('.')) + "-dev.jar");
-
-			try {
-				com.google.common.io.Files.move(destination, source);
-			} catch (IOException e) {
-				throw new RuntimeException("Could not rename " + destination.getName() + "!", e);
-			}
-		}
-
-		Path srcPath = source.toPath();
-		boolean isSrcTmp = false;
-
-		if (!source.isDirectory()) {
-			// create tmp directory
-			isSrcTmp = true;
-			srcPath = Files.createTempDirectory("fabric-loom-src");
-			ZipUtil.unpack(source, srcPath.toFile());
-		}
-
-		if (!destination.isDirectory() && destination.exists()) {
-			if (!destination.delete()) {
-				throw new RuntimeException("Could not delete " + destination.getName() + "!");
-			}
-		}
-
-		StitchUtil.FileSystemDelegate dstFs = destination.isDirectory() ? null : StitchUtil.getJarFileSystem(destination, true);
-		Path dstPath = dstFs != null ? dstFs.get().getPath("/") : destination.toPath();
-
 		try {
-			mercury.rewrite(srcPath, dstPath);
+			File source = remapTask.source();
-		} catch (Exception e) {
+			final File destination = remapTask.destination();
-			project.getLogger().warn("Could not remap " + source.getName() + " fully!", e);
-		}
 
-		copyNonJavaFiles(srcPath, dstPath, project, source);
+			if (source.equals(destination)) {
+				if (source.isDirectory()) {
+					throw new RuntimeException("Directories must differ!");
+				}
 
-		if (dstFs != null) {
+				source = new File(destination.getAbsolutePath().substring(0, destination.getAbsolutePath().lastIndexOf('.')) + "-dev.jar");
-			dstFs.close();
-		}
 
-		if (isSrcTmp) {
+				try {
-			Files.walkFileTree(srcPath, new DeletingFileVisitor());
+					com.google.common.io.Files.move(destination, source);
+				} catch (IOException e) {
+					throw new RuntimeException("Could not rename " + destination.getName() + "!", e);
+				}
+			}
+
+			Path srcPath = source.toPath();
+			boolean isSrcTmp = false;
+
+			if (!source.isDirectory()) {
+				// create tmp directory
+				isSrcTmp = true;
+				srcPath = Files.createTempDirectory("fabric-loom-src");
+				ZipUtil.unpack(source, srcPath.toFile());
+			}
+
+			if (!destination.isDirectory() && destination.exists()) {
+				if (!destination.delete()) {
+					throw new RuntimeException("Could not delete " + destination.getName() + "!");
+				}
+			}
+
+			StitchUtil.FileSystemDelegate dstFs = remapTask.destination().isDirectory() ? null : StitchUtil.getJarFileSystem(remapTask.destination(), true);
+			Path dstPath = dstFs != null ? dstFs.get().getPath("/") : remapTask.destination().toPath();
+
+			return new RemapData(Objects.requireNonNull(srcPath, "srcPath"), Objects.requireNonNull(dstPath, "dstPath"), dstFs, isSrcTmp);
+		} catch (IOException e) {
+			throw new CompletionException("Failed to prepare for remap", e);
 		}
 	}
 
-	private Mercury getMercuryInstance() {
+	private RemapData doRemap(Mercury mercury, RemapData remapData) {
-		if (this.mercury != null) {
+		try {
-			return this.mercury;
+			mercury.rewrite(remapData.source(), remapData.destination());
+		} catch (Exception e) {
+			project.getLogger().warn("Could not remap " + remapData.source().toString() + " fully!", e);
+			// TODO do something more with this error? delete the dst file in complete remap?
 		}
 
+		return remapData;
+	}
+
+	private void completeRemap(RemapData remapData, RemapTask remapTask) {
+		try {
+			copyNonJavaFiles(remapData.source(), remapData.destination(), project, remapTask.source());
+
+			if (remapData.dstFs() != null) {
+				remapData.dstFs().close();
+			}
+
+			if (remapData.isSrcTmp()) {
+				Files.walkFileTree(remapData.source(), new DeletingFileVisitor());
+			}
+
+			ZipReprocessorUtil.reprocessZip(remapTask.destination(), remapTask.reproducibleFileOrder(), remapTask.preserveFileTimestamps());
+
+			// Set the remapped sources creation date to match the sources if we're likely succeeded in making it
+			remapTask.destination().setLastModified(remapTask.source().lastModified());
+		} catch (IOException e) {
+			throw new CompletionException("Failed to complete remap", e);
+		}
+	}
+
+	private SourceProcessor createMercuryRemapper() {
 		LoomGradleExtension extension = LoomGradleExtension.get(project);
 		MappingsProviderImpl mappingsProvider = extension.getMappingsProvider();
 
@@ -172,35 +235,7 @@ public class SourceRemapper {
 			}
 		});
 
-		Mercury mercury = extension.getOrCreateSrcMercuryCache(toNamed ? 1 : 0, () -> {
+		return MercuryRemapper.create(mappings);
-			Mercury m = createMercuryWithClassPath(project, toNamed);
-
-			for (File file : extension.getUnmappedModCollection()) {
-				Path path = file.toPath();
-
-				if (Files.isRegularFile(path)) {
-					m.getClassPath().add(path);
-				}
-			}
-
-			m.getClassPath().add(extension.getMinecraftMappedProvider().getMappedJar().toPath());
-			m.getClassPath().add(extension.getMinecraftMappedProvider().getIntermediaryJar().toPath());
-
-			Set<File> files = project.getConfigurations()
-					.detachedConfiguration(project.getDependencies().create(Constants.Dependencies.JETBRAINS_ANNOTATIONS + Constants.Dependencies.Versions.JETBRAINS_ANNOTATIONS))
-					.resolve();
-
-			for (File file : files) {
-				m.getClassPath().add(file.toPath());
-			}
-
-			m.getProcessors().add(MercuryRemapper.create(mappings));
-
-			return m;
-		});
-
-		this.mercury = mercury;
-		return mercury;
 	}
 
 	private static void copyNonJavaFiles(Path from, Path to, Project project, File source) throws IOException {
@@ -218,26 +253,62 @@ public class SourceRemapper {
 	}
 
 	public static Mercury createMercuryWithClassPath(Project project, boolean toNamed) {
-		Mercury m = new Mercury();
+		var mercury = new Mercury();
-		m.setGracefulClasspathChecks(true);
+		mercury.setGracefulClasspathChecks(true);
+
+		var classpath = mercury.getClassPath();
+		classpath.addAll(getCompileClasspath(project, toNamed));
+
+		return mercury;
+	}
+
+	private static Set<Path> getCompileClasspath(Project project, boolean toNamed) {
+		var classpath = new HashSet<Path>();
 
 		for (File file : project.getConfigurations().getByName(Constants.Configurations.LOADER_DEPENDENCIES).getFiles()) {
-			m.getClassPath().add(file.toPath());
+			classpath.add(file.toPath());
 		}
 
 		if (!toNamed) {
-			for (File file : project.getConfigurations().getByName("compileClasspath").getFiles()) {
+			for (File file : project.getConfigurations().getByName(JavaPlugin.COMPILE_CLASSPATH_CONFIGURATION_NAME).getFiles()) {
-				m.getClassPath().add(file.toPath());
+				classpath.add(file.toPath());
 			}
 		} else {
 			for (RemappedConfigurationEntry entry : Constants.MOD_COMPILE_ENTRIES) {
 				for (File inputFile : project.getConfigurations().getByName(entry.sourceConfiguration()).getFiles()) {
-					m.getClassPath().add(inputFile.toPath());
+					classpath.add(inputFile.toPath());
 				}
 			}
 		}
 
-		return m;
+		return classpath;
+	}
+
+	private static Set<Path> getInputClasspath(Project project) {
+		LoomGradleExtension extension = LoomGradleExtension.get(project);
+
+		var classpath = new HashSet<Path>();
+
+		for (File file : extension.getUnmappedModCollection()) {
+			Path path = file.toPath();
+
+			if (Files.isRegularFile(path)) {
+				classpath.add(path);
+			}
+		}
+
+		classpath.add(extension.getMinecraftMappedProvider().getMappedJar().toPath());
+		classpath.add(extension.getMinecraftMappedProvider().getIntermediaryJar().toPath());
+
+		Set<File> files = project.getConfigurations()
+				.detachedConfiguration(project.getDependencies().create(Constants.Dependencies.JETBRAINS_ANNOTATIONS + Constants.Dependencies.Versions.JETBRAINS_ANNOTATIONS))
+				.resolve();
+
+		for (File file : files) {
+			classpath.add(file.toPath());
+		}
+
+		return classpath;
 	}
 
 	private static boolean isJavaFile(Path path) {
@@ -245,4 +316,10 @@ public class SourceRemapper {
 		// ".java" is not a valid java file
 		return name.endsWith(".java") && name.length() != 5;
 	}
+
+	public static record RemapTask(File source, File destination, boolean reproducibleFileOrder, boolean preserveFileTimestamps) {
+	}
+
+	public static record RemapData(Path source, Path destination, StitchUtil.FileSystemDelegate dstFs, boolean isSrcTmp) {
+	}
 }