content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License. You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.kafka.clients.producer.internals;
  
  import org.apache.kafka.clients.ApiVersions;
  import org.apache.kafka.clients.CommonClientConfigs;
  import org.apache.kafka.clients.MetadataSnapshot;
  import org.apache.kafka.clients.NodeApiVersions;
  import org.apache.kafka.clients.producer.Callback;
  import org.apache.kafka.clients.producer.RecordMetadata;
  import org.apache.kafka.common.Cluster;
  import org.apache.kafka.common.KafkaException;
  import org.apache.kafka.common.Node;
  import org.apache.kafka.common.PartitionInfo;
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.compress.Compression;
  import org.apache.kafka.common.metrics.Metrics;
  import org.apache.kafka.common.protocol.ApiKeys;
  import org.apache.kafka.common.protocol.Errors;
  import org.apache.kafka.common.record.CompressionRatioEstimator;
  import org.apache.kafka.common.record.CompressionType;
  import org.apache.kafka.common.record.DefaultRecord;
  import org.apache.kafka.common.record.DefaultRecordBatch;
  import org.apache.kafka.common.record.MemoryRecords;
  import org.apache.kafka.common.record.MemoryRecordsBuilder;
  import org.apache.kafka.common.record.MutableRecordBatch;
  import org.apache.kafka.common.record.Record;
  import org.apache.kafka.common.record.TimestampType;
  import org.apache.kafka.common.requests.MetadataResponse;
  import org.apache.kafka.common.requests.MetadataResponse.PartitionMetadata;
  import org.apache.kafka.common.utils.LogContext;
  import org.apache.kafka.common.utils.MockTime;
  import org.apache.kafka.common.utils.ProducerIdAndEpoch;
  import org.apache.kafka.common.utils.Time;
  import org.apache.kafka.test.TestUtils;
  
  import org.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  import org.mockito.Mockito;
  
  import java.nio.ByteBuffer;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.Deque;
  import java.util.HashSet;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.OptionalInt;
  import java.util.Random;
  import java.util.Set;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.Future;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.function.Function;
  import java.util.stream.Collectors;
  import java.util.stream.Stream;
  
  import static java.util.Arrays.asList;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  import static org.junit.jupiter.api.Assertions.fail;
  
  public class RecordAccumulatorTest {
  
      private final String topic = "test";
      private final int partition1 = 0;
      private final int partition2 = 1;
      private final int partition3 = 2;
      private final Node node1 = new Node(0, "localhost", 1111);
      private final Node node2 = new Node(1, "localhost", 1112);
  
      private final TopicPartition tp1 = new TopicPartition(topic, partition1);
      private final TopicPartition tp2 = new TopicPartition(topic, partition2);
      private final TopicPartition tp3 = new TopicPartition(topic, partition3);
  
      private final PartitionMetadata partMetadata1 = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()), Optional.empty(), null, null, null);
      private final PartitionMetadata partMetadata2 = new PartitionMetadata(Errors.NONE, tp2, Optional.of(node1.id()), Optional.empty(), null, null, null);
      private final PartitionMetadata partMetadata3 = new PartitionMetadata(Errors.NONE, tp3, Optional.of(node2.id()), Optional.empty(), null, null, null);
      private final List<PartitionMetadata> partMetadatas = new ArrayList<>(Arrays.asList(partMetadata1, partMetadata2, partMetadata3));
  
      private final Map<Integer, Node> nodes = Stream.of(node1, node2).collect(Collectors.toMap(Node::id, Function.identity()));
      private MetadataSnapshot metadataCache = new MetadataSnapshot(null,
          nodes,
          partMetadatas,
          Collections.emptySet(),
          Collections.emptySet(),
          Collections.emptySet(),
          null,
          Collections.emptyMap());
  
      private final Cluster cluster = metadataCache.cluster();
  
      private final MockTime time = new MockTime();
      private final byte[] key = "key".getBytes();
      private final byte[] value = "value".getBytes();
      private final int msgSize = DefaultRecord.sizeInBytes(0, 0, key.length, value.length, Record.EMPTY_HEADERS);
  
      private final Metrics metrics = new Metrics(time);
      private final long maxBlockTimeMs = 1000;
      private final LogContext logContext = new LogContext();
      private AtomicInteger mockRandom = null;
  
      @BeforeEach void setup() {}
  
      @AfterEach
      public void teardown() {
          this.metrics.close();
      }
  
      @Test
      public void testDrainBatches() throws Exception {
          // test case: node1(tp1,tp2) , node2(tp3,tp4)
          // add tp-4
          int partition4 = 3;
          TopicPartition tp4 = new TopicPartition(topic, partition4);
          PartitionMetadata partMetadata4 = new PartitionMetadata(Errors.NONE, tp4, Optional.of(node2.id()), Optional.empty(), null, null, null);
          partMetadatas.add(partMetadata4);
  
          // This test requires that partitions to be drained in order for each node i.e.
          // node1 -> tp1, tp3, and node2 -> tp2, tp4.
          // So setup cluster with this order, and pass this cluster to MetadataCache to preserve this order.
          PartitionInfo part1 = MetadataResponse.toPartitionInfo(partMetadata1, nodes);
          PartitionInfo part2 = MetadataResponse.toPartitionInfo(partMetadata2, nodes);
          PartitionInfo part3 = MetadataResponse.toPartitionInfo(partMetadata3, nodes);
          PartitionInfo part4 = MetadataResponse.toPartitionInfo(partMetadata4, nodes);
          Cluster cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1, part2, part3, part4),
              Collections.emptySet(), Collections.emptySet());
  
          metadataCache = new MetadataSnapshot(null,
              nodes,
              partMetadatas,
              Collections.emptySet(),
              Collections.emptySet(),
              Collections.emptySet(),
              null,
              Collections.emptyMap(),
              cluster);
          long batchSize = value.length + DefaultRecordBatch.RECORD_BATCH_OVERHEAD;
          RecordAccumulator accum = createTestRecordAccumulator((int) batchSize, Integer.MAX_VALUE, Compression.NONE, 10);
  
  
          //  initial data
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          accum.append(topic, partition2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          accum.append(topic, partition3, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          accum.append(topic, partition4, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
  
          // drain batches from 2 nodes: node1 => tp1, node2 => tp3, because the max request size is full after the first batch drained
          Map<Integer, List<ProducerBatch>> batches1 = accum.drain(metadataCache, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
          verifyTopicPartitionInBatches(batches1, tp1, tp3);
  
          // add record for tp1, tp3
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          accum.append(topic, partition3, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
  
          // drain batches from 2 nodes: node1 => tp2, node2 => tp4, because the max request size is full after the first batch drained
          // The drain index should start from next topic partition, that is, node1 => tp2, node2 => tp4
          Map<Integer, List<ProducerBatch>> batches2 = accum.drain(metadataCache, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
          verifyTopicPartitionInBatches(batches2, tp2, tp4);
  
          // make sure in next run, the drain index will start from the beginning
          Map<Integer, List<ProducerBatch>> batches3 = accum.drain(metadataCache, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
          verifyTopicPartitionInBatches(batches3, tp1, tp3);
  
          // add record for tp2, tp3, tp4 and mute the tp4
          accum.append(topic, partition2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          accum.append(topic, partition3, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          accum.append(topic, partition4, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          accum.mutePartition(tp4);
          // drain batches from 2 nodes: node1 => tp2, node2 => tp3 (because tp4 is muted)
          Map<Integer, List<ProducerBatch>> batches4 = accum.drain(metadataCache, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
          verifyTopicPartitionInBatches(batches4, tp2, tp3);
  
          // add record for tp1, tp2, tp3, and unmute tp4
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          accum.append(topic, partition2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          accum.append(topic, partition3, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          accum.unmutePartition(tp4);
          // set maxSize as a max value, so that the all partitions in 2 nodes should be drained: node1 => [tp1, tp2], node2 => [tp3, tp4]
          Map<Integer, List<ProducerBatch>> batches5 = accum.drain(metadataCache, new HashSet<>(Arrays.asList(node1, node2)), Integer.MAX_VALUE, 0);
          verifyTopicPartitionInBatches(batches5, tp1, tp2, tp3, tp4);
      }
  
      private void verifyTopicPartitionInBatches(Map<Integer, List<ProducerBatch>> nodeBatches, TopicPartition... tp) {
          int allTpBatchCount = (int) nodeBatches.values().stream().mapToLong(Collection::size).sum();
          assertEquals(tp.length, allTpBatchCount);
          List<TopicPartition> topicPartitionsInBatch = new ArrayList<>();
          for (Map.Entry<Integer, List<ProducerBatch>> entry : nodeBatches.entrySet()) {
              List<ProducerBatch> tpBatchList = entry.getValue();
              List<TopicPartition> tpList = tpBatchList.stream().map(producerBatch -> producerBatch.topicPartition).collect(Collectors.toList());
              topicPartitionsInBatch.addAll(tpList);
          }
  
          for (int i = 0; i < tp.length; i++) {
              assertEquals(tp[i], topicPartitionsInBatch.get(i));
          }
      }
  
      @Test
      public void testFull() throws Exception {
          long now = time.milliseconds();
  
          // test case assumes that the records do not fill the batch completely
          int batchSize = 1025;
  
          RecordAccumulator accum = createTestRecordAccumulator(
                  batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10L * batchSize, Compression.NONE, 10);
          int appends = expectedNumAppends(batchSize);
          for (int i = 0; i < appends; i++) {
              // append to the first batch
              accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), metadataCache.cluster());
              Deque<ProducerBatch> partitionBatches = accum.getDeque(tp1);
              assertEquals(1, partitionBatches.size());
  
              ProducerBatch batch = partitionBatches.peekFirst();
              assertTrue(batch.isWritable());
              assertEquals(0, accum.ready(metadataCache, now).readyNodes.size(), "No partitions should be ready.");
          }
  
          // this append doesn't fit in the first batch, so a new batch is created and the first batch is closed
  
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), metadataCache.cluster());
          Deque<ProducerBatch> partitionBatches = accum.getDeque(tp1);
          assertEquals(2, partitionBatches.size());
          Iterator<ProducerBatch> partitionBatchesIterator = partitionBatches.iterator();
          assertTrue(partitionBatchesIterator.next().isWritable());
          assertEquals(Collections.singleton(node1), accum.ready(metadataCache, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
  
          List<ProducerBatch> batches = accum.drain(metadataCache, Collections.singleton(node1), Integer.MAX_VALUE, 0).get(node1.id());
          assertEquals(1, batches.size());
          ProducerBatch batch = batches.get(0);
  
          Iterator<Record> iter = batch.records().records().iterator();
          for (int i = 0; i < appends; i++) {
              Record record = iter.next();
              assertEquals(ByteBuffer.wrap(key), record.key(), "Keys should match");
              assertEquals(ByteBuffer.wrap(value), record.value(), "Values should match");
          }
          assertFalse(iter.hasNext(), "No more records");
      }
  
      @Test
      public void testAppendLargeCompressed() throws Exception {
          testAppendLarge(Compression.gzip().build());
      }
  
      @Test
      public void testAppendLargeNonCompressed() throws Exception {
          testAppendLarge(Compression.NONE);
      }
  
      private void testAppendLarge(Compression compression) throws Exception {
          int batchSize = 512;
          byte[] value = new byte[2 * batchSize];
          RecordAccumulator accum = createTestRecordAccumulator(
                  batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, compression, 0);
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), metadataCache.cluster());
          assertEquals(Collections.singleton(node1), accum.ready(metadataCache, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
  
          Deque<ProducerBatch> batches = accum.getDeque(tp1);
          assertEquals(1, batches.size());
          ProducerBatch producerBatch = batches.peek();
          List<MutableRecordBatch> recordBatches = TestUtils.toList(producerBatch.records().batches());
          assertEquals(1, recordBatches.size());
          MutableRecordBatch recordBatch = recordBatches.get(0);
          assertEquals(0L, recordBatch.baseOffset());
          List<Record> records = TestUtils.toList(recordBatch);
          assertEquals(1, records.size());
          Record record = records.get(0);
          assertEquals(0L, record.offset());
          assertEquals(ByteBuffer.wrap(key), record.key());
          assertEquals(ByteBuffer.wrap(value), record.value());
          assertEquals(0L, record.timestamp());
      }
  
      @Test
      public void testAppendLargeOldMessageFormatCompressed() throws Exception {
          testAppendLargeOldMessageFormat(Compression.gzip().build());
      }
  
      @Test
      public void testAppendLargeOldMessageFormatNonCompressed() throws Exception {
          testAppendLargeOldMessageFormat(Compression.NONE);
      }
  
      private void testAppendLargeOldMessageFormat(Compression compression) throws Exception {
          int batchSize = 512;
          byte[] value = new byte[2 * batchSize];
  
          ApiVersions apiVersions = new ApiVersions();
          apiVersions.update(node1.idString(), NodeApiVersions.create(ApiKeys.PRODUCE.id, (short) 0, (short) 2));
  
          RecordAccumulator accum = createTestRecordAccumulator(
                  batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, compression, 0);
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), metadataCache.cluster());
          assertEquals(Collections.singleton(node1), accum.ready(metadataCache, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
  
          Deque<ProducerBatch> batches = accum.getDeque(tp1);
          assertEquals(1, batches.size());
          ProducerBatch producerBatch = batches.peek();
          List<MutableRecordBatch> recordBatches = TestUtils.toList(producerBatch.records().batches());
          assertEquals(1, recordBatches.size());
          MutableRecordBatch recordBatch = recordBatches.get(0);
          assertEquals(0L, recordBatch.baseOffset());
          List<Record> records = TestUtils.toList(recordBatch);
          assertEquals(1, records.size());
          Record record = records.get(0);
          assertEquals(0L, record.offset());
          assertEquals(ByteBuffer.wrap(key), record.key());
          assertEquals(ByteBuffer.wrap(value), record.value());
          assertEquals(0L, record.timestamp());
      }
  
      @Test
      public void testLinger() throws Exception {
          int lingerMs = 10;
          RecordAccumulator accum = createTestRecordAccumulator(
                  1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, Compression.NONE, lingerMs);
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready");
          time.sleep(10);
          assertEquals(Collections.singleton(node1), accum.ready(metadataCache, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
          List<ProducerBatch> batches = accum.drain(metadataCache, Collections.singleton(node1), Integer.MAX_VALUE, 0).get(node1.id());
          assertEquals(1, batches.size());
          ProducerBatch batch = batches.get(0);
  
          Iterator<Record> iter = batch.records().records().iterator();
          Record record = iter.next();
          assertEquals(ByteBuffer.wrap(key), record.key(), "Keys should match");
          assertEquals(ByteBuffer.wrap(value), record.value(), "Values should match");
          assertFalse(iter.hasNext(), "No more records");
      }
  
      @Test
      public void testPartialDrain() throws Exception {
          RecordAccumulator accum = createTestRecordAccumulator(
                  1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, Compression.NONE, 10);
          int appends = 1024 / msgSize + 1;
          List<TopicPartition> partitions = asList(tp1, tp2);
          for (TopicPartition tp : partitions) {
              for (int i = 0; i < appends; i++)
                  accum.append(tp.topic(), tp.partition(), 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          }
          assertEquals(Collections.singleton(node1), accum.ready(metadataCache, time.milliseconds()).readyNodes, "Partition's leader should be ready");
  
          List<ProducerBatch> batches = accum.drain(metadataCache, Collections.singleton(node1), 1024, 0).get(node1.id());
          assertEquals(1, batches.size(), "But due to size bound only one partition should have been retrieved");
      }
  
      @SuppressWarnings("unused")
      @Test
      public void testStressfulSituation() throws Exception {
          final int numThreads = 5;
          final int msgs = 10000;
          final int numParts = 2;
          final RecordAccumulator accum = createTestRecordAccumulator(
              1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, Compression.NONE, 0);
          List<Thread> threads = new ArrayList<>();
          for (int i = 0; i < numThreads; i++) {
              threads.add(new Thread(() -> {
                  for (int j = 0; j < msgs; j++) {
                      try {
                          accum.append(topic, j % numParts, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
                      } catch (Exception e) {
                          e.printStackTrace();
                      }
                  }
              }));
          }
          for (Thread t : threads)
              t.start();
          int read = 0;
          long now = time.milliseconds();
          while (read < numThreads * msgs) {
              Set<Node> nodes = accum.ready(metadataCache, now).readyNodes;
              List<ProducerBatch> batches = accum.drain(metadataCache, nodes, 5 * 1024, 0).get(node1.id());
              if (batches != null) {
                  for (ProducerBatch batch : batches) {
                      for (@SuppressWarnings("UnusedLocalVariable") Record ignored : batch.records().records())
                          read++;
                      accum.deallocate(batch);
                  }
              }
          }
  
          for (Thread t : threads)
              t.join();
      }
  
  
      @Test
      public void testNextReadyCheckDelay() throws Exception {
          // Next check time will use lingerMs since this test won't trigger any retries/backoff
          int lingerMs = 10;
  
          // test case assumes that the records do not fill the batch completely
          int batchSize = 1025;
  
          RecordAccumulator accum = createTestRecordAccumulator(batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD,
                  10 * batchSize, Compression.NONE, lingerMs);
          // Just short of going over the limit so we trigger linger time
          int appends = expectedNumAppends(batchSize);
  
          // Partition on node1 only
          for (int i = 0; i < appends; i++)
              accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
          assertEquals(0, result.readyNodes.size(), "No nodes should be ready.");
          assertEquals(lingerMs, result.nextReadyCheckDelayMs, "Next check time should be the linger time");
  
          time.sleep(lingerMs / 2);
  
          // Add partition on node2 only
          for (int i = 0; i < appends; i++)
              accum.append(topic, partition3, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          result = accum.ready(metadataCache, time.milliseconds());
          assertEquals(0, result.readyNodes.size(), "No nodes should be ready.");
          assertEquals(lingerMs / 2, result.nextReadyCheckDelayMs, "Next check time should be defined by node1, half remaining linger time");
  
          // Add data for another partition on node1, enough to make data sendable immediately
          for (int i = 0; i < appends + 1; i++)
              accum.append(topic, partition2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          result = accum.ready(metadataCache, time.milliseconds());
          assertEquals(Collections.singleton(node1), result.readyNodes, "Node1 should be ready");
          // Note this can actually be < linger time because it may use delays from partitions that aren't sendable
          // but have leaders with other sendable data.
          assertTrue(result.nextReadyCheckDelayMs <= lingerMs, "Next check time should be defined by node2, at most linger time");
      }
  
      @Test
      public void testRetryBackoff() throws Exception {
          int lingerMs = Integer.MAX_VALUE / 16;
          long retryBackoffMs = Integer.MAX_VALUE / 8;
          long retryBackoffMaxMs = retryBackoffMs * 10;
          int deliveryTimeoutMs = Integer.MAX_VALUE;
          long totalSize = 10 * 1024;
          int batchSize = 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD;
          String metricGrpName = "producer-metrics";
  
          final RecordAccumulator accum = new RecordAccumulator(logContext, batchSize,
              Compression.NONE, lingerMs, retryBackoffMs, retryBackoffMaxMs,
              deliveryTimeoutMs, metrics, metricGrpName, time, null,
              new BufferPool(totalSize, batchSize, metrics, time, metricGrpName));
  
          long now = time.milliseconds();
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now + lingerMs + 1);
          assertEquals(Collections.singleton(node1), result.readyNodes, "Node1 should be ready");
          Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, now + lingerMs + 1);
          assertEquals(1, batches.size(), "Node1 should be the only ready node.");
          assertEquals(1, batches.get(0).size(), "Partition 0 should only have one batch drained.");
  
          // Reenqueue the batch
          now = time.milliseconds();
          accum.reenqueue(batches.get(0).get(0), now);
  
          // Put message for partition 1 into accumulator
          accum.append(topic, partition2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          result = accum.ready(metadataCache, now + lingerMs + 1);
          assertEquals(Collections.singleton(node1), result.readyNodes, "Node1 should be ready");
  
          // tp1 should backoff while tp2 should not
          batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, now + lingerMs + 1);
          assertEquals(1, batches.size(), "Node1 should be the only ready node.");
          assertEquals(1, batches.get(0).size(), "Node1 should only have one batch drained.");
          assertEquals(tp2, batches.get(0).get(0).topicPartition, "Node1 should only have one batch for partition 1.");
  
          // Partition 0 can be drained after retry backoff
          long upperBoundBackoffMs = (long) (retryBackoffMs * (1 + CommonClientConfigs.RETRY_BACKOFF_JITTER));
          result = accum.ready(metadataCache, now + upperBoundBackoffMs + 1);
          assertEquals(Collections.singleton(node1), result.readyNodes, "Node1 should be ready");
          batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, now + upperBoundBackoffMs + 1);
          assertEquals(1, batches.size(), "Node1 should be the only ready node.");
          assertEquals(1, batches.get(0).size(), "Node1 should only have one batch drained.");
          assertEquals(tp1, batches.get(0).get(0).topicPartition, "Node1 should only have one batch for partition 0.");
      }
  
      private Map<Integer, List<ProducerBatch>> drainAndCheckBatchAmount(
          MetadataSnapshot metadataCache, Node leader, RecordAccumulator accum, long now, int expected) {
          RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
          if (expected > 0) {
              assertEquals(Collections.singleton(leader), result.readyNodes, "Leader should be ready");
              Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, now);
              assertEquals(expected, batches.size(), "Leader should be the only ready node.");
              assertEquals(expected, batches.get(leader.id()).size(), "Partition should only have " + expected + " batch drained.");
              return batches;
          } else {
              assertEquals(0, result.readyNodes.size(), "Leader should not be ready");
              Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, now);
              assertEquals(0, batches.size(), "Leader should not be drained.");
              return null;
          }
      }
  
      @Test
      public void testExponentialRetryBackoff() throws Exception {
          int lingerMs = Integer.MAX_VALUE / 16;
          long retryBackoffMs = 100;
          long retryBackoffMaxMs = 1000;
          int deliveryTimeoutMs = Integer.MAX_VALUE;
          long totalSize = 10 * 1024;
          int batchSize = 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD;
          String metricGrpName = "producer-metrics";
  
          final RecordAccumulator accum = new RecordAccumulator(logContext, batchSize,
                  Compression.NONE, lingerMs, retryBackoffMs, retryBackoffMaxMs,
                  deliveryTimeoutMs, metrics, metricGrpName, time, null,
                  new BufferPool(totalSize, batchSize, metrics, time, metricGrpName));
  
          long now = time.milliseconds();
          long initial = now;
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
  
          // No backoff for initial attempt
          Map<Integer, List<ProducerBatch>> batches = drainAndCheckBatchAmount(metadataCache, node1, accum, now + lingerMs + 1, 1);
          ProducerBatch batch = batches.get(0).get(0);
          long currentRetryBackoffMs = 0;
  
          for (int i = 0; currentRetryBackoffMs < retryBackoffMaxMs * (1 - CommonClientConfigs.RETRY_BACKOFF_JITTER); i++) {
              // Re-enqueue the batch
              now = time.milliseconds();
              accum.reenqueue(batch, now);
              long lowerBoundBackoffMs = (long) (retryBackoffMs * Math.pow(CommonClientConfigs.RETRY_BACKOFF_EXP_BASE, i) * (1 - CommonClientConfigs.RETRY_BACKOFF_JITTER));
              long upperBoundBackoffMs = (long) (retryBackoffMs * Math.pow(CommonClientConfigs.RETRY_BACKOFF_EXP_BASE, i) * (1 + CommonClientConfigs.RETRY_BACKOFF_JITTER));
              currentRetryBackoffMs = upperBoundBackoffMs;
              // Should back off
              drainAndCheckBatchAmount(metadataCache, node1, accum, initial + lowerBoundBackoffMs - 1, 0);
              // Should not back off
              drainAndCheckBatchAmount(metadataCache, node1, accum, initial + upperBoundBackoffMs + 1, 1);
          }
      }
  
      @Test
      public void testExponentialRetryBackoffLeaderChange() throws Exception {
          int lingerMs = Integer.MAX_VALUE / 16;
          long retryBackoffMs = 100;
          long retryBackoffMaxMs = 1000;
          int deliveryTimeoutMs = Integer.MAX_VALUE;
          long totalSize = 10 * 1024;
          int batchSize = 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD;
          String metricGrpName = "producer-metrics";
  
          PartitionMetadata part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()), Optional.empty(), null, null, null);
          PartitionMetadata part1MetadataChange = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node2.id()), Optional.empty(), null, null, null);
          PartitionMetadata part2Metadata = new PartitionMetadata(Errors.NONE, tp2, Optional.of(node1.id()), Optional.empty(), null, null, null);
          PartitionMetadata part3Metadata = new PartitionMetadata(Errors.NONE, tp3, Optional.of(node2.id()), Optional.empty(), null, null, null);
  
          MetadataSnapshot metadataCache = new MetadataSnapshot(null,
              nodes,
              Arrays.asList(part1Metadata, part2Metadata, part3Metadata),
              Collections.emptySet(),
              Collections.emptySet(),
              Collections.emptySet(),
              null,
              Collections.emptyMap());
  
          MetadataSnapshot metadataCacheChange = new MetadataSnapshot(null,
              nodes,
              Arrays.asList(part1MetadataChange, part2Metadata, part3Metadata),
              Collections.emptySet(),
              Collections.emptySet(),
              Collections.emptySet(),
              null,
              Collections.emptyMap());
  
          final RecordAccumulator accum = new RecordAccumulator(logContext, batchSize,
                  Compression.NONE, lingerMs, retryBackoffMs, retryBackoffMaxMs,
                  deliveryTimeoutMs, metrics, metricGrpName, time, null,
                  new BufferPool(totalSize, batchSize, metrics, time, metricGrpName));
  
          long now = time.milliseconds();
          long initial = now;
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
  
          // No backoff for initial attempt
          Map<Integer, List<ProducerBatch>> batches = drainAndCheckBatchAmount(metadataCache, node1, accum, now + lingerMs + 1, 1);
          ProducerBatch batch = batches.get(0).get(0);
  
          long lowerBoundBackoffMs;
          long upperBoundBackoffMs;
  
          // Retry 1 - delay by retryBackoffMs +/- jitter
          now = time.milliseconds();
          accum.reenqueue(batch, now);
          lowerBoundBackoffMs = (long) (retryBackoffMs * (1 - CommonClientConfigs.RETRY_BACKOFF_JITTER));
          upperBoundBackoffMs = (long) (retryBackoffMs * (1 + CommonClientConfigs.RETRY_BACKOFF_JITTER));
          // Should back off
          drainAndCheckBatchAmount(metadataCache, node1, accum, initial + lowerBoundBackoffMs - 1, 0);
          // Should not back off
          drainAndCheckBatchAmount(metadataCache, node1, accum, initial + upperBoundBackoffMs + 1, 1);
  
          // Retry 2 - delay by retryBackoffMs * 2 +/- jitter
          now = time.milliseconds();
          accum.reenqueue(batch, now);
          lowerBoundBackoffMs = (long) (retryBackoffMs * CommonClientConfigs.RETRY_BACKOFF_EXP_BASE * (1 - CommonClientConfigs.RETRY_BACKOFF_JITTER));
          upperBoundBackoffMs = (long) (retryBackoffMs * CommonClientConfigs.RETRY_BACKOFF_EXP_BASE * (1 + CommonClientConfigs.RETRY_BACKOFF_JITTER));
          // Should back off
          drainAndCheckBatchAmount(metadataCache, node1, accum, initial + lowerBoundBackoffMs - 1, 0);
          // Should not back off
          drainAndCheckBatchAmount(metadataCache, node1, accum, initial + upperBoundBackoffMs + 1, 1);
  
          // Retry 3 - after a leader change, delay by retryBackoffMs * 2^2 +/- jitter (could optimise to do not delay at all)
          now = time.milliseconds();
          accum.reenqueue(batch, now);
          lowerBoundBackoffMs = (long) (retryBackoffMs * Math.pow(CommonClientConfigs.RETRY_BACKOFF_EXP_BASE, 2) * (1 - CommonClientConfigs.RETRY_BACKOFF_JITTER));
          upperBoundBackoffMs = (long) (retryBackoffMs * Math.pow(CommonClientConfigs.RETRY_BACKOFF_EXP_BASE, 2) * (1 + CommonClientConfigs.RETRY_BACKOFF_JITTER));
          // Should back off
          drainAndCheckBatchAmount(metadataCacheChange, node2, accum, initial + lowerBoundBackoffMs - 1, 0);
          // Should not back off
          drainAndCheckBatchAmount(metadataCacheChange, node2, accum, initial + upperBoundBackoffMs + 1, 1);
  
          // Retry 4 - delay by retryBackoffMs * 2^3 +/- jitter (capped to retryBackoffMaxMs)
          now = time.milliseconds();
          accum.reenqueue(batch, now);
          lowerBoundBackoffMs = (long) (retryBackoffMs * Math.pow(CommonClientConfigs.RETRY_BACKOFF_EXP_BASE, 3) * (1 - CommonClientConfigs.RETRY_BACKOFF_JITTER));
          upperBoundBackoffMs = retryBackoffMaxMs;
          // Should back off
          drainAndCheckBatchAmount(metadataCacheChange, node2, accum, initial + lowerBoundBackoffMs - 1, 0);
          // Should not back off
          drainAndCheckBatchAmount(metadataCacheChange, node2, accum, initial + upperBoundBackoffMs + 1, 1);
      }
  
      @Test
      public void testFlush() throws Exception {
          int lingerMs = Integer.MAX_VALUE;
          final RecordAccumulator accum = createTestRecordAccumulator(
                  4 * 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, Compression.NONE, lingerMs);
  
          for (int i = 0; i < 100; i++) {
              accum.append(topic, i % 3, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
              assertTrue(accum.hasIncomplete());
          }
          RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
          assertEquals(0, result.readyNodes.size(), "No nodes should be ready.");
  
          accum.beginFlush();
          result = accum.ready(metadataCache, time.milliseconds());
  
          // drain and deallocate all batches
          Map<Integer, List<ProducerBatch>> results = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertTrue(accum.hasIncomplete());
  
          for (List<ProducerBatch> batches: results.values())
              for (ProducerBatch batch: batches)
                  accum.deallocate(batch);
  
          // should be complete with no unsent records.
          accum.awaitFlushCompletion();
          assertFalse(accum.hasUndrained());
          assertFalse(accum.hasIncomplete());
      }
  
  
      private void delayedInterrupt(final Thread thread, final long delayMs) {
          Thread t = new Thread(() -> {
              Time.SYSTEM.sleep(delayMs);
              thread.interrupt();
          });
          t.start();
      }
  
      @Test
      public void testAwaitFlushComplete() throws Exception {
          RecordAccumulator accum = createTestRecordAccumulator(
              4 * 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, Compression.NONE, Integer.MAX_VALUE);
          accum.append(topic, 0, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
  
          accum.beginFlush();
          assertTrue(accum.flushInProgress());
          delayedInterrupt(Thread.currentThread(), 1000L);
          try {
              accum.awaitFlushCompletion();
              fail("awaitFlushCompletion should throw InterruptException");
          } catch (InterruptedException e) {
              assertFalse(accum.flushInProgress(), "flushInProgress count should be decremented even if thread is interrupted");
          }
      }
  
      @Test
      public void testAbortIncompleteBatches() throws Exception {
          int lingerMs = Integer.MAX_VALUE;
          int numRecords = 100;
  
          final AtomicInteger numExceptionReceivedInCallback = new AtomicInteger(0);
          final RecordAccumulator accum = createTestRecordAccumulator(
              128 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, Compression.NONE, lingerMs);
          class TestCallback implements RecordAccumulator.AppendCallbacks {
              @Override
              public void onCompletion(RecordMetadata metadata, Exception exception) {
                  assertEquals("Producer is closed forcefully.", exception.getMessage());
                  numExceptionReceivedInCallback.incrementAndGet();
              }
  
              @Override
              public void setPartition(int partition) {
              }
          }
          for (int i = 0; i < numRecords; i++)
              accum.append(topic, i % 3, 0L, key, value, null, new TestCallback(), maxBlockTimeMs, time.milliseconds(), cluster);
          RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
          assertFalse(result.readyNodes.isEmpty());
          Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertTrue(accum.hasUndrained());
          assertTrue(accum.hasIncomplete());
  
          int numDrainedRecords = 0;
          for (Map.Entry<Integer, List<ProducerBatch>> drainedEntry : drained.entrySet()) {
              for (ProducerBatch batch : drainedEntry.getValue()) {
                  assertTrue(batch.isClosed());
                  assertFalse(batch.produceFuture.completed());
                  numDrainedRecords += batch.recordCount;
              }
          }
  
          assertTrue(numDrainedRecords > 0 && numDrainedRecords < numRecords);
          accum.abortIncompleteBatches();
          assertEquals(numRecords, numExceptionReceivedInCallback.get());
          assertFalse(accum.hasUndrained());
          assertFalse(accum.hasIncomplete());
      }
  
      @Test
      public void testAbortUnsentBatches() throws Exception {
          int lingerMs = Integer.MAX_VALUE;
          int numRecords = 100;
  
          final AtomicInteger numExceptionReceivedInCallback = new AtomicInteger(0);
          final RecordAccumulator accum = createTestRecordAccumulator(
                  128 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, Compression.NONE, lingerMs);
          final KafkaException cause = new KafkaException();
  
          class TestCallback implements RecordAccumulator.AppendCallbacks {
              @Override
              public void onCompletion(RecordMetadata metadata, Exception exception) {
                  assertEquals(cause, exception);
                  numExceptionReceivedInCallback.incrementAndGet();
              }
  
              @Override
              public void setPartition(int partition) {
              }
          }
          for (int i = 0; i < numRecords; i++)
              accum.append(topic, i % 3, 0L, key, value, null, new TestCallback(), maxBlockTimeMs, time.milliseconds(), cluster);
          RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
          assertFalse(result.readyNodes.isEmpty());
          Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE,
                  time.milliseconds());
          assertTrue(accum.hasUndrained());
          assertTrue(accum.hasIncomplete());
  
          accum.abortUndrainedBatches(cause);
          int numDrainedRecords = 0;
          for (Map.Entry<Integer, List<ProducerBatch>> drainedEntry : drained.entrySet()) {
              for (ProducerBatch batch : drainedEntry.getValue()) {
                  assertTrue(batch.isClosed());
                  assertFalse(batch.produceFuture.completed());
                  numDrainedRecords += batch.recordCount;
              }
          }
  
          assertTrue(numDrainedRecords > 0);
          assertTrue(numExceptionReceivedInCallback.get() > 0);
          assertEquals(numRecords, numExceptionReceivedInCallback.get() + numDrainedRecords);
          assertFalse(accum.hasUndrained());
          assertTrue(accum.hasIncomplete());
      }
  
      private void doExpireBatchSingle(int deliveryTimeoutMs) throws InterruptedException {
          int lingerMs = 300;
          List<Boolean> muteStates = Arrays.asList(false, true);
          Set<Node> readyNodes;
          List<ProducerBatch> expiredBatches;
          // test case assumes that the records do not fill the batch completely
          int batchSize = 1025;
          RecordAccumulator accum = createTestRecordAccumulator(deliveryTimeoutMs,
              batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, Compression.NONE, lingerMs);
  
          // Make the batches ready due to linger. These batches are not in retry
          for (Boolean mute: muteStates) {
              if (time.milliseconds() < System.currentTimeMillis())
                  time.setCurrentTimeMs(System.currentTimeMillis());
              accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
              assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partition should be ready.");
  
              time.sleep(lingerMs);
              readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
              assertEquals(Collections.singleton(node1), readyNodes, "Our partition's leader should be ready");
  
              expiredBatches = accum.expiredBatches(time.milliseconds());
              assertEquals(0, expiredBatches.size(), "The batch should not expire when just linger has passed");
  
              if (mute)
                  accum.mutePartition(tp1);
              else
                  accum.unmutePartition(tp1);
  
              // Advance the clock to expire the batch.
              time.sleep(deliveryTimeoutMs - lingerMs);
              expiredBatches = accum.expiredBatches(time.milliseconds());
              assertEquals(1, expiredBatches.size(), "The batch may expire when the partition is muted");
              assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
          }
      }
  
      @Test
      public void testExpiredBatchSingle() throws InterruptedException {
          doExpireBatchSingle(3200);
      }
  
      @Test
      public void testExpiredBatchSingleMaxValue() throws InterruptedException {
          doExpireBatchSingle(Integer.MAX_VALUE);
      }
  
      @Test
      public void testExpiredBatches() throws InterruptedException {
          long retryBackoffMs = 100L;
          int lingerMs = 30;
          int requestTimeout = 60;
          int deliveryTimeoutMs = 3200;
  
          // test case assumes that the records do not fill the batch completely
          int batchSize = 1025;
  
          RecordAccumulator accum = createTestRecordAccumulator(
              deliveryTimeoutMs, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, Compression.NONE, lingerMs);
          int appends = expectedNumAppends(batchSize);
  
          // Test batches not in retry
          for (int i = 0; i < appends; i++) {
              accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
              assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
          }
          // Make the batches ready due to batch full
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, time.milliseconds(), cluster);
          Set<Node> readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
          assertEquals(Collections.singleton(node1), readyNodes, "Our partition's leader should be ready");
          // Advance the clock to expire the batch.
          time.sleep(deliveryTimeoutMs + 1);
          accum.mutePartition(tp1);
          List<ProducerBatch> expiredBatches = accum.expiredBatches(time.milliseconds());
          assertEquals(2, expiredBatches.size(), "The batches will be muted no matter if the partition is muted or not");
  
          accum.unmutePartition(tp1);
          expiredBatches = accum.expiredBatches(time.milliseconds());
          assertEquals(0, expiredBatches.size(), "All batches should have been expired earlier");
          assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
  
          // Advance the clock to make the next batch ready due to linger.ms
          time.sleep(lingerMs);
          assertEquals(Collections.singleton(node1), readyNodes, "Our partition's leader should be ready");
          time.sleep(requestTimeout + 1);
  
          accum.mutePartition(tp1);
          expiredBatches = accum.expiredBatches(time.milliseconds());
          assertEquals(0, expiredBatches.size(), "The batch should not be expired when metadata is still available and partition is muted");
  
          accum.unmutePartition(tp1);
          expiredBatches = accum.expiredBatches(time.milliseconds());
          assertEquals(0, expiredBatches.size(), "All batches should have been expired");
          assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
  
          // Test batches in retry.
          // Create a retried batch
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, time.milliseconds(), cluster);
          time.sleep(lingerMs);
          readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
          assertEquals(Collections.singleton(node1), readyNodes, "Our partition's leader should be ready");
          Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertEquals(drained.get(node1.id()).size(), 1, "There should be only one batch.");
          time.sleep(1000L);
          accum.reenqueue(drained.get(node1.id()).get(0), time.milliseconds());
  
          // test expiration.
          time.sleep(requestTimeout + retryBackoffMs);
          expiredBatches = accum.expiredBatches(time.milliseconds());
          assertEquals(0, expiredBatches.size(), "The batch should not be expired.");
          time.sleep(1L);
  
          accum.mutePartition(tp1);
          expiredBatches = accum.expiredBatches(time.milliseconds());
          assertEquals(0, expiredBatches.size(), "The batch should not be expired when the partition is muted");
  
          accum.unmutePartition(tp1);
          expiredBatches = accum.expiredBatches(time.milliseconds());
          assertEquals(0, expiredBatches.size(), "All batches should have been expired.");
  
          // Test that when being throttled muted batches are expired before the throttle time is over.
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, time.milliseconds(), cluster);
          time.sleep(lingerMs);
          readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
          assertEquals(Collections.singleton(node1), readyNodes, "Our partition's leader should be ready");
          // Advance the clock to expire the batch.
          time.sleep(requestTimeout + 1);
          accum.mutePartition(tp1);
          expiredBatches = accum.expiredBatches(time.milliseconds());
          assertEquals(0, expiredBatches.size(), "The batch should not be expired when the partition is muted");
  
          long throttleTimeMs = 100L;
          accum.unmutePartition(tp1);
          // The batch shouldn't be expired yet.
          expiredBatches = accum.expiredBatches(time.milliseconds());
          assertEquals(0, expiredBatches.size(), "The batch should not be expired when the partition is muted");
  
          // Once the throttle time is over, the batch can be expired.
          time.sleep(throttleTimeMs);
          expiredBatches = accum.expiredBatches(time.milliseconds());
          assertEquals(0, expiredBatches.size(), "All batches should have been expired earlier");
          assertEquals(1, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
      }
  
      @Test
      public void testMutedPartitions() throws InterruptedException {
          long now = time.milliseconds();
          // test case assumes that the records do not fill the batch completely
          int batchSize = 1025;
  
          RecordAccumulator accum = createTestRecordAccumulator(
                  batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, Compression.NONE, 10);
          int appends = expectedNumAppends(batchSize);
          for (int i = 0; i < appends; i++) {
              accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
              assertEquals(0, accum.ready(metadataCache, now).readyNodes.size(), "No partitions should be ready.");
          }
          time.sleep(2000);
  
          // Test ready with muted partition
          accum.mutePartition(tp1);
          RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
          assertEquals(0, result.readyNodes.size(), "No node should be ready");
  
          // Test ready without muted partition
          accum.unmutePartition(tp1);
          result = accum.ready(metadataCache, time.milliseconds());
          assertFalse(result.readyNodes.isEmpty(), "The batch should be ready");
  
          // Test drain with muted partition
          accum.mutePartition(tp1);
          Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertEquals(0, drained.get(node1.id()).size(), "No batch should have been drained");
  
          // Test drain without muted partition.
          accum.unmutePartition(tp1);
          drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertFalse(drained.get(node1.id()).isEmpty(), "The batch should have been drained.");
      }
  
      @Test
      public void testRecordsDrainedWhenTransactionCompleting() throws Exception {
          int batchSize = 1025;
          int deliveryTimeoutMs = 3200;
          int lingerMs = 10;
          long totalSize = 10 * batchSize;
  
          TransactionManager transactionManager = Mockito.mock(TransactionManager.class);
          RecordAccumulator accumulator = createTestRecordAccumulator(transactionManager, deliveryTimeoutMs,
              batchSize, totalSize, Compression.NONE, lingerMs);
  
          ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(12345L, (short) 5);
          Mockito.when(transactionManager.producerIdAndEpoch()).thenReturn(producerIdAndEpoch);
          Mockito.when(transactionManager.isSendToPartitionAllowed(tp1)).thenReturn(true);
          Mockito.when(transactionManager.transactionContainsPartition(tp1)).thenReturn(true);
          Mockito.when(transactionManager.firstInFlightSequence(tp1)).thenReturn(0);
  
          // Initially, the transaction is still in progress, so we should respect the linger.
          Mockito.when(transactionManager.isCompleting()).thenReturn(false);
  
          accumulator.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs,
              time.milliseconds(), cluster);
          accumulator.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs,
              time.milliseconds(), cluster);
          assertTrue(accumulator.hasUndrained());
  
          RecordAccumulator.ReadyCheckResult firstResult = accumulator.ready(metadataCache, time.milliseconds());
          assertEquals(0, firstResult.readyNodes.size());
          Map<Integer, List<ProducerBatch>> firstDrained = accumulator.drain(metadataCache, firstResult.readyNodes,
              Integer.MAX_VALUE, time.milliseconds());
          assertEquals(0, firstDrained.size());
  
          // Once the transaction begins completion, then the batch should be drained immediately.
          Mockito.when(transactionManager.isCompleting()).thenReturn(true);
  
          RecordAccumulator.ReadyCheckResult secondResult = accumulator.ready(metadataCache, time.milliseconds());
          assertEquals(1, secondResult.readyNodes.size());
          Node readyNode = secondResult.readyNodes.iterator().next();
  
          Map<Integer, List<ProducerBatch>> secondDrained = accumulator.drain(metadataCache, secondResult.readyNodes,
              Integer.MAX_VALUE, time.milliseconds());
          assertEquals(Collections.singleton(readyNode.id()), secondDrained.keySet());
          List<ProducerBatch> batches = secondDrained.get(readyNode.id());
          assertEquals(1, batches.size());
      }
  
      @Test
      public void testSplitAndReenqueue() throws ExecutionException, InterruptedException {
          long now = time.milliseconds();
          RecordAccumulator accum = createTestRecordAccumulator(1024, 10 * 1024, Compression.gzip().build(), 10);
  
          // Create a big batch
          ByteBuffer buffer = ByteBuffer.allocate(4096);
          MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, Compression.NONE, TimestampType.CREATE_TIME, 0L);
          ProducerBatch batch = new ProducerBatch(tp1, builder, now, true);
  
          byte[] value = new byte[1024];
          final AtomicInteger acked = new AtomicInteger(0);
          Callback cb = (metadata, exception) -> acked.incrementAndGet();
          // Append two messages so the batch is too big.
          Future<RecordMetadata> future1 = batch.tryAppend(now, null, value, Record.EMPTY_HEADERS, cb, now);
          Future<RecordMetadata> future2 = batch.tryAppend(now, null, value, Record.EMPTY_HEADERS, cb, now);
          assertNotNull(future1);
          assertNotNull(future2);
          batch.close();
          // Enqueue the batch to the accumulator as if the batch was created by the accumulator.
          accum.reenqueue(batch, now);
          // Re-enqueuing counts as a second attempt, so the delay with jitter is 100 * (1 + 0.2) + 1
          time.sleep(121L);
          // Drain the batch.
          RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
          assertFalse(result.readyNodes.isEmpty(), "The batch should be ready");
          Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertEquals(1, drained.size(), "Only node1 should be drained");
          assertEquals(1, drained.get(node1.id()).size(), "Only one batch should be drained");
          // Split and reenqueue the batch.
          accum.splitAndReenqueue(drained.get(node1.id()).get(0));
          time.sleep(101L);
  
          drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertFalse(drained.isEmpty());
          assertFalse(drained.get(node1.id()).isEmpty());
          drained.get(node1.id()).get(0).complete(acked.get(), 100L);
          assertEquals(1, acked.get(), "The first message should have been acked.");
          assertTrue(future1.isDone());
          assertEquals(0, future1.get().offset());
  
          drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertFalse(drained.isEmpty());
          assertFalse(drained.get(node1.id()).isEmpty());
          drained.get(node1.id()).get(0).complete(acked.get(), 100L);
          assertEquals(2, acked.get(), "Both message should have been acked.");
          assertTrue(future2.isDone());
          assertEquals(1, future2.get().offset());
      }
  
      @Test
      public void testSplitBatchOffAccumulator() throws InterruptedException {
          long seed = System.currentTimeMillis();
          final int batchSize = 1024;
          final int bufferCapacity = 3 * 1024;
  
          // First set the compression ratio estimation to be good.
          CompressionRatioEstimator.setEstimation(tp1.topic(), CompressionType.GZIP, 0.1f);
          RecordAccumulator accum = createTestRecordAccumulator(batchSize, bufferCapacity, Compression.gzip().build(), 0);
          int numSplitBatches = prepareSplitBatches(accum, seed, 100, 20);
          assertTrue(numSplitBatches > 0, "There should be some split batches");
          // Drain all the split batches.
          RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
          for (int i = 0; i < numSplitBatches; i++) {
              Map<Integer, List<ProducerBatch>> drained =
                  accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
              assertFalse(drained.isEmpty());
              assertFalse(drained.get(node1.id()).isEmpty());
          }
          assertTrue(accum.ready(metadataCache, time.milliseconds()).readyNodes.isEmpty(), "All the batches should have been drained.");
          assertEquals(bufferCapacity, accum.bufferPoolAvailableMemory(),
              "The split batches should be allocated off the accumulator");
      }
  
      @Test
      public void testSplitFrequency() throws InterruptedException {
          long seed = System.currentTimeMillis();
          Random random = new Random();
          random.setSeed(seed);
          final int batchSize = 1024;
          final int numMessages = 1000;
  
          RecordAccumulator accum = createTestRecordAccumulator(batchSize, 3 * 1024, Compression.gzip().build(), 10);
          // Adjust the high and low compression ratio message percentage
          for (int goodCompRatioPercentage = 1; goodCompRatioPercentage < 100; goodCompRatioPercentage++) {
              int numSplit = 0;
              int numBatches = 0;
              CompressionRatioEstimator.resetEstimation(topic);
              for (int i = 0; i < numMessages; i++) {
                  int dice = random.nextInt(100);
                  byte[] value = (dice < goodCompRatioPercentage) ?
                          bytesWithGoodCompression(random) : bytesWithPoorCompression(random, 100);
                  accum.append(topic, partition1, 0L, null, value, Record.EMPTY_HEADERS, null, 0, time.milliseconds(), cluster);
                  BatchDrainedResult result = completeOrSplitBatches(accum, batchSize);
                  numSplit += result.numSplit;
                  numBatches += result.numBatches;
              }
              time.sleep(10);
              BatchDrainedResult result = completeOrSplitBatches(accum, batchSize);
              numSplit += result.numSplit;
              numBatches += result.numBatches;
              assertTrue((double) numSplit / numBatches < 0.1f, String.format("Total num batches = %d, split batches = %d, more than 10%% of the batch splits. "
                      + "Random seed is " + seed,
                  numBatches, numSplit));
          }
      }
  
      @Test
      public void testSoonToExpireBatchesArePickedUpForExpiry() throws InterruptedException {
          int lingerMs = 500;
          int batchSize = 1025;
  
          RecordAccumulator accum = createTestRecordAccumulator(
              batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, Compression.NONE, lingerMs);
  
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          Set<Node> readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
          Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertTrue(drained.isEmpty());
          //assertTrue(accum.soonToExpireInFlightBatches().isEmpty());
  
          // advanced clock and send one batch out but it should not be included in soon to expire inflight
          // batches because batch's expiry is quite far.
          time.sleep(lingerMs + 1);
          readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
          drained = accum.drain(metadataCache, readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertEquals(1, drained.size(), "A batch did not drain after linger");
          //assertTrue(accum.soonToExpireInFlightBatches().isEmpty());
  
          // Queue another batch and advance clock such that batch expiry time is earlier than request timeout.
          accum.append(topic, partition2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster);
          time.sleep(lingerMs * 4);
  
          // Now drain and check that accumulator picked up the drained batch because its expiry is soon.
          readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
          drained = accum.drain(metadataCache, readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertEquals(1, drained.size(), "A batch did not drain after linger");
      }
  
      @Test
      public void testExpiredBatchesRetry() throws InterruptedException {
          int lingerMs = 3000;
          int rtt = 1000;
          int deliveryTimeoutMs = 3200;
          Set<Node> readyNodes;
          List<ProducerBatch> expiredBatches;
          List<Boolean> muteStates = Arrays.asList(false, true);
  
          // test case assumes that the records do not fill the batch completely
          int batchSize = 1025;
          RecordAccumulator accum = createTestRecordAccumulator(
              batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, Compression.NONE, lingerMs);
  
          // Test batches in retry.
          for (Boolean mute : muteStates) {
              accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, time.milliseconds(), cluster);
              time.sleep(lingerMs);
              readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
              assertEquals(Collections.singleton(node1), readyNodes, "Our partition's leader should be ready");
              Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, readyNodes, Integer.MAX_VALUE, time.milliseconds());
              assertEquals(1, drained.get(node1.id()).size(), "There should be only one batch.");
              time.sleep(rtt);
              accum.reenqueue(drained.get(node1.id()).get(0), time.milliseconds());
  
              if (mute)
                  accum.mutePartition(tp1);
              else
                  accum.unmutePartition(tp1);
  
              // test expiration
              time.sleep(deliveryTimeoutMs - rtt);
              accum.drain(metadataCache, Collections.singleton(node1), Integer.MAX_VALUE, time.milliseconds());
              expiredBatches = accum.expiredBatches(time.milliseconds());
              assertEquals(mute ? 1 : 0, expiredBatches.size(), "RecordAccumulator has expired batches if the partition is not muted");
          }
      }
  
      @Test
      public void testUniformBuiltInPartitioner() throws Exception {
  
          mockRandom = new AtomicInteger();
  
          long totalSize = 1024 * 1024;
          int batchSize = 1024;  // note that this is also a "sticky" limit for the partitioner
          RecordAccumulator accum = createTestRecordAccumulator(batchSize, totalSize, Compression.NONE, 0);
  
          // Set up callbacks so that we know what partition is chosen.
          final AtomicInteger partition = new AtomicInteger(RecordMetadata.UNKNOWN_PARTITION);
          RecordAccumulator.AppendCallbacks callbacks = new RecordAccumulator.AppendCallbacks() {
              @Override
              public void setPartition(int p) {
                  partition.set(p);
              }
  
              @Override
              public void onCompletion(RecordMetadata metadata, Exception exception) {
  
              }
          };
  
          PartitionInfo part1 = MetadataResponse.toPartitionInfo(partMetadata1, nodes);
          PartitionInfo part2 = MetadataResponse.toPartitionInfo(partMetadata2, nodes);
          PartitionInfo part3 = MetadataResponse.toPartitionInfo(partMetadata3, nodes);
          Cluster cluster = new Cluster(null, asList(node1, node2), asList(part1, part2, part3),
                  Collections.emptySet(), Collections.emptySet());
  
          // Produce small record, we should switch to first partition.
          accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0L, null, value, Record.EMPTY_HEADERS,
                  callbacks, maxBlockTimeMs, time.milliseconds(), cluster);
          assertEquals(partition1, partition.get());
          assertEquals(1, mockRandom.get());
  
          // Produce large record, we should exceed "sticky" limit, but produce to this partition
          // as we try to switch after the "sticky" limit is exceeded.  The switch is disabled
          // because of incomplete batch.
          byte[] largeValue = new byte[batchSize];
          accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0L, null, largeValue, Record.EMPTY_HEADERS,
                  callbacks, maxBlockTimeMs, time.milliseconds(), cluster);
          assertEquals(partition1, partition.get());
          assertEquals(1, mockRandom.get());
  
          // Produce large record, we should switch to next partition as we complete
          // previous batch and exceeded sticky limit.
          accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0L, null, largeValue, Record.EMPTY_HEADERS,
                  callbacks, maxBlockTimeMs, time.milliseconds(), cluster);
          assertEquals(partition2, partition.get());
          assertEquals(2, mockRandom.get());
  
          // Produce large record, we should switch to next partition as we complete
          // previous batch and exceeded sticky limit.
          accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0L, null, largeValue, Record.EMPTY_HEADERS,
                  callbacks, maxBlockTimeMs, time.milliseconds(), cluster);
          assertEquals(partition3, partition.get());
          assertEquals(3, mockRandom.get());
  
          // Produce large record, we should switch to next partition as we complete
          // previous batch and exceeded sticky limit.
          accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0L, null, largeValue, Record.EMPTY_HEADERS,
                  callbacks, maxBlockTimeMs, time.milliseconds(), cluster);
          assertEquals(partition1, partition.get());
          assertEquals(4, mockRandom.get());
      }
  
      @Test
      public void testAdaptiveBuiltInPartitioner() throws Exception {
          // Mock random number generator with just sequential integer.
          mockRandom = new AtomicInteger();
  
          // Create accumulator with partitioner config to enable adaptive partitioning.
          RecordAccumulator.PartitionerConfig config = new RecordAccumulator.PartitionerConfig(true, 100);
          long totalSize = 1024 * 1024;
          int batchSize = 128;
          RecordAccumulator accum = new RecordAccumulator(logContext, batchSize, Compression.NONE, 0, 0L, 0L,
                  3200, config, metrics, "producer-metrics", time, null,
                  new BufferPool(totalSize, batchSize, metrics, time, "producer-internal-metrics")) {
              @Override
              BuiltInPartitioner createBuiltInPartitioner(LogContext logContext, String topic,
                                                                    int stickyBatchSize) {
                  return new SequentialPartitioner(logContext, topic, stickyBatchSize);
              }
          };
  
          byte[] largeValue = new byte[batchSize];
          int[] queueSizes = {1, 7, 2};
          int[] expectedFrequencies = new int[queueSizes.length];
          for (int i = 0; i < queueSizes.length; i++) {
              expectedFrequencies[i] = 8 - queueSizes[i];  // 8 is max(queueSizes) + 1
              for (int c = queueSizes[i]; c-- > 0; ) {
                  // Add large records to each partition, so that each record creates a batch.
                  accum.append(topic, i, 0L, null, largeValue, Record.EMPTY_HEADERS,
                          null, maxBlockTimeMs, time.milliseconds(), cluster);
              }
              assertEquals(queueSizes[i], accum.getDeque(new TopicPartition(topic, i)).size());
          }
  
          // Let the accumulator generate the probability tables.
          accum.ready(metadataCache, time.milliseconds());
  
          // Set up callbacks so that we know what partition is chosen.
          final AtomicInteger partition = new AtomicInteger(RecordMetadata.UNKNOWN_PARTITION);
          RecordAccumulator.AppendCallbacks callbacks = new RecordAccumulator.AppendCallbacks() {
              @Override
              public void setPartition(int p) {
                  partition.set(p);
              }
  
              @Override
              public void onCompletion(RecordMetadata metadata, Exception exception) {
  
              }
          };
  
          // Prime built-in partitioner so that it'd switch on every record, as switching only
          // happens after the "sticky" limit is exceeded.
          accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0L, null, largeValue, Record.EMPTY_HEADERS,
                  callbacks, maxBlockTimeMs, time.milliseconds(), cluster);
  
          // Issue a certain number of partition calls to validate that the partitions would be
          // distributed with frequencies that are reciprocal to the queue sizes.  The number of
          // iterations is defined by the last element of the cumulative frequency table which is
          // the sum of all frequencies.  We do 2 cycles, just so it's more than 1.
          final int numberOfCycles = 2;
          int numberOfIterations = accum.getBuiltInPartitioner(topic).loadStatsRangeEnd() * numberOfCycles;
          int[] frequencies = new int[queueSizes.length];
  
          for (int i = 0; i < numberOfIterations; i++) {
              accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0L, null, largeValue, Record.EMPTY_HEADERS,
                      callbacks, maxBlockTimeMs, time.milliseconds(), cluster);
              ++frequencies[partition.get()];
          }
  
          // Verify that frequencies are reciprocal of queue sizes.
          for (int i = 0; i < frequencies.length; i++) {
              assertEquals(expectedFrequencies[i] * numberOfCycles, frequencies[i],
                      "Partition " + i + " was chosen " + frequencies[i] + " times");
          }
  
          // Test that partitions residing on high-latency nodes don't get switched to.
          accum.updateNodeLatencyStats(0, time.milliseconds() - 200, true);
          accum.updateNodeLatencyStats(0, time.milliseconds(), false);
          accum.ready(metadataCache, time.milliseconds());
  
          // Do one append, because partition gets switched after append.
          accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0L, null, largeValue, Record.EMPTY_HEADERS,
                  callbacks, maxBlockTimeMs, time.milliseconds(), cluster);
  
          for (int c = 10; c-- > 0; ) {
              accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0L, null, largeValue, Record.EMPTY_HEADERS,
                      callbacks, maxBlockTimeMs, time.milliseconds(), cluster);
              assertEquals(partition3, partition.get());
          }
  
      }
  
      @Test
      public void testBuiltInPartitionerFractionalBatches() throws Exception {
          // Test how we avoid creating fractional batches with high linger.ms (see
          // BuiltInPartitioner.updatePartitionInfo).
          long totalSize = 1024 * 1024;
          int batchSize = 512;  // note that this is also a "sticky" limit for the partitioner
          int valSize = 32;
          RecordAccumulator accum = createTestRecordAccumulator(batchSize, totalSize, Compression.NONE, 10);
          byte[] value = new byte[valSize];
  
          for (int c = 10; c-- > 0; ) {
              // Produce about 2/3 of the batch size.
              for (int recCount = batchSize * 2 / 3 / valSize; recCount-- > 0; ) {
                  accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0, null, value, Record.EMPTY_HEADERS,
                      null, maxBlockTimeMs, time.milliseconds(), cluster);
              }
  
              // Advance the time to make the batch ready.
              time.sleep(10);
  
              // We should have one batch ready.
              Set<Node> nodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
              assertEquals(1, nodes.size(), "Should have 1 leader ready");
              List<ProducerBatch> batches = accum.drain(metadataCache, nodes, Integer.MAX_VALUE, 0).entrySet().iterator().next().getValue();
              assertEquals(1, batches.size(), "Should have 1 batch ready");
              int actualBatchSize = batches.get(0).records().sizeInBytes();
              assertTrue(actualBatchSize > batchSize / 2, "Batch must be greater than half batch.size");
              assertTrue(actualBatchSize < batchSize, "Batch must be less than batch.size");
          }
      }
  
      /**
       * For a batch being retried, this validates ready() and drain() whether a batch should skip-backoff(retries-immediately), or backoff, based on -
       * 1. how long it has waited between retry attempts.
       * 2. change in leader hosting the partition.
       */
      @Test
      public void testReadyAndDrainWhenABatchIsBeingRetried() throws InterruptedException {
          int part1LeaderEpoch = 100;
          // Create cluster metadata, partition1 being hosted by node1
          PartitionMetadata part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()),  Optional.of(part1LeaderEpoch), null, null, null);
          MetadataSnapshot metadataCache = new MetadataSnapshot(null, nodes, Collections.singletonList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
  
          int batchSize = 10;
          int lingerMs = 10;
          int retryBackoffMs = 100;
          int retryBackoffMaxMs = 1000;
          int deliveryTimeoutMs = Integer.MAX_VALUE;
          long totalSize = 10 * 1024;
          String metricGrpName = "producer-metrics";
          final RecordAccumulator accum = new RecordAccumulator(logContext, batchSize,
              Compression.NONE, lingerMs, retryBackoffMs, retryBackoffMaxMs,
              deliveryTimeoutMs, metrics, metricGrpName, time, null,
              new BufferPool(totalSize, batchSize, metrics, time, metricGrpName));
  
          // Create 1 batch(batchA) to be produced to partition1.
          long now = time.milliseconds();
          accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, now, cluster);
  
          // 1st attempt(not a retry) to produce batchA, it should be ready & drained to be produced.
          {
              now += lingerMs + 1;
              RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
              assertTrue(result.readyNodes.contains(node1), "Node1 is ready");
  
              Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
                  result.readyNodes, 999999 /* maxSize */, now);
              assertTrue(batches.containsKey(node1.id()) && batches.get(node1.id()).size() == 1, "Node1 has 1 batch ready & drained");
              ProducerBatch batch = batches.get(node1.id()).get(0);
              assertEquals(OptionalInt.of(part1LeaderEpoch), batch.currentLeaderEpoch());
              assertEquals(0, batch.attemptsWhenLeaderLastChanged());
              // Re-enqueue batch for subsequent retries & test-cases
              accum.reenqueue(batch, now);
          }
  
          // In this retry of batchA, wait-time between retries is less than configured and no leader change, so should backoff.
          {
              now += 1;
              RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
              assertFalse(result.readyNodes.contains(node1), "Node1 is not ready");
  
              // Try to drain from node1, it should return no batches.
              Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
                  new HashSet<>(Collections.singletonList(node1)), 999999 /* maxSize */, now);
              assertTrue(batches.containsKey(node1.id()) && batches.get(node1.id()).isEmpty(),
                  "No batches ready to be drained on Node1");
          }
  
          // In this retry of batchA, wait-time between retries is less than configured and leader has changed, so should not backoff.
          {
              now += 1;
              part1LeaderEpoch++;
              // Create cluster metadata, with new leader epoch.
              part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()),  Optional.of(part1LeaderEpoch), null, null, null);
              metadataCache = new MetadataSnapshot(null, nodes, Collections.singletonList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
              RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
              assertTrue(result.readyNodes.contains(node1), "Node1 is ready");
  
              Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
                  result.readyNodes, 999999 /* maxSize */, now);
              assertTrue(batches.containsKey(node1.id()) && batches.get(node1.id()).size() == 1, "Node1 has 1 batch ready & drained");
              ProducerBatch batch = batches.get(node1.id()).get(0);
              assertEquals(OptionalInt.of(part1LeaderEpoch), batch.currentLeaderEpoch());
              assertEquals(1, batch.attemptsWhenLeaderLastChanged());
  
              // Re-enqueue batch for subsequent retries/test-cases.
              accum.reenqueue(batch, now);
          }
  
          // In this retry of batchA, wait-time between retries is more than configured and no leader change, so should not backoff.
          {
              now += 2 * retryBackoffMaxMs;
              // Create cluster metadata, with new leader epoch.
              part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()),  Optional.of(part1LeaderEpoch), null, null, null);
              metadataCache = new MetadataSnapshot(null, nodes, Collections.singletonList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
              RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
              assertTrue(result.readyNodes.contains(node1), "Node1 is ready");
  
              Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
                  result.readyNodes, 999999 /* maxSize */, now);
              assertTrue(batches.containsKey(node1.id()) && batches.get(node1.id()).size() == 1, "Node1 has 1 batch ready & drained");
              ProducerBatch batch = batches.get(node1.id()).get(0);
              assertEquals(OptionalInt.of(part1LeaderEpoch), batch.currentLeaderEpoch());
              assertEquals(1, batch.attemptsWhenLeaderLastChanged());
  
              // Re-enqueue batch for subsequent retries/test-cases.
              accum.reenqueue(batch, now);
          }
  
          // In this retry of batchA, wait-time between retries is more than configured and leader has changed, so should not backoff.
          {
              now += 2 * retryBackoffMaxMs;
              part1LeaderEpoch++;
              // Create cluster metadata, with new leader epoch.
              part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()),  Optional.of(part1LeaderEpoch), null, null, null);
              metadataCache = new MetadataSnapshot(null, nodes, Collections.singletonList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
              RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
              assertTrue(result.readyNodes.contains(node1), "Node1 is ready");
  
              Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
                  result.readyNodes, 999999 /* maxSize */, now);
              assertTrue(batches.containsKey(node1.id()) && batches.get(node1.id()).size() == 1, "Node1 has 1 batch ready & drained");
              ProducerBatch batch = batches.get(node1.id()).get(0);
              assertEquals(OptionalInt.of(part1LeaderEpoch), batch.currentLeaderEpoch());
              assertEquals(3, batch.attemptsWhenLeaderLastChanged());
  
              // Re-enqueue batch for subsequent retries/test-cases.
              accum.reenqueue(batch, now);
          }
      }
  
      @Test
      public void testDrainWithANodeThatDoesntHostAnyPartitions() {
          int batchSize = 10;
          int lingerMs = 10;
          long totalSize = 10 * 1024;
          RecordAccumulator accum = createTestRecordAccumulator(batchSize, totalSize, Compression.NONE, lingerMs);
  
          // Create cluster metadata, node2 doesn't host any partitions.
          PartitionMetadata part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()), Optional.empty(), null, null, null);
          MetadataSnapshot metadataCache = new MetadataSnapshot(null, nodes, Collections.singletonList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
  
          // Drain for node2, it should return 0 batches,
          Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
              new HashSet<>(Collections.singletonList(node2)), 999999 /* maxSize */, time.milliseconds());
          assertTrue(batches.get(node2.id()).isEmpty());
      }
  
      private int prepareSplitBatches(RecordAccumulator accum, long seed, int recordSize, int numRecords)
          throws InterruptedException {
          Random random = new Random();
          random.setSeed(seed);
  
          // First set the compression ratio estimation to be good.
          CompressionRatioEstimator.setEstimation(tp1.topic(), CompressionType.GZIP, 0.1f);
          // Append 20 records of 100 bytes size with poor compression ratio should make the batch too big.
          for (int i = 0; i < numRecords; i++) {
              accum.append(topic, partition1, 0L, null, bytesWithPoorCompression(random, recordSize), Record.EMPTY_HEADERS, null, 0, time.milliseconds(), cluster);
          }
  
          RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
          assertFalse(result.readyNodes.isEmpty());
          Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
          assertEquals(1, batches.size());
          assertEquals(1, batches.values().iterator().next().size());
          ProducerBatch batch = batches.values().iterator().next().get(0);
          int numSplitBatches = accum.splitAndReenqueue(batch);
          accum.deallocate(batch);
  
          return numSplitBatches;
      }
  
      private BatchDrainedResult completeOrSplitBatches(RecordAccumulator accum, int batchSize) {
          int numSplit = 0;
          int numBatches = 0;
          boolean batchDrained;
          do {
              batchDrained = false;
              RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
              Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
              for (List<ProducerBatch> batchList : batches.values()) {
                  for (ProducerBatch batch : batchList) {
                      batchDrained = true;
                      numBatches++;
                      if (batch.estimatedSizeInBytes() > batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD) {
                          accum.splitAndReenqueue(batch);
                          // release the resource of the original big batch.
                          numSplit++;
                      } else {
                          batch.complete(0L, 0L);
                      }
                      accum.deallocate(batch);
                  }
              }
          } while (batchDrained);
          return new BatchDrainedResult(numSplit, numBatches);
      }
  
      /**
       * Generates the compression ratio at about 0.6
       */
      private byte[] bytesWithGoodCompression(Random random) {
          byte[] value = new byte[100];
          ByteBuffer buffer = ByteBuffer.wrap(value);
          while (buffer.remaining() > 0)
              buffer.putInt(random.nextInt(1000));
          return value;
      }
  
      /**
       * Generates the compression ratio at about 0.9
       */
      private byte[] bytesWithPoorCompression(Random random, int size) {
          byte[] value = new byte[size];
          random.nextBytes(value);
          return value;
      }
  
      private static class BatchDrainedResult {
          final int numSplit;
          final int numBatches;
          BatchDrainedResult(int numSplit, int numBatches) {
              this.numBatches = numBatches;
              this.numSplit = numSplit;
          }
      }
  
      /**
       * Return the offset delta.
       */
      private int expectedNumAppends(int batchSize) {
          int size = 0;
          int offsetDelta = 0;
          while (true) {
              int recordSize = DefaultRecord.sizeInBytes(offsetDelta, 0, key.length, value.length,
                  Record.EMPTY_HEADERS);
              if (size + recordSize > batchSize)
                  return offsetDelta;
              offsetDelta += 1;
              size += recordSize;
          }
      }
  
      private RecordAccumulator createTestRecordAccumulator(int batchSize, long totalSize, Compression compression, int lingerMs) {
          int deliveryTimeoutMs = 3200;
          return createTestRecordAccumulator(deliveryTimeoutMs, batchSize, totalSize, compression, lingerMs);
      }
  
      private RecordAccumulator createTestRecordAccumulator(int deliveryTimeoutMs, int batchSize, long totalSize, Compression compression, int lingerMs) {
          return createTestRecordAccumulator(null, deliveryTimeoutMs, batchSize, totalSize, compression, lingerMs);
      }
  
      /**
       * Return a test RecordAccumulator instance
       */
      private RecordAccumulator createTestRecordAccumulator(
          TransactionManager txnManager,
          int deliveryTimeoutMs,
          int batchSize,
          long totalSize,
          Compression compression,
          int lingerMs
      ) {
          long retryBackoffMs = 100L;
          long retryBackoffMaxMs = 1000L;
          String metricGrpName = "producer-metrics";
  
          return new RecordAccumulator(
              logContext,
              batchSize,
              compression,
              lingerMs,
              retryBackoffMs,
              retryBackoffMaxMs,
              deliveryTimeoutMs,
              metrics,
              metricGrpName,
              time,
              txnManager,
              new BufferPool(totalSize, batchSize, metrics, time, metricGrpName)) {
              @Override
              BuiltInPartitioner createBuiltInPartitioner(LogContext logContext, String topic,
                                                          int stickyBatchSize) {
                  return new SequentialPartitioner(logContext, topic, stickyBatchSize);
              }
          };
      }
  
      private class SequentialPartitioner extends BuiltInPartitioner {
  
          public SequentialPartitioner(LogContext logContext, String topic, int stickyBatchSize) {
              super(logContext, topic, stickyBatchSize);
          }
  
          @Override
          int randomPartition() {
              return mockRandom == null ? super.randomPartition() : mockRandom.getAndIncrement();
          }
      }
  }
  