All EventHandler remove extension callback when the extension unloads

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

@@ -6,6 +6,7 @@ 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.IExtensionObserver;
 import net.minestom.server.extras.selfmodification.MinestomRootClassLoader;
 import net.minestom.server.instance.Instance;
 import org.jetbrains.annotations.NotNull;
@@ -19,7 +20,7 @@ import java.util.stream.Stream;
 /**
  * Represents an element which can have {@link Event} listeners assigned to it.
  */
-public interface EventHandler {
+public interface EventHandler extends IExtensionObserver {
 
     /**
      * Gets a {@link Map} containing all the listeners assigned to a specific {@link Event} type.
@@ -48,7 +49,10 @@ public interface EventHandler {
      */
     default <E extends Event> boolean addEventCallback(@NotNull Class<E> eventClass, @NotNull EventCallback<E> eventCallback) {
         Optional<String> extensionSource = MinestomRootClassLoader.findExtensionObjectOwner(eventCallback);
-        extensionSource.ifPresent(s -> getExtensionCallbacks(s).add(eventCallback));
+        extensionSource.ifPresent(name -> {
+            MinecraftServer.getExtensionManager().getExtension(name).observe(this);
+            getExtensionCallbacks(name).add(eventCallback);
+        });
 
         Collection<EventCallback> callbacks = getEventCallbacks(eventClass);
         return callbacks.add(eventCallback);
@@ -166,4 +170,9 @@ public interface EventHandler {
         }
     }
 
+    @Override
+    default void onExtensionUnload(String extensionName) {
+        removeCallbacksOwnedByExtension(extensionName);
+    }
+
 }

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

@@ -1,10 +1,18 @@
 package net.minestom.server.extensions;
 
+import net.minestom.server.event.handler.EventHandler;
 import org.jetbrains.annotations.NotNull;
 import org.slf4j.Logger;
 
+import java.lang.ref.WeakReference;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.function.Consumer;
 
 public abstract class Extension {
     // Set by reflection
@@ -14,6 +22,14 @@ public abstract class Extension {
     @SuppressWarnings("unused")
     private Logger logger;
 
+    /**
+     * Observers that will be notified of events related to this extension.
+     * Kept as WeakReference because entities can be observer, but could become candidate to be garbage-collected while
+     * this extension holds a reference to it. A WeakReference makes sure this extension does not prevent the memory
+     * from being cleaned up.
+     */
+    private Set<WeakReference<IExtensionObserver>> observers = Collections.newSetFromMap(new ConcurrentHashMap<>());
+
     protected Extension() {
 
     }
@@ -55,6 +71,27 @@ public abstract class Extension {
         return logger;
     }
 
+    /**
+     * Adds a new observer to this extension.
+     * Will be kept as a WeakReference.
+     * @param observer
+     */
+    public void observe(IExtensionObserver observer) {
+        observers.add(new WeakReference<>(observer));
+    }
+
+    /**
+     * Calls some action on all valid observers of this extension
+     * @param action code to execute on each observer
+     */
+    public void triggerChange(Consumer<IExtensionObserver> action) {
+        for(WeakReference<IExtensionObserver> weakObserver : observers) {
+            IExtensionObserver observer = weakObserver.get();
+            if(observer != null) {
+                action.accept(observer);
+            }
+        }
+    }
 
     public static class ExtensionDescription {
         private final String name;

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

@@ -525,8 +525,7 @@ public class ExtensionManager {
         ext.terminate();
         // remove callbacks for this extension
         String extensionName = ext.getDescription().getName();
-        MinecraftServer.getGlobalEventHandler().removeCallbacksOwnedByExtension(extensionName);
-        MinecraftServer.getSchedulerManager().removeExtensionTasksOnUnload(extensionName);
+        ext.triggerChange(observer -> observer.onExtensionUnload(extensionName));
         // TODO: more callback types
 
         ext.postTerminate();

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

@@ -0,0 +1,14 @@
+package net.minestom.server.extensions;
+
+/**
+ * Observes events related to extensions
+ */
+public interface IExtensionObserver {
+
+    /**
+     * Called before unloading an extension (that is, right after Extension#terminate and right before Extension#unload)
+     * @param extensionName the name of the extension that is being unloaded
+     */
+    void onExtensionUnload(String extensionName);
+
+}

src/main/java/net/minestom/server/timer/SchedulerManager.java

@@ -4,6 +4,8 @@ import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
 import it.unimi.dsi.fastutil.objects.ObjectCollection;
 import net.minestom.server.MinecraftServer;
+import net.minestom.server.extensions.Extension;
+import net.minestom.server.extensions.IExtensionObserver;
 import net.minestom.server.utils.thread.MinestomThread;
 import org.jetbrains.annotations.NotNull;
 
@@ -24,7 +26,7 @@ import java.util.stream.Collectors;
  * <p>
  * Shutdown tasks are built with {@link #buildShutdownTask(Runnable)} and are executed, as the name implies, when the server stops.
  */
-public final class SchedulerManager {
+public final class SchedulerManager implements IExtensionObserver {
 
     private static boolean instanced;
     // A counter for all normal tasks
@@ -188,6 +190,9 @@ public final class SchedulerManager {
     void onScheduleFromExtension(String owningExtension, Task task) {
         List<Task> scheduledForThisExtension = extensionTasks.computeIfAbsent(owningExtension, s -> new CopyOnWriteArrayList<>());
         scheduledForThisExtension.add(task);
+
+        Extension ext = MinecraftServer.getExtensionManager().getExtension(owningExtension);
+        ext.observe(this);
     }
 
     /**
@@ -198,6 +203,9 @@ public final class SchedulerManager {
     void onScheduleShutdownFromExtension(String owningExtension, Task task) {
         List<Task> scheduledForThisExtension = extensionShutdownTasks.computeIfAbsent(owningExtension, s -> new CopyOnWriteArrayList<>());
         scheduledForThisExtension.add(task);
+
+        Extension ext = MinecraftServer.getExtensionManager().getExtension(owningExtension);
+        ext.observe(this);
     }
 
     /**
@@ -226,4 +234,9 @@ public final class SchedulerManager {
             shutdownTasks.values().removeAll(toCancel);
         }
     }
+
+    @Override
+    public void onExtensionUnload(String extensionName) {
+        removeExtensionTasksOnUnload(extensionName);
+    }
 }

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

@@ -1,11 +1,15 @@
 package improveextensions.unloadcallbacks;
 
 import net.minestom.server.MinecraftServer;
+import net.minestom.server.entity.type.monster.EntityZombie;
 import net.minestom.server.event.EventCallback;
 import net.minestom.server.event.GlobalEventHandler;
+import net.minestom.server.event.entity.EntityTickEvent;
 import net.minestom.server.event.instance.InstanceTickEvent;
 import net.minestom.server.extensions.Extension;
 import net.minestom.server.extras.selfmodification.MinestomRootClassLoader;
+import net.minestom.server.instance.Instance;
+import net.minestom.server.utils.Position;
 import net.minestom.server.utils.time.TimeUnit;
 import org.junit.jupiter.api.Assertions;
 import org.opentest4j.AssertionFailedError;
@@ -18,6 +22,8 @@ public class UnloadCallbacksExtension extends Extension {
     private boolean ticked2 = false;
     private boolean tickedScheduledNonTransient = false;
     private boolean tickedScheduledTransient = false;
+    private boolean zombieTicked = false;
+    private boolean instanceTicked = false;
     private final EventCallback<InstanceTickEvent> callback = this::onTick;
 
     private void onTick(InstanceTickEvent e) {
@@ -32,6 +38,19 @@ public class UnloadCallbacksExtension extends Extension {
         // this one too
         globalEvents.addEventCallback(InstanceTickEvent.class, e -> ticked2 = true);
 
+        Instance instance = MinecraftServer.getInstanceManager().getInstances().stream().findFirst().orElseThrow();
+
+        // add an event callback on an instance
+        instance.addEventCallback(InstanceTickEvent.class, e -> instanceTicked = true);
+        instance.loadChunk(0,0);
+
+        // add an event callback on an entity
+        EntityZombie zombie = new EntityZombie(new Position(8,64,8) /* middle of chunk */);
+        zombie.addEventCallback(EntityTickEvent.class, e -> {
+            zombieTicked = true;
+        });
+        zombie.setInstance(instance);
+
         // this callback will be cancelled
         MinecraftServer.getSchedulerManager().buildTask(() -> {
             tickedScheduledNonTransient = true;
@@ -58,10 +77,20 @@ public class UnloadCallbacksExtension extends Extension {
 
     @Override
     public void terminate() {
-        ticked1 = false;
-        ticked2 = false;
-        tickedScheduledNonTransient = false;
-        tickedScheduledTransient = false;
+        new Thread(() -> {
+            try {
+                // wait for complete termination of this extension
+                Thread.sleep(10);
+                ticked1 = false;
+                ticked2 = false;
+                tickedScheduledNonTransient = false;
+                tickedScheduledTransient = false;
+                instanceTicked = false;
+                zombieTicked = false;
+            } catch (InterruptedException e) {
+                e.printStackTrace();
+            }
+        }).start();
 
         AtomicBoolean executedDelayTaskAfterTerminate = new AtomicBoolean(false);
         // because terminate is called just before unscheduling and removing event callbacks,
@@ -79,6 +108,8 @@ public class UnloadCallbacksExtension extends Extension {
                 Assertions.assertFalse(ticked1, "ticked1 should be false because the callback has been unloaded");
                 Assertions.assertFalse(ticked2, "ticked2 should be false because the callback has been unloaded");
                 Assertions.assertFalse(tickedScheduledNonTransient, "tickedScheduledNonTransient should be false because the callback has been unloaded");
+                Assertions.assertFalse(zombieTicked, "zombieTicked should be false because the callback has been unloaded");
+                Assertions.assertFalse(instanceTicked, "instanceTicked should be false because the callback has been unloaded");
                 Assertions.assertTrue(tickedScheduledTransient, "tickedScheduledNonTransient should be true because the callback has NOT been unloaded");
                 Assertions.assertFalse(executedDelayTaskAfterTerminate.get(), "executedDelayTaskAfterTerminate should be false because the callback has been unloaded before executing");
                 System.out.println("All tests passed.");