KAFKA-17645 Enable warmup in producer performance test (KIP-1052) (#17340)

In order to better analyze steady-state performance of Kafka, this PR
enables a warmup in the Producer Performance test. The warmup duration
is specified as a number of records that are a subset of the total
numRecords. If warmup records is greater than 0, the warmup is
represented by a second Stats object which holds warmup results. Once
warmup records have been exhausted, the test switches to using the
existing Stats object. At end of test, if warmup was enabled, the
summary of the whole test (warump + steady state) is printed followed by
the summary of the steady-state portion of the test.  If no warmup is
used, summary prints don't change from existing behavior. This
contribution is an original work and is licensed to the Kafka project
under the Apache license

Testing strategy comprises new Java unit tests added to
ProducerPerformanceTests.java.

Reviewers: Kirk True <kirk@kirktrue.pro>, Federico Valeri
 <fedevaleri@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>

Author: matt-welch

tools/src/main/java/org/apache/kafka/tools/ProducerPerformance.java

@@ -76,7 +76,12 @@ void start(String[] args) throws IOException {
             // not thread-safe, do not share with other threads
             SplittableRandom random = new SplittableRandom(0);
             ProducerRecord<byte[], byte[]> record;
-            stats = new Stats(config.numRecords, 5000);
+
+            if (config.warmupRecords > 0) {
+                System.out.println("Warmup first " + config.warmupRecords + " records. Steady state results will print after the complete test summary.");
+            }
+            boolean isSteadyState = false;
+            stats = new Stats(config.numRecords, isSteadyState);
             long startMs = System.currentTimeMillis();
 
             ThroughputThrottler throttler = new ThroughputThrottler(config.throughput, startMs);
@@ -95,7 +100,11 @@ void start(String[] args) throws IOException {
                 record = new ProducerRecord<>(config.topicName, payload);
 
                 long sendStartMs = System.currentTimeMillis();
-                cb = new PerfCallback(sendStartMs, payload.length, stats);
+                if ((isSteadyState = config.warmupRecords > 0) && i == config.warmupRecords) {
+                    steadyStateStats = new Stats(config.numRecords - config.warmupRecords, isSteadyState);
+                    stats.suppressPrinting();
+                }
+                cb = new PerfCallback(sendStartMs, payload.length, stats, steadyStateStats);
                 producer.send(record, cb);
 
                 currentTransactionSize++;
@@ -117,6 +126,10 @@ record = new ProducerRecord<>(config.topicName, payload);
 
                 /* print final results */
                 stats.printTotal();
+                /* print steady-state stats if relevant */
+                if (steadyStateStats != null) {
+                    steadyStateStats.printTotal();
+                }
             } else {
                 // Make sure all messages are sent before printing out the stats and the metrics
                 // We need to do this in a different branch for now since tests/kafkatest/sanity_checks/test_performance_services.py
@@ -125,6 +138,10 @@ record = new ProducerRecord<>(config.topicName, payload);
 
                 /* print final results */
                 stats.printTotal();
+                /* print steady-state stats if relevant */
+                if (steadyStateStats != null) {
+                    steadyStateStats.printTotal();
+                }
 
                 /* print out metrics */
                 ToolsUtils.printMetrics(producer.metrics());
@@ -147,8 +164,8 @@ KafkaProducer<byte[], byte[]> createKafkaProducer(Properties props) {
     }
 
     Callback cb;
-
     Stats stats;
+    Stats steadyStateStats;
 
     static byte[] generateRandomPayload(Integer recordSize, List<byte[]> payloadByteList, byte[] payload,
             SplittableRandom random, boolean payloadMonotonic, long recordValue) {
@@ -164,7 +181,7 @@ static byte[] generateRandomPayload(Integer recordSize, List<byte[]> payloadByte
         }
         return payload;
     }
-    
+
     static Properties readProps(List<String> producerProps, String producerConfig) throws IOException {
         Properties props = new Properties();
         if (producerConfig != null) {
@@ -331,6 +348,16 @@ static ArgumentParser argParser() {
                        "--producer.config, or --transactional-id but --transaction-duration-ms is not specified, " +
                        "the default value will be 3000.");
 
+        parser.addArgument("--warmup-records")
+               .action(store())
+               .required(false)
+               .type(Long.class)
+               .metavar("WARMUP-RECORDS")
+               .dest("warmupRecords")
+               .setDefault(0L)
+               .help("The number of records to treat as warmup; these initial records will not be included in steady-state statistics. " +
+                       "An additional summary line will be printed describing the steady-state statistics. (default: 0).");
+
         return parser;
     }
 
@@ -351,8 +378,10 @@ static class Stats {
         private long windowTotalLatency;
         private long windowBytes;
         private long windowStart;
+        private final boolean isSteadyState;
+        private boolean suppressPrint;
 
-        public Stats(long numRecords, int reportingInterval) {
+        public Stats(long numRecords, boolean isSteadyState) {
             this.start = System.currentTimeMillis();
             this.windowStart = System.currentTimeMillis();
             this.iteration = 0;
@@ -365,7 +394,9 @@ public Stats(long numRecords, int reportingInterval) {
             this.windowTotalLatency = 0;
             this.windowBytes = 0;
             this.totalLatency = 0;
-            this.reportingInterval = reportingInterval;
+            this.reportingInterval = 5000;
+            this.isSteadyState = isSteadyState;
+            this.suppressPrint = false;
         }
 
         public void record(int latency, int bytes, long time) {
@@ -383,9 +414,15 @@ public void record(int latency, int bytes, long time) {
             }
             /* maybe report the recent perf */
             if (time - windowStart >= reportingInterval) {
-                printWindow();
+                if (this.isSteadyState && count == windowCount) {
+                    System.out.println("In steady state.");
+                }
+                if (!this.suppressPrint) {
+                    printWindow();
+                }
                 newWindow();
             }
+            this.iteration++;
         }
 
         public long totalCount() {
@@ -433,8 +470,9 @@ public void printTotal() {
             double recsPerSec = 1000.0 * count / (double) elapsed;
             double mbPerSec = 1000.0 * this.bytes / (double) elapsed / (1024.0 * 1024.0);
             int[] percs = percentiles(this.latencies, index, 0.5, 0.95, 0.99, 0.999);
-            System.out.printf("%d records sent, %.1f records/sec (%.2f MB/sec), %.2f ms avg latency, %.2f ms max latency, %d ms 50th, %d ms 95th, %d ms 99th, %d ms 99.9th.%n",
+            System.out.printf("%d%s records sent, %f records/sec (%.2f MB/sec), %.2f ms avg latency, %.2f ms max latency, %d ms 50th, %d ms 95th, %d ms 99th, %d ms 99.9th.%n",
                               count,
+                              this.isSteadyState ? " steady state" : "",
                               recsPerSec,
                               mbPerSec,
                               totalLatency / (double) count,
@@ -455,16 +493,22 @@ private static int[] percentiles(int[] latencies, int count, double... percentil
             }
             return values;
         }
+
+        public void suppressPrinting() {
+            this.suppressPrint = true;
+        }
     }
 
     static final class PerfCallback implements Callback {
         private final long start;
         private final int bytes;
         private final Stats stats;
+        private final Stats steadyStateStats;
 
-        public PerfCallback(long start, int bytes, Stats stats) {
+        public PerfCallback(long start, int bytes, Stats stats, Stats steadyStateStats) {
             this.start = start;
             this.stats = stats;
+            this.steadyStateStats = steadyStateStats;
             this.bytes = bytes;
         }
 
@@ -475,7 +519,9 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
             // magically printed when the sending fails.
             if (exception == null) {
                 this.stats.record(latency, bytes, now);
-                this.stats.iteration++;
+                if (steadyStateStats != null) {
+                    this.steadyStateStats.record(latency, bytes, now);
+                }
             }
             if (exception != null)
                 exception.printStackTrace();
@@ -484,7 +530,8 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
 
     static final class ConfigPostProcessor {
         final String topicName;
-        final Long numRecords;
+        final long numRecords;
+        final long warmupRecords;
         final Integer recordSize;
         final double throughput;
         final boolean payloadMonotonic;
@@ -498,6 +545,7 @@ public ConfigPostProcessor(ArgumentParser parser, String[] args) throws IOExcept
             Namespace namespace = parser.parseArgs(args);
             this.topicName = namespace.getString("topic");
             this.numRecords = namespace.getLong("numRecords");
+            this.warmupRecords = Math.max(namespace.getLong("warmupRecords"), 0);
             this.recordSize = namespace.getInt("recordSize");
             this.throughput = namespace.getDouble("throughput");
             this.payloadMonotonic = namespace.getBoolean("payloadMonotonic");
@@ -508,9 +556,12 @@ public ConfigPostProcessor(ArgumentParser parser, String[] args) throws IOExcept
             String payloadFilePath = namespace.getString("payloadFile");
             Long transactionDurationMsArg = namespace.getLong("transactionDurationMs");
             String transactionIdArg = namespace.getString("transactionalId");
-            if (numRecords != null && numRecords <= 0) {
+            if (numRecords <= 0) {
                 throw new ArgumentParserException("--num-records should be greater than zero", parser);
             }
+            if (warmupRecords >= numRecords) {
+                throw new ArgumentParserException("The value for --warmup-records must be strictly fewer than the number of records in the test, --num-records.", parser);
+            }
             if (recordSize != null && recordSize <= 0) {
                 throw new ArgumentParserException("--record-size should be greater than zero", parser);
             }

tools/src/test/java/org/apache/kafka/tools/ProducerPerformanceTest.java

@@ -381,16 +381,16 @@ public void testDefaultClientId() throws Exception {
     @Test
     public void testStatsInitializationWithLargeNumRecords() {
         long numRecords = Long.MAX_VALUE;
-        assertDoesNotThrow(() -> new ProducerPerformance.Stats(numRecords, 5000));
+        assertDoesNotThrow(() -> new ProducerPerformance.Stats(numRecords, false));
     }
 
     @Test
     public void testStatsCorrectness() throws Exception {
         ExecutorService singleThreaded = Executors.newSingleThreadExecutor();
         final long numRecords = 1000000;
-        ProducerPerformance.Stats stats = new ProducerPerformance.Stats(numRecords, 5000);
+        ProducerPerformance.Stats stats = new ProducerPerformance.Stats(numRecords, false);
         for (long i = 0; i < numRecords; i++) {
-            final Callback callback = new ProducerPerformance.PerfCallback(0, 100, stats);
+            final Callback callback = new ProducerPerformance.PerfCallback(0, 100, stats, null);
             CompletableFuture.runAsync(() -> {
                 callback.onCompletion(null, null);
             }, singleThreaded);
@@ -567,4 +567,77 @@ public void testEnableTransactionByTransactionDurationMs() throws IOException, A
         assertTrue(configs.producerProps.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG).toString()
                 .startsWith(ProducerPerformance.DEFAULT_TRANSACTION_ID_PREFIX));
     }
+
+    @Test
+    public void testWarmupRecordsFractionalValue() throws Exception {
+        String[] args = new String[] {
+            "--topic", "Hello-Kafka",
+            "--num-records", "10",
+            "--warmup-records", "1.5",
+            "--throughput", "100",
+            "--record-size", "100",
+            "--producer-props", "bootstrap.servers=localhost:9000"};
+        ArgumentParser parser = ProducerPerformance.argParser();
+        ArgumentParserException thrown = assertThrows(ArgumentParserException.class, () -> parser.parseArgs(args));
+        thrown.printStackTrace();
+    }
+
+    @Test
+    public void testWarmupRecordsString() throws Exception {
+        String[] args = new String[] {
+            "--topic", "Hello-Kafka",
+            "--num-records", "10",
+            "--warmup-records", "foo",
+            "--throughput", "100",
+            "--record-size", "100",
+            "--producer-props", "bootstrap.servers=localhost:9000"};
+        ArgumentParser parser = ProducerPerformance.argParser();
+        ArgumentParserException thrown = assertThrows(ArgumentParserException.class, () -> parser.parseArgs(args));
+        thrown.printStackTrace();
+    }
+
+    @Test
+    public void testWarmupNumberOfSuccessfulSendAndClose() throws IOException {
+        doReturn(producerMock).when(producerPerformanceSpy).createKafkaProducer(any(Properties.class));
+        doAnswer(invocation -> {
+            producerPerformanceSpy.cb.onCompletion(null, null);
+            return null;
+        }).when(producerMock).send(any(), any());
+
+        String[] args = new String[] {
+            "--topic", "Hello-Kafka",
+            "--num-records", "10",
+            "--warmup-records", "2",
+            "--throughput", "1",
+            "--record-size", "100",
+            "--producer-props", "bootstrap.servers=localhost:9000"};
+        producerPerformanceSpy.start(args);
+
+        verify(producerMock, times(10)).send(any(), any());
+        assertEquals(10, producerPerformanceSpy.stats.totalCount());
+        assertEquals(10 - 2, producerPerformanceSpy.steadyStateStats.totalCount());
+        verify(producerMock, times(1)).close();
+    }
+
+    @Test
+    public void testWarmupNegativeRecordsNormalTest() throws IOException {
+        doReturn(producerMock).when(producerPerformanceSpy).createKafkaProducer(any(Properties.class));
+        doAnswer(invocation -> {
+            producerPerformanceSpy.cb.onCompletion(null, null);
+            return null;
+        }).when(producerMock).send(any(), any());
+
+        String[] args = new String[] {
+            "--topic", "Hello-Kafka",
+            "--num-records", "10",
+            "--warmup-records", "-1",
+            "--throughput", "1",
+            "--record-size", "100",
+            "--producer-props", "bootstrap.servers=localhost:9000"};
+        producerPerformanceSpy.start(args);
+
+        verify(producerMock, times(10)).send(any(), any());
+        assertEquals(10, producerPerformanceSpy.stats.totalCount());
+        verify(producerMock, times(1)).close();
+    }
 }