Search over protobuf + jrt

4 files changed, 152 insertions, 17 deletions

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 708caafa3f5..0cc58801298 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
@@ -2,8 +2,10 @@
 package com.yahoo.search.dispatch;
 
 import com.yahoo.prelude.fastsearch.FS4InvokerFactory;
+import com.yahoo.prelude.fastsearch.VespaBackEndSearcher;
 import com.yahoo.processing.request.CompoundName;
 import com.yahoo.search.Query;
+import com.yahoo.search.dispatch.rpc.RpcInvokerFactory;
 import com.yahoo.search.dispatch.searchcluster.Node;
 import com.yahoo.search.dispatch.searchcluster.SearchCluster;
 import com.yahoo.vespa.config.search.DispatchConfig;
@@ -36,7 +38,7 @@ public class DispatcherTest {
         builder.useMultilevelDispatch(true);
         DispatchConfig dc = new DispatchConfig(builder);
 
-        Dispatcher disp = new Dispatcher(cl, dc);
+        Dispatcher disp = new Dispatcher(cl, dc, new MockFS4InvokerFactory(cl), new MockRpcInvokerFactory());
         assertThat(disp.getSearchInvoker(query(), null).isPresent(), is(false));
     }
 
@@ -45,13 +47,13 @@ public class DispatcherTest {
         SearchCluster cl = new MockSearchCluster("1", 2, 2);
         Query q = query();
         q.getModel().setSearchPath("1/0"); // second node in first group
-        Dispatcher disp = new Dispatcher(cl, createDispatchConfig());
         MockFS4InvokerFactory invokerFactory = new MockFS4InvokerFactory(cl, (nodes, a) -> {
             assertThat(nodes.size(), is(1));
             assertThat(nodes.get(0).key(), is(2));
             return true;
         });
-        Optional<SearchInvoker> invoker = disp.getSearchInvoker(q, invokerFactory);
+        Dispatcher disp = new Dispatcher(cl, createDispatchConfig(), invokerFactory, new MockRpcInvokerFactory());
+        Optional<SearchInvoker> invoker = disp.getSearchInvoker(q, null);
         assertThat(invoker.isPresent(), is(true));
         invokerFactory.verifyAllEventsProcessed();
     }
@@ -64,9 +66,9 @@ public class DispatcherTest {
                 return Optional.of(new Node(1, "test", 123, 1));
             }
         };
-        Dispatcher disp = new Dispatcher(cl, createDispatchConfig());
         MockFS4InvokerFactory invokerFactory = new MockFS4InvokerFactory(cl, (n, a) -> true);
-        Optional<SearchInvoker> invoker = disp.getSearchInvoker(query(), invokerFactory);
+        Dispatcher disp = new Dispatcher(cl, createDispatchConfig(), invokerFactory, new MockRpcInvokerFactory());
+        Optional<SearchInvoker> invoker = disp.getSearchInvoker(query(), null);
         assertThat(invoker.isPresent(), is(true));
         invokerFactory.verifyAllEventsProcessed();
     }
