Group affinity through searchPath; remove Query from LB assignment

3 files changed, 13 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 708d80d4969..e95255b31e6 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
@@ -123,6 +123,7 @@ public class Dispatcher extends AbstractComponent {
             if (nodes.isEmpty()) {
                 return Optional.empty();
             } else {
+                query.trace(false, 2, "Dispatching internally with search path ", searchPath);
                 return invokerFactory.supply(query, -1, nodes, true);
             }
         } catch (InvalidSearchPathException e) {
@@ -141,7 +142,7 @@ public class Dispatcher extends AbstractComponent {
         int max = Integer.min(searchCluster.orderedGroups().size(), MAX_GROUP_SELECTION_ATTEMPTS);
         Set<Integer> rejected = null;
         for (int i = 0; i < max; i++) {
-            Optional<Group> groupInCluster = loadBalancer.takeGroupForQuery(query, rejected);
+            Optional<Group> groupInCluster = loadBalancer.takeGroupForQuery(rejected);
             if (!groupInCluster.isPresent()) {
                 // No groups available
                 break;
@@ -150,7 +151,8 @@ public class Dispatcher extends AbstractComponent {
             boolean acceptIncompleteCoverage = (i == max - 1);
             Optional<SearchInvoker> invoker = invokerFactory.supply(query, group.id(), group.nodes(), acceptIncompleteCoverage);
             if (invoker.isPresent()) {
-                query.trace(false, 2, "Dispatching internally to ", group);
+                query.trace(false, 2, "Dispatching internally to search group ", group.id());
+                query.getModel().setSearchPath("/" + group.id());
                 invoker.get().teardown(() -> loadBalancer.releaseGroup(group));
                 return invoker;
             } else {
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 22573fac2d9..0bc7a9e47d7 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
@@ -52,7 +52,7 @@ public class LoadBalancer {
      * @param rejectedGroups if not null, the load balancer will only return groups with IDs not in the set
      * @return The node group to target, or <i>empty</i> if the internal dispatch logic cannot be used
      */
-    public Optional<Group> takeGroupForQuery(Query query, Set<Integer> rejectedGroups) {
+    public Optional<Group> takeGroupForQuery(Set<Integer> rejectedGroups) {
         if (scoreboard == null) {
             return Optional.empty();
         }
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 698cee743e4..13b75cf9c5a 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
@@ -1,7 +1,6 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.search.dispatch;
 
-import com.yahoo.search.Query;
 import com.yahoo.search.dispatch.searchcluster.Group;
 import com.yahoo.search.dispatch.searchcluster.Node;
 import com.yahoo.search.dispatch.searchcluster.SearchCluster;
@@ -26,7 +25,7 @@ public class LoadBalancerTest {
         SearchCluster cluster = new SearchCluster(88.0, 99.0, 0, Arrays.asList(n1), null, 1, null);
         LoadBalancer lb = new LoadBalancer(cluster, true);
 
-        Optional<Group> grp = lb.takeGroupForQuery(new Query(), null);
+        Optional<Group> grp = lb.takeGroupForQuery(null);
         Group group = grp.orElseGet(() -> {
             throw new AssertionFailedError("Expected a SearchCluster.Group");
         });
@@ -40,7 +39,7 @@ public class LoadBalancerTest {
         SearchCluster cluster = new SearchCluster(88.0, 99.0, 0, Arrays.asList(n1, n2), null, 1, null);
         LoadBalancer lb = new LoadBalancer(cluster, true);
 
-        Optional<Group> grp = lb.takeGroupForQuery(new Query(), null);
+        Optional<Group> grp = lb.takeGroupForQuery(null);
         Group group = grp.orElseGet(() -> {
             throw new AssertionFailedError("Expected a SearchCluster.Group");
         });
@@ -56,7 +55,7 @@ public class LoadBalancerTest {
         SearchCluster cluster = new SearchCluster(88.0, 99.0, 0, Arrays.asList(n1, n2, n3, n4), null, 2, null);
         LoadBalancer lb = new LoadBalancer(cluster, true);
 
-        Optional<Group> grp = lb.takeGroupForQuery(new Query(), null);
+        Optional<Group> grp = lb.takeGroupForQuery(null);
         assertThat(grp.isPresent(), is(true));
     }
 
@@ -68,14 +67,14 @@ public class LoadBalancerTest {
         LoadBalancer lb = new LoadBalancer(cluster, true);
 
         // get first group
-        Optional<Group> grp = lb.takeGroupForQuery(new Query(), null);
+        Optional<Group> grp = lb.takeGroupForQuery(null);
         Group group = grp.get();
         int id1 = group.id();
         // release allocation
         lb.releaseGroup(group);
 
         // get second group
-        grp = lb.takeGroupForQuery(new Query(), null);
+        grp = lb.takeGroupForQuery(null);
         group = grp.get();
         assertThat(group.id(), not(equalTo(id1)));
     }
@@ -88,12 +87,12 @@ public class LoadBalancerTest {
         LoadBalancer lb = new LoadBalancer(cluster, true);
 
         // get first group
-        Optional<Group> grp = lb.takeGroupForQuery(new Query(), null);
+        Optional<Group> grp = lb.takeGroupForQuery(null);
         Group group = grp.get();
         int id1 = group.id();
 
         // get second group
-        grp = lb.takeGroupForQuery(new Query(), null);
+        grp = lb.takeGroupForQuery(null);
         group = grp.get();
         int id2 = group.id();
         assertThat(id2, not(equalTo(id1)));
@@ -101,7 +100,7 @@ public class LoadBalancerTest {
         lb.releaseGroup(group);
 
         // get third group
-        grp = lb.takeGroupForQuery(new Query(), null);
+        grp = lb.takeGroupForQuery(null);
         group = grp.get();
         assertThat(group.id(), equalTo(id2));
     }