Backport annotation scanner parallelization

Related: neoforged/FancyModLoader#32
TauMC · Nov 2, 2023 · f0b29cf · f0b29cf
1 parent 9b5ef9e
commit f0b29cf
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 1 deletion.
diff --git a/main/src/main/java/org/embeddedt/blacksmith/impl/TransformerCore.java b/main/src/main/java/org/embeddedt/blacksmith/impl/TransformerCore.java
@@ -37,6 +37,7 @@ public class TransformerCore {
         TRANSFORMERS.add(new FinalFieldHelperTransformer());
         TRANSFORMERS.add(new ModuleLayerHandlerTransformer());
         TRANSFORMERS.add(new FileWatcherTransformer());
+        TRANSFORMERS.add(new BackgroundScanHandlerTransformer());
     }
 
     public static void log(String s) {

diff --git a/main/src/main/java/org/embeddedt/blacksmith/impl/hooks/Hooks.java b/main/src/main/java/org/embeddedt/blacksmith/impl/hooks/Hooks.java
@@ -16,6 +16,10 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Predicate;
 import java.util.jar.JarFile;
 import java.util.jar.Manifest;
@@ -44,6 +48,17 @@ public static <T> Stream<T> dummyFilter(Stream<T> stream, Predicate<T> filter) {
         return stream;
     }
 
+    public static ExecutorService makeScanningExecutor(ThreadFactory factory) {
+        int maxScanThreads = Math.max(Runtime.getRuntime().availableProcessors() - 1, 1);
+        System.out.println("Using " + maxScanThreads + " threads for scanner");
+        AtomicInteger tCount = new AtomicInteger();
+        return Executors.newFixedThreadPool(maxScanThreads, r -> {
+            Thread t = factory.newThread(r);
+            t.setName("Scan-Handler-" + tCount.getAndIncrement());
+            return t;
+        });
+    }
+
     public static Stream<Path> getPackagesSkippingAssets(Path basePath, FileVisitOption[] options) throws IOException {
         List<Path> validPaths = new ArrayList<>();
         Files.walkFileTree(basePath, new SimpleFileVisitor<Path>() {

diff --git a/...ain/java/org/embeddedt/blacksmith/impl/transformers/BackgroundScanHandlerTransformer.java b/...ain/java/org/embeddedt/blacksmith/impl/transformers/BackgroundScanHandlerTransformer.java
@@ -0,0 +1,37 @@
+package org.embeddedt.blacksmith.impl.transformers;
+
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.tree.*;
+
+import java.lang.instrument.IllegalClassFormatException;
+import java.util.Collections;
+import java.util.List;
+
+public class BackgroundScanHandlerTransformer implements RuntimeTransformer {
+    @Override
+    public List<String> getTransformedClasses() {
+        return Collections.singletonList("net/minecraftforge/fml/loading/moddiscovery/BackgroundScanHandler");
+    }
+
+    @Override
+    public void transformClass(ClassNode data) throws IllegalClassFormatException {
+        for(MethodNode m : data.methods) {
+            if(m.name.equals("<init>")) {
+                for(AbstractInsnNode n : m.instructions) {
+                    if(n.getOpcode() == Opcodes.INVOKESTATIC) {
+                        MethodInsnNode mNode = (MethodInsnNode)n;
+                        if(mNode.name.equals("newSingleThreadExecutor")) {
+                            mNode.name = "makeScanningExecutor";
+                            mNode.owner = RuntimeTransformer.HOOK_CLASS;
+                            System.out.println("Parallelized scanner");
+                            break;
+                        }
+                    }
+                }
+            } else if(m.name.equals("addCompletedFile")) {
+                /* synchronize */
+                m.access |= Opcodes.ACC_SYNCHRONIZED;
+            }
+        }
+    }
+}
diff --git a/main/src/main/java/org/embeddedt/blacksmith/impl/transformers/RuntimeTransformer.java b/main/src/main/java/org/embeddedt/blacksmith/impl/transformers/RuntimeTransformer.java
@@ -10,13 +10,14 @@
 import java.util.List;
 
 public interface RuntimeTransformer {
+    String HOOK_CLASS = "org/embeddedt/blacksmith/impl/hooks/Hooks";
     List<String> getTransformedClasses();
     void transformClass(ClassNode data) throws IllegalClassFormatException;
 
     default int getWriteFlags() { return 0; }
 
     static MethodInsnNode redirectToStaticHook(String hookName, String hookDesc) {
-        return new MethodInsnNode(Opcodes.INVOKESTATIC, "org/embeddedt/blacksmith/impl/hooks/Hooks", hookName, hookDesc, false);
+        return new MethodInsnNode(Opcodes.INVOKESTATIC, HOOK_CLASS, hookName, hookDesc, false);
     }
 
     static <T extends AbstractInsnNode> T swapInstruction(InsnList list, AbstractInsnNode oldInsn, T newInsn) {