Move vespaclient-java to Vespa open-source

1 files changed, 162 insertions, 0 deletions

diff --git a/vespaclient-java/src/main/java/com/yahoo/vespasummarybenchmark/VespaSummaryBenchmark.java b/vespaclient-java/src/main/java/com/yahoo/vespasummarybenchmark/VespaSummaryBenchmark.java
new file mode 100644
index 00000000000..803445d16f5
--- /dev/null
+++ b/vespaclient-java/src/main/java/com/yahoo/vespasummarybenchmark/VespaSummaryBenchmark.java
@@ -0,0 +1,162 @@
+// Copyright 2017 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespasummarybenchmark;
+
+import com.yahoo.compress.CompressionType;
+import com.yahoo.document.GlobalId;
+import com.yahoo.document.idstring.IdString;
+import com.yahoo.document.serialization.DeserializationException;
+import com.yahoo.jrt.*;
+import com.yahoo.log.LogSetup;
+import com.yahoo.slime.*;
+import net.jpountz.lz4.LZ4Factory;
+import net.jpountz.lz4.LZ4FastDecompressor;
+
+import java.io.*;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ *
+ * This is used for testing and benchmarking rpc docsum interface.
+ * time vespa-summary-benchmark file-containing-docids connectionspec summary-class repetitions threads
+ * fx ' time vespa-summary-benchmark feed.xml tcp/localhost:19115 keyvaluesummary 10000 32'
+ *
+ * @author baldersheim
+ */
+public class VespaSummaryBenchmark {
+
+    private final Supervisor supervisor = new Supervisor(new Transport());
+
+    private VespaSummaryBenchmark() { }
+
+    private static List<String> getDocIds(String fileName) {
+        try {
+            FileInputStream fstream = new FileInputStream(fileName);
+            DataInputStream in = new DataInputStream(fstream);
+            BufferedReader br = new BufferedReader(new InputStreamReader(in));
+            String strLine;
+
+            List<String> docIds = new ArrayList<>();
+            while ((strLine = br.readLine()) != null) {
+                docIds.add(strLine);
+            }
+            in.close();
+            return docIds;
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private List<Target> getTargets(String connectionSpec, int numTargets) {
+        List<Target> targets = new ArrayList<>(numTargets);
+        for ( int i=0; i < numTargets; i++) {
+            targets.add(supervisor.connect(new Spec(connectionSpec)));
+        }
+        return targets;
+    }
+
+    private static Slime createDocsumRequest(String summaryClass, List<GlobalId> gids) {
+        Slime docsumRequest = new Slime();
+        Cursor root = docsumRequest.setObject();
+        root.setString("class", summaryClass);
+        Cursor gidCursor = root.setArray("gids");
+        for (GlobalId gid : gids) {
+            gidCursor.addData(gid.getRawId());
+        }
+        return docsumRequest;
+    }
+
+    private static class Waiter implements RequestWaiter {
+
+        int waitingFor;
+        boolean dump;
+
+        Waiter(int expect, boolean dump) {
+            waitingFor = expect;
+            this.dump = dump;
+        }
+
+        private void print(Request request) {
+            Values ret = request.returnValues();
+            CompressionType type = CompressionType.valueOf(ret.get(0).asInt8());
+            int uncompressedSize = ret.get(1).asInt32();
+            byte [] blob = ret.get(2).asData();
+            if (type == CompressionType.LZ4) {
+                LZ4Factory factory = LZ4Factory.fastestInstance();
+                LZ4FastDecompressor decompressor = factory.fastDecompressor();
+                byte [] uncompressed = new byte [uncompressedSize];
+                int compressedLength = decompressor.decompress(blob, 0, uncompressed, 0, uncompressedSize);
+                if (compressedLength != blob.length) {
+                    throw new DeserializationException("LZ4 decompression failed. compressed size does not match. Expected " + blob.length + ". Got " + compressedLength);
+                }
+                blob = uncompressed;
+            }
+            Slime slime = BinaryFormat.decode(blob);
+            try {
+                new JsonFormat(true).encode(System.out, slime);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        @Override
+        public void handleRequestDone(Request request) {
+            synchronized (this) {
+                if (dump) {
+                    print(request);
+                    dump = false;
+                }
+                waitingFor--;
+                if (waitingFor == 0) {
+                    this.notifyAll();
+                }
+            }
+        }
+        void waitForReplies() throws InterruptedException {
+            synchronized (this) {
+                while (waitingFor > 0) {
+                    this.wait();
+                }
+            }
+        }
+    }
+
+    private static void fetchDocIds(String summaryClass, List<Target> targets, List<GlobalId> gids, boolean dump) {
+        Slime docsumRequest = createDocsumRequest(summaryClass, gids);
+        byte [] blob = BinaryFormat.encode(docsumRequest);
+        Waiter waiter = new Waiter(targets.size(), dump);
+        for (Target target : targets) {
+            Request r = new Request("proton.getDocsums");
+            r.parameters().add(new Int8Value(CompressionType.NONE.getCode()));
+            r.parameters().add(new Int32Value(blob.length));
+            r.parameters().add(new DataValue(blob));
+            target.invokeAsync(r, 100.0, waiter);
+        }
+        try {
+            waiter.waitForReplies();
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+
+    }
+
+    public static void main(String[] args) {
+        LogSetup.initVespaLogging("vespasummarybenchmark");
+        String docidFileName = args[0];
+        String connectionSpec = args[1];
+        String summaryClass = args[2];
+        int numRuns = Integer.parseInt(args[3]);
+        int numTargets = Integer.parseInt(args[4]);
+        VespaSummaryBenchmark benchmark = new VespaSummaryBenchmark();
+        List<String> docIds = getDocIds(docidFileName);
+        List<GlobalId> gids = new ArrayList<>(docIds.size());
+        for (String docid : docIds) {
+            GlobalId gid = new GlobalId(IdString.createIdString(docid));
+            gids.add(gid);
+        }
+        List<Target> targets = benchmark.getTargets(connectionSpec, numTargets);
+        for (int i = 0; i < numRuns; i++) {
+            fetchDocIds(summaryClass, targets, gids, i==0);
+        }
+    }
+}