Revert "Revert "Disable topk optimisation on dispatch when content distribution is se…""

6 files changed, 102 insertions, 36 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/InterleavedSearchInvoker.java b/container-search/src/main/java/com/yahoo/search/dispatch/InterleavedSearchInvoker.java
index d8fb7b46440..c60e1bf39cb 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/InterleavedSearchInvoker.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/InterleavedSearchInvoker.java
@@ -43,6 +43,7 @@ public class InterleavedSearchInvoker extends SearchInvoker implements ResponseM
     private final SearchCluster searchCluster;
     private final LinkedBlockingQueue<SearchInvoker> availableForProcessing;
     private final Set<Integer> alreadyFailedNodes;
+    private final boolean isContentWellBalanced;
     private Query query;
 
     private boolean adaptiveTimeoutCalculated = false;
@@ -59,13 +60,14 @@ public class InterleavedSearchInvoker extends SearchInvoker implements ResponseM
     private boolean timedOut = false;
     private boolean degradedByMatchPhase = false;
 
-    public InterleavedSearchInvoker(Collection<SearchInvoker> invokers, SearchCluster searchCluster, Set<Integer> alreadyFailedNodes) {
+    public InterleavedSearchInvoker(Collection<SearchInvoker> invokers, boolean isContentWellBalanced, SearchCluster searchCluster, Set<Integer> alreadyFailedNodes) {
         super(Optional.empty());
         this.invokers = Collections.newSetFromMap(new IdentityHashMap<>());
         this.invokers.addAll(invokers);
         this.searchCluster = searchCluster;
         this.availableForProcessing = newQueue();
         this.alreadyFailedNodes = alreadyFailedNodes;
+        this.isContentWellBalanced = isContentWellBalanced;
     }
 
     /**
@@ -82,10 +84,13 @@ public class InterleavedSearchInvoker extends SearchInvoker implements ResponseM
         int originalHits = query.getHits();
         int originalOffset = query.getOffset();
         int neededHits = originalHits + originalOffset;
-        Double topkProbabilityOverrride = query.properties().getDouble(Dispatcher.topKProbability);
-        int q = (topkProbabilityOverrride != null)
-                ? searchCluster.estimateHitsToFetch(neededHits, invokers.size(), topkProbabilityOverrride)
-                : searchCluster.estimateHitsToFetch(neededHits, invokers.size());
+        int q = neededHits;
+        if (isContentWellBalanced) {
+            Double topkProbabilityOverrride = query.properties().getDouble(Dispatcher.topKProbability);
+            q = (topkProbabilityOverrride != null)
+                    ? searchCluster.estimateHitsToFetch(neededHits, invokers.size(), topkProbabilityOverrride)
+                    : searchCluster.estimateHitsToFetch(neededHits, invokers.size());
+        }
         query.setHits(q);
         query.setOffset(0);
 
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 03160e6c9c7..f65e0e43757 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
@@ -46,12 +46,12 @@ public abstract class InvokerFactory {
      * @return the invoker or empty if some node in the
      *         list is invalid and the remaining coverage is not sufficient
      */
-    public Optional<SearchInvoker> createSearchInvoker(VespaBackEndSearcher searcher,
-                                                       Query query,
-                                                       OptionalInt groupId,
-                                                       List<Node> nodes,
-                                                       boolean acceptIncompleteCoverage,
-                                                       int maxHits) {
+    Optional<SearchInvoker> createSearchInvoker(VespaBackEndSearcher searcher,
+                                                Query query,
+                                                OptionalInt groupId,
+                                                List<Node> nodes,
+                                                boolean acceptIncompleteCoverage,
+                                                int maxHits) {
         List<SearchInvoker> invokers = new ArrayList<>(nodes.size());
         Set<Integer> failed = null;
         for (Node node : nodes) {
@@ -90,7 +90,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, failed));
+            return Optional.of(new InterleavedSearchInvoker(invokers, searchCluster.isGroupWellBalanced(groupId), searchCluster, failed));
         }
     }
 
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 ec616a18e09..23255fff330 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
@@ -22,6 +22,8 @@ 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 static double MAX_UNBALANCE = 0.10; // If documents on a node is more than 10% off from the average the group is unbalanced
 
     public Group(int id, List<Node> nodes) {
         this.id = id;
@@ -53,25 +55,30 @@ public class Group {
     }
 
     public int workingNodes() {
-        int nodesUp = 0;
-        for (Node node : nodes) {
-            if (node.isWorking() == Boolean.TRUE) {
-                nodesUp++;
-            }
-        }
-        return nodesUp;
+        return (int) nodes.stream().filter(node -> node.isWorking() == Boolean.TRUE).count();
     }
 
     void aggregateNodeValues() {
-        activeDocuments.set(nodes.stream().filter(node -> node.isWorking() == Boolean.TRUE).mapToLong(Node::getActiveDocuments).sum());
-        isBlockingWrites.set(nodes.stream().anyMatch(node -> node.isBlockingWrites()));
+        long activeDocs = nodes.stream().filter(node -> node.isWorking() == Boolean.TRUE).mapToLong(Node::getActiveDocuments).sum();
+        activeDocuments.set(activeDocs);
+        isBlockingWrites.set(nodes.stream().anyMatch(Node::isBlockingWrites));
+        int numWorkingNodes = workingNodes();
+        if (numWorkingNodes > 0) {
+            long average = activeDocs / numWorkingNodes;
+            long deviation = nodes.stream().filter(node -> node.isWorking() == Boolean.TRUE).mapToLong(node -> Math.abs(node.getActiveDocuments() - average)).sum();
+            isContentWellBalanced.set(deviation <= (activeDocs * MAX_UNBALANCE));
+        } else {
+            isContentWellBalanced.set(true);
+        }
+
     }
 
-    /** Returns the active documents on this node. If unknown, 0 is returned. */
+    /** Returns the active documents on this group. If unknown, 0 is returned. */
     long getActiveDocuments() { return activeDocuments.get(); }
 
     /** Returns whether any node in this group is currently blocking write operations */
     public boolean isBlockingWrites() { return isBlockingWrites.get(); }
+    public boolean isContentWellBalanced() { return isContentWellBalanced.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 2f62b07ac04..56165ec150b 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
@@ -368,6 +368,12 @@ public class SearchCluster implements NodeManager<Node> {
         return workingNodes + nodesAllowedDown >= nodesInGroup;
     }
 
+    public boolean isGroupWellBalanced(OptionalInt groupId) {
+        if (groupId.isEmpty()) return false;
+        Group group = groups().get(groupId);
+        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/InterleavedSearchInvokerTest.java b/container-search/src/test/java/com/yahoo/search/dispatch/InterleavedSearchInvokerTest.java
index 2bfa778a2ba..8cab7884152 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/InterleavedSearchInvokerTest.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/InterleavedSearchInvokerTest.java
@@ -204,8 +204,8 @@ public class InterleavedSearchInvokerTest {
     private static final List<Double> A5Aux = Arrays.asList(-1.0,11.0,8.5,7.5,-7.0,3.0,2.0);
     private static final List<Double> B5Aux = Arrays.asList(9.0,8.0,-3.0,7.0,6.0,1.0, -1.0);
 
-    private void validateThatTopKProbabilityOverrideTakesEffect(Double topKProbability, int expectedK) throws IOException {
-        InterleavedSearchInvoker invoker = createInterLeavedTestInvoker(A5, B5);
+    private void validateThatTopKProbabilityOverrideTakesEffect(Double topKProbability, int expectedK, boolean isContentWellBalanced) throws IOException {
+        InterleavedSearchInvoker invoker = createInterLeavedTestInvoker(A5, B5, isContentWellBalanced);
         query.setHits(8);
         query.properties().set(Dispatcher.topKProbability, topKProbability);
         SearchInvoker [] invokers = invoker.invokers().toArray(new SearchInvoker[0]);
@@ -227,13 +227,17 @@ public class InterleavedSearchInvokerTest {
 
     @Test
     public void requireThatTopKProbabilityOverrideTakesEffect() throws IOException {
-        validateThatTopKProbabilityOverrideTakesEffect(null, 8);
-        validateThatTopKProbabilityOverrideTakesEffect(0.8, 7);
+        validateThatTopKProbabilityOverrideTakesEffect(null, 8, true);
+        validateThatTopKProbabilityOverrideTakesEffect(0.8, 7, true);
+    }
+    @Test
+    public void requireThatTopKProbabilityOverrideIsDisabledOnContentSkew() throws IOException {
+        validateThatTopKProbabilityOverrideTakesEffect(0.8, 8, false);
     }
 
     @Test
     public void requireThatMergeOfConcreteHitsObeySorting() throws IOException {
-        InterleavedSearchInvoker invoker = createInterLeavedTestInvoker(A5, B5);
+        InterleavedSearchInvoker invoker = createInterLeavedTestInvoker(A5, B5, true);
         query.setHits(12);
         Result result = invoker.search(query, null);
         assertEquals(10, result.hits().size());
@@ -242,7 +246,7 @@ public class InterleavedSearchInvokerTest {
         assertEquals(0, result.getQuery().getOffset());
         assertEquals(12, result.getQuery().getHits());
 
-        invoker = createInterLeavedTestInvoker(B5, A5);
+        invoker = createInterLeavedTestInvoker(B5, A5, true);
         result = invoker.search(query, null);
         assertEquals(10, result.hits().size());
         assertEquals(11.0, result.hits().get(0).getRelevance().getScore(), DELTA);
@@ -253,7 +257,7 @@ public class InterleavedSearchInvokerTest {
 
     @Test
     public void requireThatMergeOfConcreteHitsObeyOffset() throws IOException {
-        InterleavedSearchInvoker invoker = createInterLeavedTestInvoker(A5, B5);
+        InterleavedSearchInvoker invoker = createInterLeavedTestInvoker(A5, B5, true);
         query.setHits(3);
         query.setOffset(5);
         Result result = invoker.search(query, null);
@@ -263,7 +267,7 @@ public class InterleavedSearchInvokerTest {
         assertEquals(0, result.getQuery().getOffset());
         assertEquals(3, result.getQuery().getHits());
 
-        invoker = createInterLeavedTestInvoker(B5, A5);
+        invoker = createInterLeavedTestInvoker(B5, A5, true);
         query.setOffset(5);
         result = invoker.search(query, null);
         assertEquals(3, result.hits().size());
@@ -275,7 +279,7 @@ public class InterleavedSearchInvokerTest {
 
     @Test
     public void requireThatMergeOfConcreteHitsObeyOffsetWithAuxilliaryStuff() throws IOException {
-        InterleavedSearchInvoker invoker = createInterLeavedTestInvoker(A5Aux, B5Aux);
+        InterleavedSearchInvoker invoker = createInterLeavedTestInvoker(A5Aux, B5Aux, true);
         query.setHits(3);
         query.setOffset(5);
         Result result = invoker.search(query, null);
@@ -286,7 +290,7 @@ public class InterleavedSearchInvokerTest {
         assertEquals(0, result.getQuery().getOffset());
         assertEquals(3, result.getQuery().getHits());
 
-        invoker = createInterLeavedTestInvoker(B5Aux, A5Aux);
+        invoker = createInterLeavedTestInvoker(B5Aux, A5Aux, true);
         query.setOffset(5);
         result = invoker.search(query, null);
         assertEquals(7, result.hits().size());
@@ -297,12 +301,13 @@ public class InterleavedSearchInvokerTest {
         assertEquals(3, result.getQuery().getHits());
     }
 
-    private static InterleavedSearchInvoker createInterLeavedTestInvoker(List<Double> a, List<Double> b) {
+    private static InterleavedSearchInvoker createInterLeavedTestInvoker(List<Double> a, List<Double> b,
+                                                                         boolean isContentWellBalanced) {
         SearchCluster cluster = new MockSearchCluster("!", 1, 2);
         List<SearchInvoker> invokers = new ArrayList<>();
         invokers.add(createInvoker(a, 0));
         invokers.add(createInvoker(b, 1));
-        InterleavedSearchInvoker invoker = new InterleavedSearchInvoker(invokers, cluster, Collections.emptySet());
+        InterleavedSearchInvoker invoker = new InterleavedSearchInvoker(invokers, isContentWellBalanced, cluster, Collections.emptySet());
         invoker.responseAvailable(invokers.get(0));
         invoker.responseAvailable(invokers.get(1));
         return invoker;
@@ -353,7 +358,7 @@ public class InterleavedSearchInvokerTest {
             invokers.add(new MockInvoker(i));
         }
 
-        return new InterleavedSearchInvoker(invokers, searchCluster, null) {
+        return new InterleavedSearchInvoker(invokers, false, searchCluster, null) {
             @Override
             protected long currentTime() {
                 return clock.millis();
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 09024150a9a..c6fd48836fe 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
@@ -8,7 +8,6 @@ import com.yahoo.net.HostName;
 import com.yahoo.prelude.Pong;
 import com.yahoo.search.cluster.ClusterMonitor;
 import com.yahoo.search.dispatch.MockSearchCluster;
-import com.yahoo.search.dispatch.TopKEstimator;
 import com.yahoo.search.result.ErrorMessage;
 import org.junit.Test;
 
@@ -335,4 +334,48 @@ public class SearchClusterTest {
         assertEquals(3, node.getLastReceivedPongId());
     }
 
+    @Test
+    public void requireThatEmptyGroupIsInBalance() {
+        Group group = new Group(0, new ArrayList<>());
+        assertTrue(group.isContentWellBalanced());
+        group.aggregateNodeValues();
+        assertTrue(group.isContentWellBalanced());
+    }
+
+    @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());
+        group.aggregateNodeValues();
+        assertTrue(group.isContentWellBalanced());
+        group.nodes().get(0).setActiveDocuments(1000);
+        group.aggregateNodeValues();
+        assertTrue(group.isContentWellBalanced());
+    }
+
+    @Test
+    public void requireThatMultiNodeGroupDetectsBalance() {
+        Group group = new Group(0, Arrays.asList(new Node(1, "n1", 1), new Node(2, "n2", 1)));
+        assertTrue(group.isContentWellBalanced());
+        group.nodes().forEach(node -> node.setWorking(true));
+        assertTrue(group.isContentWellBalanced());
+        group.aggregateNodeValues();
+        assertTrue(group.isContentWellBalanced());
+        group.nodes().get(0).setActiveDocuments(1000);
+        group.aggregateNodeValues();
+        assertFalse(group.isContentWellBalanced());
+        group.nodes().get(1).setActiveDocuments(100);
+        group.aggregateNodeValues();
+        assertFalse(group.isContentWellBalanced());
+        group.nodes().get(1).setActiveDocuments(800);
+        group.aggregateNodeValues();
+        assertFalse(group.isContentWellBalanced());
+        group.nodes().get(1).setActiveDocuments(818);
+        group.aggregateNodeValues();
+        assertFalse(group.isContentWellBalanced());
+        group.nodes().get(1).setActiveDocuments(819);
+        group.aggregateNodeValues();
+        assertTrue(group.isContentWellBalanced());
+    }
 }