@@ -75,7 +77,6 @@ public class DispatcherTest {
     public void requireThatInvokerConstructionIsRetriedAndLastAcceptsAnyCoverage() {
         SearchCluster cl = new MockSearchCluster("1", 2, 1);
 
-        Dispatcher disp = new Dispatcher(cl, createDispatchConfig());
         MockFS4InvokerFactory invokerFactory = new MockFS4InvokerFactory(cl, (n, acceptIncompleteCoverage) -> {
             assertThat(acceptIncompleteCoverage, is(false));
             return false;
@@ -83,7 +84,8 @@ public class DispatcherTest {
             assertThat(acceptIncompleteCoverage, is(true));
             return true;
         });
-        Optional<SearchInvoker> invoker = disp.getSearchInvoker(query(), invokerFactory);
+        Dispatcher disp = new Dispatcher(cl, createDispatchConfig(), invokerFactory, new MockRpcInvokerFactory());
+        Optional<SearchInvoker> invoker = disp.getSearchInvoker(query(), null);
         assertThat(invoker.isPresent(), is(true));
         invokerFactory.verifyAllEventsProcessed();
     }
@@ -92,9 +94,9 @@ public class DispatcherTest {
     public void requireThatInvokerConstructionDoesNotRepeatGroups() {
         SearchCluster cl = new MockSearchCluster("1", 2, 1);
 
-        Dispatcher disp = new Dispatcher(cl, createDispatchConfig());
         MockFS4InvokerFactory invokerFactory = new MockFS4InvokerFactory(cl, (n, a) -> false, (n, a) -> false);
-        Optional<SearchInvoker> invoker = disp.getSearchInvoker(query(), invokerFactory);
+        Dispatcher disp = new Dispatcher(cl, createDispatchConfig(), invokerFactory, null);
+        Optional<SearchInvoker> invoker = disp.getSearchInvoker(query(), null);
         assertThat(invoker.isPresent(), is(false));
         invokerFactory.verifyAllEventsProcessed();
     }
@@ -108,12 +110,13 @@ public class DispatcherTest {
         private int step = 0;
 
         public MockFS4InvokerFactory(SearchCluster cl, FactoryStep... events) {
-            super(null, cl, null);
+            super(null, cl);
             this.events = events;
         }
 
         @Override
-        public Optional<SearchInvoker> getSearchInvoker(Query query, OptionalInt groupId, List<Node> nodes, boolean acceptIncompleteCoverage) {
+        public Optional<SearchInvoker> createSearchInvoker(VespaBackEndSearcher searcher, Query query, OptionalInt groupId,
+                List<Node> nodes, boolean acceptIncompleteCoverage) {
             if (step >= events.length) {
                 throw new RuntimeException("Was not expecting more calls to getSearchInvoker");
             }
@@ -130,4 +133,14 @@ public class DispatcherTest {
             assertThat(step, is(events.length));
         }
     }
+
+    public class MockRpcInvokerFactory extends RpcInvokerFactory {
+        public MockRpcInvokerFactory() {
+            super(null, null);
+        }
+
+        @Override
+        public void release() {
+        }
+    }
 }
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/FillTestCase.java b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/FillTestCase.java
index 173b2de494f..2adbd12a2aa 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/FillTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/FillTestCase.java
@@ -1,5 +1,5 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.search.dispatch;
+package com.yahoo.search.dispatch.rpc;
 
 import com.yahoo.prelude.fastsearch.DocsumDefinition;
 import com.yahoo.prelude.fastsearch.DocsumDefinitionSet;
@@ -8,6 +8,9 @@ import com.yahoo.prelude.fastsearch.DocumentDatabase;
 import com.yahoo.prelude.fastsearch.FastHit;
 import com.yahoo.search.Query;
 import com.yahoo.search.Result;
+import com.yahoo.search.dispatch.rpc.Client;
+import com.yahoo.search.dispatch.rpc.RpcInvokerFactory;
+import com.yahoo.search.dispatch.rpc.RpcResourcePool;
 import org.junit.Test;
 
 import java.util.ArrayList;
@@ -36,6 +39,7 @@ public class FillTestCase {
         nodes.put(1, client.createConnection("host1", 123));
         nodes.put(2, client.createConnection("host2", 123));
         RpcResourcePool rpcResourcePool = new RpcResourcePool(client, nodes);
+        RpcInvokerFactory factory = new RpcInvokerFactory(rpcResourcePool, null);
 
         Query query = new Query();
         Result result = new Result(query);
@@ -51,7 +55,7 @@ public class FillTestCase {
         client.setDocsumReponse("host2", 3, "summaryClass1", map("field1", "s.2.3", "field2", 3));
         client.setDocsumReponse("host0", 4, "summaryClass1", map("field1", "s.0.4", "field2", 4));
 
-        rpcResourcePool.getFillInvoker(db()).fill(result, "summaryClass1");
+        factory.createFillInvoker(db()).fill(result, "summaryClass1");
 
         assertEquals("s.0.0", result.hits().get("hit:0").getField("field1").toString());
         assertEquals("s.2.1", result.hits().get("hit:1").getField("field1").toString());
@@ -72,6 +76,7 @@ public class FillTestCase {
         nodes.put(1, client.createConnection("host1", 123));
         nodes.put(2, client.createConnection("host2", 123));
         RpcResourcePool rpcResourcePool = new RpcResourcePool(client, nodes);
+        RpcInvokerFactory factory = new RpcInvokerFactory(rpcResourcePool, null);
 
         Query query = new Query();
         Result result = new Result(query);
@@ -87,7 +92,7 @@ public class FillTestCase {
         client.setDocsumReponse("host2", 3, "summaryClass1", map("field1", "s.2.3", "field2", 3));
         client.setDocsumReponse("host0", 4, "summaryClass1",new HashMap<>());
 
-        rpcResourcePool.getFillInvoker(db()).fill(result, "summaryClass1");
+        factory.createFillInvoker(db()).fill(result, "summaryClass1");
 
         assertEquals("s.0.0", result.hits().get("hit:0").getField("field1").toString());
         assertEquals("s.2.1", result.hits().get("hit:1").getField("field1").toString());
@@ -111,12 +116,13 @@ public class FillTestCase {
         Map<Integer, Client.NodeConnection> nodes = new HashMap<>();
         nodes.put(0, client.createConnection("host0", 123));
         RpcResourcePool rpcResourcePool = new RpcResourcePool(client, nodes);
+        RpcInvokerFactory factory = new RpcInvokerFactory(rpcResourcePool, null);
 
         Query query = new Query();
         Result result = new Result(query);
         result.hits().add(createHit(0, 0));
 
-        rpcResourcePool.getFillInvoker(db()).fill(result, "summaryClass1");
+        factory.createFillInvoker(db()).fill(result, "summaryClass1");
 
         assertEquals("Malfunctioning", result.hits().getError().getDetailedMessage());
     }
@@ -128,6 +134,7 @@ public class FillTestCase {
         Map<Integer, Client.NodeConnection> nodes = new HashMap<>();
         nodes.put(0, client.createConnection("host0", 123));
         RpcResourcePool rpcResourcePool = new RpcResourcePool(client, nodes);
+        RpcInvokerFactory factory = new RpcInvokerFactory(rpcResourcePool, null);
 
         Query query = new Query();
         Result result = new Result(query);
@@ -135,7 +142,7 @@ public class FillTestCase {
         result.hits().add(createHit(1, 1));
 
 
-        rpcResourcePool.getFillInvoker(db()).fill(result, "summaryClass1");
+        factory.createFillInvoker(db()).fill(result, "summaryClass1");
 
         assertEquals("Could not fill hits from unknown node 1", result.hits().getError().getDetailedMessage());
     }
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/MockClient.java b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/MockClient.java
index a4cb8ae641c..f9b628e594a 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/MockClient.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/MockClient.java
@@ -1,11 +1,15 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.search.dispatch;
+package com.yahoo.search.dispatch.rpc;
 
 import com.yahoo.compress.CompressionType;
 import com.yahoo.compress.Compressor;
 import com.yahoo.document.GlobalId;
 import com.yahoo.document.idstring.IdIdString;
 import com.yahoo.prelude.fastsearch.FastHit;
+import com.yahoo.search.Result;
+import com.yahoo.search.dispatch.rpc.Client;
+import com.yahoo.search.dispatch.rpc.RpcFillInvoker;
+import com.yahoo.search.dispatch.rpc.RpcSearchInvoker;
 import com.yahoo.slime.ArrayTraverser;
 import com.yahoo.slime.BinaryFormat;
 import com.yahoo.slime.Cursor;
@@ -25,6 +29,7 @@ public class MockClient implements Client {
     private final Map<DocsumKey, Map<String, Object>> docsums = new HashMap<>();
     private final Compressor compressor = new Compressor();
     private boolean malfunctioning = false;
+    private Result searchResult;
 
     /** Set to true to cause this to produce an error instead of a regular response */
     public void setMalfunctioning(boolean malfunctioning) { this.malfunctioning = malfunctioning; }
@@ -72,6 +77,24 @@ public class MockClient implements Client {
         responseReceiver.receive(GetDocsumsResponseOrError.fromResponse(response));
     }
 
+    @Override
+    public void search(NodeConnection node, CompressionType compression, int uncompressedLength, byte[] compressedPayload,
+            RpcSearchInvoker responseReceiver, double timeoutSeconds) {
+        if (malfunctioning) {
+            responseReceiver.receive(SearchResponseOrError.fromError("Malfunctioning"));
+            return;
+        }
+
+        if(searchResult == null) {
+            responseReceiver.receive(SearchResponseOrError.fromError("No result defined"));
+            return;
+        }
+        var payload = ProtobufSerialization.serializeResult(searchResult);
+        var compressionResult = compressor.compress(compression, payload);
+        var response = new SearchResponse(compressionResult.type().getCode(), payload.length, compressionResult.data());
+        responseReceiver.receive(SearchResponseOrError.fromResponse(response));
+    }
+
     public void setDocsumReponse(String nodeId, int docId, String docsumClass, Map<String, Object> docsumValues) {
         docsums.put(new DocsumKey(nodeId, globalIdFrom(docId), docsumClass), docsumValues);
     }
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/rpc/RpcSearchInvokerTest.java b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/RpcSearchInvokerTest.java
new file mode 100644
index 00000000000..689be53de23
--- /dev/null
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/RpcSearchInvokerTest.java
@@ -0,0 +1,92 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+package com.yahoo.search.dispatch.rpc;
+
+import ai.vespa.searchlib.searchprotocol.protobuf.SearchProtocol;
+import com.google.common.collect.ImmutableMap;
+import com.yahoo.compress.CompressionType;
+import com.yahoo.fs4.QueryPacket;
+import com.yahoo.prelude.fastsearch.FastHit;
+import com.yahoo.prelude.fastsearch.VespaBackEndSearcher;
+import com.yahoo.search.Query;
+import com.yahoo.search.Result;
+import com.yahoo.search.dispatch.rpc.RpcFillInvoker.GetDocsumsResponseReceiver;
+import com.yahoo.search.dispatch.searchcluster.Node;
+import com.yahoo.search.searchchain.Execution;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThan;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.fail;
+
+/**
+ * @author ollivir
+ */
+public class RpcSearchInvokerTest {
+    @Test
+    public void testProtobufSerialization() throws IOException {
+        var compressionTypeHolder = new AtomicReference<CompressionType>();
+        var payloadHolder = new AtomicReference<byte[]>();
+        var lengthHolder = new AtomicInteger();
+        var mockClient = parameterCollectorClient(compressionTypeHolder, payloadHolder, lengthHolder);
+        var mockPool = new RpcResourcePool(mockClient, ImmutableMap.of(7, () -> {}));
+        @SuppressWarnings("resource")
+        var invoker = new RpcSearchInvoker(mockSearcher(), new Node(7, "seven", 77, 1), mockPool);
+
+        Query q = new Query("search/?query=test&hits=10&offset=3");
+        invoker.sendSearchRequest(q, null);
+
+        var bytes = mockPool.compressor().decompress(payloadHolder.get(), compressionTypeHolder.get(), lengthHolder.get());
+        var request = SearchProtocol.SearchRequest.newBuilder().mergeFrom(bytes).build();
+
+        assertThat(request.getHits(), equalTo(10));
+        assertThat(request.getOffset(), equalTo(3));
+        assertThat(request.getQueryTreeBlob().size(), greaterThan(0));
+    }
+
+    private Client parameterCollectorClient(AtomicReference<CompressionType> compressionTypeHolder, AtomicReference<byte[]> payloadHolder,
+            AtomicInteger lengthHolder) {
+        return new Client() {
+            @Override
+            public void search(NodeConnection node, CompressionType compression, int uncompressedLength, byte[] compressedPayload,
+                    RpcSearchInvoker responseReceiver, double timeoutSeconds) {
+                compressionTypeHolder.set(compression);
+                payloadHolder.set(compressedPayload);
+                lengthHolder.set(uncompressedLength);
+            }
+
+            @Override
+            public void getDocsums(List<FastHit> hits, NodeConnection node, CompressionType compression, int uncompressedLength,
+                    byte[] compressedSlime, GetDocsumsResponseReceiver responseReceiver, double timeoutSeconds) {
+                fail("Unexpected call");
+            }
+
+            @Override
+            public NodeConnection createConnection(String hostname, int port) {
+                fail("Unexpected call");
+                return null;
+            }
+        };
+    }
+
+    private VespaBackEndSearcher mockSearcher() {
+        return new VespaBackEndSearcher() {
+            @Override
+            protected Result doSearch2(Query query, QueryPacket queryPacket, Execution execution) {
+                fail("Unexpected call");
+                return null;
+            }
+
+            @Override
+            protected void doPartialFill(Result result, String summaryClass) {
+                fail("Unexpected call");
+            }
+        };
+    }
+}