Test that distribution key is carried from direct search to rpc summary

10 files changed, 68 insertions, 37 deletions

diff --git a/container-search/src/main/java/com/yahoo/fs4/mplex/Backend.java b/container-search/src/main/java/com/yahoo/fs4/mplex/Backend.java
index 68a13c06e32..3eabc3c6a6c 100644
--- a/container-search/src/main/java/com/yahoo/fs4/mplex/Backend.java
+++ b/container-search/src/main/java/com/yahoo/fs4/mplex/Backend.java
@@ -66,7 +66,7 @@ public class Backend implements ConnectionFactory {
     /**
      * For unit testing.  do not use
      */
-    protected Backend() {
+    protected Backend(Optional<Integer> distributionKey) {
         listeners = null;
         host = null;
         port = 0;
@@ -74,7 +74,7 @@ public class Backend implements ConnectionFactory {
         packetDumper = null;
         address = null;
         connectionPool = new ConnectionPool();
-        distributionKey = Optional.empty();
+        this.distributionKey = distributionKey;
     }
 
     public Backend(String host,
@@ -197,7 +197,7 @@ public class Backend implements ConnectionFactory {
     //============================================================
 
     /** Opens a new channel to fdispatch.  Analogous to the "Channel" concept as used in FS4. */
-    public FS4Channel openChannel () {
+    public FS4Channel openChannel() {
         int cachedChannelId;
         synchronized (this) {
             if (channelId >= ((1 << 31) - 2)) {
@@ -214,7 +214,7 @@ public class Backend implements ConnectionFactory {
         return chan;
     }
 
-    public FS4Channel openPingChannel () {
+    public FS4Channel openPingChannel() {
         FS4Channel chan = FS4Channel.createPingChannel(this);
         synchronized (pingChannels) {
             pingChannels.add(chan);
@@ -251,7 +251,7 @@ public class Backend implements ConnectionFactory {
      * Return the first channel in the queue waiting for pings or
      * <code>null</code> if none.
      */
-    public FS4Channel getPingChannel () {
+    public FS4Channel getPingChannel() {
         synchronized (pingChannels) {
             return (pingChannels.isEmpty()) ? null : pingChannels.getFirst();
         }
@@ -266,7 +266,7 @@ public class Backend implements ConnectionFactory {
      *         or <code>null</code> if the channel is not in the
      *         set of active channels.
      */
-    public FS4Channel getChannel (int id) {
+    public FS4Channel getChannel(int id) {
         return getChannel(Integer.valueOf(id));
     }
 
@@ -280,7 +280,7 @@ public class Backend implements ConnectionFactory {
      *         with this id or <code>null</code> if the channel is
      *         not in the set of active channels.
      */
-    protected FS4Channel removeChannel (Integer id) {
+    protected FS4Channel removeChannel(Integer id) {
         synchronized (activeChannels) {
             return activeChannels.remove(id);
         }
@@ -295,7 +295,7 @@ public class Backend implements ConnectionFactory {
      *         the queue of ping channels or <code>null</code>
      *         if there are no active ping channels.
      */
-    protected FS4Channel removePingChannel () {
+    protected FS4Channel removePingChannel() {
         synchronized (pingChannels) {
             if (pingChannels.isEmpty())
                 return null;
diff --git a/container-search/src/main/java/com/yahoo/prelude/cluster/ClusterSearcher.java b/container-search/src/main/java/com/yahoo/prelude/cluster/ClusterSearcher.java
index 0d9534457e9..116db906755 100644
--- a/container-search/src/main/java/com/yahoo/prelude/cluster/ClusterSearcher.java
+++ b/container-search/src/main/java/com/yahoo/prelude/cluster/ClusterSearcher.java
@@ -152,11 +152,11 @@ public class ClusterSearcher extends Searcher {
         } else {
             for (int dispatcherIndex = 0; dispatcherIndex < searchClusterConfig.dispatcher().size(); dispatcherIndex++) {
                 try {
-                    if (! isRemote(searchClusterConfig.dispatcher(dispatcherIndex).host())) {
-                        Backend b = createBackend(searchClusterConfig.dispatcher(dispatcherIndex));
+                    if ( ! isRemote(searchClusterConfig.dispatcher(dispatcherIndex).host())) {
+                        Backend dispatchBackend = createBackend(searchClusterConfig.dispatcher(dispatcherIndex));
                         FastSearcher searcher = searchDispatch(searchClusterIndex, fs4ResourcePool,
                                                                cacheParams, emulationConfig, docSumParams,
-                                                               documentDbConfig, b, dispatcher, dispatcherIndex);
+                                                               documentDbConfig, dispatchBackend, dispatcher, dispatcherIndex);
                         addBackendSearcher(searcher);
                     }
                 } catch (UnknownHostException e) {
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 3eb010da555..5ef81403f26 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
@@ -103,15 +103,15 @@ public class Dispatcher extends AbstractComponent {
 
     /** Return a map of hits by their search node (partition) id */
     private static ListMap<Integer, FastHit> hitsByNode(Result result) {
-        ListMap<Integer, FastHit> hitsByPartition = new ListMap<>();
+        ListMap<Integer, FastHit> hitsByNode = new ListMap<>();
         for (Iterator<Hit> i = result.hits().unorderedDeepIterator() ; i.hasNext(); ) {
             Hit h = i.next();
             if ( ! (h instanceof FastHit)) continue;
             FastHit hit = (FastHit)h;
 
-            hitsByPartition.put(hit.getDistributionKey(), hit);
+            hitsByNode.put(hit.getDistributionKey(), hit);
         }
-        return hitsByPartition;
+        return hitsByNode;
     }
 
     /** Send a getDocsums request to a node. Responses will be added to the given receiver. */
diff --git a/container-search/src/test/java/com/yahoo/fs4/PacketQueryTracerTestCase.java b/container-search/src/test/java/com/yahoo/fs4/PacketQueryTracerTestCase.java
index 27d5e677871..20d9b61c177 100644
--- a/container-search/src/test/java/com/yahoo/fs4/PacketQueryTracerTestCase.java
+++ b/container-search/src/test/java/com/yahoo/fs4/PacketQueryTracerTestCase.java
@@ -16,11 +16,12 @@ import com.yahoo.fs4.mplex.InvalidChannelException;
 import com.yahoo.search.Query;
 
 /**
- * Ensure hex dumping of packets seems to work.
+ * Tests hex dumping of packets.
  *
- * @author <a href="mailto:steinar@yahoo-inc.com">Steinar Knutsen</a>
+ * @author Steinar Knutsen
  */
 public class PacketQueryTracerTestCase {
+
     FS4Channel channel;
     BasicPacket packet;
     PacketListener tracer;
@@ -47,26 +48,22 @@ public class PacketQueryTracerTestCase {
         }
 
         @Override
-        public boolean sendPacket(BasicPacket packet)
-                throws InvalidChannelException, IOException {
+        public boolean sendPacket(BasicPacket packet) {
             return true;
         }
 
         @Override
-        public BasicPacket[] receivePackets(long timeout, int packetCount)
-                throws InvalidChannelException, ChannelTimeoutException {
+        public BasicPacket[] receivePackets(long timeout, int packetCount) {
             return null;
         }
 
         @Override
-        public BasicPacket nextPacket(long timeout)
-                throws InterruptedException, InvalidChannelException {
+        public BasicPacket nextPacket(long timeout) {
             return null;
         }
 
         @Override
-        protected void addPacket(BasicPacket packet)
-                throws InterruptedException, InvalidChannelException {
+        protected void addPacket(BasicPacket packet) {
         }
 
         @Override
diff --git a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/DirectSearchTestCase.java b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/DirectSearchTestCase.java
index ca3a25d3bfa..d1be5c109d9 100644
--- a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/DirectSearchTestCase.java
+++ b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/DirectSearchTestCase.java
@@ -1,6 +1,8 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.prelude.fastsearch.test;
 
+import com.yahoo.prelude.fastsearch.FastHit;
+import com.yahoo.search.Result;
 import org.junit.Test;
 
 import static org.junit.Assert.assertEquals;
@@ -50,6 +52,17 @@ public class DirectSearchTestCase {
     }
 
     @Test
+    public void testDirectSearchSummaryFetchGoToLocalNode() {
+        FastSearcherTester tester = new FastSearcherTester(2, "otherhost:9999:1", FastSearcherTester.selfHostname + ":9999:0");
+        int localDistributionKey = tester.dispatcher().searchCluster().nodesByHost().get(FastSearcherTester.selfHostname).asList().get(0).key();
+        assertEquals(1, localDistributionKey);
+        Result result = tester.search("?query=test&dispatch.direct=true");
+        assertEquals(1, tester.requestCount(FastSearcherTester.selfHostname, 9999));
+        FastHit hit = (FastHit)result.hits().get(0);
+        assertEquals(localDistributionKey, hit.getDistributionKey());
+    }
+
+    @Test
     public void testNoDirectSearchWhenMultipleNodesPerGroup() {
         FastSearcherTester tester = new FastSearcherTester(2, FastSearcherTester.selfHostname + ":9999:0", "otherhost:9999:0");
         tester.search("?query=test&dispatch.direct=true");
diff --git a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/FastSearcherTester.java b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/FastSearcherTester.java
index 035710c612c..c70e1138954 100644
--- a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/FastSearcherTester.java
+++ b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/FastSearcherTester.java
@@ -20,6 +20,7 @@ import com.yahoo.search.searchchain.Execution;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
@@ -47,7 +48,7 @@ class FastSearcherTester {
     public FastSearcherTester(int containerClusterSize, List<SearchCluster.Node> searchNodes) {
         mockFS4ResourcePool = new MockFS4ResourcePool();
         mockDispatcher = new MockDispatcher(searchNodes, mockFS4ResourcePool, containerClusterSize, vipStatus);
-        fastSearcher = new FastSearcher(new MockBackend(selfHostname, MockFSChannel::new),
+        fastSearcher = new FastSearcher(new MockBackend(Optional.empty(), selfHostname, 0L, true),
                                         mockFS4ResourcePool,
                                         mockDispatcher,
                                         new SummaryParameters(null),
@@ -77,6 +78,8 @@ class FastSearcherTester {
         return mockFS4ResourcePool.requestCount(hostname, port);
     }
 
+    public MockDispatcher dispatcher() { return mockDispatcher; }
+
     /** Sets the response status of a node and ping it to update the monitor status */
     public void setResponding(String hostname, boolean responding) {
         // Start/stop returning a failing backend
diff --git a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockBackend.java b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockBackend.java
index c01df2604c7..01ae9aa8f33 100644
--- a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockBackend.java
+++ b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockBackend.java
@@ -4,6 +4,7 @@ package com.yahoo.prelude.fastsearch.test.fs4mock;
 import com.yahoo.fs4.mplex.Backend;
 import com.yahoo.fs4.mplex.FS4Channel;
 
+import java.util.Optional;
 import java.util.function.Supplier;
 
 /**
@@ -11,27 +12,39 @@ import java.util.function.Supplier;
  */
 public class MockBackend extends Backend {
 
-    private MockFSChannel channel;
     private String hostname;
+    private final long activeDocumentsInBackend;
+    private final boolean working;
+
+    /** Created lazily as we want to have just one but it depends on the channel */
+    private MockFSChannel channel = null;
 
     public MockBackend() {
-        this("", MockFSChannel::new);
+        this(Optional.empty(), "", 0L, true);
     }
     
-    public MockBackend(String hostname, Supplier<MockFSChannel> channelSupplier) {
+    public MockBackend(Optional<Integer> distributionKey, String hostname, long activeDocumentsInBackend, boolean working) {
+        super(distributionKey);
         this.hostname = hostname;
-        channel = channelSupplier.get();
+        this.activeDocumentsInBackend = activeDocumentsInBackend;
+        this.working = working;
     }
 
     @Override
-    public FS4Channel openChannel() { return channel; }
+    public FS4Channel openChannel() {
+        if (channel == null)
+            channel = working ? new MockFSChannel(activeDocumentsInBackend, this)
+                              : new NonWorkingMockFSChannel(this);
+        return channel;
+    }
 
     @Override
-    public FS4Channel openPingChannel() { return channel; }
+    public FS4Channel openPingChannel() { return openChannel(); }
 
     @Override
     public String getHost() { return hostname; }
 
+    /** Returns the channel in use or null if no channel has been used yet */
     public MockFSChannel getChannel() { return channel; }
 
     public void shutdown() {}
diff --git a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockFS4ResourcePool.java b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockFS4ResourcePool.java
index 704dcd31c0d..9b5f4b99f20 100644
--- a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockFS4ResourcePool.java
+++ b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockFS4ResourcePool.java
@@ -29,10 +29,9 @@ public class MockFS4ResourcePool extends FS4ResourcePool {
     public Backend getBackend(String hostname, int port, Optional<Integer> distributionKey) {
         countRequest(hostname + ":" + port);
         if (nonRespondingBackends.contains(hostname))
-            return new MockBackend(hostname, NonWorkingMockFSChannel::new);
+            return new MockBackend(distributionKey, hostname, 0L, false);
         else
-            return new MockBackend(hostname, 
-                                   () -> new MockFSChannel(activeDocumentsInBackend.getOrDefault(hostname, 0L)));
+            return new MockBackend(distributionKey, hostname, activeDocumentsInBackend.getOrDefault(hostname, 0L), true);
     }
 
     /** 
diff --git a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockFSChannel.java b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockFSChannel.java
index 37a74831e56..d3d58bbcc0e 100644
--- a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockFSChannel.java
+++ b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/MockFSChannel.java
@@ -28,11 +28,12 @@ public class MockFSChannel extends FS4Channel {
     /** The number of active documents this should report in ping reponses */
     private final long activeDocuments;
     
-    public MockFSChannel() { 
-        this(0); 
+    public MockFSChannel(Backend backend) {
+        this(0, backend);
     }
 
-    public MockFSChannel(long activeDocuments) {
+    public MockFSChannel(long activeDocuments, Backend backend) {
+        super(backend, 0);
         this.activeDocuments = activeDocuments;
     }
 
diff --git a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/NonWorkingMockFSChannel.java b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/NonWorkingMockFSChannel.java
index ea58cccd96c..c7425afd611 100644
--- a/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/NonWorkingMockFSChannel.java
+++ b/container-search/src/test/java/com/yahoo/prelude/fastsearch/test/fs4mock/NonWorkingMockFSChannel.java
@@ -2,12 +2,17 @@
 package com.yahoo.prelude.fastsearch.test.fs4mock;
 
 import com.yahoo.fs4.BasicPacket;
+import com.yahoo.fs4.mplex.Backend;
 
 /**
  * @author bratseth
  */
 public class NonWorkingMockFSChannel extends MockFSChannel {
 
+    public NonWorkingMockFSChannel(Backend backend) {
+        super(backend);
+    }
+
     @Override
     public synchronized boolean sendPacket(BasicPacket bPacket) {
         return false;