Early loading of Mixin and code modifiers

+ System property to disable early loading if necessary

src/main/java/net/minestom/server/Bootstrap.java

@@ -1,5 +1,6 @@
 package net.minestom.server;
 
+import net.minestom.server.extensions.ExtensionManager;
 import net.minestom.server.extras.selfmodification.MinestomRootClassLoader;
 import net.minestom.server.extras.selfmodification.mixins.MixinCodeModifier;
 import net.minestom.server.extras.selfmodification.mixins.MixinServiceMinestom;
@@ -22,6 +23,8 @@ public final class Bootstrap {
             startMixin(args);
             MinestomRootClassLoader.getInstance().addCodeModifier(new MixinCodeModifier());
 
+            ExtensionManager.loadCodeModifiersEarly();
+
             MixinServiceMinestom.gotoPreinitPhase();
             // ensure extensions are loaded when starting the server
             Class<?> serverClass = classLoader.loadClass("net.minestom.server.MinecraftServer");

src/main/java/net/minestom/server/extensions/ExtensionManager.java

@@ -27,6 +27,8 @@ import java.util.zip.ZipFile;
 
 public class ExtensionManager {
 
+    public final static String DISABLE_EARLY_LOAD_SYSTEM_KEY = "minestom.extension.disable_early_load";
+
     public final static Logger LOGGER = LoggerFactory.getLogger(ExtensionManager.class);
 
     private final static String INDEV_CLASSES_FOLDER = "minestom.extension.indevfolder.classes";
@@ -469,7 +471,7 @@ public class ExtensionManager {
     }
 
     @NotNull
-    public Map<String, URLClassLoader> getExtensionLoaders() {
+    public Map<String, MinestomExtensionClassLoader> getExtensionLoaders() {
         return new HashMap<>(extensionLoaders);
     }
 
@@ -485,6 +487,10 @@ public class ExtensionManager {
             return;
         }
         MinestomRootClassLoader modifiableClassLoader = (MinestomRootClassLoader) cl;
+        setupCodeModifiers(extensions, modifiableClassLoader);
+    }
+
+    private void setupCodeModifiers(@NotNull List<DiscoveredExtension> extensions, MinestomRootClassLoader modifiableClassLoader) {
         LOGGER.info("Start loading code modifiers...");
         for (DiscoveredExtension extension : extensions) {
             try {
@@ -664,4 +670,33 @@ public class ExtensionManager {
     public void shutdown() {
         this.extensionList.forEach(this::unload);
     }
+
+    /**
+     * Loads code modifiers early, that is before <code>MinecraftServer.init()</code> is called.
+     */
+    public static void loadCodeModifiersEarly() {
+        // allow users to disable early code modifier load
+        if("true".equalsIgnoreCase(System.getProperty(DISABLE_EARLY_LOAD_SYSTEM_KEY))) {
+            return;
+        }
+        LOGGER.info("Early load of code modifiers from extensions.");
+        ExtensionManager manager = new ExtensionManager();
+
+        // discover extensions that are present
+        List<DiscoveredExtension> discovered = manager.discoverExtensions();
+
+        // setup extension class loaders, so that Mixin can load the json configuration file correctly
+        for(DiscoveredExtension e : discovered) {
+            manager.setupClassLoader(e);
+        }
+
+        // setup code modifiers and mixins
+        manager.setupCodeModifiers(discovered, MinestomRootClassLoader.getInstance());
+
+        // setup is done, remove all extension classloaders
+        for(MinestomExtensionClassLoader extensionLoader : manager.getExtensionLoaders().values()) {
+            MinestomRootClassLoader.getInstance().removeChildInHierarchy(extensionLoader);
+        }
+        LOGGER.info("Early load of code modifiers from extensions done!");
+    }
 }

src/main/java/net/minestom/server/extras/selfmodification/MinestomRootClassLoader.java

@@ -65,6 +65,11 @@ public class MinestomRootClassLoader extends HierarchyClassLoader {
     // TODO: priorities?
     private final List<CodeModifier> modifiers = new LinkedList<>();
 
+    /**
+     * List of already loaded code modifier class names. This prevents loading the same class twice.
+     */
+    private final Set<String> alreadyLoadedCodeModifiers = new HashSet<>();
+
     private MinestomRootClassLoader(ClassLoader parent) {
         super("Minestom Root ClassLoader", extractURLsFromClasspath(), parent);
         asmClassLoader = newChild(new URL[0]);
@@ -246,6 +251,9 @@ public class MinestomRootClassLoader extends HierarchyClassLoader {
     }
 
     public void loadModifier(File[] originFiles, String codeModifierClass) {
+        if(alreadyLoadedCodeModifiers.contains(codeModifierClass)) {
+            return;
+        }
         URL[] urls = new URL[originFiles.length];
         try {
             for (int i = 0; i < originFiles.length; i++) {
@@ -258,6 +266,7 @@ public class MinestomRootClassLoader extends HierarchyClassLoader {
                 synchronized (modifiers) {
                     LOGGER.warn("Added Code modifier: {}", modifier);
                     addCodeModifier(modifier);
+                    alreadyLoadedCodeModifiers.add(codeModifierClass);
                 }
             }
         } catch (MalformedURLException | ClassNotFoundException | InvocationTargetException | InstantiationException | IllegalAccessException | NoSuchMethodException e) {

src/test/java/improveextensions/DisableEarlyLoad.java

@@ -0,0 +1,34 @@
+package improveextensions;
+
+import net.minestom.server.MinecraftServer;
+import net.minestom.server.extensions.Extension;
+import net.minestom.server.instance.InstanceContainer;
+import net.minestom.server.world.DimensionType;
+import org.junit.jupiter.api.Assertions;
+import org.opentest4j.AssertionFailedError;
+
+import java.util.UUID;
+
+/**
+ * Extensions should be able to use Mixins for classes loaded very early by Minestom (InstanceContainer for instance)
+ */
+public class DisableEarlyLoad extends Extension {
+
+    @Override
+    public void initialize() {
+        // force load of InstanceContainer class
+        InstanceContainer c = new InstanceContainer(UUID.randomUUID(), DimensionType.OVERWORLD, null);
+        System.out.println(c.toString());
+        try {
+            Assertions.assertFalse(MixinIntoMinestomCore.success, "InstanceContainer must NOT have been mixed in with improveextensions.InstanceContainerMixin, because early loading has been disabled");
+        } catch (AssertionFailedError e) {
+            e.printStackTrace();
+        }
+        MinecraftServer.stopCleanly();
+    }
+
+    @Override
+    public void terminate() {
+        getLogger().info("Terminate extension");
+    }
+}

src/test/java/improveextensions/MixinIntoMinestomCore.java

@@ -0,0 +1,37 @@
+package improveextensions;
+
+import net.minestom.server.MinecraftServer;
+import net.minestom.server.extensions.Extension;
+import net.minestom.server.instance.InstanceContainer;
+import net.minestom.server.world.DimensionType;
+import org.junit.jupiter.api.Assertions;
+import org.opentest4j.AssertionFailedError;
+
+import java.util.UUID;
+
+/**
+ * Extensions should be able to use Mixins for classes loaded very early by Minestom (InstanceContainer for instance)
+ */
+public class MixinIntoMinestomCore extends Extension {
+
+    public static boolean success = false;
+
+    @Override
+    public void initialize() {
+        // force load of InstanceContainer class
+        InstanceContainer c = new InstanceContainer(UUID.randomUUID(), DimensionType.OVERWORLD, null);
+        System.out.println(c.toString());
+        try {
+            Assertions.assertTrue(success, "InstanceContainer must have been mixed in with improveextensions.InstanceContainerMixin");
+            Assertions.assertEquals(1, MinecraftServer.getExtensionManager().getExtensionLoaders().size(), "Only one extension classloader (this extension's) must be active.");
+        } catch (AssertionFailedError e) {
+            e.printStackTrace();
+        }
+        MinecraftServer.stopCleanly();
+    }
+
+    @Override
+    public void terminate() {
+        getLogger().info("Terminate extension");
+    }
+}

src/test/java/improveextensions/MixinIntoMinestomCoreLauncher.java

@@ -0,0 +1,15 @@
+package improveextensions;
+
+import net.minestom.server.Bootstrap;
+
+// To launch with VM arguments:
+
+// To test early Mixin injections:
+// -Dminestom.extension.indevfolder.classes=build/classes/java/test/ -Dminestom.extension.indevfolder.resources=build/resources/test/improveextensions
+// To test disabling early Mixin injections:
+// -Dminestom.extension.disable_early_load=true -Dminestom.extension.indevfolder.classes=build/classes/java/test/ -Dminestom.extension.indevfolder.resources=build/resources/test/improveextensions/disableearlyload
+public class MixinIntoMinestomCoreLauncher {
+     public static void main(String[] args) {
+            Bootstrap.bootstrap("demo.MainDemo", args);
+        }
+}

src/test/java/improveextensions/mixins/InstanceContainerMixin.java

@@ -0,0 +1,19 @@
+package improveextensions.mixins;
+
+import improveextensions.MixinIntoMinestomCore;
+import net.minestom.server.instance.InstanceContainer;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+@Mixin(InstanceContainer.class)
+public class InstanceContainerMixin {
+
+    @Inject(method = "<init>", at = @At("RETURN"), require = 1)
+    private void constructorHead(CallbackInfo ci) {
+        System.out.println("Mixin into InstanceContainerMixin");
+        MixinIntoMinestomCore.success = true;
+    }
+
+}

src/test/resources/improveextensions/disableearlyload/extension.json

@@ -0,0 +1,6 @@
+{
+  "entrypoint": "improveextensions.DisableEarlyLoad",
+  "name": "DisableEarlyLoad",
+  "codeModifiers": [],
+  "mixinConfig": "mixins.minestomcore.json"
+}

src/test/resources/improveextensions/disableearlyload/mixins.minestomcore.json

@@ -0,0 +1,10 @@
+{
+  "required": true,
+  "minVersion": "0.8",
+  "package": "improveextensions.mixins",
+  "target": "@env(DEFAULT)",
+  "compatibilityLevel": "JAVA_11",
+  "mixins": [
+    "InstanceContainerMixin"
+  ]
+}

src/test/resources/improveextensions/extension.json

@@ -0,0 +1,6 @@
+{
+  "entrypoint": "improveextensions.MixinIntoMinestomCore",
+  "name": "MixinIntoMinestomCore",
+  "codeModifiers": [],
+  "mixinConfig": "mixins.minestomcore.json"
+}

src/test/resources/improveextensions/mixins.minestomcore.json

@@ -0,0 +1,10 @@
+{
+  "required": true,
+  "minVersion": "0.8",
+  "package": "improveextensions.mixins",
+  "target": "@env(DEFAULT)",
+  "compatibilityLevel": "JAVA_11",
+  "mixins": [
+    "InstanceContainerMixin"
+  ]
+}