diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java index 386fa75636ad7..655f99914870b 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java @@ -21,6 +21,7 @@ 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.ProducerConfig; import org.apache.kafka.clients.producer.RecordMetadata; import org.apache.kafka.common.Cluster; import org.apache.kafka.common.KafkaException; @@ -52,6 +53,8 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; import org.mockito.Mockito; import java.nio.ByteBuffer; @@ -80,6 +83,8 @@ import java.util.stream.Stream; import static java.util.Arrays.asList; +import static org.apache.kafka.clients.producer.ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL; +import static org.apache.kafka.clients.producer.ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -134,8 +139,9 @@ public void teardown() { this.metrics.close(); } - @Test - public void testDrainBatches() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testDrainBatches(String allocationStrategy) throws Exception { // test case: node1(tp1,tp2) , node2(tp3,tp4) // add tp-4 int partition4 = 3; @@ -162,8 +168,9 @@ public void testDrainBatches() throws Exception { null, Collections.emptyMap(), cluster); + byte[] value = new byte[512]; long batchSize = value.length + DefaultRecordBatch.RECORD_BATCH_OVERHEAD; - RecordAccumulator accum = createTestRecordAccumulator((int) batchSize, Integer.MAX_VALUE, Compression.NONE, 10); + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, (int) batchSize, Integer.MAX_VALUE, Compression.NONE, 10); // initial data @@ -223,14 +230,15 @@ private void verifyTopicPartitionInBatches(Map> nod } } - @Test - public void testFull() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testFull(String allocationStrategy) throws Exception { long now = time.milliseconds(); // test case assumes that the records do not fill the batch completely int batchSize = 1025; - RecordAccumulator accum = createTestRecordAccumulator( + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10L * batchSize, Compression.NONE, 10); int appends = expectedNumAppends(batchSize); for (int i = 0; i < appends; i++) { @@ -268,18 +276,19 @@ public void testFull() throws Exception { @Test public void testAppendLargeCompressed() throws Exception { - testAppendLarge(Compression.gzip().build()); + testAppendLarge(BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, Compression.gzip().build()); } - @Test - public void testAppendLargeNonCompressed() throws Exception { - testAppendLarge(Compression.NONE); + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testAppendLargeNonCompressed(String allocationStrategy) throws Exception { + testAppendLarge(allocationStrategy, Compression.NONE); } - private void testAppendLarge(Compression compression) throws Exception { + private void testAppendLarge(String allocationStrategy, Compression compression) throws Exception { int batchSize = 512; byte[] value = new byte[2 * batchSize]; - RecordAccumulator accum = createTestRecordAccumulator( + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 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"); @@ -302,22 +311,23 @@ private void testAppendLarge(Compression compression) throws Exception { @Test public void testAppendLargeOldMessageFormatCompressed() throws Exception { - testAppendLargeOldMessageFormat(Compression.gzip().build()); + testAppendLargeOldMessageFormat(BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, Compression.gzip().build()); } - @Test - public void testAppendLargeOldMessageFormatNonCompressed() throws Exception { - testAppendLargeOldMessageFormat(Compression.NONE); + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testAppendLargeOldMessageFormatNonCompressed(String allocationStrategy) throws Exception { + testAppendLargeOldMessageFormat(allocationStrategy, Compression.NONE); } - private void testAppendLargeOldMessageFormat(Compression compression) throws Exception { + private void testAppendLargeOldMessageFormat(String allocationStrategy, 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( + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 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"); @@ -338,10 +348,11 @@ private void testAppendLargeOldMessageFormat(Compression compression) throws Exc assertEquals(0L, record.timestamp()); } - @Test - public void testLinger() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testLinger(String allocationStrategy) throws Exception { int lingerMs = 10; - RecordAccumulator accum = createTestRecordAccumulator( + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 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"); @@ -358,9 +369,10 @@ public void testLinger() throws Exception { assertFalse(iter.hasNext(), "No more records"); } - @Test - public void testPartialDrain() throws Exception { - RecordAccumulator accum = createTestRecordAccumulator( + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testPartialDrain(String allocationStrategy) throws Exception { + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, Compression.NONE, 10); int appends = 1024 / msgSize + 1; List partitions = asList(tp1, tp2); @@ -375,12 +387,13 @@ public void testPartialDrain() throws Exception { } @SuppressWarnings("unused") - @Test - public void testStressfulSituation() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testStressfulSituation(String allocationStrategy) throws Exception { final int numThreads = 5; final int msgs = 10000; final int numParts = 2; - final RecordAccumulator accum = createTestRecordAccumulator( + final RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, Compression.NONE, 0); List threads = new ArrayList<>(); for (int i = 0; i < numThreads; i++) { @@ -415,16 +428,17 @@ public void testStressfulSituation() throws Exception { } - @Test - public void testNextReadyCheckDelay() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testNextReadyCheckDelay(String allocationStrategy) 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); + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 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); @@ -454,8 +468,9 @@ public void testNextReadyCheckDelay() throws Exception { assertTrue(result.nextReadyCheckDelayMs <= lingerMs, "Next check time should be defined by node2, at most linger time"); } - @Test - public void testRetryBackoff() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testRetryBackoff(String allocationStrategy) throws Exception { int lingerMs = Integer.MAX_VALUE / 16; long retryBackoffMs = Integer.MAX_VALUE / 8; long retryBackoffMaxMs = retryBackoffMs * 10; @@ -464,10 +479,8 @@ public void testRetryBackoff() throws Exception { 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)); + final RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, batchSize, totalSize, + lingerMs, retryBackoffMs, retryBackoffMaxMs, deliveryTimeoutMs, metricGrpName); long now = time.milliseconds(); accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); @@ -519,8 +532,9 @@ private Map> drainAndCheckBatchAmount( } } - @Test - public void testExponentialRetryBackoff() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testExponentialRetryBackoff(String allocationStrategy) throws Exception { int lingerMs = Integer.MAX_VALUE / 16; long retryBackoffMs = 100; long retryBackoffMaxMs = 1000; @@ -529,10 +543,8 @@ public void testExponentialRetryBackoff() throws Exception { 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)); + final RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, batchSize, totalSize, + lingerMs, retryBackoffMs, retryBackoffMaxMs, deliveryTimeoutMs, metricGrpName); long now = time.milliseconds(); long initial = now; @@ -557,8 +569,9 @@ public void testExponentialRetryBackoff() throws Exception { } } - @Test - public void testExponentialRetryBackoffLeaderChange() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testExponentialRetryBackoffLeaderChange(String allocationStrategy) throws Exception { int lingerMs = Integer.MAX_VALUE / 16; long retryBackoffMs = 100; long retryBackoffMaxMs = 1000; @@ -590,10 +603,8 @@ public void testExponentialRetryBackoffLeaderChange() throws Exception { 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)); + final RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, batchSize, totalSize, + lingerMs, retryBackoffMs, retryBackoffMaxMs, deliveryTimeoutMs, metricGrpName); long now = time.milliseconds(); long initial = now; @@ -647,10 +658,11 @@ public void testExponentialRetryBackoffLeaderChange() throws Exception { drainAndCheckBatchAmount(metadataCacheChange, node2, accum, initial + upperBoundBackoffMs + 1, 1); } - @Test - public void testFlush() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testFlush(String allocationStrategy) throws Exception { int lingerMs = Integer.MAX_VALUE; - final RecordAccumulator accum = createTestRecordAccumulator( + final RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 4 * 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, Compression.NONE, lingerMs); for (int i = 0; i < 100; i++) { @@ -686,9 +698,10 @@ private void delayedInterrupt(final Thread thread, final long delayMs) { t.start(); } - @Test - public void testAwaitFlushComplete() throws Exception { - RecordAccumulator accum = createTestRecordAccumulator( + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testAwaitFlushComplete(String allocationStrategy) throws Exception { + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 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); @@ -703,13 +716,14 @@ public void testAwaitFlushComplete() throws Exception { } } - @Test - public void testAbortIncompleteBatches() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testAbortIncompleteBatches(String allocationStrategy) throws Exception { int lingerMs = Integer.MAX_VALUE; int numRecords = 100; final AtomicInteger numExceptionReceivedInCallback = new AtomicInteger(0); - final RecordAccumulator accum = createTestRecordAccumulator( + final RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 128 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, Compression.NONE, lingerMs); class TestCallback implements RecordAccumulator.AppendCallbacks { @Override @@ -746,13 +760,14 @@ public void setPartition(int partition) { assertFalse(accum.hasIncomplete()); } - @Test - public void testAbortUnsentBatches() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testAbortUnsentBatches(String allocationStrategy) throws Exception { int lingerMs = Integer.MAX_VALUE; int numRecords = 100; final AtomicInteger numExceptionReceivedInCallback = new AtomicInteger(0); - final RecordAccumulator accum = createTestRecordAccumulator( + final RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 128 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, Compression.NONE, lingerMs); final KafkaException cause = new KafkaException(); @@ -793,14 +808,14 @@ public void setPartition(int partition) { assertTrue(accum.hasIncomplete()); } - private void doExpireBatchSingle(int deliveryTimeoutMs) throws InterruptedException { + private void doExpireBatchSingle(String allocationStrategy, int deliveryTimeoutMs) throws InterruptedException { int lingerMs = 300; List muteStates = Arrays.asList(false, true); Set readyNodes; List expiredBatches; // test case assumes that the records do not fill the batch completely int batchSize = 1025; - RecordAccumulator accum = createTestRecordAccumulator(deliveryTimeoutMs, + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, deliveryTimeoutMs, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, Compression.NONE, lingerMs); // Make the batches ready due to linger. These batches are not in retry @@ -830,18 +845,21 @@ private void doExpireBatchSingle(int deliveryTimeoutMs) throws InterruptedExcept } } - @Test - public void testExpiredBatchSingle() throws InterruptedException { - doExpireBatchSingle(3200); + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testExpiredBatchSingle(String allocationStrategy) throws InterruptedException { + doExpireBatchSingle(allocationStrategy, 3200); } - @Test - public void testExpiredBatchSingleMaxValue() throws InterruptedException { - doExpireBatchSingle(Integer.MAX_VALUE); + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testExpiredBatchSingleMaxValue(String allocationStrategy) throws InterruptedException { + doExpireBatchSingle(allocationStrategy, Integer.MAX_VALUE); } - @Test - public void testExpiredBatches() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testExpiredBatches(String allocationStrategy) throws InterruptedException { long retryBackoffMs = 100L; int lingerMs = 30; int requestTimeout = 60; @@ -850,7 +868,7 @@ public void testExpiredBatches() throws InterruptedException { // test case assumes that the records do not fill the batch completely int batchSize = 1025; - RecordAccumulator accum = createTestRecordAccumulator( + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, deliveryTimeoutMs, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, Compression.NONE, lingerMs); int appends = expectedNumAppends(batchSize); @@ -937,13 +955,14 @@ public void testExpiredBatches() throws InterruptedException { assertEquals(1, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready."); } - @Test - public void testMutedPartitions() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testMutedPartitions(String allocationStrategy) throws InterruptedException { long now = time.milliseconds(); // test case assumes that the records do not fill the batch completely int batchSize = 1025; - RecordAccumulator accum = createTestRecordAccumulator( + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, Compression.NONE, 10); int appends = expectedNumAppends(batchSize); for (int i = 0; i < appends; i++) { @@ -973,16 +992,17 @@ public void testMutedPartitions() throws InterruptedException { assertFalse(drained.get(node1.id()).isEmpty(), "The batch should have been drained."); } - @Test - public void testRecordsDrainedWhenTransactionCompleting() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testRecordsDrainedWhenTransactionCompleting(String allocationStrategy) 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); + RecordAccumulator accumulator = createTestRecordAccumulator(allocationStrategy, transactionManager, + deliveryTimeoutMs, batchSize, totalSize, Compression.NONE, lingerMs); ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(12345L, (short) 5); Mockito.when(transactionManager.producerIdAndEpoch()).thenReturn(producerIdAndEpoch); @@ -1022,7 +1042,7 @@ public void testRecordsDrainedWhenTransactionCompleting() throws Exception { @Test public void testSplitAndReenqueue() throws ExecutionException, InterruptedException { long now = time.milliseconds(); - RecordAccumulator accum = createTestRecordAccumulator(1024, 10 * 1024, Compression.gzip().build(), 10); + RecordAccumulator accum = createTestRecordAccumulator(BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, 1024, 10 * 1024, Compression.gzip().build(), 10); // Create a big batch ByteBuffer buffer = ByteBuffer.allocate(4096); @@ -1112,7 +1132,7 @@ public void testSplitBatchOffAccumulator() throws InterruptedException { // 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); + RecordAccumulator accum = createTestRecordAccumulator(BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, 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. @@ -1136,7 +1156,7 @@ public void testSplitFrequency() throws InterruptedException { final int batchSize = 1024; final int numMessages = 1000; - RecordAccumulator accum = createTestRecordAccumulator(batchSize, 3 * 1024, Compression.gzip().build(), 10); + RecordAccumulator accum = createTestRecordAccumulator(BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, 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; @@ -1161,12 +1181,13 @@ public void testSplitFrequency() throws InterruptedException { } } - @Test - public void testSoonToExpireBatchesArePickedUpForExpiry() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testSoonToExpireBatchesArePickedUpForExpiry(String allocationStrategy) throws InterruptedException { int lingerMs = 500; int batchSize = 1025; - RecordAccumulator accum = createTestRecordAccumulator( + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 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); @@ -1193,8 +1214,9 @@ public void testSoonToExpireBatchesArePickedUpForExpiry() throws InterruptedExce assertEquals(1, drained.size(), "A batch did not drain after linger"); } - @Test - public void testExpiredBatchesRetry() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testExpiredBatchesRetry(String allocationStrategy) throws InterruptedException { int lingerMs = 3000; int rtt = 1000; int deliveryTimeoutMs = 3200; @@ -1204,7 +1226,7 @@ public void testExpiredBatchesRetry() throws InterruptedException { // test case assumes that the records do not fill the batch completely int batchSize = 1025; - RecordAccumulator accum = createTestRecordAccumulator( + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, Compression.NONE, lingerMs); // Test batches in retry. @@ -1231,14 +1253,15 @@ public void testExpiredBatchesRetry() throws InterruptedException { } } - @Test - public void testUniformBuiltInPartitioner() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testUniformBuiltInPartitioner(String allocationStrategy) 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); + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, batchSize, totalSize, Compression.NONE, 0); // Set up callbacks so that we know what partition is chosen. final AtomicInteger partition = new AtomicInteger(RecordMetadata.UNKNOWN_PARTITION); @@ -1297,8 +1320,9 @@ public void onCompletion(RecordMetadata metadata, Exception exception) { assertEquals(4, mockRandom.get()); } - @Test - public void testAdaptiveBuiltInPartitioner() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testAdaptiveBuiltInPartitioner(String allocationStrategy) throws Exception { // Mock random number generator with just sequential integer. mockRandom = new AtomicInteger(); @@ -1306,15 +1330,24 @@ public void testAdaptiveBuiltInPartitioner() throws Exception { RecordAccumulator.PartitionerConfig config = new RecordAccumulator.PartitionerConfig(true, 100, false, ""); 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, boolean rackAware, String rack) { - return new SequentialPartitioner(logContext, topic, stickyBatchSize, rackAware, rack); + BufferPool pool = createTestBufferPool(allocationStrategy, totalSize, batchSize, "producer-internal-metrics"); + RecordAccumulator accum = allocationStrategy.equals(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL) + ? new ChunkedRecordAccumulator(logContext, batchSize, Compression.NONE, 0, 0L, 0L, + 3200, config, metrics, "producer-metrics", time, null, pool) { + @Override + BuiltInPartitioner createBuiltInPartitioner(LogContext logContext, String topic, + int stickyBatchSize, boolean rackAware, String rack) { + return new SequentialPartitioner(logContext, topic, stickyBatchSize, rackAware, rack); + } } - }; + : new RecordAccumulator(logContext, batchSize, Compression.NONE, 0, 0L, 0L, + 3200, config, metrics, "producer-metrics", time, null, pool) { + @Override + BuiltInPartitioner createBuiltInPartitioner(LogContext logContext, String topic, + int stickyBatchSize, boolean rackAware, String rack) { + return new SequentialPartitioner(logContext, topic, stickyBatchSize, rackAware, rack); + } + }; byte[] largeValue = new byte[batchSize]; int[] queueSizes = {1, 7, 2}; @@ -1388,14 +1421,15 @@ public void onCompletion(RecordMetadata metadata, Exception exception) { } - @Test - public void testBuiltInPartitionerFractionalBatches() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testBuiltInPartitionerFractionalBatches(String allocationStrategy) 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); + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, batchSize, totalSize, Compression.NONE, 10); byte[] value = new byte[valSize]; for (int c = 10; c-- > 0; ) { @@ -1424,24 +1458,25 @@ public void testBuiltInPartitionerFractionalBatches() throws Exception { * 1. how long it has waited between retry attempts. * 2. change in leader hosting the partition. */ - @Test - public void testReadyAndDrainWhenABatchIsBeingRetried() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testReadyAndDrainWhenABatchIsBeingRetried(String allocationStrategy) 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; + // Comfortably above the chunk size, so a chunked batch here spans several chunks. Readiness + // comes from linger below, not from the batch filling up, so the exact value doesn't matter. + int batchSize = 1024; 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)); + final RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, batchSize, totalSize, + lingerMs, retryBackoffMs, retryBackoffMaxMs, deliveryTimeoutMs, metricGrpName); // Create 1 batch(batchA) to be produced to partition1. long now = time.milliseconds(); @@ -1539,12 +1574,14 @@ public void testReadyAndDrainWhenABatchIsBeingRetried() throws InterruptedExcept } } - @Test - public void testDrainWithANodeThatDoesntHostAnyPartitions() { - int batchSize = 10; + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testDrainWithANodeThatDoesntHostAnyPartitions(String allocationStrategy) { + // Nothing is ever appended here, so batch.size only has to stay above the chunk size. + int batchSize = 1024; int lingerMs = 10; long totalSize = 10 * 1024; - RecordAccumulator accum = createTestRecordAccumulator(batchSize, totalSize, Compression.NONE, lingerMs); + RecordAccumulator accum = createTestRecordAccumulator(allocationStrategy, 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); @@ -1651,19 +1688,31 @@ private int expectedNumAppends(int batchSize) { } } - private RecordAccumulator createTestRecordAccumulator(int batchSize, long totalSize, Compression compression, int lingerMs) { + private static Stream allocationStrategies() { + return Stream.of(BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL); + } + + private BufferPool createTestBufferPool(String allocationStrategy, long totalSize, int batchSize, String metricGrpName) { + return allocationStrategy.equals(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL) + ? new BufferPool(totalSize, 128, metrics, time, metricGrpName, BufferPool.AllocationMode.INCREMENTAL) + : new BufferPool(totalSize, batchSize, metrics, time, metricGrpName); + } + + private RecordAccumulator createTestRecordAccumulator(String allocationStrategy, int batchSize, long totalSize, Compression compression, int lingerMs) { int deliveryTimeoutMs = 3200; - return createTestRecordAccumulator(deliveryTimeoutMs, batchSize, totalSize, compression, lingerMs); + return createTestRecordAccumulator(allocationStrategy, 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); + private RecordAccumulator createTestRecordAccumulator(String allocationStrategy, int deliveryTimeoutMs, int batchSize, long totalSize, Compression compression, int lingerMs) { + return createTestRecordAccumulator(allocationStrategy, null, deliveryTimeoutMs, batchSize, totalSize, compression, lingerMs); } /** - * Return a test RecordAccumulator instance + * Return a test RecordAccumulator instance serving the given allocation strategy, with the + * built-in partitioner replaced by a sequential one. */ private RecordAccumulator createTestRecordAccumulator( + String allocationStrategy, TransactionManager txnManager, int deliveryTimeoutMs, int batchSize, @@ -1674,6 +1723,29 @@ private RecordAccumulator createTestRecordAccumulator( long retryBackoffMs = 100L; long retryBackoffMaxMs = 1000L; String metricGrpName = "producer-metrics"; + BufferPool pool = createTestBufferPool(allocationStrategy, totalSize, batchSize, metricGrpName); + + if (allocationStrategy.equals(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL)) { + return new ChunkedRecordAccumulator( + logContext, + batchSize, + compression, + lingerMs, + retryBackoffMs, + retryBackoffMaxMs, + deliveryTimeoutMs, + metrics, + metricGrpName, + time, + txnManager, + pool) { + @Override + BuiltInPartitioner createBuiltInPartitioner(LogContext logContext, String topic, + int stickyBatchSize, boolean rackAware, String rack) { + return new SequentialPartitioner(logContext, topic, stickyBatchSize, rackAware, rack); + } + }; + } return new RecordAccumulator( logContext, @@ -1687,7 +1759,7 @@ private RecordAccumulator createTestRecordAccumulator( metricGrpName, time, txnManager, - new BufferPool(totalSize, batchSize, metrics, time, metricGrpName)) { + pool) { @Override BuiltInPartitioner createBuiltInPartitioner(LogContext logContext, String topic, int stickyBatchSize, boolean rackAware, String rack) { @@ -1696,6 +1768,21 @@ BuiltInPartitioner createBuiltInPartitioner(LogContext logContext, String topic, }; } + /** + * Return an uncompressed test RecordAccumulator instance serving the given allocation strategy, + * with an explicit retry/timeout configuration and the production built-in partitioner. + */ + private RecordAccumulator createTestRecordAccumulator(String allocationStrategy, int batchSize, long totalSize, + int lingerMs, long retryBackoffMs, long retryBackoffMaxMs, + int deliveryTimeoutMs, String metricGrpName) { + BufferPool pool = createTestBufferPool(allocationStrategy, totalSize, batchSize, metricGrpName); + return allocationStrategy.equals(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL) + ? new ChunkedRecordAccumulator(logContext, batchSize, Compression.NONE, lingerMs, retryBackoffMs, + retryBackoffMaxMs, deliveryTimeoutMs, metrics, metricGrpName, time, null, pool) + : new RecordAccumulator(logContext, batchSize, Compression.NONE, lingerMs, retryBackoffMs, + retryBackoffMaxMs, deliveryTimeoutMs, metrics, metricGrpName, time, null, pool); + } + private class SequentialPartitioner extends BuiltInPartitioner { public SequentialPartitioner(LogContext logContext, String topic, int stickyBatchSize, boolean rackAware, String rack) { @@ -1722,8 +1809,8 @@ public void testSplitAndReenqueuePreventInfiniteRecursion() throws InterruptedEx // Initialize test environment with a large batch size long now = time.milliseconds(); int batchSize = 1024 * 1024; // 1MB batch size - RecordAccumulator accum = createTestRecordAccumulator(batchSize, 10 * batchSize, Compression.gzip().build(), - 10); + RecordAccumulator accum = createTestRecordAccumulator(BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, batchSize, + 10 * batchSize, Compression.gzip().build(), 10); // Create a large producer batch manually (bypassing the accumulator's normal // append process) diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java index 61dcab8ef07de..c2c17b4f23e8e 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java @@ -26,6 +26,7 @@ import org.apache.kafka.clients.MockClient; import org.apache.kafka.clients.NetworkClient; import org.apache.kafka.clients.NodeApiVersions; +import org.apache.kafka.clients.producer.ProducerConfig; import org.apache.kafka.clients.producer.RecordMetadata; import org.apache.kafka.common.Cluster; import org.apache.kafka.common.InvalidRecordException; @@ -98,6 +99,7 @@ import org.junit.jupiter.api.Timeout; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; import org.mockito.InOrder; import java.nio.ByteBuffer; @@ -119,7 +121,10 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Stream; +import static org.apache.kafka.clients.producer.ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL; +import static org.apache.kafka.clients.producer.ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL; import static org.apache.kafka.clients.producer.internals.ProducerTestUtils.runUntil; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -546,16 +551,16 @@ public void senderThreadShouldNotGetStuckWhenThrottledAndAddingPartitionsToTxn() } } - @Test - public void testNodeLatencyStats() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testNodeLatencyStats(String allocationStrategy) throws Exception { try (Metrics m = new Metrics()) { // Create a new record accumulator with non-0 partitionAvailabilityTimeoutMs // otherwise it wouldn't update the stats. RecordAccumulator.PartitionerConfig config = new RecordAccumulator.PartitionerConfig(false, 42, false, ""); long totalSize = 1024 * 1024; - accumulator = new RecordAccumulator(logContext, batchSize, Compression.NONE, 0, 0L, 0L, - DELIVERY_TIMEOUT_MS, config, m, "producer-metrics", time, null, - new BufferPool(totalSize, batchSize, m, time, "producer-internal-metrics")); + accumulator = createAccumulator(allocationStrategy, config, 0L, 0L, m, "producer-metrics", + createBufferPool(allocationStrategy, totalSize, m, "producer-internal-metrics")); SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m); apiVersions.update("0", NodeApiVersions.create(ApiKeys.PRODUCE.id, ApiKeys.PRODUCE.oldestVersion(), ApiKeys.PRODUCE.latestVersion())); @@ -3308,21 +3313,19 @@ public void testAbortableErrorIsConvertedToFatalErrorDuringAbort() throws Except } } - @Test - public void testProducerBatchRetriesWhenPartitionLeaderChanges() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testProducerBatchRetriesWhenPartitionLeaderChanges(String allocationStrategy) throws Exception { Metrics m = new Metrics(); SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m); try { // SETUP String metricGrpName = "producer-metrics-test-stats-1"; long totalSize = 1024 * 1024; - BufferPool pool = new BufferPool(totalSize, batchSize, metrics, time, - metricGrpName); + BufferPool pool = createBufferPool(allocationStrategy, totalSize, metrics, metricGrpName); long retryBackoffMaxMs = 100L; - // lingerMs is 0 to send batch as soon as any records are available on it. - this.accumulator = new RecordAccumulator(logContext, batchSize, - Compression.NONE, 0, 10L, retryBackoffMaxMs, - DELIVERY_TIMEOUT_MS, metrics, metricGrpName, time, null, pool); + this.accumulator = createAccumulator(allocationStrategy, new RecordAccumulator.PartitionerConfig(), + 10L, retryBackoffMaxMs, metrics, metricGrpName, pool); Sender sender = new Sender(logContext, client, metadata, this.accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL, 10, senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, null); @@ -3421,8 +3424,9 @@ public void testSenderShouldCloseWhenTransactionManagerInErrorState() { * Test the scenario that FetchResponse returns NOT_LEADER_OR_FOLLOWER, indicating change in leadership, but it * does not contain new leader info(defined in KIP-951). */ - @Test - public void testWhenProduceResponseReturnsWithALeaderShipChangeErrorButNoNewLeaderInformation() + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testWhenProduceResponseReturnsWithALeaderShipChangeErrorButNoNewLeaderInformation(String allocationStrategy) throws InterruptedException { // Setup 3 partitions, tp0 & tp1 return with NOT_LEADER_OR_FOLLOWER, tp2 doesn't return an error. Metrics m = new Metrics(); @@ -3431,13 +3435,10 @@ public void testWhenProduceResponseReturnsWithALeaderShipChangeErrorButNoNewLead // SETUP String metricGrpName = "producer-metrics-test-stats-1"; long totalSize = 1024 * 1024; - BufferPool pool = new BufferPool(totalSize, batchSize, metrics, time, - metricGrpName); + BufferPool pool = createBufferPool(allocationStrategy, totalSize, metrics, metricGrpName); long retryBackoffMaxMs = 100L; - // lingerMs is 0 to send batch as soon as any records are available on it. - this.accumulator = new RecordAccumulator(logContext, batchSize, - Compression.NONE, 0, 10L, retryBackoffMaxMs, - DELIVERY_TIMEOUT_MS, metrics, metricGrpName, time, null, pool); + this.accumulator = createAccumulator(allocationStrategy, new RecordAccumulator.PartitionerConfig(), + 10L, retryBackoffMaxMs, metrics, metricGrpName, pool); Sender sender = new Sender(logContext, client, metadata, this.accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL, 10, senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, null); @@ -3501,8 +3502,9 @@ public void testWhenProduceResponseReturnsWithALeaderShipChangeErrorButNoNewLead * Test the scenario that FetchResponse returns NOT_LEADER_OR_FOLLOWER, indicating change in leadership, along with * new leader info(defined in KIP-951). */ - @Test - public void testWhenProduceResponseReturnsWithALeaderShipChangeErrorAndNewLeaderInformation() + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testWhenProduceResponseReturnsWithALeaderShipChangeErrorAndNewLeaderInformation(String allocationStrategy) throws InterruptedException { // Setup 3 partitions, tp0 & tp1 return with NOT_LEADER_OR_FOLLOWER, tp2 doesn't return an error. Metrics m = new Metrics(); @@ -3511,13 +3513,10 @@ public void testWhenProduceResponseReturnsWithALeaderShipChangeErrorAndNewLeader // SETUP String metricGrpName = "producer-metrics-test-stats-1"; long totalSize = 1024 * 1024; - BufferPool pool = new BufferPool(totalSize, batchSize, metrics, time, - metricGrpName); + BufferPool pool = createBufferPool(allocationStrategy, totalSize, metrics, metricGrpName); long retryBackoffMaxMs = 100L; - // lingerMs is 0 to send batch as soon as any records are available on it. - this.accumulator = new RecordAccumulator(logContext, batchSize, - Compression.NONE, 0, 10L, retryBackoffMaxMs, - DELIVERY_TIMEOUT_MS, metrics, metricGrpName, time, null, pool); + this.accumulator = createAccumulator(allocationStrategy, new RecordAccumulator.PartitionerConfig(), + 10L, retryBackoffMaxMs, metrics, metricGrpName, pool); Sender sender = new Sender(logContext, client, metadata, this.accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL, 10, senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, null); @@ -3818,6 +3817,26 @@ private ProduceResponse produceResponse(TopicPartition tp, long offset, Errors e return produceResponse(tp, offset, error, throttleTimeMs, -1L, null); } + private static Stream allocationStrategies() { + return Stream.of(BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL); + } + + private BufferPool createBufferPool(String allocationStrategy, long totalSize, Metrics metrics, String metricGrpName) { + return allocationStrategy.equals(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL) + ? new BufferPool(totalSize, 128, metrics, time, metricGrpName, BufferPool.AllocationMode.INCREMENTAL) + : new BufferPool(totalSize, batchSize, metrics, time, metricGrpName); + } + + private RecordAccumulator createAccumulator(String allocationStrategy, RecordAccumulator.PartitionerConfig partitionerConfig, + long retryBackoffMs, long retryBackoffMaxMs, Metrics metrics, + String metricGrpName, BufferPool pool) { + return allocationStrategy.equals(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL) + ? new ChunkedRecordAccumulator(logContext, batchSize, Compression.NONE, 0, retryBackoffMs, retryBackoffMaxMs, + DELIVERY_TIMEOUT_MS, partitionerConfig, metrics, metricGrpName, time, null, pool) + : new RecordAccumulator(logContext, batchSize, Compression.NONE, 0, retryBackoffMs, retryBackoffMaxMs, + DELIVERY_TIMEOUT_MS, partitionerConfig, metrics, metricGrpName, time, null, pool); + } + private TransactionManager createTransactionManager() { return new TransactionManager(new LogContext(), null, 0, RETRY_BACKOFF_MS, new ApiVersions(), metadata, false); } diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java index 97fe1913fa837..0dc4d2660dc84 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java @@ -24,6 +24,7 @@ import org.apache.kafka.clients.consumer.CommitFailedException; import org.apache.kafka.clients.consumer.ConsumerGroupMetadata; import org.apache.kafka.clients.consumer.OffsetAndMetadata; +import org.apache.kafka.clients.producer.ProducerConfig; import org.apache.kafka.clients.producer.RecordMetadata; import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.Node; @@ -89,7 +90,9 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; import java.nio.ByteBuffer; @@ -108,10 +111,13 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Supplier; +import java.util.stream.Stream; import static java.util.Collections.singleton; import static java.util.Collections.singletonList; import static java.util.Collections.singletonMap; +import static org.apache.kafka.clients.producer.ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL; +import static org.apache.kafka.clients.producer.ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; @@ -173,6 +179,28 @@ public void setup() { initializeTransactionManager(Optional.of(transactionalId), false, false); } + private static Stream allocationStrategies() { + return Stream.of(BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL); + } + + private static Stream transactionV2AndAllocationStrategies() { + return Stream.of(true, false).flatMap(v2 -> allocationStrategies().map(strategy -> Arguments.of(v2, strategy))); + } + + private RecordAccumulator createAccumulator(String allocationStrategy, int batchSize, long totalSize, + int deliveryTimeoutMs, Metrics metrics, String metricGrpName, + TransactionManager transactionManager) { + boolean incremental = allocationStrategy.equals(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL); + BufferPool pool = incremental + ? new BufferPool(totalSize, 128, metrics, time, metricGrpName, BufferPool.AllocationMode.INCREMENTAL) + : new BufferPool(totalSize, batchSize, metrics, time, metricGrpName); + return incremental + ? new ChunkedRecordAccumulator(logContext, batchSize, Compression.NONE, 0, 0L, 0L, + deliveryTimeoutMs, metrics, metricGrpName, time, transactionManager, pool) + : new RecordAccumulator(logContext, batchSize, Compression.NONE, 0, 0L, 0L, + deliveryTimeoutMs, metrics, metricGrpName, time, transactionManager, pool); + } + private void initializeTransactionManager( Optional transactionalId, boolean transactionV2Enabled @@ -184,6 +212,15 @@ private void initializeTransactionManager( Optional transactionalId, boolean transactionV2Enabled, boolean enable2pc + ) { + initializeTransactionManager(transactionalId, transactionV2Enabled, enable2pc, BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL); + } + + private void initializeTransactionManager( + Optional transactionalId, + boolean transactionV2Enabled, + boolean enable2pc, + String allocationStrategy ) { Metrics metrics = new Metrics(time); @@ -220,9 +257,8 @@ private void initializeTransactionManager( String metricGrpName = "producer-metrics"; this.brokerNode = new Node(0, "localhost", 2211); - this.accumulator = new RecordAccumulator(logContext, batchSize, Compression.NONE, 0, 0L, 0L, - deliveryTimeoutMs, metrics, metricGrpName, time, transactionManager, - new BufferPool(totalSize, batchSize, metrics, time, metricGrpName)); + this.accumulator = createAccumulator(allocationStrategy, batchSize, totalSize, deliveryTimeoutMs, + metrics, metricGrpName, transactionManager); this.sender = new Sender(logContext, this.client, this.metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL, MAX_RETRIES, new SenderMetricsRegistry(metrics), this.time, REQUEST_TIMEOUT, @@ -750,18 +786,17 @@ public void testBatchCompletedAfterProducerReset(boolean transactionV2Enabled) { } @ParameterizedTest - @ValueSource(booleans = {true, false}) - public void testDuplicateSequenceAfterProducerReset(boolean transactionV2Enabled) throws Exception { - initializeTransactionManager(Optional.empty(), transactionV2Enabled); + @MethodSource("transactionV2AndAllocationStrategies") + public void testDuplicateSequenceAfterProducerReset(boolean transactionV2Enabled, String allocationStrategy) throws Exception { + initializeTransactionManager(Optional.empty(), transactionV2Enabled, false, allocationStrategy); initializeIdempotentProducerId(producerId, epoch); Metrics metrics = new Metrics(time); final int requestTimeout = 10000; final int deliveryTimeout = 15000; - RecordAccumulator accumulator = new RecordAccumulator(logContext, 16 * 1024, Compression.NONE, 0, 0L, 0L, - deliveryTimeout, metrics, "", time, transactionManager, - new BufferPool(1024 * 1024, 16 * 1024, metrics, time, "")); + RecordAccumulator accumulator = createAccumulator(allocationStrategy, 16 * 1024, 1024 * 1024, + deliveryTimeout, metrics, "", transactionManager); Sender sender = new Sender(logContext, this.client, this.metadata, accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL, MAX_RETRIES, new SenderMetricsRegistry(metrics), this.time, requestTimeout, @@ -882,8 +917,10 @@ public void testProducerIdReset(boolean transactionV2Enabled) { assertEquals(3, transactionManager.sequenceNumber(tp1)); } - @Test - public void testBasicTransaction() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testBasicTransaction(String allocationStrategy) throws InterruptedException { + initializeTransactionManager(Optional.of(transactionalId), false, false, allocationStrategy); doInitTransactions(); transactionManager.beginTransaction(); @@ -2421,8 +2458,10 @@ public void testFindCoordinatorAllowedInAbortableErrorState() throws Interrupted assertTrue(transactionManager.hasAbortableError()); } - @Test - public void testCancelUnsentAddPartitionsAndProduceOnAbort() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testCancelUnsentAddPartitionsAndProduceOnAbort(String allocationStrategy) throws InterruptedException { + initializeTransactionManager(Optional.of(transactionalId), false, false, allocationStrategy); doInitTransactions(); transactionManager.beginTransaction(); @@ -2468,8 +2507,10 @@ public void testAbortResendsAddPartitionErrorIfRetried() throws InterruptedExcep TestUtils.assertFutureThrows(TransactionAbortedException.class, responseFuture); } - @Test - public void testAbortResendsProduceRequestIfRetried() throws Exception { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testAbortResendsProduceRequestIfRetried(String allocationStrategy) throws Exception { + initializeTransactionManager(Optional.of(transactionalId), false, false, allocationStrategy); doInitTransactions(producerId, epoch); transactionManager.beginTransaction(); @@ -2987,8 +3028,10 @@ public void resendFailedProduceRequestAfterAbortableError() throws Exception { assertNotNull(responseFuture.get()); // should throw the exception which caused the transaction to be aborted. } - @Test - public void testTransitionToAbortableErrorOnBatchExpiry() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testTransitionToAbortableErrorOnBatchExpiry(String allocationStrategy) throws InterruptedException { + initializeTransactionManager(Optional.of(transactionalId), false, false, allocationStrategy); doInitTransactions(); transactionManager.beginTransaction(); @@ -3026,8 +3069,10 @@ public void testTransitionToAbortableErrorOnBatchExpiry() throws InterruptedExce assertTrue(transactionManager.hasAbortableError()); } - @Test - public void testTransitionToAbortableErrorOnMultipleBatchExpiry() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testTransitionToAbortableErrorOnMultipleBatchExpiry(String allocationStrategy) throws InterruptedException { + initializeTransactionManager(Optional.of(transactionalId), false, false, allocationStrategy); doInitTransactions(); transactionManager.beginTransaction(); @@ -3082,8 +3127,10 @@ public void testTransitionToAbortableErrorOnMultipleBatchExpiry() throws Interru assertTrue(transactionManager.hasAbortableError()); } - @Test - public void testDropCommitOnBatchExpiry() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testDropCommitOnBatchExpiry(String allocationStrategy) throws InterruptedException { + initializeTransactionManager(Optional.of(transactionalId), false, false, allocationStrategy); doInitTransactions(); transactionManager.beginTransaction(); @@ -3141,8 +3188,10 @@ public void testDropCommitOnBatchExpiry() throws InterruptedException { assertFalse(transactionManager.transactionContainsPartition(tp0)); } - @Test - public void testTransitionToFatalErrorWhenRetriedBatchIsExpired() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testTransitionToFatalErrorWhenRetriedBatchIsExpired(String allocationStrategy) throws InterruptedException { + initializeTransactionManager(Optional.of(transactionalId), false, false, allocationStrategy); apiVersions.update("0", new NodeApiVersions(Arrays.asList( new ApiVersion() .setApiKey(ApiKeys.INIT_PRODUCER_ID.id) @@ -3651,8 +3700,10 @@ public void testBumpTransactionalEpochOnUnknownProducerIdError() throws Interrup assertEquals(0, transactionManager.sequenceNumber(tp0)); } - @Test - public void testBumpTransactionalEpochOnTimeout() throws InterruptedException { + @ParameterizedTest + @MethodSource("allocationStrategies") + public void testBumpTransactionalEpochOnTimeout(String allocationStrategy) throws InterruptedException { + initializeTransactionManager(Optional.of(transactionalId), false, false, allocationStrategy); final short initialEpoch = 1; final short bumpedEpoch = 2; @@ -3764,10 +3815,10 @@ public void testBumpTransactionalEpochOnRecoverableAddOffsetsRequestError() thro } @ParameterizedTest - @ValueSource(booleans = {true, false}) - public void testHealthyPartitionRetriesDuringEpochBump(boolean transactionV2Enabled) throws InterruptedException { + @MethodSource("transactionV2AndAllocationStrategies") + public void testHealthyPartitionRetriesDuringEpochBump(boolean transactionV2Enabled, String allocationStrategy) throws InterruptedException { // Use a custom Sender to allow multiple inflight requests - initializeTransactionManager(Optional.empty(), transactionV2Enabled); + initializeTransactionManager(Optional.empty(), transactionV2Enabled, false, allocationStrategy); Sender sender = new Sender(logContext, this.client, this.metadata, this.accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL, MAX_RETRIES, new SenderMetricsRegistry(new Metrics(time)), this.time, REQUEST_TIMEOUT, 50, transactionManager); @@ -3889,10 +3940,10 @@ public void testNeedToTriggerEpochBumpFromClientDuringCoordinatorDisconnect() { } @ParameterizedTest - @ValueSource(booleans = {true, false}) - public void testFailedInflightBatchAfterEpochBump(boolean transactionV2Enabled) throws InterruptedException { + @MethodSource("transactionV2AndAllocationStrategies") + public void testFailedInflightBatchAfterEpochBump(boolean transactionV2Enabled, String allocationStrategy) throws InterruptedException { // Use a custom Sender to allow multiple inflight requests - initializeTransactionManager(Optional.empty(), transactionV2Enabled); + initializeTransactionManager(Optional.empty(), transactionV2Enabled, false, allocationStrategy); Sender sender = new Sender(logContext, this.client, this.metadata, this.accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL, MAX_RETRIES, new SenderMetricsRegistry(new Metrics(time)), this.time, REQUEST_TIMEOUT, 50, transactionManager);