Add benchmark mode to JsonStreamFeeder + CliClient wiring

2 files changed, 35 insertions, 7 deletions

diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/JsonStreamFeeder.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/JsonStreamFeeder.java
index 72eb96c2356..17162f19d3f 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/JsonStreamFeeder.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/JsonStreamFeeder.java
@@ -12,7 +12,9 @@ import java.io.InputStream;
 import java.io.InterruptedIOException;
 import java.io.UncheckedIOException;
 import java.time.Duration;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
 import static ai.vespa.feed.client.FeedClient.OperationType.put;
@@ -67,24 +69,36 @@ public class JsonStreamFeeder implements Closeable {
      * Note that {@code "id"} is an alias for the document put operation.
      */
     public void feed(InputStream jsonStream) throws IOException {
-        feed(jsonStream, 1 << 26);
+        feed(jsonStream, 1 << 26, false);
     }
 
-    void feed(InputStream jsonStream, int size) throws IOException {
+    BenchmarkResult benchmark(InputStream jsonStream) throws IOException {
+        return feed(jsonStream, 1 << 26, true).get();
+    }
+
+    Optional<BenchmarkResult> feed(InputStream jsonStream, int size, boolean benchmark) throws IOException {
         RingBufferStream buffer = new RingBufferStream(jsonStream, size);
         buffer.expect(JsonToken.START_ARRAY);
+        AtomicInteger okCount = new AtomicInteger();
+        AtomicInteger failedCount = new AtomicInteger();
+        long startTime = System.nanoTime();
         CompletableFuture<Result> result;
         AtomicReference<Throwable> thrown = new AtomicReference<>();
         while ((result = buffer.next()) != null) {
             result.whenComplete((r, t) -> {
-                if (t != null)
-                    thrown.set(t);
-                else
-                    ; // Aggregate stats.
+                if (t != null) {
+                    failedCount.incrementAndGet();
+                    if (!benchmark) thrown.set(t);
+                } else
+                    okCount.incrementAndGet();
             });
             if (thrown.get() != null)
                 sneakyThrow(thrown.get());
         }
+        if (!benchmark) return Optional.empty();
+        Duration duration = Duration.ofNanos(System.nanoTime() - startTime);
+        double throughPut = (double)okCount.get() / duration.toMillis() * 1000D;
+        return Optional.of(new BenchmarkResult(okCount.get(), failedCount.get(), duration, throughPut));
     }
 
     @SuppressWarnings("unchecked")
@@ -333,4 +347,18 @@ public class JsonStreamFeeder implements Closeable {
 
     }
 
+    static class BenchmarkResult {
+        final int okCount;
+        final int errorCount;
+        final Duration duration;
+        final double throughput;
+
+        BenchmarkResult(int okCount, int errorCount, Duration duration, double throughput) {
+            this.okCount = okCount;
+            this.errorCount = errorCount;
+            this.duration = duration;
+            this.throughput = throughput;
+        }
+    }
+
 }
diff --git a/vespa-feed-client/src/test/java/ai/vespa/feed/client/JsonStreamFeederTest.java b/vespa-feed-client/src/test/java/ai/vespa/feed/client/JsonStreamFeederTest.java
index 25f64e3c98a..8ef8ae57f5e 100644
--- a/vespa-feed-client/src/test/java/ai/vespa/feed/client/JsonStreamFeederTest.java
+++ b/vespa-feed-client/src/test/java/ai/vespa/feed/client/JsonStreamFeederTest.java
@@ -60,7 +60,7 @@ class JsonStreamFeederTest {
             public void close() throws IOException {
 
             }
-        }).build().feed(in, 1 << 7); // TODO: hangs on 1 << 6.
+        }).build().feed(in, 1 << 7, false); // TODO: hangs on 1 << 6.
         assertEquals(docs + 1, ids.size());
     }