implement round robbing worker assignemnt

littlehorse-enterprises · May 17, 2024 · c4dfd55 · c4dfd55
1 parent d770f0b
commit c4dfd55
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 100 deletions.
diff --git a/...java/io/littlehorse/common/model/corecommand/subcommand/internals/RoundRobinAssignor.java b/...java/io/littlehorse/common/model/corecommand/subcommand/internals/RoundRobinAssignor.java
@@ -1,12 +1,9 @@
 package io.littlehorse.common.model.corecommand.subcommand.internals;
 
-import static io.littlehorse.common.LHConstants.MIN_WORKER_ASSIGNMENT_BY_SERVER;
-
 import com.google.common.collect.Iterables;
 import io.littlehorse.common.model.getable.core.taskworkergroup.HostModel;
 import io.littlehorse.common.model.getable.core.taskworkergroup.TaskWorkerMetadataModel;
 import java.util.*;
-import java.util.stream.IntStream;
 
 public class RoundRobinAssignor implements TaskWorkerAssignor {
 
@@ -28,11 +25,9 @@ public void assign(Collection<HostModel> hosts, Collection<TaskWorkerMetadataMod
 
         // Assigning N workers to a server
         for (HostModel host : sortedHosts) {
-            IntStream.range(0, MIN_WORKER_ASSIGNMENT_BY_SERVER).forEach(i -> {
-                TaskWorkerMetadataModel worker = roundRobinWorkers.next();
-                remainingWorkers.remove(worker);
-                worker.hosts.add(host);
-            });
+            TaskWorkerMetadataModel worker = roundRobinWorkers.next();
+            remainingWorkers.remove(worker);
+            worker.hosts.add(host);
         }
 
         // Assign remaining workers to a server

diff --git a/.../io/littlehorse/common/model/corecommand/subcommand/internals/RoundRobinAssignorTest.java b/.../io/littlehorse/common/model/corecommand/subcommand/internals/RoundRobinAssignorTest.java
@@ -6,6 +6,7 @@
 import io.littlehorse.common.model.getable.core.taskworkergroup.TaskWorkerMetadataModel;
 import java.util.*;
 import net.datafaker.Faker;
+import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 public class RoundRobinAssignorTest {
@@ -24,61 +25,35 @@ void assignWithOneHostAndTwoWorkers() {
         assertTrue(taskWorkersMetadata.get(1).hosts.contains(hosts.get(0)));
     }
 
-    @Test
-    void assignWithTwoHostsAndTwoWorkers() {
-        List<HostModel> hosts = generateHosts(2);
-        List<TaskWorkerMetadataModel> taskWorkersMetadata = generateTaskWorkersMetadata(2);
-
-        robinAssignor.assign(hosts, taskWorkersMetadata);
-
-        for (HostModel host : hosts) {
-            assertTrue(taskWorkersMetadata.get(0).hosts.contains(host));
-            assertTrue(taskWorkersMetadata.get(1).hosts.contains(host));
-        }
-    }
-
-    @Test
-    void assignWithTwoHostsAndFourWorkers() {
-        List<HostModel> hosts = generateHosts(2);
-        List<TaskWorkerMetadataModel> taskWorkersMetadata = generateTaskWorkersMetadata(4);
-
-        robinAssignor.assign(hosts, taskWorkersMetadata);
-
-        assertTrue(taskWorkersMetadata.get(0).hosts.contains(hosts.get(0))
-                && !taskWorkersMetadata.get(0).hosts.contains(hosts.get(1)));
-        assertTrue(taskWorkersMetadata.get(1).hosts.contains(hosts.get(0))
-                && !taskWorkersMetadata.get(1).hosts.contains(hosts.get(1)));
-        assertTrue(taskWorkersMetadata.get(2).hosts.contains(hosts.get(1))
-                && !taskWorkersMetadata.get(2).hosts.contains(hosts.get(0)));
-        assertTrue(taskWorkersMetadata.get(3).hosts.contains(hosts.get(1))
-                && !taskWorkersMetadata.get(3).hosts.contains(hosts.get(0)));
-    }
-
     @Test
     void assignWithTwoHostsAndFiveWorkers() {
-        List<HostModel> hosts = generateHosts(2);
+        HostModel hostA = new HostModel("a", 0);
+        HostModel hostB = new HostModel("b", 0);
+        Set<HostModel> hosts = new TreeSet<>(List.of(hostA, hostB));
         List<TaskWorkerMetadataModel> taskWorkersMetadata = generateTaskWorkersMetadata(5);
 
         robinAssignor.assign(hosts, taskWorkersMetadata);
 
-        assertTrue(taskWorkersMetadata.get(0).hosts.contains(hosts.get(0)));
-        assertTrue(taskWorkersMetadata.get(1).hosts.contains(hosts.get(0)));
-        assertTrue(taskWorkersMetadata.get(2).hosts.contains(hosts.get(1)));
-        assertTrue(taskWorkersMetadata.get(3).hosts.contains(hosts.get(1)));
-        assertTrue(taskWorkersMetadata.get(4).hosts.contains(hosts.get(0)));
+        Assertions.assertThat(taskWorkersMetadata.get(0).hosts).containsExactly(hostA);
+        Assertions.assertThat(taskWorkersMetadata.get(1).hosts).containsExactly(hostB);
+        Assertions.assertThat(taskWorkersMetadata.get(2).hosts).containsExactly(hostA);
+        Assertions.assertThat(taskWorkersMetadata.get(3).hosts).containsExactly(hostB);
+        Assertions.assertThat(taskWorkersMetadata.get(4).hosts).containsExactly(hostA);
     }
 
     @Test
     void assignWithFiveHostsAndTwoWorkers() {
-        List<HostModel> hosts = generateHosts(5);
+        HostModel hostA = new HostModel("a", 0);
+        HostModel hostB = new HostModel("b", 0);
+        HostModel hostC = new HostModel("c", 0);
+        HostModel hostD = new HostModel("d", 0);
+        HostModel hostE = new HostModel("e", 0);
+        Set<HostModel> hosts = new TreeSet<>(List.of(hostA, hostB, hostC, hostD, hostE));
         List<TaskWorkerMetadataModel> taskWorkersMetadata = generateTaskWorkersMetadata(2);
 
         robinAssignor.assign(hosts, taskWorkersMetadata);
-
-        for (HostModel host : hosts) {
-            assertTrue(taskWorkersMetadata.get(0).hosts.contains(host));
-            assertTrue(taskWorkersMetadata.get(1).hosts.contains(host));
-        }
+        Assertions.assertThat(taskWorkersMetadata.get(0).hosts).containsExactly(hostA, hostC, hostE);
+        Assertions.assertThat(taskWorkersMetadata.get(1).hosts).containsExactly(hostB, hostD);
     }
 
     @Test
@@ -87,21 +62,11 @@ void assignWithFiveHostsAndFiveWorkers() {
         List<TaskWorkerMetadataModel> taskWorkersMetadata = generateTaskWorkersMetadata(5);
         robinAssignor.assign(hosts, taskWorkersMetadata);
 
-        assertTrue(taskWorkersMetadata.get(0).hosts.contains(hosts.get(0))
-                && taskWorkersMetadata.get(0).hosts.contains(hosts.get(2))
-                && taskWorkersMetadata.get(0).hosts.size() == 2);
-        assertTrue(taskWorkersMetadata.get(1).hosts.contains(hosts.get(0))
-                && taskWorkersMetadata.get(1).hosts.contains(hosts.get(3))
-                && taskWorkersMetadata.get(1).hosts.size() == 2);
-        assertTrue(taskWorkersMetadata.get(2).hosts.contains(hosts.get(1))
-                && taskWorkersMetadata.get(2).hosts.contains(hosts.get(3))
-                && taskWorkersMetadata.get(2).hosts.size() == 2);
-        assertTrue(taskWorkersMetadata.get(3).hosts.contains(hosts.get(1))
-                && taskWorkersMetadata.get(3).hosts.contains(hosts.get(4))
-                && taskWorkersMetadata.get(3).hosts.size() == 2);
-        assertTrue(taskWorkersMetadata.get(4).hosts.contains(hosts.get(2))
-                && taskWorkersMetadata.get(4).hosts.contains(hosts.get(4))
-                && taskWorkersMetadata.get(4).hosts.size() == 2);
+        Assertions.assertThat(taskWorkersMetadata.get(0).hosts).containsExactly(hosts.get(0));
+        Assertions.assertThat(taskWorkersMetadata.get(1).hosts).containsExactly(hosts.get(1));
+        Assertions.assertThat(taskWorkersMetadata.get(2).hosts).containsExactly(hosts.get(2));
+        Assertions.assertThat(taskWorkersMetadata.get(3).hosts).containsExactly(hosts.get(3));
+        Assertions.assertThat(taskWorkersMetadata.get(4).hosts).containsExactly(hosts.get(4));
     }
 
     @Test
@@ -112,17 +77,13 @@ void assignWithThreeHostsAndFourWorkers() {
         robinAssignor.assign(hosts, taskWorkersMetadata);
 
         assertTrue(taskWorkersMetadata.get(0).hosts.contains(hosts.get(0))
-                && taskWorkersMetadata.get(0).hosts.contains(hosts.get(2))
-                && taskWorkersMetadata.get(0).hosts.size() == 2);
-        assertTrue(taskWorkersMetadata.get(1).hosts.contains(hosts.get(0))
-                && taskWorkersMetadata.get(1).hosts.contains(hosts.get(2))
-                && taskWorkersMetadata.get(1).hosts.size() == 2);
-        assertTrue(taskWorkersMetadata.get(2).hosts.contains(hosts.get(1))
-                && !taskWorkersMetadata.get(2).hosts.contains(hosts.get(2))
-                && taskWorkersMetadata.get(2).hosts.size() == 1);
-        assertTrue(taskWorkersMetadata.get(3).hosts.contains(hosts.get(1))
-                && !taskWorkersMetadata.get(3).hosts.contains(hosts.get(2))
-                && taskWorkersMetadata.get(3).hosts.size() == 1);
+                && taskWorkersMetadata.get(0).hosts.size() == 1);
+        assertTrue(taskWorkersMetadata.get(1).hosts.contains(hosts.get(1))
+                && taskWorkersMetadata.get(1).hosts.size() == 1);
+        assertTrue(taskWorkersMetadata.get(2).hosts.contains(hosts.get(2))
+                && taskWorkersMetadata.get(1).hosts.size() == 1);
+        assertTrue(taskWorkersMetadata.get(3).hosts.contains(hosts.get(0))
+                && taskWorkersMetadata.get(1).hosts.size() == 1);
     }
 
     @Test
@@ -134,24 +95,15 @@ void assignWithThreeHostsAndNineWorkers() {
         robinAssignor.assign(hosts, taskWorkersMetadata);
         robinAssignor.assign(hosts, taskWorkersMetadata);
 
-        assertTrue(taskWorkersMetadata.get(0).hosts.contains(hosts.get(0))
-                && taskWorkersMetadata.get(0).hosts.size() == 1);
-        assertTrue(taskWorkersMetadata.get(1).hosts.contains(hosts.get(0))
-                && taskWorkersMetadata.get(1).hosts.size() == 1);
-        assertTrue(taskWorkersMetadata.get(2).hosts.contains(hosts.get(1))
-                && taskWorkersMetadata.get(2).hosts.size() == 1);
-        assertTrue(taskWorkersMetadata.get(3).hosts.contains(hosts.get(1))
-                && taskWorkersMetadata.get(3).hosts.size() == 1);
-        assertTrue(taskWorkersMetadata.get(4).hosts.contains(hosts.get(2))
-                && taskWorkersMetadata.get(4).hosts.size() == 1);
-        assertTrue(taskWorkersMetadata.get(5).hosts.contains(hosts.get(2))
-                && taskWorkersMetadata.get(5).hosts.size() == 1);
-        assertTrue(taskWorkersMetadata.get(6).hosts.contains(hosts.get(0))
-                && taskWorkersMetadata.get(6).hosts.size() == 1);
-        assertTrue(taskWorkersMetadata.get(7).hosts.contains(hosts.get(1))
-                && taskWorkersMetadata.get(7).hosts.size() == 1);
-        assertTrue(taskWorkersMetadata.get(8).hosts.contains(hosts.get(2))
-                && taskWorkersMetadata.get(8).hosts.size() == 1);
+        Assertions.assertThat(taskWorkersMetadata.get(0).hosts).containsExactly(hosts.get(0));
+        Assertions.assertThat(taskWorkersMetadata.get(1).hosts).containsExactly(hosts.get(1));
+        Assertions.assertThat(taskWorkersMetadata.get(2).hosts).containsExactly(hosts.get(2));
+        Assertions.assertThat(taskWorkersMetadata.get(3).hosts).containsExactly(hosts.get(0));
+        Assertions.assertThat(taskWorkersMetadata.get(4).hosts).containsExactly(hosts.get(1));
+        Assertions.assertThat(taskWorkersMetadata.get(5).hosts).containsExactly(hosts.get(2));
+        Assertions.assertThat(taskWorkersMetadata.get(6).hosts).containsExactly(hosts.get(0));
+        Assertions.assertThat(taskWorkersMetadata.get(7).hosts).containsExactly(hosts.get(1));
+        Assertions.assertThat(taskWorkersMetadata.get(8).hosts).containsExactly(hosts.get(2));
     }
 
     @Test
@@ -163,17 +115,17 @@ void assignWithNineHostsAndSixWorkers() {
         robinAssignor.assign(hosts, taskWorkersMetadata);
         robinAssignor.assign(hosts, taskWorkersMetadata);
         org.assertj.core.api.Assertions.assertThat(taskWorkersMetadata.get(0).hosts)
-                .containsExactlyInAnyOrder((hosts.get(0)), (hosts.get(3)), (hosts.get(6)));
+                .containsExactlyInAnyOrder((hosts.get(0)), (hosts.get(6)));
         org.assertj.core.api.Assertions.assertThat(taskWorkersMetadata.get(1).hosts)
-                .containsExactlyInAnyOrder((hosts.get(0)), (hosts.get(3)), (hosts.get(6)));
+                .containsExactlyInAnyOrder((hosts.get(1)), (hosts.get(7)));
         org.assertj.core.api.Assertions.assertThat(taskWorkersMetadata.get(2).hosts)
-                .containsExactlyInAnyOrder((hosts.get(1)), (hosts.get(4)), (hosts.get(7)));
+                .containsExactlyInAnyOrder((hosts.get(2)), (hosts.get(8)));
         org.assertj.core.api.Assertions.assertThat(taskWorkersMetadata.get(3).hosts)
-                .containsExactlyInAnyOrder((hosts.get(1)), (hosts.get(4)), (hosts.get(7)));
+                .containsExactlyInAnyOrder((hosts.get(3)));
         org.assertj.core.api.Assertions.assertThat(taskWorkersMetadata.get(4).hosts)
-                .containsExactlyInAnyOrder((hosts.get(2)), (hosts.get(5)), (hosts.get(8)));
+                .containsExactlyInAnyOrder((hosts.get(4)));
         org.assertj.core.api.Assertions.assertThat(taskWorkersMetadata.get(5).hosts)
-                .containsExactlyInAnyOrder((hosts.get(2)), (hosts.get(5)), (hosts.get(8)));
+                .containsExactlyInAnyOrder((hosts.get(5)));
     }
 
     public List<HostModel> generateHosts(int q) {