From 62fff7a56ead98e38028cffcd56da65b5b22889d Mon Sep 17 00:00:00 2001
From: Julian Krings <juliankrings6@gmail.com>
Date: Sun, 8 Sep 2024 13:41:08 +0200
Subject: [PATCH] take server size into account when picking a server

---
 .../iris/server/master/IrisMasterClient.java  |  9 +++++
 .../iris/server/master/IrisMasterServer.java  | 36 ++++++++++++++-----
 .../iris/server/master/IrisMasterSession.java |  2 +-
 3 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/core/src/main/java/com/volmit/iris/server/master/IrisMasterClient.java b/core/src/main/java/com/volmit/iris/server/master/IrisMasterClient.java
index af4ed90aa..83d48276c 100644
--- a/core/src/main/java/com/volmit/iris/server/master/IrisMasterClient.java
+++ b/core/src/main/java/com/volmit/iris/server/master/IrisMasterClient.java
@@ -8,6 +8,7 @@ import com.volmit.iris.server.packet.init.PingPacket;
 import com.volmit.iris.server.util.ConnectionHolder;
 import com.volmit.iris.server.util.PacketListener;
 import com.volmit.iris.server.util.PacketSendListener;
+import com.volmit.iris.util.collection.KList;
 import lombok.Getter;
 import org.jetbrains.annotations.Nullable;
 
@@ -55,6 +56,14 @@ public class IrisMasterClient implements ConnectionHolder, PacketListener {
         session.onClientPacket(this, raw);
     }
 
+    public KList<String> getVersions() {
+        try {
+            return pingResponse.get().getVersion();
+        } catch (InterruptedException | ExecutionException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     public int getNodeCount() {
         try {
             return nodeCount.get();
diff --git a/core/src/main/java/com/volmit/iris/server/master/IrisMasterServer.java b/core/src/main/java/com/volmit/iris/server/master/IrisMasterServer.java
index e1346005f..f4ff713fd 100644
--- a/core/src/main/java/com/volmit/iris/server/master/IrisMasterServer.java
+++ b/core/src/main/java/com/volmit/iris/server/master/IrisMasterServer.java
@@ -6,9 +6,11 @@ import com.volmit.iris.server.util.PregenHolder;
 import com.volmit.iris.util.collection.KList;
 import com.volmit.iris.util.collection.KMap;
 import com.volmit.iris.util.collection.KSet;
+import com.volmit.iris.util.parallel.MultiBurst;
 import lombok.extern.java.Log;
 
 import java.net.InetSocketAddress;
+import java.util.Collection;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
@@ -45,13 +47,16 @@ public class IrisMasterServer extends IrisServer {
     private void addNode(InetSocketAddress address) throws InterruptedException {
         var ping = new PingConnection(address);
         try {
-            for (String version : ping.getVersion().get())
-                addNode(address, version);
+            addNode(address, ping.getVersion().get());
         } catch (ExecutionException e) {
             throw new RuntimeException(e);
         }
     }
 
+    public void addNode(InetSocketAddress address, Collection<String> versions) {
+        versions.forEach(v -> addNode(address, v));
+    }
+
     public void addNode(InetSocketAddress address, String version) {
         nodes.computeIfAbsent(version, v -> new KSet<>()).add(address);
     }
@@ -80,14 +85,27 @@ public class IrisMasterServer extends IrisServer {
 
         master.getLogger().info("Requesting nodes for session " + uuid);
         var map = new KMap<IrisMasterClient, KMap<UUID, PregenHolder>>();
-        for (var address : master.nodes.getOrDefault(version, new KSet<>())) {
-            try {
-                map.put(IrisConnection.connect(address, new IrisMasterClient(version, session)), new KMap<>());
-            } catch (Throwable e) {
-                master.getLogger().log(Level.WARNING, "Failed to connect to server " + address, e);
-                master.removeNode(address);
-            }
+        if (!master.nodes.containsKey(version)) {
+            master.getLogger().warning("No nodes found for version " + version);
+            return map;
         }
+        var nodes = master.nodes.get(version);
+        var remove = new KList<InetSocketAddress>();
+        var burst = MultiBurst.burst.burst(nodes.size());
+        for (var address : nodes) {
+            burst.queue(() -> {
+                try {
+                    var client = new IrisMasterClient(version, session);
+                    master.addNode(address, client.getVersions());
+                    map.put(IrisConnection.connect(address, client), new KMap<>());
+                } catch (Throwable e) {
+                    master.getLogger().log(Level.WARNING, "Failed to connect to server " + address, e);
+                    remove.add(address);
+                }
+            });
+        }
+        burst.complete();
+        remove.forEach(nodes::remove);
 
         master.sessions.put(uuid, map);
         return map;
diff --git a/core/src/main/java/com/volmit/iris/server/master/IrisMasterSession.java b/core/src/main/java/com/volmit/iris/server/master/IrisMasterSession.java
index 985555c6b..cc42b67c7 100644
--- a/core/src/main/java/com/volmit/iris/server/master/IrisMasterSession.java
+++ b/core/src/main/java/com/volmit/iris/server/master/IrisMasterSession.java
@@ -116,7 +116,7 @@ public class IrisMasterSession implements ConnectionHolder, PacketListener {
     private IrisMasterClient pick() throws RejectedException {
         return clients.keySet()
                 .stream()
-                .min(Comparator.comparingInt(c -> clients.get(c).size()))
+                .min(Comparator.comparingDouble(c -> (double) clients.get(c).size() / c.getNodeCount()))
                 .orElseThrow(() -> new RejectedException("No clients available"));
     }