Allow to register main class as listener after unregistering it

proxy/src/main/java/com/velocitypowered/proxy/event/VelocityEventManager.java

@@ -37,6 +37,7 @@
 import com.velocitypowered.proxy.event.UntargetedEventHandler.EventTaskHandler;
 import com.velocitypowered.proxy.event.UntargetedEventHandler.VoidHandler;
 import com.velocitypowered.proxy.event.UntargetedEventHandler.WithContinuationHandler;
+import com.velocitypowered.proxy.plugin.loader.VelocityPluginContainer;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.VarHandle;
@@ -132,27 +133,16 @@ public <F> void registerHandlerAdapter(
 
   /**
    * Represents the registration of a single {@link EventHandler}.
+   *
+   * @param instance The instance of the {@link EventHandler} or the listener instance that was registered.
    */
-  static final class HandlerRegistration {
-
-    final PluginContainer plugin;
-    final short order;
-    final Class<?> eventType;
-    final EventHandler<Object> handler;
-
-    /**
-     * The instance of the {@link EventHandler} or the listener instance that was registered.
-     */
-    final Object instance;
-
-    public HandlerRegistration(final PluginContainer plugin, final short order,
-        final Class<?> eventType, final Object instance, final EventHandler<Object> handler) {
-      this.plugin = plugin;
-      this.order = order;
-      this.eventType = eventType;
-      this.instance = instance;
-      this.handler = handler;
-    }
+  record HandlerRegistration(
+          PluginContainer plugin,
+          short order,
+          Class<?> eventType,
+          Object instance,
+          EventHandler<Object> handler
+  ) {
   }
 
   enum AsyncType {
@@ -166,13 +156,7 @@ enum AsyncType {
     NEVER
   }
 
-  static final class HandlersCache {
-
-    final HandlerRegistration[] handlers;
-
-    HandlersCache(final HandlerRegistration[] handlers) {
-      this.handlers = handlers;
-    }
+  record HandlersCache(HandlerRegistration[] handlers) {
   }
 
   private @Nullable HandlersCache bakeHandlers(final Class<?> eventType) {
@@ -226,23 +210,13 @@ private UntargetedEventHandler buildUntargetedMethodHandler(final Method method)
     return LambdaFactory.create(type.defineClassesWith(lookup), methodHandle);
   }
 
-  static final class MethodHandlerInfo {
-
-    final Method method;
-    final @Nullable Class<?> eventType;
-    final short order;
-    final @Nullable String errors;
-    final @Nullable Class<?> continuationType;
-
-    private MethodHandlerInfo(final Method method, final @Nullable Class<?> eventType,
-        final short order, final @Nullable String errors,
-        final @Nullable Class<?> continuationType) {
-      this.method = method;
-      this.eventType = eventType;
-      this.order = order;
-      this.errors = errors;
-      this.continuationType = continuationType;
-    }
+  record MethodHandlerInfo(
+          Method method,
+          @Nullable Class<?> eventType,
+          short order,
+          @Nullable String errors,
+          @Nullable Class<?> continuationType
+  ) {
   }
 
   private void collectMethods(final Class<?> targetClass,
@@ -339,9 +313,15 @@ private void register(final List<HandlerRegistration> registrations) {
   @Override
   public void register(final Object plugin, final Object listener) {
     requireNonNull(listener, "listener");
-    final PluginContainer pluginContainer = pluginManager.ensurePluginContainer(plugin);
+    final VelocityPluginContainer pluginContainer
+            = (VelocityPluginContainer) pluginManager.ensurePluginContainer(plugin);
     if (plugin == listener) {
-      throw new IllegalArgumentException("The plugin main instance is automatically registered.");
+      // Allow to register main class as listener after unregistering it
+      if (pluginContainer.unregisteredMainClassListener()) {
+        pluginContainer.unregisteredMainClassListener(false);
+      } else {
+        throw new IllegalArgumentException("The plugin main instance is automatically registered.");
+      }
     }
     registerInternally(pluginContainer, listener);
   }
@@ -393,14 +373,24 @@ public void registerInternally(final PluginContainer pluginContainer, final Obje
 
   @Override
   public void unregisterListeners(final Object plugin) {
-    final PluginContainer pluginContainer = pluginManager.ensurePluginContainer(plugin);
+    final VelocityPluginContainer pluginContainer
+            = (VelocityPluginContainer) pluginManager.ensurePluginContainer(plugin);
+    // If the main class is unregistered, it is marked as such,
+    // to allow it to be re-registered
+    pluginContainer.unregisteredMainClassListener(true);
     unregisterIf(registration -> registration.plugin == pluginContainer);
   }
 
   @Override
   public void unregisterListener(final Object plugin, final Object handler) {
-    final PluginContainer pluginContainer = pluginManager.ensurePluginContainer(plugin);
     requireNonNull(handler, "handler");
+    final VelocityPluginContainer pluginContainer
+            = (VelocityPluginContainer) pluginManager.ensurePluginContainer(plugin);
+    // If the main class is unregistered, it is marked as such,
+    // to allow it to be re-registered
+    if (plugin == handler) {
+      pluginContainer.unregisteredMainClassListener(true);
+    }
     unregisterIf(registration ->
         registration.plugin == pluginContainer && registration.instance == handler);
   }

proxy/src/main/java/com/velocitypowered/proxy/plugin/loader/VelocityPluginContainer.java

@@ -32,6 +32,7 @@ public class VelocityPluginContainer implements PluginContainer {
   private final PluginDescription description;
   private Object instance;
   private volatile ExecutorService service;
+  private boolean unregisteredMainClassListener = false;
 
   public VelocityPluginContainer(PluginDescription description) {
     this.description = description;
@@ -75,4 +76,12 @@ public ExecutorService getExecutorService() {
   public boolean hasExecutorService() {
     return this.service != null;
   }
+
+  public void unregisteredMainClassListener(final boolean unregistered) {
+    this.unregisteredMainClassListener = unregistered;
+  }
+
+  public boolean unregisteredMainClassListener() {
+    return this.unregisteredMainClassListener;
+  }
 }

proxy/src/test/java/com/velocitypowered/proxy/event/EventTest.java

@@ -289,13 +289,7 @@ interface FancyContinuation {
     void resumeWithError(Exception exception);
   }
 
-  private static final class FancyContinuationImpl implements FancyContinuation {
-
-    private final Continuation continuation;
-
-    private FancyContinuationImpl(final Continuation continuation) {
-      this.continuation = continuation;
-    }
+  private record FancyContinuationImpl(Continuation continuation) implements FancyContinuation {
 
     @Override
     public void resume() {

proxy/src/test/java/com/velocitypowered/proxy/testutil/FakePluginManager.java

@@ -21,6 +21,7 @@
 import com.velocitypowered.api.plugin.PluginContainer;
 import com.velocitypowered.api.plugin.PluginDescription;
 import com.velocitypowered.api.plugin.PluginManager;
+import com.velocitypowered.proxy.plugin.loader.VelocityPluginContainer;
 import java.nio.file.Path;
 import java.util.Collection;
 import java.util.Optional;
@@ -52,14 +53,11 @@ public class FakePluginManager implements PluginManager {
 
   @Override
   public @NonNull Optional<PluginContainer> getPlugin(@NonNull String id) {
-    switch (id) {
-      case "a":
-        return Optional.of(PC_A);
-      case "b":
-        return Optional.of(PC_B);
-      default:
-        return Optional.empty();
-    }
+    return switch (id) {
+      case "a" -> Optional.of(PC_A);
+      case "b" -> Optional.of(PC_B);
+      default -> Optional.empty();
+    };
   }
 
   @Override
@@ -77,13 +75,14 @@ public void addToClasspath(@NonNull Object plugin, @NonNull Path path) {
     throw new UnsupportedOperationException();
   }
 
-  private static class FakePluginContainer implements PluginContainer {
+  private static class FakePluginContainer extends VelocityPluginContainer {
 
     private final String id;
     private final Object instance;
     private final ExecutorService service;
 
     private FakePluginContainer(String id, Object instance) {
+      super(null);
       this.id = id;
       this.instance = instance;
       this.service = ForkJoinPool.commonPool();