From 942e2f2e1ad489561ce8ea2eb6ead5b4dea88f6b Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Mon, 25 May 2020 11:56:56 -0400
Subject: [PATCH] Better generic cleanup.

---
 .../proxy/connection/client/ConnectedPlayer.java    | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java
index 4c5954cf6..e646a2179 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java
@@ -757,11 +757,18 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
                   VelocityServerConnection con = new VelocityServerConnection(vrs,
                       ConnectedPlayer.this, server);
                   connectionInFlight = con;
-                  return con.connect();
+                  return con.connect().whenCompleteAsync((result, throwable) ->
+                      this.cleanupIfRequired(con));
                 }, connection.eventLoop());
           });
     }
 
+    private void cleanupIfRequired(VelocityServerConnection establishedConnection) {
+      if (establishedConnection == connectionInFlight) {
+        resetInFlightConnection();
+      }
+    }
+
     @Override
     public CompletableFuture<Result> connect() {
       return this.internalConnect()
@@ -769,11 +776,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
             if (status != null && !status.isSuccessful()) {
               if (!status.isSafe()) {
                 handleConnectionException(status.getAttemptedConnection(), throwable, false);
-              } else if (status.getStatus() == Status.SERVER_DISCONNECTED) {
-                resetInFlightConnection();
               }
-            } else if (throwable != null) {
-              resetInFlightConnection();
             }
           }, connection.eventLoop())
           .thenApply(x -> x);