Automatically unload GlobalEventHandler callbacks

src/main/java/net/minestom/server/entity/Entity.java

@@ -120,6 +120,7 @@ public abstract class Entity implements Viewable, EventHandler, DataContainer, P
 
     // Events
     private final Map<Class<? extends Event>, Collection<EventCallback>> eventCallbacks = new ConcurrentHashMap<>();
+    private final Map<String, Collection<EventCallback<?>>> extensionCallbacks = new ConcurrentHashMap<>();
 
     protected Metadata metadata = new Metadata(this);
 
@@ -640,6 +641,12 @@ public abstract class Entity implements Viewable, EventHandler, DataContainer, P
         return eventCallbacks;
     }
 
+    @NotNull
+    @Override
+    public Collection<EventCallback<?>> getExtensionCallbacks(String extension) {
+        return extensionCallbacks.computeIfAbsent(extension, e -> new CopyOnWriteArrayList<>());
+    }
+
     /**
      * Each entity has an unique id (server-wide) which will change after a restart.
      *

src/main/java/net/minestom/server/event/GlobalEventHandler.java

@@ -6,6 +6,7 @@ import org.jetbrains.annotations.NotNull;
 import java.util.Collection;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
 
 /**
  * Object containing all the global event listeners.
@@ -14,10 +15,17 @@ public final class GlobalEventHandler implements EventHandler {
 
     // Events
     private final Map<Class<? extends Event>, Collection<EventCallback>> eventCallbacks = new ConcurrentHashMap<>();
+    private final Map<String, Collection<EventCallback<?>>> extensionCallbacks = new ConcurrentHashMap<>();
 
     @NotNull
     @Override
     public Map<Class<? extends Event>, Collection<EventCallback>> getEventCallbacksMap() {
         return eventCallbacks;
     }
+
+    @NotNull
+    @Override
+    public Collection<EventCallback<?>> getExtensionCallbacks(String extension) {
+        return extensionCallbacks.computeIfAbsent(extension, e -> new CopyOnWriteArrayList<>());
+    }
 }

src/main/java/net/minestom/server/event/handler/EventHandler.java

@@ -6,12 +6,16 @@ import net.minestom.server.event.CancellableEvent;
 import net.minestom.server.event.Event;
 import net.minestom.server.event.EventCallback;
 import net.minestom.server.event.GlobalEventHandler;
+import net.minestom.server.extensions.ExtensionManager;
+import net.minestom.server.extras.selfmodification.MinestomExtensionClassLoader;
+import net.minestom.server.extras.selfmodification.MinestomRootClassLoader;
 import net.minestom.server.instance.Instance;
 import net.minestom.server.utils.validate.Check;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.Collection;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.CopyOnWriteArraySet;
 import java.util.stream.Stream;
 
@@ -28,6 +32,38 @@ public interface EventHandler {
     @NotNull
     Map<Class<? extends Event>, Collection<EventCallback>> getEventCallbacksMap();
 
+    /**
+     * Gets a {@link Collection} containing all the listeners assigned to a specific extension (represented by its name).
+     * Used to unload all callbacks when the extension is unloaded
+     *
+     * @return a {@link Collection} with all the listeners
+     */
+    @NotNull
+    Collection<EventCallback<?>> getExtensionCallbacks(String extension);
+
+    /**
+     * Tries to know which extension created this callback, based on the classloader of the callback
+     * @param callback the callback to get the extension of
+     * @return <code>Optional.empty()</code> if no extension has been found, <code>Optional.of(&lt;name&gt;)</code> with 'name' being the extension name
+     */
+    static Optional<String> getExtensionOwningCallback(@NotNull EventCallback<?> callback) {
+        ClassLoader cl = callback.getClass().getClassLoader();
+        if(cl instanceof MinestomExtensionClassLoader) {
+            return Optional.of(((MinestomExtensionClassLoader) cl).getExtensionName());
+        } else if(System.getProperty(ExtensionManager.INDEV_CLASSES_FOLDER) != null) { // in a dev environment, the extension will always be loaded with the root classloader
+            StackTraceElement[] stackTrace = Thread.currentThread().getStackTrace();
+            // 0 -> getStackTrace
+            // 1 -> getExtensionOwningCallback
+            // 2 -> add/remove EventCallback
+            // 3 -> Potentially the extension
+            if(stackTrace.length >= 4) {
+                StackTraceElement potentialExtensionCall = stackTrace[3];
+                System.out.println(potentialExtensionCall.getClassLoaderName());
+            }
+        }
+        return Optional.empty();
+    }
+
     /**
      * Adds a new event callback for the specified type {@code eventClass}.
      *
@@ -37,6 +73,9 @@ public interface EventHandler {
      * @return true if the callback collection changed as a result of the call
      */
     default <E extends Event> boolean addEventCallback(@NotNull Class<E> eventClass, @NotNull EventCallback<E> eventCallback) {
+        Optional<String> extensionSource = getExtensionOwningCallback(eventCallback);
+        extensionSource.ifPresent(s -> getExtensionCallbacks(s).add(eventCallback));
+
         Collection<EventCallback> callbacks = getEventCallbacks(eventClass);
         return callbacks.add(eventCallback);
     }
