Fixed not reconnecting after one of multiple nodes die

Earleys · Sep 17, 2017 · 383db10 · 383db10
1 parent ed78af9
commit 383db10
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git 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
@@ -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{" +