Add best-of-random-2 dispatch policy.

3 files changed, 74 insertions, 12 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/Dispatcher.java b/container-search/src/main/java/com/yahoo/search/dispatch/Dispatcher.java
index 68a8e351b34..8bec018f8f8 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/Dispatcher.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/Dispatcher.java
@@ -96,6 +96,20 @@ public class Dispatcher extends AbstractComponent {
         this(new ClusterMonitor<>(searchCluster, true), searchCluster, dispatchConfig, new RpcInvokerFactory(resourcePool, searchCluster), metric);
     }
 
+    private static LoadBalancer.Policy toLoadBalancerPolicy(DispatchConfig.DistributionPolicy.Enum policy) {
+        if (policy == DispatchConfig.DistributionPolicy.ROUNDROBIN) {
+            return LoadBalancer.Policy.ROUNDROBIN;
+        } else if (policy == DispatchConfig.DistributionPolicy.BEST_OF_RANDOM_2) {
+            return LoadBalancer.Policy.BEST_OF_RANDOM_2;
+        } else if (policy == DispatchConfig.DistributionPolicy.LATENCY_AMORTIZED_OVER_REQUESTS) {
+            return LoadBalancer.Policy.LATENCY_AMORTIZED_OVER_REQUESTS;
+        } else if (policy == DispatchConfig.DistributionPolicy.LATENCY_AMORTIZED_OVER_TIME) {
+            return LoadBalancer.Policy.LATENCY_AMORTIZED_OVER_TIME;
+        } else {
+            return LoadBalancer.Policy.LATENCY_AMORTIZED_OVER_REQUESTS;
+        }
+    }
+
     /* Protected for simple mocking in tests. Beware that searchCluster is shutdown on in deconstruct() */
     protected Dispatcher(ClusterMonitor<Node> clusterMonitor,
                          SearchCluster searchCluster,
@@ -107,8 +121,7 @@ public class Dispatcher extends AbstractComponent {
 
         this.searchCluster = searchCluster;
         this.clusterMonitor = clusterMonitor;
-        this.loadBalancer = new LoadBalancer(searchCluster,
-                                  dispatchConfig.distributionPolicy() == DispatchConfig.DistributionPolicy.ROUNDROBIN);
+        this.loadBalancer = new LoadBalancer(searchCluster, toLoadBalancerPolicy(dispatchConfig.distributionPolicy()));
         this.invokerFactory = invokerFactory;
         this.metric = metric;
         this.metricContext = metric.createContext(null);
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/LoadBalancer.java b/container-search/src/main/java/com/yahoo/search/dispatch/LoadBalancer.java
index 4c0bcb38d15..4c2af5b885f 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/LoadBalancer.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/LoadBalancer.java
@@ -31,16 +31,22 @@ public class LoadBalancer {
     private final List<GroupStatus> scoreboard;
     private final GroupScheduler scheduler;
 
-    public LoadBalancer(SearchCluster searchCluster, boolean roundRobin) {
+    public enum Policy { ROUNDROBIN, LATENCY_AMORTIZED_OVER_REQUESTS, LATENCY_AMORTIZED_OVER_TIME, BEST_OF_RANDOM_2}
+
+    public LoadBalancer(SearchCluster searchCluster, Policy policy) {
         this.scoreboard = new ArrayList<>(searchCluster.groups().size());
         for (Group group : searchCluster.orderedGroups()) {
             scoreboard.add(new GroupStatus(group));
         }
-        if (roundRobin || scoreboard.size() == 1) {
-            this.scheduler = new RoundRobinScheduler(scoreboard);
-        } else {
-            this.scheduler = new AdaptiveScheduler(new Random(), scoreboard);
-        }
+        if (scoreboard.size() == 1)
+            policy = Policy.ROUNDROBIN;
+
+        this.scheduler = switch (policy) {
+            case ROUNDROBIN: yield new RoundRobinScheduler(scoreboard);
+            case BEST_OF_RANDOM_2: yield new BestOfRandom2(new Random(), scoreboard);
+            case LATENCY_AMORTIZED_OVER_REQUESTS: yield new AdaptiveScheduler(new Random(), scoreboard);
+            case LATENCY_AMORTIZED_OVER_TIME: yield new AdaptiveScheduler(new Random(), scoreboard);
+        };
     }
 
     /**
@@ -239,4 +245,47 @@ public class LoadBalancer {
         }
     }
 
+    static class BestOfRandom2 implements GroupScheduler {
+        private final Random random;
+        private final List<GroupStatus> scoreboard;
+        public BestOfRandom2(Random random, List<GroupStatus> scoreboard) {
+            this.random = random;
+            this.scoreboard = scoreboard;
+        }
+        @Override
+        public Optional<GroupStatus> takeNextGroup(Set<Integer> rejectedGroups) {
+            GroupStatus gs = selectBestOf2(rejectedGroups, true);
+            return (gs != null)
+                    ? Optional.of(gs)
+                    : Optional.ofNullable(selectBestOf2(rejectedGroups, false));
+        }
+
+        private GroupStatus selectBestOf2(Set<Integer> rejectedGroups, boolean requireCoverage) {
+            List<Integer> candidates = new ArrayList<>(scoreboard.size());
+            for (int i=0; i < scoreboard.size(); i++) {
+                GroupStatus gs = scoreboard.get(i);
+                if (rejectedGroups == null || !rejectedGroups.contains(gs.group.id())) {
+                    if (!requireCoverage || gs.group.hasSufficientCoverage()) {
+                        candidates.add(i);
+                    }
+                }
+            }
+            GroupStatus candA = selectRandom(candidates);
+            GroupStatus candB = selectRandom(candidates);
+            if (candA == null) return candB;
+            if (candB == null) return candA;
+            if (candB.allocations < candA.allocations) return candB;
+            return candA;
+        }
+        private GroupStatus selectRandom(List<Integer> candidates) {
+            if ( ! candidates.isEmpty()) {
+                int index = random.nextInt(candidates.size());
+                Integer groupIndex = candidates.remove(index);
+                return scoreboard.get(groupIndex);
+            }
+            return null;
+        }
+
+    }
+
 }
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/LoadBalancerTest.java b/container-search/src/test/java/com/yahoo/search/dispatch/LoadBalancerTest.java
index e9ed1c48302..350ede820eb 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/LoadBalancerTest.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/LoadBalancerTest.java
@@ -29,7 +29,7 @@ public class LoadBalancerTest {
     void requireThatLoadBalancerServesSingleNodeSetups() {
         Node n1 = new Node(0, "test-node1", 0);
         SearchCluster cluster = new SearchCluster("a", createDispatchConfig(n1), null, null);
-        LoadBalancer lb = new LoadBalancer(cluster, true);
+        LoadBalancer lb = new LoadBalancer(cluster, LoadBalancer.Policy.ROUNDROBIN);
 
         Optional<Group> grp = lb.takeGroup(null);
         Group group = grp.orElseGet(() -> {
@@ -43,7 +43,7 @@ public class LoadBalancerTest {
         Node n1 = new Node(0, "test-node1", 0);
         Node n2 = new Node(1, "test-node2", 1);
         SearchCluster cluster = new SearchCluster("a", createDispatchConfig(n1, n2), null, null);
-        LoadBalancer lb = new LoadBalancer(cluster, true);
+        LoadBalancer lb = new LoadBalancer(cluster, LoadBalancer.Policy.ROUNDROBIN);
 
         Optional<Group> grp = lb.takeGroup(null);
         Group group = grp.orElseGet(() -> {
@@ -59,7 +59,7 @@ public class LoadBalancerTest {
         Node n3 = new Node(0, "test-node3", 1);
         Node n4 = new Node(1, "test-node4", 1);
         SearchCluster cluster = new SearchCluster("a", createDispatchConfig(n1, n2, n3, n4), null, null);
-        LoadBalancer lb = new LoadBalancer(cluster, true);
+        LoadBalancer lb = new LoadBalancer(cluster, LoadBalancer.Policy.ROUNDROBIN);
 
         Optional<Group> grp = lb.takeGroup(null);
         assertTrue(grp.isPresent());
@@ -70,7 +70,7 @@ public class LoadBalancerTest {
         Node n1 = new Node(0, "test-node1", 0);
         Node n2 = new Node(1, "test-node2", 1);
         SearchCluster cluster = new SearchCluster("a", createDispatchConfig(n1, n2), null, null);
-        LoadBalancer lb = new LoadBalancer(cluster, true);
+        LoadBalancer lb = new LoadBalancer(cluster, LoadBalancer.Policy.ROUNDROBIN);
 
         // get first group
         Optional<Group> grp = lb.takeGroup(null);