From c5860d6ce2032013ba2dbf3278d7824244e8aef4 Mon Sep 17 00:00:00 2001
From: chenhang <chenhang@apache.org>
Date: Fri, 10 Mar 2023 23:33:48 +0800
Subject: [PATCH] add V2 protocal and warmupMessages support for benchMark

---
 .../benchmark/BenchThroughputLatency.java           | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/bookkeeper-benchmark/src/main/java/org/apache/bookkeeper/benchmark/BenchThroughputLatency.java b/bookkeeper-benchmark/src/main/java/org/apache/bookkeeper/benchmark/BenchThroughputLatency.java
index 84573eec7ca..83ebe4229cf 100644
--- a/bookkeeper-benchmark/src/main/java/org/apache/bookkeeper/benchmark/BenchThroughputLatency.java
+++ b/bookkeeper-benchmark/src/main/java/org/apache/bookkeeper/benchmark/BenchThroughputLatency.java
@@ -257,6 +257,8 @@ public static void main(String[] args)
         options.addOption("skipwarmup", false, "Skip warm up, default false");
         options.addOption("sendlimit", true, "Max number of entries to send. Default 20000000");
         options.addOption("latencyFile", true, "File to dump latencies. Default is latencyDump.dat");
+        options.addOption("useV2", false, "Whether use V2 protocol to send requests to the bookie server");
+        options.addOption("warmupMessages", true, "Number of messages to warm up. Default 10000");
         options.addOption("help", false, "This message");
 
         CommandLineParser parser = new PosixParser();
@@ -281,6 +283,7 @@ public static void main(String[] args)
         }
         int throttle = Integer.parseInt(cmd.getOptionValue("throttle", "10000"));
         int sendLimit = Integer.parseInt(cmd.getOptionValue("sendlimit", "20000000"));
+        int warmupMessages = Integer.parseInt(cmd.getOptionValue("warmupMessages", "10000"));
 
         final int sockTimeout = Integer.parseInt(cmd.getOptionValue("sockettimeout", "5"));
 
@@ -321,11 +324,15 @@ public void run() {
         ClientConfiguration conf = new ClientConfiguration();
         conf.setThrottleValue(throttle).setReadTimeout(sockTimeout).setZkServers(servers);
 
+        if (cmd.hasOption("useV2")) {
+            conf.setUseV2WireProtocol(true);
+        }
+
         if (!cmd.hasOption("skipwarmup")) {
             long throughput;
             LOG.info("Starting warmup");
 
-            throughput = warmUp(data, ledgers, ensemble, quorum, passwd, conf);
+            throughput = warmUp(data, ledgers, ensemble, quorum, passwd, warmupMessages, conf);
             LOG.info("Warmup tp: " + throughput);
             LOG.info("Warmup phase finished");
         }
@@ -438,7 +445,7 @@ private static double percentile(long[] latency, int percentile) {
      * <p>TODO: update benchmark to use metadata service uri {@link https://github.com/apache/bookkeeper/issues/1331}
      */
     private static long warmUp(byte[] data, int ledgers, int ensemble, int qSize,
-                               byte[] passwd, ClientConfiguration conf)
+                               byte[] passwd, int warmupMessages, ClientConfiguration conf)
             throws KeeperException, IOException, InterruptedException, BKException {
         final CountDownLatch connectLatch = new CountDownLatch(1);
         final int bookies;
@@ -465,7 +472,7 @@ public void process(WatchedEvent event) {
         }
 
         BenchThroughputLatency warmup = new BenchThroughputLatency(bookies, bookies, bookies, passwd,
-                                                                   ledgers, 10000, conf);
+                                                                   ledgers, warmupMessages, conf);
         warmup.setEntryData(data);
         Thread thread = new Thread(warmup);
         thread.start();