Merge upstream with my fork

src/main/java/lavalink/client/io/Lavalink.java

@@ -34,23 +34,30 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.*;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
-@SuppressWarnings("unused")
 public abstract class Lavalink<T extends Link> {
 
     private static final Logger log = LoggerFactory.getLogger(Lavalink.class);
 
+    @SuppressWarnings("WeakerAccess")
     protected final int numShards;
-    private final String userId;
+    /** User id may be set at a later time */
+    @Nullable
+    private String userId;
     private boolean holdEvents = false;
     private final ConcurrentHashMap<String, T> links = new ConcurrentHashMap<>();
-    private final List<LavalinkSocket> nodes = new CopyOnWriteArrayList<>();
+    final List<LavalinkSocket> nodes = new CopyOnWriteArrayList<>();
     final LavalinkLoadBalancer loadBalancer = new LavalinkLoadBalancer(this);
+
     private final ScheduledExecutorService reconnectService;
 
-    public Lavalink(String userId, int numShards) {
+    public Lavalink(@Nullable String userId, int numShards) {
         this.userId = userId;
         this.numShards = numShards;
 
@@ -62,48 +69,48 @@ public Lavalink(String userId, int numShards) {
         reconnectService.scheduleWithFixedDelay(new ReconnectTask(this), 0, 500, TimeUnit.MILLISECONDS);
     }
 
-    private static final AtomicInteger nodeCounter = new AtomicInteger(0);
-
     /**
-     * @param serverUri uri of the node to be added
-     * @param password  password of the node to be added
+     * Creates a Lavalink instance.
+     * N.B: You must set the user ID before adding a node
      */
-    @SuppressWarnings("UnusedReturnValue")
-    public LavalinkSocket addNode(@NonNull URI serverUri, @NonNull String password) {
-        return addNode("Lavalink_Node_#" + nodeCounter.getAndIncrement(), serverUri, password, null);
+    @SuppressWarnings("unused")
+    public Lavalink(int numShards) {
+        this(null, numShards);
     }
 
-    /**
-     * @param serverUri uri of the node to be added
-     * @param password  password of the node to be added
-     * @param resumeKey key to use when resuming
-     */
-    public LavalinkSocket addNode(@NonNull URI serverUri, @NonNull String password, @Nullable String resumeKey) {
-        return addNode("Lavalink_Node_#" + nodeCounter.getAndIncrement(), serverUri, password, resumeKey);
-    }
+    private static final AtomicInteger nodeCounter = new AtomicInteger(0);
 
-    /**
-     * @param name      A name to identify this node. May show up in metrics and other places.
-     * @param serverUri uri of the node to be added
-     * @param password  password of the node to be added
-     */
-    public LavalinkSocket addNode(@NonNull String name, @NonNull URI serverUri, @NonNull String password) {
-        return addNode(name, serverUri, password, null);
+    public LavalinkSocket addNode(@NonNull URI serverUri, @NonNull String password) {
+        return addNode("Lavalink_Node_#" + nodeCounter.getAndIncrement(), serverUri, password, null);
     }
 
     /**
-     * @param name      A name to identify this node. May show up in metrics and other places.
-     * @param serverUri uri of the node to be added
-     * @param password  password of the node to be added
-     * @param resumeKey key to use when resuming
+     *
+     * @param name
+     *         A name to identify this node. May show up in metrics and other places.
+     * @param serverUri
+     *         uri of the node to be added
+     * @param password
+     *         password of the node to be added
+     * @throws IllegalStateException if no userId has been set.
+     * @throws IllegalArgumentException if a node with that name already exists.
+     * @see #setUserId(String)
      */
     @SuppressWarnings("WeakerAccess")
-    public LavalinkSocket addNode(@NonNull String name, @NonNull URI serverUri,
-                                  @NonNull String password, @Nullable String resumeKey) {
+    public LavalinkSocket addNode(@NonNull String name, @NonNull URI serverUri, @NonNull String password, @Nullable String resumeKey) {
+        if (userId == null) {
+            throw new IllegalStateException("We need a userId to connect to Lavalink");
+        }
+
+        if (nodes.stream().anyMatch(sock -> sock.getName().equals(name))) {
+            throw new IllegalArgumentException("A node with the name " + name + " already exists.");
+        }
+
         HashMap<String, String> headers = new HashMap<>();
         headers.put("Authorization", password);
         headers.put("Num-Shards", Integer.toString(numShards));
         headers.put("User-Id", userId);
+        headers.put("Client-Name", "Lavalink-Client");
 
         LavalinkSocket socket = new LavalinkSocket(name, this, serverUri, new Draft_6455(), headers, resumeKey);
         socket.connect();
@@ -124,6 +131,7 @@ public LavalinkLoadBalancer getLoadBalancer() {
         return loadBalancer;
     }
 
+    @SuppressWarnings("WeakerAccess")
     @NonNull
     public T getLink(@NonNull String guildId) {
         return links.computeIfAbsent(guildId, __ -> buildNewLink(guildId));
@@ -144,6 +152,7 @@ public T getExistingLink(@NonNull String guildId) {
      */
     protected abstract T buildNewLink(String guildId);
 
+    @SuppressWarnings({"WeakerAccess", "unused"})
     public int getNumShards() {
         return numShards;
     }
@@ -154,11 +163,24 @@ public Collection<T> getLinks() {
         return links.values();
     }
 
+    @SuppressWarnings("WeakerAccess")
     @NonNull
     public List<LavalinkSocket> getNodes() {
         return nodes;
     }
 
+    /**
+     * The user id of this bot.
+     * @throws IllegalStateException if any nodes are registered.
+     */
+    @SuppressWarnings("unused")
+    public void setUserId(@Nullable String userId) {
+        if (!nodes.isEmpty()) {
+            throw new IllegalStateException("Can't set userId if we already have nodes registered!");
+        }
+        this.userId = userId;
+    }
+
     public void shutdown() {
         reconnectService.shutdown();
         nodes.forEach(ReusableWebSocket::close);
@@ -169,6 +191,7 @@ void removeDestroyedLink(Link link) {
         links.remove(link.getGuildId());
     }
 
+    @SuppressWarnings("WeakerAccess")
     protected Map<String, T> getLinksMap() {
         return links;
     }

src/main/java/lavalink/client/io/LavalinkLoadBalancer.java

@@ -76,7 +76,7 @@ void onNodeDisconnect(LavalinkSocket disconnected) {
         Collection<Link> links = lavalink.getLinks();
         links.forEach(link -> {
             if (disconnected.equals(link.getNode(false)))
-                link.changeNode(determineBestSocket(link.getGuildIdLong()));
+                link.changeNode(lavalink.loadBalancer.determineBestSocket(link.getGuildIdLong()));
         });
     }
 
@@ -93,9 +93,8 @@ void onNodeConnect(LavalinkSocket connected) {
         @SuppressWarnings("unchecked")
         Collection<Link> links = lavalink.getLinks();
         links.forEach(link -> {
-            LavalinkSocket current = link.getNode(false);
-            if (current == null) link.changeNode(connected);
-            else if (current == connected) link.onNodeConnected();
+            if (link.getNode(false) == null)
+                link.changeNode(connected);
         });
     }
 
@@ -208,5 +207,4 @@ public String toString() {
                     '}';
         }
     }
-
 }