Protobuf object coversion moved to separate class. RPC classes moved to subpackage

3 files changed, 413 insertions, 0 deletions

diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/rpc/FillTestCase.java b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/FillTestCase.java
new file mode 100644
index 00000000000..2adbd12a2aa
--- /dev/null
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/FillTestCase.java
@@ -0,0 +1,174 @@
+// 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.rpc;
+
+import com.yahoo.prelude.fastsearch.DocsumDefinition;
+import com.yahoo.prelude.fastsearch.DocsumDefinitionSet;
+import com.yahoo.prelude.fastsearch.DocsumField;
+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;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+
+/**
+ * Tests using a dispatcher to fill a result
+ *
+ * @author bratseth
+ */
+public class FillTestCase {
+
+    private MockClient client = new MockClient();
+
+    @Test
+    public void testFilling() {
+        Map<Integer, Client.NodeConnection> nodes = new HashMap<>();
+        nodes.put(0, client.createConnection("host0", 123));
+        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);
+        result.hits().add(createHit(0, 0));
+        result.hits().add(createHit(2, 1));
+        result.hits().add(createHit(1, 2));
+        result.hits().add(createHit(2, 3));
+        result.hits().add(createHit(0, 4));
+
+        client.setDocsumReponse("host0", 0, "summaryClass1", map("field1", "s.0.0", "field2", 0));
+        client.setDocsumReponse("host2", 1, "summaryClass1", map("field1", "s.2.1", "field2", 1));
+        client.setDocsumReponse("host1", 2, "summaryClass1", map("field1", "s.1.2", "field2", 2));
+        client.setDocsumReponse("host2", 3, "summaryClass1", map("field1", "s.2.3", "field2", 3));
+        client.setDocsumReponse("host0", 4, "summaryClass1", map("field1", "s.0.4", "field2", 4));
+
+        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());
+        assertEquals("s.1.2", result.hits().get("hit:2").getField("field1").toString());
+        assertEquals("s.2.3", result.hits().get("hit:3").getField("field1").toString());
+        assertEquals("s.0.4", result.hits().get("hit:4").getField("field1").toString());
+        assertEquals(0L, result.hits().get("hit:0").getField("field2"));
+        assertEquals(1L, result.hits().get("hit:1").getField("field2"));
+        assertEquals(2L, result.hits().get("hit:2").getField("field2"));
+        assertEquals(3L, result.hits().get("hit:3").getField("field2"));
+        assertEquals(4L, result.hits().get("hit:4").getField("field2"));
+    }
+
+    @Test
+    public void testEmptyHits() {
+        Map<Integer, Client.NodeConnection> nodes = new HashMap<>();
+        nodes.put(0, client.createConnection("host0", 123));
+        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);
+        result.hits().add(createHit(0, 0));
+        result.hits().add(createHit(2, 1));
+        result.hits().add(createHit(1, 2));
+        result.hits().add(createHit(2, 3));
+        result.hits().add(createHit(0, 4));
+
+        client.setDocsumReponse("host0", 0, "summaryClass1", map("field1", "s.0.0", "field2", 0));
+        client.setDocsumReponse("host2", 1, "summaryClass1", map("field1", "s.2.1", "field2", 1));
+        client.setDocsumReponse("host1", 2, "summaryClass1", new HashMap<>());
+        client.setDocsumReponse("host2", 3, "summaryClass1", map("field1", "s.2.3", "field2", 3));
+        client.setDocsumReponse("host0", 4, "summaryClass1",new HashMap<>());
+
+        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());
+        assertNull(result.hits().get("hit:2").getField("field1"));
+        assertEquals("s.2.3", result.hits().get("hit:3").getField("field1").toString());
+        assertNull(result.hits().get("hit:4").getField("field1"));
+
+        assertEquals(0L, result.hits().get("hit:0").getField("field2"));
+        assertEquals(1L, result.hits().get("hit:1").getField("field2"));
+        assertNull(result.hits().get("hit:2").getField("field2"));
+        assertEquals(3L, result.hits().get("hit:3").getField("field2"));
+        assertNull(result.hits().get("hit:4").getField("field2"));
+
+        assertEquals("Missing hit summary data for summary summaryClass1 for 2 hits", result.hits().getError().getDetailedMessage());
+    }
+
+    @Test
+    public void testErrorHandling() {
+        client.setMalfunctioning(true);
+
+        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));
+
+        factory.createFillInvoker(db()).fill(result, "summaryClass1");
+
+        assertEquals("Malfunctioning", result.hits().getError().getDetailedMessage());
+    }
+
+    @Test
+    public void testSendingFill2UnknownNode() {
+        client.setMalfunctioning(true);
+
+        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));
+        result.hits().add(createHit(1, 1));
+
+
+        factory.createFillInvoker(db()).fill(result, "summaryClass1");
+
+        assertEquals("Could not fill hits from unknown node 1", result.hits().getError().getDetailedMessage());
+    }
+
+    private DocumentDatabase db() {
+        List<DocsumField> fields = new ArrayList<>();
+        fields.add(DocsumField.create("field1", "string"));
+        fields.add(DocsumField.create("field2", "int64"));
+        DocsumDefinitionSet docsums = new DocsumDefinitionSet(Collections.singleton(new DocsumDefinition("summaryClass1",
+                                                                                                         fields)));
+        return new DocumentDatabase("default", docsums, Collections.emptySet());
+    }
+
+    private FastHit createHit(int sourceNodeId, int hitId) {
+        FastHit hit = new FastHit("hit:" + hitId, 1.0);
+        hit.setPartId(sourceNodeId);
+        hit.setDistributionKey(sourceNodeId);
+        hit.setGlobalId(client.globalIdFrom(hitId));
+        return hit;
+    }
+
+    private Map<String, Object> map(String stringKey, String stringValue, String intKey, int intValue) {
+        Map<String, Object> map = new HashMap<>();
+        map.put(stringKey, stringValue);
+        map.put(intKey, intValue);
+        return map;
+    }
+
+}
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/rpc/MockClient.java b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/MockClient.java
new file mode 100644
index 00000000000..f9b628e594a
--- /dev/null
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/MockClient.java
@@ -0,0 +1,144 @@
+// 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.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;
+import com.yahoo.slime.Inspector;
+import com.yahoo.slime.Slime;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * @author bratseth
+ */
+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; }
+
+    @Override
+    public NodeConnection createConnection(String hostname, int port) {
+        return new MockNodeConnection(hostname, port);
+    }
+
+    @Override
+    public void getDocsums(List<FastHit> hitsContext, NodeConnection node, CompressionType compression,
+                           int uncompressedSize, byte[] compressedSlime, RpcFillInvoker.GetDocsumsResponseReceiver responseReceiver,
+                           double timeoutSeconds) {
+        if (malfunctioning) {
+            responseReceiver.receive(GetDocsumsResponseOrError.fromError("Malfunctioning"));
+            return;
+        }
+
+        Inspector request = BinaryFormat.decode(compressor.decompress(compressedSlime, compression, uncompressedSize)).get();
+        String docsumClass = request.field("class").asString();
+        List<Map<String, Object>> docsumsToReturn = new ArrayList<>();
+        request.field("gids").traverse((ArrayTraverser)(index, gid) -> {
+            GlobalId docId = new GlobalId(gid.asData());
+            docsumsToReturn.add(docsums.get(new DocsumKey(node.toString(), docId, docsumClass)));
+        });
+        Slime responseSlime = new Slime();
+        Cursor root = responseSlime.setObject();
+        Cursor docsums = root.setArray("docsums");
+        for (Map<String, Object> docsumFields : docsumsToReturn) {
+            Cursor docsumItem = docsums.addObject();
+            Cursor docsum = docsumItem.setObject("docsum");
+            for (Map.Entry<String, Object> field : docsumFields.entrySet()) {
+                if (field.getValue() instanceof Integer)
+                    docsum.setLong(field.getKey(), (Integer)field.getValue());
+                else if (field.getValue() instanceof String)
+                    docsum.setString(field.getKey(), (String)field.getValue());
+                else
+                    throw new RuntimeException();
+            }
+        }
+        byte[] slimeBytes = BinaryFormat.encode(responseSlime);
+        Compressor.Compression compressionResult = compressor.compress(compression, slimeBytes);
+        GetDocsumsResponse response = new GetDocsumsResponse(compressionResult.type().getCode(), slimeBytes.length,
+                                                             compressionResult.data(), hitsContext);
+        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);
+    }
+
+    public GlobalId globalIdFrom(int hitId) {
+        return new GlobalId(new IdIdString("", "test", "", String.valueOf(hitId)));
+    }
+
+    private static class MockNodeConnection implements Client.NodeConnection {
+
+        private final String hostname;
+
+        public MockNodeConnection(String hostname, int port) {
+            this.hostname = hostname;
+        }
+
+        @Override
+        public void close() { }
+
+        @Override
+        public String toString() { return hostname; }
+
+    }
+
+    private static class DocsumKey {
+
+        private final String internalKey;
+
+        public DocsumKey(String nodeId, GlobalId docId, String docsumClass) {
+            internalKey = docsumClass + "." + nodeId + "." + docId;
+        }
+
+        @Override
+        public int hashCode() { return internalKey.hashCode(); }
+
+        @Override
+        public boolean equals(Object other) {
+            if ( ! (other instanceof DocsumKey)) return false;
+            return ((DocsumKey)other).internalKey.equals(this.internalKey);
+        }
+
+        @Override
+        public String toString() { return internalKey; }
+
+    }
+
+}
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..4792b45281f
--- /dev/null
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/RpcSearchInvokerTest.java
@@ -0,0 +1,95 @@
+// 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.Search;
+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.Client;
+import com.yahoo.search.dispatch.rpc.RpcResourcePool;
+import com.yahoo.search.dispatch.rpc.RpcSearchInvoker;
+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 = Search.Request.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");
+            }
+        };
+    }
+}