From b33d18af2b5cd54bf8e55fdc11f106298def9311 Mon Sep 17 00:00:00 2001
From: Owen <23108066+Owen1212055@users.noreply.github.com>
Date: Fri, 27 Oct 2023 01:39:03 -0400
Subject: [PATCH] Allow closing active proxy listeners (#1109)

---
 .../api/proxy/ProxyServer.java                |  6 ++++
 .../velocitypowered/proxy/VelocityServer.java |  5 ++++
 .../proxy/network/ConnectionManager.java      | 30 ++++++++++++++-----
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/api/src/main/java/com/velocitypowered/api/proxy/ProxyServer.java b/api/src/main/java/com/velocitypowered/api/proxy/ProxyServer.java
index 8f4ec2f65..f043b0c6d 100644
--- a/api/src/main/java/com/velocitypowered/api/proxy/ProxyServer.java
+++ b/api/src/main/java/com/velocitypowered/api/proxy/ProxyServer.java
@@ -41,6 +41,12 @@ public interface ProxyServer extends Audience {
    */
   void shutdown();
 
+  /**
+   * Closes all listening endpoints for this server.
+   * This includes the main minecraft listener and query channel.
+   */
+  void closeListeners();
+
   /**
    * Retrieves the player currently connected to this proxy by their Minecraft username. The search
    * is case-insensitive.
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java b/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
index f82fd45d5..e5a4828e8 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
@@ -572,6 +572,11 @@ public class VelocityServer implements ProxyServer, ForwardingAudience {
     shutdown(true);
   }
 
+  @Override
+  public void closeListeners() {
+    this.cm.closeEndpoints(false);
+  }
+
   public AsyncHttpClient getAsyncHttpClient() {
     return cm.getHttpClient();
   }
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/network/ConnectionManager.java b/proxy/src/main/java/com/velocitypowered/proxy/network/ConnectionManager.java
index c5668db2a..bf596e8e3 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/network/ConnectionManager.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/network/ConnectionManager.java
@@ -212,9 +212,11 @@ public final class ConnectionManager {
   }
 
   /**
-   * Closes all endpoints.
+   * Closes all the currently registered endpoints.
+   *
+   * @param interrupt should closing forward interruptions
    */
-  public void shutdown() {
+  public void closeEndpoints(boolean interrupt) {
     for (final Map.Entry<InetSocketAddress, Endpoint> entry : this.endpoints.entrySet()) {
       final InetSocketAddress address = entry.getKey();
       final Endpoint endpoint = entry.getValue();
@@ -223,14 +225,26 @@ public final class ConnectionManager {
       // should have a chance to be notified before the server stops accepting connections.
       server.getEventManager().fire(new ListenerCloseEvent(address, endpoint.getType())).join();
 
-      try {
-        LOGGER.info("Closing endpoint {}", address);
-        endpoint.getChannel().close().sync();
-      } catch (final InterruptedException e) {
-        LOGGER.info("Interrupted whilst closing endpoint", e);
-        Thread.currentThread().interrupt();
+      LOGGER.info("Closing endpoint {}", address);
+      if (interrupt) {
+        try {
+          endpoint.getChannel().close().sync();
+        } catch (final InterruptedException e) {
+          LOGGER.info("Interrupted whilst closing endpoint", e);
+          Thread.currentThread().interrupt();
+        }
+      } else {
+        endpoint.getChannel().close().syncUninterruptibly();
       }
     }
+    this.endpoints.clear();
+  }
+
+  /**
+   * Closes all endpoints.
+   */
+  public void shutdown() {
+    this.closeEndpoints(true);
 
     this.resolver.shutdown();
   }