Merge branch 'master' into bratseth/balanced

7 files changed, 118 insertions, 142 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/Query.java b/container-search/src/main/java/com/yahoo/search/Query.java
index fb7281e1f24..83fa18d847f 100644
--- a/container-search/src/main/java/com/yahoo/search/Query.java
+++ b/container-search/src/main/java/com/yahoo/search/Query.java
@@ -726,8 +726,7 @@ public class Query extends com.yahoo.processing.Request implements Cloneable {
      * if the trace level of the query is sufficiently high.
      *
      * @param message      the message to add
-     * @param includeQuery true to append the query root stringValue
-     *        at the end of the message
+     * @param includeQuery true to append the query root stringValue at the end of the message
      * @param traceLevel   the context level of the message, this method will do nothing
      *                     if the traceLevel of the query is lower than this value
      */
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 d7c9f1dce53..5e38f6b4bdd 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
@@ -116,8 +116,6 @@ public class InterleavedSearchInvoker extends SearchInvoker implements ResponseM
         InvokerResult result = new InvokerResult(query, query.getHits());
         List<LeanHit> merged = Collections.emptyList();
         long nextTimeout = query.getTimeLeft();
-        boolean extraDebug = (query.getOffset() == 0) && (query.getHits() == 7) && log.isLoggable(java.util.logging.Level.FINE);
-        List<InvokerResult> processed = new ArrayList<>();
         var groupingResultAggregator = new GroupingResultAggregator();
         try {
             while (!invokers.isEmpty() && nextTimeout >= 0) {
@@ -127,9 +125,6 @@ public class InterleavedSearchInvoker extends SearchInvoker implements ResponseM
                     break;
                 } else {
                     InvokerResult toMerge = invoker.getSearchResult(execution);
-                    if (extraDebug) {
-                        processed.add(toMerge);
-                    }
                     merged = mergeResult(result.getResult(), toMerge, merged, groupingResultAggregator);
                     ejectInvoker(invoker);
                 }
@@ -143,32 +138,6 @@ public class InterleavedSearchInvoker extends SearchInvoker implements ResponseM
         insertNetworkErrors(result.getResult());
         result.getResult().setCoverage(createCoverage());
 
-        if (extraDebug && merged.size() > 0) {
-            int firstPartId = merged.get(0).getPartId();
-            for (int index = 1; index < merged.size(); index++) {
-                if (merged.get(index).getPartId() != firstPartId) {
-                    extraDebug = false;
-                    log.fine("merged["+index+"/"+merged.size()+"] from partId "+merged.get(index).getPartId()+", first "+firstPartId);
-                    break;
-                }
-            }
-        }
-        if (extraDebug) {
-            log.fine("Interleaved "+processed.size()+" results");
-            for (int pIdx = 0; pIdx < processed.size(); ++pIdx) {
-                var p = processed.get(pIdx);
-                log.fine("InvokerResult "+pIdx+" total hits "+p.getResult().getTotalHitCount());
-                var lean = p.getLeanHits();
-                for (int idx = 0; idx < lean.size(); ++idx) {
-                    var hit = lean.get(idx);
-                    log.fine("lean hit "+idx+" relevance "+hit.getRelevance()+" partid "+hit.getPartId());
-                }
-            }
-            for (int mIdx = 0; mIdx < merged.size(); ++mIdx) {
-                var hit = merged.get(mIdx);
-                log.fine("merged hit "+mIdx+" relevance "+hit.getRelevance()+" partid "+hit.getPartId());                
-            }
-        }
         int needed = query.getOffset() + query.getHits();
         for (int index = query.getOffset(); (index < merged.size()) && (index < needed); index++) {
             result.getLeanHits().add(merged.get(index));
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 7f4d8fc4739..4c0bcb38d15 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
@@ -12,13 +12,14 @@ import java.util.Set;
 import java.util.logging.Logger;
 
 /**
- * LoadBalancer determines which group of content nodes should be accessed next for each search query when the internal java dispatcher is
- * used.
+ * LoadBalancer determines which group of content nodes should be accessed next for each search query when the
+ * internal java dispatcher is used.
+ *
+ * The implementation here is a simplistic least queries in flight + round-robin load balancer
  *
  * @author ollivir
  */
 public class LoadBalancer {
-    // The implementation here is a simplistic least queries in flight + round-robin load balancer
 
     private static final Logger log = Logger.getLogger(LoadBalancer.class.getName());
 
@@ -84,6 +85,7 @@ public class LoadBalancer {
     }
 
     static class GroupStatus {
+
         private final Group group;
         private int allocations = 0;
         private long queries = 0;
@@ -174,24 +176,10 @@ public class LoadBalancer {
          * @return the better of the two
          */
         private static GroupStatus betterGroup(GroupStatus first, GroupStatus second) {
-            if (second == null) {
-                return first;
-            }
-            if (first == null) {
-                return second;
-            }
-
-            // different coverage
-            if (first.group.hasSufficientCoverage() != second.group.hasSufficientCoverage()) {
-                if (!first.group.hasSufficientCoverage()) {
-                    // first doesn't have coverage, second does
-                    return second;
-                } else {
-                    // second doesn't have coverage, first does
-                    return first;
-                }
-            }
-
+            if (second == null) return first;
+            if (first == null) return second;
+            if (first.group.hasSufficientCoverage() != second.group.hasSufficientCoverage())
+                return first.group.hasSufficientCoverage() ? first : second;
             return first;
         }
 
@@ -246,11 +234,8 @@ public class LoadBalancer {
         public Optional<GroupStatus> takeNextGroup(Set<Integer> rejectedGroups) {
             double needle = random.nextDouble();
             Optional<GroupStatus> gs = selectGroup(needle, true, rejectedGroups);
-            if (gs.isPresent()) {
-                return gs;
-            }
-            // fallback - any coverage better than none
-            return selectGroup(needle, false, rejectedGroups);
+            if (gs.isPresent()) return gs;
+            return selectGroup(needle, false, rejectedGroups); // any coverage better than none
         }
     }
 
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 5e12fb550c8..af9834e282a 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
@@ -17,6 +17,7 @@ import java.util.logging.Logger;
 public class Group {
 
     private static final Logger log = Logger.getLogger(Group.class.getName());
+
     private final static double maxContentSkew = 0.10;
     private final static int minDocsPerNodeToRequireLowSkew = 100;
 
@@ -41,8 +42,7 @@ public class Group {
 
     /**
      * Returns the unique identity of this group.
-     * NOTE: This is a contiguous index from 0, NOT necessarily the group id assigned
-     * by the user or node repo.
+     * NOTE: This is a contiguous index from 0, NOT necessarily the group id assigned by the user or node repo.
      */
     public int id() { return id; }
 
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 36d7e7a85a9..f8c4627473d 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
@@ -1,10 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.search.dispatch.searchcluster;
 
-import com.google.common.collect.ImmutableCollection;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableMultimap;
 import com.google.common.math.Quantiles;
 import com.yahoo.container.handler.VipStatus;
 import com.yahoo.net.HostName;
@@ -32,11 +29,10 @@ public class SearchCluster implements NodeManager<Node> {
     private static final Logger log = Logger.getLogger(SearchCluster.class.getName());
 
     private final DispatchConfig dispatchConfig;
-    private final int size;
     private final String clusterId;
-    private final ImmutableMap<Integer, Group> groups;
-    private final ImmutableMultimap<String, Node> nodesByHost;
-    private final ImmutableList<Group> orderedGroups;
+    private final Map<Integer, Group> groups;
+    private final List<Group> orderedGroups;
+    private final List<Node> nodes;
     private final VipStatus vipStatus;
     private final PingFactory pingFactory;
     private final TopKEstimator hitEstimator;
@@ -60,8 +56,7 @@ public class SearchCluster implements NodeManager<Node> {
         this.vipStatus = vipStatus;
         this.pingFactory = pingFactory;
 
-        List<Node> nodes = toNodes(dispatchConfig);
-        this.size = nodes.size();
+        this.nodes = toNodes(dispatchConfig);
 
         // Create groups
         ImmutableMap.Builder<Integer, Group> groupsBuilder = new ImmutableMap.Builder<>();
@@ -72,16 +67,10 @@ public class SearchCluster implements NodeManager<Node> {
         this.groups = groupsBuilder.build();
         LinkedHashMap<Integer, Group> groupIntroductionOrder = new LinkedHashMap<>();
         nodes.forEach(node -> groupIntroductionOrder.put(node.group(), groups.get(node.group())));
-        this.orderedGroups = ImmutableList.<Group>builder().addAll(groupIntroductionOrder.values()).build();
+        this.orderedGroups = List.copyOf(groupIntroductionOrder.values());
 
-        // Index nodes by host
-        ImmutableMultimap.Builder<String, Node> nodesByHostBuilder = new ImmutableMultimap.Builder<>();
-        for (Node node : nodes)
-            nodesByHostBuilder.put(node.hostname(), node);
-        this.nodesByHost = nodesByHostBuilder.build();
         hitEstimator = new TopKEstimator(30.0, dispatchConfig.topKProbability(), SKEW_FACTOR);
-
-        this.localCorpusDispatchTarget = findLocalCorpusDispatchTarget(HostName.getLocalhost(), nodesByHost, groups);
+        this.localCorpusDispatchTarget = findLocalCorpusDispatchTarget(HostName.getLocalhost(), nodes, groups);
     }
 
     @Override
@@ -95,13 +84,15 @@ public class SearchCluster implements NodeManager<Node> {
     }
 
     private static Optional<Node> findLocalCorpusDispatchTarget(String selfHostname,
-                                                                ImmutableMultimap<String, Node> nodesByHost,
-                                                                ImmutableMap<Integer, Group> groups) {
+                                                                List<Node> nodes,
+                                                                Map<Integer, Group> groups) {
         // A search node in the search cluster in question is configured on the same host as the currently running container.
         // It has all the data <==> No other nodes in the search cluster have the same group id as this node.
         //         That local search node responds.
         // The search cluster to be searched has at least as many nodes as the container cluster we're running in.
-        ImmutableCollection<Node> localSearchNodes = nodesByHost.get(selfHostname);
+        List<Node> localSearchNodes = nodes.stream()
+                                           .filter(node -> node.hostname().equals(selfHostname))
+                                           .collect(Collectors.toList());
         // Only use direct dispatch if we have exactly 1 search node on the same machine:
         if (localSearchNodes.size() != 1) return Optional.empty();
 
@@ -114,25 +105,24 @@ public class SearchCluster implements NodeManager<Node> {
         return Optional.of(localSearchNode);
     }
 
-    private static ImmutableList<Node> toNodes(DispatchConfig dispatchConfig) {
-        ImmutableList.Builder<Node> nodesBuilder = new ImmutableList.Builder<>();
-        for (DispatchConfig.Node node : dispatchConfig.node())
-            nodesBuilder.add(new Node(node.key(), node.host(), node.group()));
-        return nodesBuilder.build();
+    private static List<Node> toNodes(DispatchConfig dispatchConfig) {
+        return dispatchConfig.node().stream()
+                             .map(n -> new Node(n.key(), n.host(), n.group()))
+                             .collect(Collectors.toUnmodifiableList());
     }
 
     public DispatchConfig dispatchConfig() {
         return dispatchConfig;
     }
 
-    /** Returns the number of nodes in this cluster (across all groups) */
-    public int size() { return size; }
+    /** Returns an immutable list of all nodes in this. */
+    public List<Node> nodes() { return nodes; }
 
     /** Returns the groups of this cluster as an immutable map indexed by group id */
-    public ImmutableMap<Integer, Group> groups() { return groups; }
+    public Map<Integer, Group> groups() { return groups; }
 
     /** Returns the groups of this cluster as an immutable list in introduction order */
-    public ImmutableList<Group> orderedGroups() { return orderedGroups; }
+    public List<Group> orderedGroups() { return orderedGroups; }
 
     /** Returns the n'th (zero-indexed) group in the cluster if possible */
     public Optional<Group> group(int n) {
@@ -143,23 +133,12 @@ public class SearchCluster implements NodeManager<Node> {
         }
     }
 
-    /**
-     * Returns the wanted number of nodes per group - size()/groups.size().
-     * The actual node count for a given group may differ due to node retirements.
-     */
-    public int wantedGroupSize() {
-        if (groups().size() == 0) return size();
-        return size() / groups().size();
+    public boolean allGroupsHaveSize1() {
+        return nodes.size() == groups.size();
     }
 
     public int groupsWithSufficientCoverage() {
-        int covered = 0;
-        for (Group g : orderedGroups()) {
-            if (g.hasSufficientCoverage()) {
-                covered++;
-            }
-        }
-        return covered;
+        return (int)groups.values().stream().filter(g -> g.hasSufficientCoverage()).count();
     }
 
     /**
@@ -210,8 +189,8 @@ public class SearchCluster implements NodeManager<Node> {
         }
         else if (usesLocalCorpusIn(node)) { // follow the status of this node
             // Do not take this out of rotation if we're a combined cluster of size 1,
-            // as that can't be helpful, and leads to a deadlock where this node is never taken back in servic e
-            if (nodeIsWorking || size() > 1)
+            // as that can't be helpful, and leads to a deadlock where this node is never set back in service
+            if (nodeIsWorking || nodes.size() > 1)
                 setInRotationOnlyIf(nodeIsWorking);
         }
     }
@@ -240,11 +219,11 @@ public class SearchCluster implements NodeManager<Node> {
     }
 
     public boolean hasInformationAboutAllNodes() {
-        return nodesByHost.values().stream().allMatch(node -> node.isWorking() != null);
+        return nodes.stream().allMatch(node -> node.isWorking() != null);
     }
 
     private boolean hasWorkingNodes() {
-        return nodesByHost.values().stream().anyMatch(node -> node.isWorking() != Boolean.FALSE );
+        return nodes.stream().anyMatch(node -> node.isWorking() != Boolean.FALSE );
     }
 
     private boolean usesLocalCorpusIn(Node node) {
@@ -255,31 +234,6 @@ public class SearchCluster implements NodeManager<Node> {
         return localCorpusDispatchTarget.isPresent() && localCorpusDispatchTarget.get().group() == group.id();
     }
 
-    private static class PongCallback implements PongHandler {
-
-        private final ClusterMonitor<Node> clusterMonitor;
-        private final Node node;
-
-        PongCallback(Node node, ClusterMonitor<Node> clusterMonitor) {
-            this.node = node;
-            this.clusterMonitor = clusterMonitor;
-        }
-
-        @Override
-        public void handle(Pong pong) {
-            if (pong.badResponse()) {
-                clusterMonitor.failed(node, pong.error().get());
-            } else {
-                if (pong.activeDocuments().isPresent()) {
-                    node.setActiveDocuments(pong.activeDocuments().get());
-                    node.setBlockingWrites(pong.isBlockingWrites());
-                }
-                clusterMonitor.responded(node);
-            }
-        }
-
-    }
-
     /** Used by the cluster monitor to manage node status */
     @Override
     public void ping(ClusterMonitor clusterMonitor, Node node, Executor executor) {
@@ -293,19 +247,15 @@ public class SearchCluster implements NodeManager<Node> {
         // With just one group sufficient coverage may not be the same as full coverage, as the
         // group will always be marked sufficient for use.
         updateSufficientCoverage(group, true);
-        boolean sufficientCoverage = isGroupCoverageSufficient(group.activeDocuments(),
-                                                               group.activeDocuments());
+        boolean sufficientCoverage = isGroupCoverageSufficient(group.activeDocuments(), group.activeDocuments());
         trackGroupCoverageChanges(group, sufficientCoverage, group.activeDocuments());
     }
 
     private void pingIterationCompletedMultipleGroups() {
         orderedGroups().forEach(Group::aggregateNodeValues);
         long medianDocuments = medianDocumentsPerGroup();
-        boolean anyGroupsSufficientCoverage = false;
         for (Group group : orderedGroups()) {
-            boolean sufficientCoverage = isGroupCoverageSufficient(group.activeDocuments(),
-                                                                   medianDocuments);
-            anyGroupsSufficientCoverage = anyGroupsSufficientCoverage || sufficientCoverage;
+            boolean sufficientCoverage = isGroupCoverageSufficient(group.activeDocuments(), medianDocuments);
             updateSufficientCoverage(group, sufficientCoverage);
             trackGroupCoverageChanges(group, sufficientCoverage, medianDocuments);
         }
@@ -372,4 +322,29 @@ public class SearchCluster implements NodeManager<Node> {
         }
     }
 
+    private static class PongCallback implements PongHandler {
+
+        private final ClusterMonitor<Node> clusterMonitor;
+        private final Node node;
+
+        PongCallback(Node node, ClusterMonitor<Node> clusterMonitor) {
+            this.node = node;
+            this.clusterMonitor = clusterMonitor;
+        }
+
+        @Override
+        public void handle(Pong pong) {
+            if (pong.badResponse()) {
+                clusterMonitor.failed(node, pong.error().get());
+            } else {
+                if (pong.activeDocuments().isPresent()) {
+                    node.setActiveDocuments(pong.activeDocuments().get());
+                    node.setBlockingWrites(pong.isBlockingWrites());
+                }
+                clusterMonitor.responded(node);
+            }
+        }
+
+    }
+
 }
diff --git a/container-search/src/main/java/com/yahoo/search/grouping/result/FlatteningSearcher.java b/container-search/src/main/java/com/yahoo/search/grouping/result/FlatteningSearcher.java
new file mode 100644
index 00000000000..321f86facd0
--- /dev/null
+++ b/container-search/src/main/java/com/yahoo/search/grouping/result/FlatteningSearcher.java
@@ -0,0 +1,50 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.search.grouping.result;
+
+import com.yahoo.component.chain.dependencies.Before;
+import com.yahoo.search.Query;
+import com.yahoo.search.Result;
+import com.yahoo.search.Searcher;
+import com.yahoo.search.grouping.vespa.GroupingExecutor;
+import com.yahoo.search.result.Hit;
+import com.yahoo.search.result.HitGroup;
+import com.yahoo.search.searchchain.Execution;
+
+import java.util.Iterator;
+
+/**
+ * Flattens a grouping result into a flat list of hits on the top level in the returned result.
+ * Useful when using grouping to create hits with diversity and similar.
+ *
+ * @author bratseth
+ */
+@Before(GroupingExecutor.COMPONENT_NAME)
+public class FlatteningSearcher extends Searcher {
+
+    @Override
+    public Result search(Query query, Execution execution) {
+        if ( ! query.properties().getBoolean("flatten", true)) return execution.search(query);
+
+        query.trace("Flattening groups", 2);
+        int originalHits = query.getHits();
+        query.setHits(0);
+        Result result = execution.search(query);
+        query.setHits(originalHits);
+        flatten(result.hits(), result);
+        return result;
+    }
+
+    public void flatten(HitGroup hits, Result result) {
+        int hitsLeft = hits.size(); // Iterate only through the initial size
+        for (Iterator<Hit> i = hits.iterator(); i.hasNext() && hitsLeft-- > 0;) {
+            Hit hit = i.next();
+            if (hit instanceof HitGroup) {
+                flatten((HitGroup)hit, result);
+                i.remove();
+            } else {
+                result.hits().add(hit);
+            }
+        }
+    }
+
+}
diff --git a/container-search/src/main/java/com/yahoo/search/result/HitGroup.java b/container-search/src/main/java/com/yahoo/search/result/HitGroup.java
index 6d09bf66175..efe25e04f2e 100644
--- a/container-search/src/main/java/com/yahoo/search/result/HitGroup.java
+++ b/container-search/src/main/java/com/yahoo/search/result/HitGroup.java
@@ -249,9 +249,7 @@ public class HitGroup extends Hit implements DataList<Hit>, Cloneable, Iterable<
         return hit;
     }
 
-    /**
-     * Adds a list of hits to this group, the same
-     */
+    /** Adds a list of hits to this group, the same as calling add for each item in the list. */
     public void addAll(List<Hit> hits) {
         for (Hit hit : hits)
             add(hit);