Infer group

9 files changed, 36 insertions, 57 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 9b92a78a7c9..af731e3ade0 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
@@ -186,8 +186,8 @@ public class Dispatcher extends AbstractComponent {
             if (nodes.isEmpty()) return Optional.empty();
 
             query.trace(false, 2, "Dispatching with search path ", searchPath);
-            return invokerFactory.createSearchInvoker(searcher, query,
-                                                      OptionalInt.empty(),
+            return invokerFactory.createSearchInvoker(searcher,
+                                                      query,
                                                       nodes,
                                                       true,
                                                       maxHitsPerNode);
@@ -203,7 +203,6 @@ public class Dispatcher extends AbstractComponent {
             query.trace(false, 2, "Dispatching to ", node);
             return invokerFactory.createSearchInvoker(searcher,
                                                       query,
-                                                      OptionalInt.empty(),
                                                       List.of(node),
                                                       true,
                                                       maxHitsPerNode)
@@ -222,7 +221,6 @@ public class Dispatcher extends AbstractComponent {
             boolean acceptIncompleteCoverage = (i == max - 1);
             Optional<SearchInvoker> invoker = invokerFactory.createSearchInvoker(searcher,
                                                                                  query,
-                                                                                 OptionalInt.of(group.id()),
                                                                                  group.nodes(),
                                                                                  acceptIncompleteCoverage,
                                                                                  maxHitsPerNode);
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/InvokerFactory.java b/container-search/src/main/java/com/yahoo/search/dispatch/InvokerFactory.java
index 1de274ce6cf..bd66ca88622 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/InvokerFactory.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/InvokerFactory.java
@@ -4,6 +4,7 @@ package com.yahoo.search.dispatch;
 import com.yahoo.prelude.fastsearch.VespaBackEndSearcher;
 import com.yahoo.search.Query;
 import com.yahoo.search.Result;
+import com.yahoo.search.dispatch.searchcluster.Group;
 import com.yahoo.search.dispatch.searchcluster.Node;
 import com.yahoo.search.dispatch.searchcluster.SearchCluster;
 import com.yahoo.search.result.Coverage;
@@ -13,7 +14,6 @@ import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
-import java.util.OptionalInt;
 import java.util.Set;
 
 /**
@@ -39,8 +39,7 @@ public abstract class InvokerFactory {
      *
      * @param searcher the searcher processing the query
      * @param query the search query being processed
-     * @param groupId the id of the node group to which the nodes belong
-     * @param nodes pre-selected list of content nodes
+     * @param nodes pre-selected list of content nodes, all in a group or a subset of a group
      * @param acceptIncompleteCoverage if some of the nodes are unavailable and this parameter is
      *                                 false, verify that the remaining set of nodes has sufficient coverage
      * @return the invoker or empty if some node in the
@@ -48,10 +47,10 @@ public abstract class InvokerFactory {
      */
     Optional<SearchInvoker> createSearchInvoker(VespaBackEndSearcher searcher,
                                                 Query query,
-                                                OptionalInt groupId,
                                                 List<Node> nodes,
                                                 boolean acceptIncompleteCoverage,
                                                 int maxHits) {
+        Group group = searchCluster.group(nodes.get(0).group()).get(); // Nodes must be of the same group
         List<SearchInvoker> invokers = new ArrayList<>(nodes.size());
         Set<Integer> failed = null;
         for (Node node : nodes) {
@@ -90,7 +89,7 @@ public abstract class InvokerFactory {
         if (invokers.size() == 1 && failed == null) {
             return Optional.of(invokers.get(0));
         } else {
-            return Optional.of(new InterleavedSearchInvoker(invokers, searchCluster.isGroupWellBalanced(groupId), searchCluster, failed));
+            return Optional.of(new InterleavedSearchInvoker(invokers, group.isBalanced(), searchCluster, failed));
         }
     }
 
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/SearchPath.java b/container-search/src/main/java/com/yahoo/search/dispatch/SearchPath.java
index f6480f80c01..b29c3297aea 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/SearchPath.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/SearchPath.java
@@ -41,7 +41,7 @@ public class SearchPath {
         if (sp.isPresent()) {
             return sp.get().mapToNodes(cluster);
         } else {
-            return Collections.emptyList();
+            return List.of();
         }
     }
 
@@ -75,7 +75,7 @@ public class SearchPath {
 
     private List<Node> mapToNodes(SearchCluster cluster) {
         if (cluster.groups().isEmpty()) {
-            return Collections.emptyList();
+            return List.of();
         }
 
         Group selectedGroup = selectGroup(cluster);
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/Group.java b/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/Group.java
index 7faad9d51cc..b6b0c15887f 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/Group.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/Group.java
@@ -23,7 +23,7 @@ public class Group {
     private final AtomicBoolean hasFullCoverage = new AtomicBoolean(true);
     private final AtomicLong activeDocuments = new AtomicLong(0);
     private final AtomicBoolean isBlockingWrites = new AtomicBoolean(false);
-    private final AtomicBoolean isContentWellBalanced = new AtomicBoolean(true);
+    private final AtomicBoolean isBalanced = new AtomicBoolean(true);
     private final static double MAX_UNBALANCE = 0.10; // If documents on a node is more than 10% off from the average the group is unbalanced
     private static final Logger log = Logger.getLogger(Group.class.getName());
 
@@ -69,13 +69,13 @@ public class Group {
             long average = activeDocs / numWorkingNodes;
             long deviation = nodes.stream().filter(node -> node.isWorking() == Boolean.TRUE).mapToLong(node -> Math.abs(node.getActiveDocuments() - average)).sum();
             boolean isDeviationSmall = deviation <= maxUnbalance(activeDocs);
-            if ((!isContentWellBalanced.get() || isDeviationSmall != isContentWellBalanced.get()) && (activeDocs > 0)) {
+            if ((!isBalanced.get() || isDeviationSmall != isBalanced.get()) && (activeDocs > 0)) {
                 log.info("Content is " + (isDeviationSmall ? "" : "not ") + "well balanced. Current deviation = " + deviation*100/activeDocs + " %" +
                          ". activeDocs = " + activeDocs + ", deviation = " + deviation + ", average = " + average);
-                isContentWellBalanced.set(isDeviationSmall);
+                isBalanced.set(isDeviationSmall);
             }
         } else {
-            isContentWellBalanced.set(true);
+            isBalanced.set(true);
         }
     }
 
@@ -88,7 +88,9 @@ public class Group {
 
     /** Returns whether any node in this group is currently blocking write operations */
     public boolean isBlockingWrites() { return isBlockingWrites.get(); }
-    public boolean isContentWellBalanced() { return isContentWellBalanced.get(); }
+
+    /** Returns whether the nodes in the group have about the same number of documents */
+    public boolean isBalanced() { return isBalanced.get(); }
 
     public boolean isFullCoverageStatusChanged(boolean hasFullCoverageNow) {
         boolean previousState = hasFullCoverage.getAndSet(hasFullCoverageNow);
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/SearchCluster.java b/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/SearchCluster.java
index 9ae25518969..f80fdb6c4a5 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/SearchCluster.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/SearchCluster.java
@@ -357,12 +357,6 @@ public class SearchCluster implements NodeManager<Node> {
         return true;
     }
 
-    public boolean isGroupWellBalanced(OptionalInt groupId) {
-        if (groupId.isEmpty()) return false;
-        Group group = groups().get(groupId.getAsInt());
-        return (group != null) && group.isContentWellBalanced();
-    }
-
     /**
      * Calculate whether a subset of nodes in a group has enough coverage
      */
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/DispatcherTest.java b/container-search/src/test/java/com/yahoo/search/dispatch/DispatcherTest.java
index 943390cb10c..be761acf2c2 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/DispatcherTest.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/DispatcherTest.java
@@ -157,7 +157,6 @@ public class DispatcherTest {
         @Override
         public Optional<SearchInvoker> createSearchInvoker(VespaBackEndSearcher searcher,
                                                            Query query,
-                                                           OptionalInt groupId,
                                                            List<Node> nodes,
                                                            boolean acceptIncompleteCoverage,
                                                            int maxHitsPerNode) {
@@ -167,7 +166,7 @@ public class DispatcherTest {
             boolean nonEmpty = events[step].returnInvoker(nodes, acceptIncompleteCoverage);
             step++;
             if (nonEmpty) {
-                return Optional.of(new MockInvoker(nodes.get(0).key(), groupId));
+                return Optional.of(new MockInvoker(nodes.get(0).key()));
             } else {
                 return Optional.empty();
             }
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/MockInvoker.java b/container-search/src/test/java/com/yahoo/search/dispatch/MockInvoker.java
index 53d1a2457d0..d86fcdfc25d 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/MockInvoker.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/MockInvoker.java
@@ -4,6 +4,7 @@ package com.yahoo.search.dispatch;
 import com.yahoo.prelude.fastsearch.FastHit;
 import com.yahoo.search.Query;
 import com.yahoo.search.Result;
+import com.yahoo.search.dispatch.searchcluster.Group;
 import com.yahoo.search.dispatch.searchcluster.Node;
 import com.yahoo.search.result.Coverage;
 import com.yahoo.search.result.Hit;
@@ -17,27 +18,17 @@ import java.util.OptionalInt;
 class MockInvoker extends SearchInvoker {
 
     private final Coverage coverage;
-    private final OptionalInt groupId;
     private Query query;
     private List<Hit> hits;
     int hitsRequested;
 
-    protected MockInvoker(int key, Coverage coverage, OptionalInt groupId) {
+    protected MockInvoker(int key, Coverage coverage) {
         super(Optional.of(new Node(key, "?", 0)));
         this.coverage = coverage;
-        this.groupId = groupId;
-    }
-
-    protected MockInvoker(int key, OptionalInt groupId) {
-        this(key, null, groupId);
-    }
-
-    protected MockInvoker(int key, Coverage coverage) {
-        this(key, coverage, OptionalInt.empty());
     }
 
     protected MockInvoker(int key) {
-        this(key, null, OptionalInt.empty());
+        this(key, null);
     }
 
     MockInvoker setHits(List<Hit> hits) {
@@ -45,18 +36,15 @@ class MockInvoker extends SearchInvoker {
         return this;
     }
 
-    /** Returns the group to be invoked, if known */
-    public OptionalInt groupId() { return groupId; }
-
     @Override
-    protected Object sendSearchRequest(Query query, Object context) throws IOException {
+    protected Object sendSearchRequest(Query query, Object context) {
         this.query = query;
         hitsRequested = query.getHits();
         return context;
     }
 
     @Override
-    protected InvokerResult getSearchResult(Execution execution) throws IOException {
+    protected InvokerResult getSearchResult(Execution execution) {
         InvokerResult ret = new InvokerResult(query, 10);
         if (coverage != null) {
             ret.getResult().setCoverage(coverage);
@@ -80,8 +68,7 @@ class MockInvoker extends SearchInvoker {
 
     @Override
     public String toString() {
-        return "invoker with key " + distributionKey() +
-               (groupId().isPresent() ? " of group " + groupId().getAsInt() : "");
+        return "invoker with key " + distributionKey();
     }
 
 }
\ No newline at end of file
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterCoverageTest.java b/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterCoverageTest.java
index c9f7469acbb..ee51b983d64 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterCoverageTest.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterCoverageTest.java
@@ -115,7 +115,7 @@ public class SearchClusterCoverageTest {
         node1.setActiveDocuments(0);
 
         tester.pingIterationCompleted();
-        assertTrue(tester.group(0).isContentWellBalanced());
+        assertTrue(tester.group(0).isBalanced());
     }
 
 }
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterTest.java b/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterTest.java
index 48134094faf..f46717ce180 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterTest.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterTest.java
@@ -337,45 +337,45 @@ public class SearchClusterTest {
     @Test
     public void requireThatEmptyGroupIsInBalance() {
         Group group = new Group(0, new ArrayList<>());
-        assertTrue(group.isContentWellBalanced());
+        assertTrue(group.isBalanced());
         group.aggregateNodeValues();
-        assertTrue(group.isContentWellBalanced());
+        assertTrue(group.isBalanced());
     }
 
     @Test
     public void requireThatSingleNodeGroupIsInBalance() {
         Group group = new Group(0, Arrays.asList(new Node(1, "n", 1)));
         group.nodes().forEach(node -> node.setWorking(true));
-        assertTrue(group.isContentWellBalanced());
+        assertTrue(group.isBalanced());
         group.aggregateNodeValues();
-        assertTrue(group.isContentWellBalanced());
+        assertTrue(group.isBalanced());
         group.nodes().get(0).setActiveDocuments(1000);
         group.aggregateNodeValues();
-        assertTrue(group.isContentWellBalanced());
+        assertTrue(group.isBalanced());
     }
 
     @Test
     public void requireThatMultiNodeGroupDetectsBalance() {
         Group group = new Group(0, Arrays.asList(new Node(1, "n1", 1), new Node(2, "n2", 1)));
-        assertTrue(group.isContentWellBalanced());
+        assertTrue(group.isBalanced());
         group.nodes().forEach(node -> node.setWorking(true));
-        assertTrue(group.isContentWellBalanced());
+        assertTrue(group.isBalanced());
         group.aggregateNodeValues();
-        assertTrue(group.isContentWellBalanced());
+        assertTrue(group.isBalanced());
         group.nodes().get(0).setActiveDocuments(1000);
         group.aggregateNodeValues();
-        assertFalse(group.isContentWellBalanced());
+        assertFalse(group.isBalanced());
         group.nodes().get(1).setActiveDocuments(100);
         group.aggregateNodeValues();
-        assertFalse(group.isContentWellBalanced());
+        assertFalse(group.isBalanced());
         group.nodes().get(1).setActiveDocuments(800);
         group.aggregateNodeValues();
-        assertFalse(group.isContentWellBalanced());
+        assertFalse(group.isBalanced());
         group.nodes().get(1).setActiveDocuments(818);
         group.aggregateNodeValues();
-        assertFalse(group.isContentWellBalanced());
+        assertFalse(group.isBalanced());
         group.nodes().get(1).setActiveDocuments(819);
         group.aggregateNodeValues();
-        assertTrue(group.isContentWellBalanced());
+        assertTrue(group.isBalanced());
     }
 }