From 383db108f39eb31aee62a9fbee8bbe5b81da05f0 Mon Sep 17 00:00:00 2001
From: Frederikam <frogkr@gmail.com>
Date: Sun, 17 Sep 2017 15:00:29 +0200
Subject: [PATCH] Fixed not reconnecting after one of multiple nodes die

---
 .../main/java/lavalink/client/io/LavalinkLoadBalancer.java | 5 +++--
 .../src/main/java/lavalink/client/io/LavalinkSocket.java   | 7 +++++++
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/LavalinkClient/src/main/java/lavalink/client/io/LavalinkLoadBalancer.java b/LavalinkClient/src/main/java/lavalink/client/io/LavalinkLoadBalancer.java
index 7fdb017b7..ff24c54b1 100644
--- a/LavalinkClient/src/main/java/lavalink/client/io/LavalinkLoadBalancer.java
+++ b/LavalinkClient/src/main/java/lavalink/client/io/LavalinkLoadBalancer.java
@@ -71,6 +71,7 @@ void onNodeDisconnect(LavalinkSocket disconnected) {
             boolean isAffected = disconnected.equals(socket.orElse(null));
 
             LavalinkSocket newSocket = isAffected ? determineBestSocket() : socket.orElse(null);
+            socketMap.put(guildId, Optional.ofNullable(newSocket)); // A bit redundant, but we need this for the setSocket call
             if (isAffected)
                 lavalink.getPlayer(guildId).setSocket(newSocket);
 
@@ -138,14 +139,14 @@ public int getNullFramePenalty() {
         }
 
         public int getTotal() {
-            if (!socket.isOpen()) return Integer.MAX_VALUE - 1;
+            if (!socket.isAvailable()) return Integer.MAX_VALUE - 1;
 
             return playerPenalty + cpuPenalty + deficitFramePenalty + nullFramePenalty;
         }
 
         @Override
         public String toString() {
-            if (!socket.isOpen()) return "Penalties{" +
+            if (!socket.isAvailable()) return "Penalties{" +
                     "unavailable=" + (Integer.MAX_VALUE - 1) +
                     '}';
 
diff --git a/LavalinkClient/src/main/java/lavalink/client/io/LavalinkSocket.java b/LavalinkClient/src/main/java/lavalink/client/io/LavalinkSocket.java
index d4ae1e997..b070037f8 100644
--- a/LavalinkClient/src/main/java/lavalink/client/io/LavalinkSocket.java
+++ b/LavalinkClient/src/main/java/lavalink/client/io/LavalinkSocket.java
@@ -58,6 +58,7 @@ public class LavalinkSocket extends ReusableWebSocket {
     long lastReconnectAttempt = 0;
     private int reconnectsAttempted = 0;
     private final URI remoteUri;
+    private boolean available = false;
 
     LavalinkSocket(Lavalink lavalink, URI serverUri, Draft protocolDraft, Map<String, String> headers) {
         super(serverUri, protocolDraft, headers, TIMEOUT_MS);
@@ -73,6 +74,7 @@ public class LavalinkSocket extends ReusableWebSocket {
     @Override
     public void onOpen(ServerHandshake handshakeData) {
         log.info("Received handshake from server");
+        available = true;
         lavalink.loadBalancer.onNodeConnect(this);
         reconnectsAttempted = 0;
     }
@@ -196,6 +198,7 @@ private void handleEvent(JSONObject json) throws IOException {
 
     @Override
     public void onClose(int code, String reason, boolean remote) {
+        available = false;
         reason = reason == null ? "<no reason given>" : reason;
         if (code == 1000) {
             log.info("Connection to " + getRemoteSocketAddress() + " closed gracefully with reason: " + reason + " :: Remote=" + remote);
@@ -245,6 +248,10 @@ public RemoteStats getStats() {
         return stats;
     }
 
+    public boolean isAvailable() {
+        return available && isOpen() && !isClosing();
+    }
+
     @Override
     public String toString() {
         return "LavalinkSocket{" +