Provide a way to execute tasks with correct context classloader

RaymondBlaze · Jun 8, 2024 · 5a30549 · 5a30549
1 parent a095c6c
commit 5a30549
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 1 deletion.
diff --git a/loader/src/main/java/net/neoforged/fml/ModContainer.java b/loader/src/main/java/net/neoforged/fml/ModContainer.java
@@ -11,6 +11,8 @@
 import java.util.IdentityHashMap;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.ForkJoinWorkerThread;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 import net.neoforged.bus.api.BusBuilder;
 import net.neoforged.bus.api.Event;
@@ -150,6 +152,24 @@ protected void constructMod() {}
     @Nullable
     public abstract IEventBus getEventBus();
 
+    /**
+     * Executes a task, ensures correct context classloader if executed in a {@link ForkJoinWorkerThread}.
+     * 
+     * @param task task to accept
+     */
+    public final void execute(Runnable task) {
+        this.execute(container -> task.run());
+    }
+
+    /**
+     * Executes a task, ensures correct context classloader if executed in a {@link ForkJoinWorkerThread}.
+     * 
+     * @param task task to accept
+     */
+    public void execute(Consumer<ModContainer> task) {
+        task.accept(this);
+    }
+
     /**
      * Accept an arbitrary event for processing by the mod. Posted to {@link #getEventBus()}.
      * 

diff --git a/loader/src/main/java/net/neoforged/fml/ModLoader.java b/loader/src/main/java/net/neoforged/fml/ModLoader.java
@@ -232,7 +232,7 @@ public static void dispatchParallelTask(String name, Executor parallelExecutor,
 
                                     try {
                                         ModLoadingContext.get().setActiveContainer(modContainer);
-                                        task.accept(modContainer);
+                                        modContainer.execute(task);
                                     } finally {
                                         progress.increment();
                                         ModLoadingContext.get().setActiveContainer(null);

diff --git a/loader/src/main/java/net/neoforged/fml/javafmlmod/FMLModContainer.java b/loader/src/main/java/net/neoforged/fml/javafmlmod/FMLModContainer.java
@@ -11,6 +11,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ForkJoinWorkerThread;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import net.neoforged.api.distmarker.Dist;
 import net.neoforged.bus.EventBusErrorMessage;
@@ -37,6 +39,7 @@ public class FMLModContainer extends ModContainer {
     private final IEventBus eventBus;
     private final List<Class<?>> modClasses;
     private final Module layer;
+    private final ClassLoader classLoader;
 
     public FMLModContainer(IModInfo info, List<String> entrypoints, ModFileScanData modFileScanResults, ModuleLayer gameLayer) {
         super(info);
@@ -48,6 +51,7 @@ public FMLModContainer(IModInfo info, List<String> entrypoints, ModFileScanData
                 .allowPerPhasePost()
                 .build();
         this.layer = gameLayer.findModule(info.getOwningFile().moduleName()).orElseThrow();
+        this.classLoader = layer.getClassLoader();
 
         final FMLJavaModLoadingContext contextExtension = new FMLJavaModLoadingContext(this);
         this.contextExtension = () -> contextExtension;
@@ -129,4 +133,18 @@ protected void constructMod() {
     public IEventBus getEventBus() {
         return this.eventBus;
     }
+
+    @Override
+    public void execute(Consumer<ModContainer> task) {
+        if (Thread.currentThread() instanceof ForkJoinWorkerThread worker) {
+            final var contextClassLoader = worker.getContextClassLoader();
+            if (classLoader != contextClassLoader) {
+                worker.setContextClassLoader(classLoader);
+                task.accept(this);
+                worker.setContextClassLoader(contextClassLoader);
+                return;
+            }
+        }
+        task.accept(this);
+    }
 }