@@ -51,6 +90,9 @@ public interface EventHandler {
      */
     default <E extends Event> boolean removeEventCallback(@NotNull Class<E> eventClass, @NotNull EventCallback<E> eventCallback) {
         Collection<EventCallback> callbacks = getEventCallbacks(eventClass);
+        Optional<String> extensionSource = getExtensionOwningCallback(eventCallback);
+        extensionSource.ifPresent(s -> getExtensionCallbacks(s).remove(eventCallback));
+
         return callbacks.remove(eventCallback);
     }
 
@@ -126,6 +168,24 @@ public interface EventHandler {
         }
     }
 
+    /**
+     * Remove all event callbacks owned by the given extension
+     * @param extension the extension to remove callbacks from
+     */
+    default void removeCallbacksOwnedByExtension(String extension) {
+        Collection<EventCallback<?>> extensionCallbacks = getExtensionCallbacks(extension);
+        for(EventCallback<?> callback : extensionCallbacks) {
+
+            // try to remove this callback from all callback collections
+            //  we do this because we do not have information about the event class at this point
+            for(Collection<EventCallback> eventCallbacks : getEventCallbacksMap().values()) {
+                eventCallbacks.remove(callback);
+            }
+        }
+
+        extensionCallbacks.clear();
+    }
+
     private <E extends Event> void runEvent(@NotNull Collection<EventCallback> eventCallbacks, @NotNull E event) {
         for (EventCallback<E> eventCallback : eventCallbacks) {
             eventCallback.run(event);

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

@@ -32,8 +32,8 @@ public class ExtensionManager {
 
     public final static Logger LOGGER = LoggerFactory.getLogger(ExtensionManager.class);
 
-    private final static String INDEV_CLASSES_FOLDER = "minestom.extension.indevfolder.classes";
+    public final static String INDEV_CLASSES_FOLDER = "minestom.extension.indevfolder.classes";
-    private final static String INDEV_RESOURCES_FOLDER = "minestom.extension.indevfolder.resources";
+    public final static String INDEV_RESOURCES_FOLDER = "minestom.extension.indevfolder.resources";
     private final static Gson GSON = new Gson();
 
     private final Map<String, MinestomExtensionClassLoader> extensionLoaders = new HashMap<>();
@@ -432,7 +432,7 @@ public class ExtensionManager {
     @NotNull
     public MinestomExtensionClassLoader newClassLoader(@NotNull DiscoveredExtension extension, @NotNull URL[] urls) {
         MinestomRootClassLoader root = MinestomRootClassLoader.getInstance();
-        MinestomExtensionClassLoader loader = new MinestomExtensionClassLoader(extension.getName(), urls, root);
+        MinestomExtensionClassLoader loader = new MinestomExtensionClassLoader(extension.getName(), extension.getEntrypoint(), urls, root);
         if (extension.getDependencies().length == 0) {
             // orphaned extension, we can insert it directly
             root.addChild(loader);
@@ -523,6 +523,10 @@ public class ExtensionManager {
     private void unload(Extension ext) {
         ext.preTerminate();
         ext.terminate();
+        // remove callbacks for this extension
+        MinecraftServer.getGlobalEventHandler().removeCallbacksOwnedByExtension(ext.getDescription().getName());
+        // TODO: more callback types
+
         ext.postTerminate();
         ext.unload();
 

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

@@ -10,9 +10,34 @@ public class MinestomExtensionClassLoader extends HierarchyClassLoader {
      */
     private final MinestomRootClassLoader root;
 
-    public MinestomExtensionClassLoader(String name, URL[] urls, MinestomRootClassLoader root) {
+    /**
-        super(name, urls, root);
+     * Main of the main class of the extension linked to this classloader
+     */
+    private final String mainClassName;
+
+    public MinestomExtensionClassLoader(String extensionName, String mainClassName, URL[] urls, MinestomRootClassLoader root) {
+        super(extensionName, urls, root);
         this.root = root;
+        this.mainClassName = mainClassName;
+    }
+
+    /**
+     * Returns the name of the extension linked to this classloader
+     * @return the name of the extension linked to this classloader
+     */
+    public String getExtensionName() {
+        // simply calls ClassLoader#getName as the extension name is used to name this classloader
+        //  this method is simply for ease-of-use
+        return getName();
+    }
+
+    /**
+     * Returns the main class name linked to the extension responsible for this classloader.
+     * Used by the root classloader to let extensions load themselves in a dev environment.
+     * @return the main class name linked to the extension responsible for this classloader
+     */
+    public String getMainClassName() {
+        return mainClassName;
     }
 
     @Override
@@ -72,4 +97,16 @@ public class MinestomExtensionClassLoader extends HierarchyClassLoader {
         super.finalize();
         System.err.println("Class loader "+getName()+" finalized.");
     }
+
+    /**
+     * Is the given class name the name of the entry point of one the extensions from this classloader chain?
+     * @param name the class name to check
+     * @return whether the given class name the name of the entry point of one the extensions from this classloader chain
+     * @see MinestomRootClassLoader#loadBytes(String, boolean) for more information
+     */
+    protected boolean isMainExtensionClass(String name) {
+        if(mainClassName.equals(name))
+            return true;
+        return children.stream().anyMatch(c -> c.isMainExtensionClass(name));
+    }
 }

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

@@ -1,6 +1,7 @@
 package net.minestom.server.extras.selfmodification;
 
 import net.minestom.server.MinecraftServer;
+import net.minestom.server.extensions.ExtensionManager;
 import org.jetbrains.annotations.NotNull;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.ClassWriter;
@@ -13,6 +14,7 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.lang.reflect.InvocationTargetException;
 import java.net.MalformedURLException;
+import java.net.URISyntaxException;
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.util.HashSet;
@@ -65,6 +67,12 @@ public class MinestomRootClassLoader extends HierarchyClassLoader {
     // TODO: priorities?
     private final List<CodeModifier> modifiers = new LinkedList<>();
 
+    /**
+     * Whether Minestom detected that it is running in a dev environment.
+     * Determined by the existence of the system property {@link ExtensionManager#INDEV_CLASSES_FOLDER}
+     */
+    private boolean inDevEnvironment = false;
+
     /**
      * List of already loaded code modifier class names. This prevents loading the same class twice.
      */
@@ -73,6 +81,7 @@ public class MinestomRootClassLoader extends HierarchyClassLoader {
     private MinestomRootClassLoader(ClassLoader parent) {
         super("Minestom Root ClassLoader", extractURLsFromClasspath(), parent);
         asmClassLoader = newChild(new URL[0]);
+        inDevEnvironment = System.getProperty(ExtensionManager.INDEV_CLASSES_FOLDER) != null;
     }
 
     public static MinestomRootClassLoader getInstance() {
@@ -184,6 +193,22 @@ public class MinestomRootClassLoader extends HierarchyClassLoader {
         if (name == null)
             throw new ClassNotFoundException();
         String path = name.replace(".", "/") + ".class";
+
+        if(inDevEnvironment) {
+            // check if the class to load is the entry point of the extension
+            boolean isMainExtensionClass = false;
+            for(MinestomExtensionClassLoader c : children) {
+                if(c.isMainExtensionClass(name)) {
+                    isMainExtensionClass = true;
+                    break;
+                }
+            }
+            if(isMainExtensionClass) { // entry point of the extension, force load through extension classloader
+                throw new ClassNotFoundException("The class "+name+" is the entry point of an extension. " +
+                        "Because we are in a dev environment, we force its load through its extension classloader, " +
+                        "even though the root classloader has access.");
+            }
+        }
         InputStream input = getResourceAsStream(path);
         if (input == null) {
             throw new ClassNotFoundException("Could not find resource " + path);

src/main/java/net/minestom/server/instance/Instance.java

@@ -41,6 +41,7 @@ import org.jetbrains.annotations.Nullable;
 
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CopyOnWriteArraySet;
 import java.util.function.Consumer;
 
@@ -79,6 +80,7 @@ public abstract class Instance implements BlockModifier, EventHandler, DataConta
     private long lastTickAge = System.currentTimeMillis();
 
     private final Map<Class<? extends Event>, Collection<EventCallback>> eventCallbacks = new ConcurrentHashMap<>();
+    private final Map<String, Collection<EventCallback<?>>> extensionCallbacks = new ConcurrentHashMap<>();
 
     // Entities present in this instance
     protected final Set<Entity> entities = new CopyOnWriteArraySet<>();
@@ -862,6 +864,12 @@ public abstract class Instance implements BlockModifier, EventHandler, DataConta
         return eventCallbacks;
     }
 
+    @NotNull
+    @Override
+    public Collection<EventCallback<?>> getExtensionCallbacks(String extension) {
+        return extensionCallbacks.computeIfAbsent(extension, e -> new CopyOnWriteArrayList<>());
+    }
+
     // UNSAFE METHODS (need most of time to be synchronized)
 
     /**

src/test/java/improveextensions/unloadcallbacks/UnloadCallbacksExtension.java

@@ -0,0 +1,57 @@
+package improveextensions.unloadcallbacks;
+
+import net.minestom.server.MinecraftServer;
+import net.minestom.server.event.EventCallback;
+import net.minestom.server.event.GlobalEventHandler;
+import net.minestom.server.event.handler.EventHandler;
+import net.minestom.server.event.instance.InstanceTickEvent;
+import net.minestom.server.extensions.Extension;
+import net.minestom.server.utils.time.TimeUnit;
+import org.junit.jupiter.api.Assertions;
+import org.opentest4j.AssertionFailedError;
+
+public class UnloadCallbacksExtension extends Extension {
+
+    private boolean ticked = false;
+    private final EventCallback<InstanceTickEvent> callback = this::onTick;
+
+    private void onTick(InstanceTickEvent e) {
+        ticked = true;
+    }
+
+    @Override
+    public void initialize() {
+        GlobalEventHandler globalEvents = MinecraftServer.getGlobalEventHandler();
+        // this callback will be automatically removed when unloading the extension
+        globalEvents.addEventCallback(InstanceTickEvent.class, callback);
+
+        try {
+            Assertions.assertTrue(EventHandler.getExtensionOwningCallback(callback).isPresent());
+            Assertions.assertEquals("UnloadCallbacksExtension", EventHandler.getExtensionOwningCallback(callback).get());
+        } catch (AssertionFailedError e) {
+            e.printStackTrace();
+            System.exit(-1);
+        }
+
+        MinecraftServer.getSchedulerManager().buildTask(() -> {
+            // unload self
+            MinecraftServer.getExtensionManager().unloadExtension(getDescription().getName());
+        }).delay(1L, TimeUnit.SECOND).schedule();
+    }
+
+    @Override
+    public void terminate() {
+        ticked = false;
+
+        // TODO: set to transient task to avoid losing the task on termination
+        MinecraftServer.getSchedulerManager().buildTask(() -> {
+            // Make sure callback is disabled
+            try {
+                Assertions.assertFalse(ticked, "ticked should be false because the callback has been unloaded");
+            } catch (AssertionFailedError e) {
+                e.printStackTrace();
+            }
+            MinecraftServer.stopCleanly();
+        }).delay(1L, TimeUnit.SECOND).schedule();
+    }
+}

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

@@ -0,0 +1,4 @@
+{
+  "entrypoint": "improveextensions.unloadcallbacks.UnloadCallbacksExtension",
+  "name": "UnloadCallbacksExtension"
+}