Fix not disconnecting on shutdown

Gaming32 · Jul 1, 2024 · ff5f5ad · ff5f5ad
1 parent f912a1e
commit ff5f5ad
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 6 deletions.
diff --git a/src/main/java/io/github/gaming32/worldhost/WorldHost.java b/src/main/java/io/github/gaming32/worldhost/WorldHost.java
@@ -76,7 +76,10 @@
 import java.util.UUID;
 import java.util.WeakHashMap;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
@@ -265,6 +268,13 @@ private static void init() {
         if (!CONFIG.isNoUPnP()) {
             scanUpnp();
         }
+
+        Runtime.getRuntime().addShutdownHook(
+            Thread.ofPlatform()
+                .name("World Host Shutdown Thread")
+                .unstarted(() -> {
+                })
+        );
     }
 
     public static void loadConfig() {
@@ -415,18 +425,36 @@ public static void scanUpnp() {
     }
 
     public static void reconnect(boolean successToast, boolean failureToast) {
+        shutdownClients();
+        LOGGER.info("Attempting to connect to WH server at {}", CONFIG.getServerIp());
+        protoClient = new ProtocolClient(CONFIG.getServerIp(), successToast, failureToast);
+        connectingFuture = protoClient.getConnectingFuture();
+        protoClient.authenticate(Minecraft.getInstance().getUser());
+    }
+
+    public static void shutdownClients() {
         if (protoClient != null) {
             protoClient.close();
-            protoClient = null;
         }
         if (proxyProtocolClient != null) {
             proxyProtocolClient.close();
-            proxyProtocolClient = null;
         }
-        LOGGER.info("Attempting to connect to WH server at {}", CONFIG.getServerIp());
-        protoClient = new ProtocolClient(CONFIG.getServerIp(), successToast, failureToast);
-        connectingFuture = protoClient.getConnectingFuture();
-        protoClient.authenticate(Minecraft.getInstance().getUser());
+        if (protoClient != null) {
+            try {
+                protoClient.getShutdownFuture().get(5L, TimeUnit.SECONDS);
+            } catch (Exception e) {
+                LOGGER.error("Failed to wait for protocol client shutdown", e);
+            }
+        }
+        protoClient = null;
+        if (proxyProtocolClient != null) {
+            try {
+                proxyProtocolClient.getShutdownFuture().get(5L, TimeUnit.SECONDS);
+            } catch (Exception e) {
+                LOGGER.error("Failed to wait for proxy protocol client shutdown", e);
+            }
+        }
+        proxyProtocolClient = null;
     }
 
     public static String getName(GameProfile profile) {

diff --git a/src/main/java/io/github/gaming32/worldhost/protocol/ProtocolClient.java b/src/main/java/io/github/gaming32/worldhost/protocol/ProtocolClient.java
@@ -48,6 +48,8 @@ public final class ProtocolClient implements AutoCloseable, ProxyPassthrough {
     final CompletableFuture<Void> connectingFuture = new CompletableFuture<>();
     private final BlockingQueue<Optional<WorldHostC2SMessage>> sendQueue = new LinkedBlockingQueue<>();
 
+    private final CompletableFuture<Void> shutdownFuture = new CompletableFuture<>();
+
     private CompletableFuture<User> authUser = new CompletableFuture<>();
 
     private boolean authenticated, closed;
@@ -184,6 +186,8 @@ public ProtocolClient(String host, boolean successToast, boolean failureToast) {
                         .show();
                 }
             }
+
+            shutdownFuture.complete(null);
         });
     }
 
@@ -314,6 +318,10 @@ public Future<Void> getConnectingFuture() {
         return connectingFuture;
     }
 
+    public CompletableFuture<Void> getShutdownFuture() {
+        return shutdownFuture;
+    }
+
     public long getConnectionId() {
         return connectionId;
     }

diff --git a/src/main/java/io/github/gaming32/worldhost/protocol/proxy/ProxyProtocolClient.java b/src/main/java/io/github/gaming32/worldhost/protocol/proxy/ProxyProtocolClient.java
@@ -10,6 +10,7 @@
 import java.net.SocketException;
 import java.util.Optional;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.LinkedBlockingQueue;
 
 public class ProxyProtocolClient implements AutoCloseable, ProxyPassthrough {
@@ -19,6 +20,8 @@ public class ProxyProtocolClient implements AutoCloseable, ProxyPassthrough {
 
     private final BlockingQueue<Optional<ProxyMessage>> sendQueue = new LinkedBlockingQueue<>();
 
+    private final CompletableFuture<Void> shutdownFuture = new CompletableFuture<>();
+
     private final String baseAddr;
     private final int mcPort;
 
@@ -108,6 +111,8 @@ public ProxyProtocolClient(String host, int port, long connectionId, String base
             } catch (IOException e) {
                 WorldHost.LOGGER.error("Failed to close WHEP socket.", e);
             }
+
+            shutdownFuture.complete(null);
         });
     }
 
@@ -141,6 +146,10 @@ public void proxyDisconnect(long connectionId) {
         close(connectionId);
     }
 
+    public CompletableFuture<Void> getShutdownFuture() {
+        return shutdownFuture;
+    }
+
     public String getBaseAddr() {
         return baseAddr;
     }