From ca31ff183a2ac063a2b9ec560039eb0defbdd82b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Armando=20Garc=C3=ADa=20Sancio?= Date: Thu, 4 Feb 2021 18:49:42 -0700 Subject: [PATCH 1/8] KAFKA-12258: Add support for splitting appending records --- .../apache/kafka/raft/KafkaRaftClient.java | 25 ++- .../org/apache/kafka/raft/RaftClient.java | 48 ++++-- .../raft/internals/BatchAccumulator.java | 102 ++++++++--- .../kafka/raft/internals/BatchBuilder.java | 9 +- .../raft/internals/BatchAccumulatorTest.java | 158 +++++++++++++----- 5 files changed, 259 insertions(+), 83 deletions(-) diff --git a/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java b/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java index 09f8672d287cd..f01852802968e 100644 --- a/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java +++ b/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java @@ -134,7 +134,7 @@ public class KafkaRaftClient implements RaftClient { private static final int RETRY_BACKOFF_BASE_MS = 100; private static final int FETCH_MAX_WAIT_MS = 1000; - static final int MAX_BATCH_SIZE = 1024 * 1024; + static final int MAX_BATCH_SIZE = 8 * 1024 * 1024; private final AtomicReference shutdown = new AtomicReference<>(); private final Logger logger; @@ -2166,13 +2166,27 @@ public void poll() throws IOException { @Override public Long scheduleAppend(int epoch, List records) { + return append(epoch, records, false); + } + + @Override + public Long scheduleAtomicAppend(int epoch, List records) { + return append(epoch, records, true); + } + + private Long append(int epoch, List records, boolean isAtomic) { BatchAccumulator accumulator = this.accumulator; if (accumulator == null) { return Long.MAX_VALUE; } boolean isFirstAppend = accumulator.isEmpty(); - Long offset = accumulator.append(epoch, records); + final Long offset; + if (isAtomic) { + offset = accumulator.appendAtomic(epoch, records); + } else { + offset = accumulator.append(epoch, records); + } // Wakeup the network channel if either this is the first append // or the accumulator is ready to drain now. Checking for the first @@ -2329,9 +2343,10 @@ public void fireHandleCommit(long baseOffset, Records records) { /** * This API is used for committed records originating from {@link #scheduleAppend(int, List)} - * on this instance. In this case, we are able to save the original record objects, - * which saves the need to read them back from disk. This is a nice optimization - * for the leader which is typically doing more work than all of the followers. + * or {@link #scheduleAtomicAppend(int, List)} on this instance. In this case, we are able to + * save the original record objects, which saves the need to read them back from disk. This is + * a nice optimization for the leader which is typically doing more work than all of the + * followers. */ public void fireHandleCommit(long baseOffset, int epoch, List records) { BatchReader.Batch batch = new BatchReader.Batch<>(baseOffset, epoch, records); diff --git a/raft/src/main/java/org/apache/kafka/raft/RaftClient.java b/raft/src/main/java/org/apache/kafka/raft/RaftClient.java index 554ce6173df98..5c60a185b3742 100644 --- a/raft/src/main/java/org/apache/kafka/raft/RaftClient.java +++ b/raft/src/main/java/org/apache/kafka/raft/RaftClient.java @@ -32,11 +32,13 @@ interface Listener { * after consuming the reader. * * Note that there is not a one-to-one correspondence between writes through - * {@link #scheduleAppend(int, List)} and this callback. The Raft implementation - * is free to batch together the records from multiple append calls provided - * that batch boundaries are respected. This means that each batch specified - * through {@link #scheduleAppend(int, List)} is guaranteed to be a subset of - * a batch provided by the {@link BatchReader}. + * {@link #scheduleAppend(int, List)} or {@link #scheduleAtomicAppend(int, List)} + * and this callback. The Raft implementation is free to batch together the records + * from multiple append calls provided that batch boundaries are respected. Records + * specified through {@link #scheduleAtomicAppend(int, List)} are guaranteed to be a + * subset of a batch provided by the {@link BatchReader}. Records specified through + * {@link #scheduleAppend(int, List)} are guaranteed to be in the same order but + * they can map to any number of batches provided by the {@link BatchReader}. * * @param reader reader instance which must be iterated and closed */ @@ -48,7 +50,7 @@ interface Listener { * {@link #handleCommit(BatchReader)}. * * After becoming a leader, the client is eligible to write to the log - * using {@link #scheduleAppend(int, List)}. + * using {@link #scheduleAppend(int, List)} or {@link #scheduleAtomicAppend(int, List)}. * * @param epoch the claimed leader epoch */ @@ -77,6 +79,29 @@ default void handleResign() {} */ void register(Listener listener); + /** + * Append a list of records to the log. The write will be scheduled for some time + * in the future. There is no guarantee that appended records will be written to + * the log and eventually committed. While the order of the records is preserve, they can + * be appended to the log using one or more batches. This means that each record could + * be committed independently. + * + * If the provided current leader epoch does not match the current epoch, which + * is possible when the state machine has yet to observe the epoch change, then + * this method will return {@link Long#MAX_VALUE} to indicate an offset which is + * not possible to become committed. The state machine is expected to discard all + * uncommitted entries after observing an epoch change. + * + * @param epoch the current leader epoch + * @param records the list of records to append + * @return the expected offset of the last record; {@link Long#MAX_VALUE} if the records could + * be committed; null if no memory could be allocated for the batch at this time + * @throws RecordBatchTooLargeException if the size of the records is greater than the maximum + * batch size; if this exception is throw none of the elements in records were + * committed + */ + Long scheduleAppend(int epoch, List records); + /** * Append a list of records to the log. The write will be scheduled for some time * in the future. There is no guarantee that appended records will be written to @@ -91,11 +116,14 @@ default void handleResign() {} * * @param epoch the current leader epoch * @param records the list of records to append - * @return the offset within the current epoch that the log entries will be appended, - * or null if the leader was unable to accept the write (e.g. due to memory - * being reached). + * @return the expected offset of the last record; {@link Long#MAX_VALUE} if the records could + * be committed; null if no memory could be allocated for the batch at this time + * @throws RecordBatchTooLargeException if the size of the records is greater than the maximum + * batch size; if this exception is throw none of the elements in records were + * committed */ - Long scheduleAppend(int epoch, List records); + Long scheduleAtomicAppend(int epoch, List records); + /** * Attempt a graceful shutdown of the client. This allows the leader to proactively diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java index 5331e4dd145d8..62650b7eede9d 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.raft.internals; +import org.apache.kafka.common.errors.RecordBatchTooLargeException; import org.apache.kafka.common.memory.MemoryPool; import org.apache.kafka.common.protocol.ObjectSerializationCache; import org.apache.kafka.common.record.CompressionType; @@ -79,24 +80,69 @@ public BatchAccumulator( } /** - * Append a list of records into an atomic batch. We guarantee all records - * are included in the same underlying record batch so that either all of - * the records become committed or none of them do. + * Append a list of records into as many batches as necessary. * - * @param epoch the expected leader epoch. If this does not match, then - * {@link Long#MAX_VALUE} will be returned as an offset which - * cannot become committed. - * @param records the list of records to include in a batch - * @return the expected offset of the last record (which will be - * {@link Long#MAX_VALUE} if the epoch does not match), or null if - * no memory could be allocated for the batch at this time + * The order of the elements in the records argument will match the order in the batches. + * This method will use as many batches as necessary to serialize all of the records. Since + * this method can split the records into multiple batches it is possible that some of the + * recors will get committed while other will not when the leader fails. + * + * @param epoch the expected leader epoch. If this does not match, then {@link Long#MAX_VALUE} + * will be returned as an offset which cannot become committed + * @param records the list of records to include in the batches + * @return the expected offset of the last record; {@link Long#MAX_VALUE} if the epoch does not + * match; null if no memory could be allocated for the batch at this time + * @throws RecordBatchTooLargeException if the size of one record T is greater than the maximum + * batch size; if this exception is throw some of the elements in records may have + * been committed */ public Long append(int epoch, List records) { if (epoch != this.epoch) { - // If the epoch does not match, then the state machine probably - // has not gotten the notification about the latest epoch change. - // In this case, ignore the append and return a large offset value - // which will never be committed. + return Long.MAX_VALUE; + } + + ObjectSerializationCache serializationCache = new ObjectSerializationCache(); + + appendLock.lock(); + try { + maybeCompleteDrain(); + + for (T record : records) { + BatchBuilder batch = maybeAllocateBatch( + serde.recordSize(record, serializationCache) + ); + if (batch == null) { + return null; + } + + batch.appendRecord(record, serializationCache); + nextOffset += 1; + } + + maybeResetLinger(); + + return nextOffset - 1; + } finally { + appendLock.unlock(); + } + } + + /** + * Append a list of records into an atomic batch. We guarantee all records are included in the + * same underlying record batch so that either all of the records become committed or none of + * them do. + * + * @param epoch the expected leader epoch. If this does not match, then {@link Long#MAX_VALUE} + * will be returned as an offset which cannot become committed + * @param records the list of records to include in a batch + * @return the expected offset of the last record; {@link Long#MAX_VALUE} if the epoch does not + * match; null if no memory could be allocated for the batch at this time + * @throws RecordBatchTooLargeException if the size of the records is greater than the maximum + * batch size; if this exception is throw none of the elements in records were + * committed + */ + public Long appendAtomic(int epoch, List records) { + if (epoch != this.epoch) { return Long.MAX_VALUE; } @@ -106,11 +152,6 @@ public Long append(int epoch, List records) { batchSize += serde.recordSize(record, serializationCache); } - if (batchSize > maxBatchSize) { - throw new IllegalArgumentException("The total size of " + records + " is " + batchSize + - ", which exceeds the maximum allowed batch size of " + maxBatchSize); - } - appendLock.lock(); try { maybeCompleteDrain(); @@ -120,24 +161,35 @@ public Long append(int epoch, List records) { return null; } - // Restart the linger timer if necessary - if (!lingerTimer.isRunning()) { - lingerTimer.reset(time.milliseconds() + lingerMs); - } - for (T record : records) { batch.appendRecord(record, serializationCache); nextOffset += 1; } + maybeResetLinger(); + return nextOffset - 1; } finally { appendLock.unlock(); } } + private void maybeResetLinger() { + if (!lingerTimer.isRunning()) { + lingerTimer.reset(time.milliseconds() + lingerMs); + } + } + private BatchBuilder maybeAllocateBatch(int batchSize) { - if (currentBatch == null) { + if (batchSize > maxBatchSize) { + throw new RecordBatchTooLargeException( + String.format( + "The total record(s) size of %s exceeds the maximum allowed batch size of %s", + batchSize, + maxBatchSize + ) + ); + } else if (currentBatch == null) { startNewBatch(); } else if (!currentBatch.hasRoomFor(batchSize)) { completeCurrentBatch(); diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java index 542bb5197c581..bac61dc85a018 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java @@ -144,15 +144,16 @@ public boolean hasRoomFor(int sizeInBytes) { sizeInBytes, DefaultRecord.EMPTY_HEADERS ); + int bytesNeeded = ByteUtils.sizeOfVarint(recordSizeInBytes) + recordSizeInBytes; + int approxUnusedSizeInBytes = maxBytes - approximateSizeInBytes(); - int unusedSizeInBytes = maxBytes - approximateSizeInBytes(); - if (unusedSizeInBytes >= recordSizeInBytes) { + if (approxUnusedSizeInBytes >= bytesNeeded) { return true; } else if (unflushedBytes > 0) { recordOutput.flush(); unflushedBytes = 0; - unusedSizeInBytes = maxBytes - flushedSizeInBytes(); - return unusedSizeInBytes >= recordSizeInBytes; + int unusedSizeInBytes = maxBytes - flushedSizeInBytes(); + return unusedSizeInBytes >= bytesNeeded; } else { return false; } diff --git a/raft/src/test/java/org/apache/kafka/raft/internals/BatchAccumulatorTest.java b/raft/src/test/java/org/apache/kafka/raft/internals/BatchAccumulatorTest.java index 24e289db2656a..0855bd80cedd3 100644 --- a/raft/src/test/java/org/apache/kafka/raft/internals/BatchAccumulatorTest.java +++ b/raft/src/test/java/org/apache/kafka/raft/internals/BatchAccumulatorTest.java @@ -19,7 +19,11 @@ import org.apache.kafka.common.memory.MemoryPool; import org.apache.kafka.common.protocol.ObjectSerializationCache; import org.apache.kafka.common.protocol.Writable; +import org.apache.kafka.common.record.AbstractRecords; import org.apache.kafka.common.record.CompressionType; +import org.apache.kafka.common.record.DefaultRecord; +import org.apache.kafka.common.record.RecordBatch; +import org.apache.kafka.common.utils.ByteUtils; import org.apache.kafka.common.utils.MockTime; import org.apache.kafka.common.utils.Utils; import org.junit.jupiter.api.Test; @@ -29,6 +33,8 @@ import java.util.Collections; import java.util.List; import java.util.concurrent.CountDownLatch; +import java.util.stream.Collectors; +import java.util.stream.Stream; import static java.util.Arrays.asList; import static java.util.Collections.singletonList; @@ -164,47 +170,84 @@ public void testUnflushedBuffersReleasedByClose() { @Test public void testSingleBatchAccumulation() { - int leaderEpoch = 17; - long baseOffset = 157; - int lingerMs = 50; - int maxBatchSize = 512; - - Mockito.when(memoryPool.tryAllocate(maxBatchSize)) - .thenReturn(ByteBuffer.allocate(maxBatchSize)); - - BatchAccumulator acc = buildAccumulator( - leaderEpoch, - baseOffset, - lingerMs, - maxBatchSize - ); - - List records = asList("a", "b", "c", "d", "e", "f", "g", "h", "i"); - assertEquals(baseOffset, acc.append(leaderEpoch, records.subList(0, 1))); - assertEquals(baseOffset + 2, acc.append(leaderEpoch, records.subList(1, 3))); - assertEquals(baseOffset + 5, acc.append(leaderEpoch, records.subList(3, 6))); - assertEquals(baseOffset + 7, acc.append(leaderEpoch, records.subList(6, 8))); - assertEquals(baseOffset + 8, acc.append(leaderEpoch, records.subList(8, 9))); - - time.sleep(lingerMs); - assertTrue(acc.needsDrain(time.milliseconds())); - - List> batches = acc.drain(); - assertEquals(1, batches.size()); - assertFalse(acc.needsDrain(time.milliseconds())); - assertEquals(Long.MAX_VALUE - time.milliseconds(), acc.timeUntilDrain(time.milliseconds())); - - BatchAccumulator.CompletedBatch batch = batches.get(0); - assertEquals(records, batch.records); - assertEquals(baseOffset, batch.baseOffset); + asList(APPEND, APPEND_ATOMIC).forEach(appender -> { + int leaderEpoch = 17; + long baseOffset = 157; + int lingerMs = 50; + int maxBatchSize = 512; + + Mockito.when(memoryPool.tryAllocate(maxBatchSize)) + .thenReturn(ByteBuffer.allocate(maxBatchSize)); + + BatchAccumulator acc = buildAccumulator( + leaderEpoch, + baseOffset, + lingerMs, + maxBatchSize + ); + + List records = asList("a", "b", "c", "d", "e", "f", "g", "h", "i"); + assertEquals(baseOffset, appender.call(acc, leaderEpoch, records.subList(0, 1))); + assertEquals(baseOffset + 2, appender.call(acc, leaderEpoch, records.subList(1, 3))); + assertEquals(baseOffset + 5, appender.call(acc, leaderEpoch, records.subList(3, 6))); + assertEquals(baseOffset + 7, appender.call(acc, leaderEpoch, records.subList(6, 8))); + assertEquals(baseOffset + 8, appender.call(acc, leaderEpoch, records.subList(8, 9))); + + time.sleep(lingerMs); + assertTrue(acc.needsDrain(time.milliseconds())); + + List> batches = acc.drain(); + assertEquals(1, batches.size()); + assertFalse(acc.needsDrain(time.milliseconds())); + assertEquals(Long.MAX_VALUE - time.milliseconds(), acc.timeUntilDrain(time.milliseconds())); + + BatchAccumulator.CompletedBatch batch = batches.get(0); + assertEquals(records, batch.records); + assertEquals(baseOffset, batch.baseOffset); + }); } @Test public void testMultipleBatchAccumulation() { + asList(APPEND, APPEND_ATOMIC).forEach(appender -> { + int leaderEpoch = 17; + long baseOffset = 157; + int lingerMs = 50; + int maxBatchSize = 256; + + Mockito.when(memoryPool.tryAllocate(maxBatchSize)) + .thenReturn(ByteBuffer.allocate(maxBatchSize)); + + BatchAccumulator acc = buildAccumulator( + leaderEpoch, + baseOffset, + lingerMs, + maxBatchSize + ); + + // Append entries until we have 4 batches to drain (3 completed, 1 building) + while (acc.numCompletedBatches() < 3) { + appender.call(acc, leaderEpoch, singletonList("foo")); + } + + List> batches = acc.drain(); + assertEquals(4, batches.size()); + assertTrue(batches.stream().allMatch(batch -> batch.data.sizeInBytes() <= maxBatchSize)); + }); + } + + @Test void testRecordsAreSplit() { int leaderEpoch = 17; long baseOffset = 157; int lingerMs = 50; - int maxBatchSize = 256; + String record = "a"; + int numberOfRecords = 9; + int recordsPerBatch = 2; + int batchHeaderSize = AbstractRecords.recordBatchHeaderSizeInBytes( + RecordBatch.MAGIC_VALUE_V2, + CompressionType.NONE + ); + int maxBatchSize = batchHeaderSize + recordsPerBatch * recordSizeInBytes(record, recordsPerBatch); Mockito.when(memoryPool.tryAllocate(maxBatchSize)) .thenReturn(ByteBuffer.allocate(maxBatchSize)); @@ -216,13 +259,19 @@ public void testMultipleBatchAccumulation() { maxBatchSize ); - // Append entries until we have 4 batches to drain (3 completed, 1 building) - while (acc.numCompletedBatches() < 3) { - acc.append(leaderEpoch, singletonList("foo")); - } + List records = Stream + .generate(() -> record) + .limit(numberOfRecords) + .collect(Collectors.toList()); + assertEquals(baseOffset + numberOfRecords - 1, acc.append(leaderEpoch, records)); + + time.sleep(lingerMs); + assertTrue(acc.needsDrain(time.milliseconds())); List> batches = acc.drain(); - assertEquals(4, batches.size()); + // ceilingDiv(records.size(), recordsPerBatch) + int expectedBatches = (records.size() + recordsPerBatch - 1) / recordsPerBatch; + assertEquals(expectedBatches, batches.size()); assertTrue(batches.stream().allMatch(batch -> batch.data.sizeInBytes() <= maxBatchSize)); } @@ -306,4 +355,35 @@ public void testDrainDoesNotBlockWithConcurrentAppend() throws Exception { }); } + int recordSizeInBytes(String record, int numberOfRecords) { + int serdeSize = serde.recordSize("a", new ObjectSerializationCache()); + + int recordSizeInBytes = DefaultRecord.sizeOfBodyInBytes( + numberOfRecords, + 0, + -1, + serdeSize, + DefaultRecord.EMPTY_HEADERS + ); + + return ByteUtils.sizeOfVarint(recordSizeInBytes) + recordSizeInBytes; + } + + static interface Appender { + Long call(BatchAccumulator acc, int epoch, List records); + } + + static final Appender APPEND_ATOMIC = new Appender() { + @Override + public Long call(BatchAccumulator acc, int epoch, List records) { + return acc.appendAtomic(epoch, records); + } + }; + + static final Appender APPEND = new Appender() { + @Override + public Long call(BatchAccumulator acc, int epoch, List records) { + return acc.append(epoch, records); + } + }; } From 76a0f8806886c3086389435fc1c7c166bab055ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Armando=20Garc=C3=ADa=20Sancio?= Date: Wed, 10 Feb 2021 20:27:16 -0700 Subject: [PATCH 2/8] Compute the needed size using the individual record sizes --- .../raft/internals/BatchAccumulator.java | 19 +++++++----- .../kafka/raft/internals/BatchBuilder.java | 30 +++++++++++-------- .../raft/internals/BatchBuilderTest.java | 4 +-- 3 files changed, 30 insertions(+), 23 deletions(-) diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java index 62650b7eede9d..5dd3f92770470 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java @@ -32,6 +32,7 @@ import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.ReentrantLock; +import java.util.stream.IntStream; public class BatchAccumulator implements Closeable { private final int epoch; @@ -109,7 +110,7 @@ public Long append(int epoch, List records) { for (T record : records) { BatchBuilder batch = maybeAllocateBatch( - serde.recordSize(record, serializationCache) + new int[] {serde.recordSize(record, serializationCache)} ); if (batch == null) { return null; @@ -147,16 +148,16 @@ public Long appendAtomic(int epoch, List records) { } ObjectSerializationCache serializationCache = new ObjectSerializationCache(); - int batchSize = 0; - for (T record : records) { - batchSize += serde.recordSize(record, serializationCache); - } + int[] recordSizes = records + .stream() + .mapToInt(record -> serde.recordSize(record, serializationCache)) + .toArray(); appendLock.lock(); try { maybeCompleteDrain(); - BatchBuilder batch = maybeAllocateBatch(batchSize); + BatchBuilder batch = maybeAllocateBatch(recordSizes); if (batch == null) { return null; } @@ -180,7 +181,9 @@ private void maybeResetLinger() { } } - private BatchBuilder maybeAllocateBatch(int batchSize) { + private BatchBuilder maybeAllocateBatch(int[] recordSizes) { + // TODO: This is incorrect. Need to look at the batch overhead + int batchSize = IntStream.of(recordSizes).sum(); if (batchSize > maxBatchSize) { throw new RecordBatchTooLargeException( String.format( @@ -191,7 +194,7 @@ private BatchBuilder maybeAllocateBatch(int batchSize) { ); } else if (currentBatch == null) { startNewBatch(); - } else if (!currentBatch.hasRoomFor(batchSize)) { + } else if (!currentBatch.hasRoomFor(recordSizes)) { completeCurrentBatch(); startNewBatch(); } diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java index bac61dc85a018..ebc7ff4ec6c8c 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java @@ -95,7 +95,7 @@ public BatchBuilder( /** * Append a record to this patch. The caller must first verify there is room for the batch - * using {@link #hasRoomFor(int)}. + * using {@link #hasRoomFor(int[])}. * * @param record the record to append * @param serializationCache serialization cache for use in {@link RecordSerde#write(Object, ObjectSerializationCache, Writable)} @@ -123,12 +123,12 @@ public long appendRecord(T record, ObjectSerializationCache serializationCache) } /** - * Check whether the batch has enough room for a record of the given size in bytes. + * Check whether the batch has enough room for all the record values. * - * @param sizeInBytes the size of the record to be appended - * @return true if there is room for the record to be appended, false otherwise + * @param recordSizes the sizes of the records to be appended + * @return true if there is room for the records to be appended, false otherwise */ - public boolean hasRoomFor(int sizeInBytes) { + public boolean hasRoomFor(int[] recordSizes) { if (!isOpenForAppends) { return false; } @@ -137,14 +137,18 @@ public boolean hasRoomFor(int sizeInBytes) { return false; } - int recordSizeInBytes = DefaultRecord.sizeOfBodyInBytes( - (int) (nextOffset - baseOffset), - 0, - -1, - sizeInBytes, - DefaultRecord.EMPTY_HEADERS - ); - int bytesNeeded = ByteUtils.sizeOfVarint(recordSizeInBytes) + recordSizeInBytes; + int bytesNeeded = 0; + for (int size: recordSizes) { + int recordSizeInBytes = DefaultRecord.sizeOfBodyInBytes( + (int) (nextOffset - baseOffset), + 0, + -1, + size, + DefaultRecord.EMPTY_HEADERS + ); + + bytesNeeded += ByteUtils.sizeOfVarint(recordSizeInBytes) + recordSizeInBytes; + } int approxUnusedSizeInBytes = maxBytes - approximateSizeInBytes(); if (approxUnusedSizeInBytes >= bytesNeeded) { diff --git a/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java b/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java index f860df7afd1cc..d817be47546c6 100644 --- a/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java +++ b/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java @@ -69,7 +69,7 @@ void testBuildBatch(CompressionType compressionType) { records.forEach(record -> builder.appendRecord(record, null)); MemoryRecords builtRecordSet = builder.build(); - assertFalse(builder.hasRoomFor(1)); + assertFalse(builder.hasRoomFor(new int[] {1})); assertThrows(IllegalArgumentException.class, () -> builder.appendRecord("a", null)); List builtBatches = Utils.toList(builtRecordSet.batchIterator()); @@ -114,7 +114,7 @@ public void testHasRoomForUncompressed(int batchSize) { String record = "i am a record"; int recordSize = serde.recordSize(record); - while (builder.hasRoomFor(recordSize)) { + while (builder.hasRoomFor(new int[] {recordSize})) { builder.appendRecord(record, null); } From 12bc1adb36107247bbac16c7fc725d24dc8cf671 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Armando=20Garc=C3=ADa=20Sancio?= Date: Mon, 15 Feb 2021 12:27:27 -0700 Subject: [PATCH 3/8] Improve size calculation to include the number of records --- .../raft/internals/BatchAccumulator.java | 8 +- .../kafka/raft/internals/BatchBuilder.java | 93 +++++++++++++++---- .../raft/internals/BatchBuilderTest.java | 22 ++++- 3 files changed, 98 insertions(+), 25 deletions(-) diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java index 5dd3f92770470..817aa35ca8c2e 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java @@ -32,7 +32,6 @@ import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.ReentrantLock; -import java.util.stream.IntStream; public class BatchAccumulator implements Closeable { private final int epoch; @@ -182,13 +181,12 @@ private void maybeResetLinger() { } private BatchBuilder maybeAllocateBatch(int[] recordSizes) { - // TODO: This is incorrect. Need to look at the batch overhead - int batchSize = IntStream.of(recordSizes).sum(); - if (batchSize > maxBatchSize) { + int bytesNeeded = BatchBuilder.batchSizeForRecordSizes(compressionType, nextOffset, recordSizes); + if (bytesNeeded > maxBatchSize) { throw new RecordBatchTooLargeException( String.format( "The total record(s) size of %s exceeds the maximum allowed batch size of %s", - batchSize, + bytesNeeded, maxBatchSize ) ); diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java index ebc7ff4ec6c8c..86f65121cbb9b 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java @@ -85,8 +85,7 @@ public BatchBuilder( this.maxBytes = maxBytes; this.records = new ArrayList<>(); - int batchHeaderSizeInBytes = AbstractRecords.recordBatchHeaderSizeInBytes( - RecordBatch.MAGIC_VALUE_V2, compressionType); + int batchHeaderSizeInBytes = batchHeaderSizeInBytes(compressionType); batchOutput.position(initialPosition + batchHeaderSizeInBytes); this.recordOutput = new DataOutputStreamWritable(new DataOutputStream( @@ -103,7 +102,7 @@ public BatchBuilder( */ public long appendRecord(T record, ObjectSerializationCache serializationCache) { if (!isOpenForAppends) { - throw new IllegalArgumentException("Cannot append new records after the batch has been built"); + throw new IllegalStateException("Cannot append new records after the batch has been built"); } if (nextOffset - baseOffset > Integer.MAX_VALUE) { @@ -133,24 +132,9 @@ public boolean hasRoomFor(int[] recordSizes) { return false; } - if (nextOffset - baseOffset >= Integer.MAX_VALUE) { - return false; - } + int bytesNeeded = bytesNeededForRecordSizes(baseOffset, nextOffset, recordSizes); - int bytesNeeded = 0; - for (int size: recordSizes) { - int recordSizeInBytes = DefaultRecord.sizeOfBodyInBytes( - (int) (nextOffset - baseOffset), - 0, - -1, - size, - DefaultRecord.EMPTY_HEADERS - ); - - bytesNeeded += ByteUtils.sizeOfVarint(recordSizeInBytes) + recordSizeInBytes; - } int approxUnusedSizeInBytes = maxBytes - approximateSizeInBytes(); - if (approxUnusedSizeInBytes >= bytesNeeded) { return true; } else if (unflushedBytes > 0) { @@ -312,4 +296,75 @@ public int writeRecord( recordOutput.writeVarint(0); return ByteUtils.sizeOfVarint(sizeInBytes) + sizeInBytes; } + + /** + * Computes the bytes required to store records with the given sizes in a new batch. + * + * @param compressionType the type of compression + * @param baseOffset the base offset of the batch + * @param recordSizes the size of the records + * @return the bytes required, {@link Integer.MAX_VALUE} when offset delta is too large. + */ + public static int batchSizeForRecordSizes( + CompressionType compressionType, + long baseOffset, + int[] recordSizes + ) { + int bytesNeeded = bytesNeededForRecordSizes(baseOffset, 0, recordSizes); + if (bytesNeeded != Integer.MAX_VALUE) { + return batchHeaderSizeInBytes(compressionType) + bytesNeeded; + } else { + return Integer.MAX_VALUE; + } + } + + /** + * Computes the bytes overhead of an empty batch. + * + * @param compressionType the type of compression + * @return the minimal bytes size of a batch + */ + static int batchHeaderSizeInBytes(CompressionType compressionType) { + return AbstractRecords.recordBatchHeaderSizeInBytes( + RecordBatch.MAGIC_VALUE_V2, + compressionType + ); + } + + /** + * Compute the number of bytes required to store the give record sizes. + * + * @param baseOffset the base offset for the batch + * @param startingOffset the offset of the first record + * @param recordSizes the sizes of all of the records + * @return the bytes required to store records with the given sizes, {@link Integer.MAX_VALUE} + * when offset delta is too large. + */ + static int bytesNeededForRecordSizes( + long baseOffset, + long startingOffset, + int[] recordSizes + ) { + long expectedNextOffset = startingOffset; + int bytesNeeded = 0; + for (int size: recordSizes) { + if (expectedNextOffset - baseOffset >= Integer.MAX_VALUE) { + return Integer.MAX_VALUE; + } + + int recordSizeInBytes = DefaultRecord.sizeOfBodyInBytes( + (int) (expectedNextOffset - baseOffset), + 0, + -1, + size, + DefaultRecord.EMPTY_HEADERS + ); + + bytesNeeded += ByteUtils.sizeOfVarint(recordSizeInBytes) + recordSizeInBytes; + + expectedNextOffset += 1; + } + + return bytesNeeded; + } } diff --git a/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java b/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java index d817be47546c6..42ba9882a57df 100644 --- a/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java +++ b/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java @@ -21,6 +21,7 @@ import org.apache.kafka.common.record.MutableRecordBatch; import org.apache.kafka.common.utils.MockTime; import org.apache.kafka.common.utils.Utils; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.junit.jupiter.params.provider.ValueSource; @@ -70,7 +71,7 @@ void testBuildBatch(CompressionType compressionType) { records.forEach(record -> builder.appendRecord(record, null)); MemoryRecords builtRecordSet = builder.build(); assertFalse(builder.hasRoomFor(new int[] {1})); - assertThrows(IllegalArgumentException.class, () -> builder.appendRecord("a", null)); + assertThrows(IllegalStateException.class, () -> builder.appendRecord("a", null)); List builtBatches = Utils.toList(builtRecordSet.batchIterator()); assertEquals(1, builtBatches.size()); @@ -126,4 +127,23 @@ public void testHasRoomForUncompressed(int batchSize) { + sizeInBytes + " is larger than max batch size " + batchSize); } + @Test + public void testSizeCalculations() { + CompressionType compressionType = CompressionType.NONE; + int batchHeaderSizeInBytes = BatchBuilder.batchHeaderSizeInBytes(compressionType); + long baseOffset = 57; + int recordSize = 1; + int bytesNeededForRecordSizes = BatchBuilder.bytesNeededForRecordSizes( + baseOffset, + baseOffset, + new int[] {recordSize} + ); + + assertEquals(61, batchHeaderSizeInBytes); + assertEquals(8, bytesNeededForRecordSizes); + assertEquals( + batchHeaderSizeInBytes + bytesNeededForRecordSizes, + BatchBuilder.batchSizeForRecordSizes(compressionType, baseOffset, new int[] {recordSize}) + ); + } } From e8d538fa6a1914b0b61f88de5dc5c49492d87d32 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Armando=20Garc=C3=ADa=20Sancio?= Date: Mon, 15 Feb 2021 12:48:12 -0700 Subject: [PATCH 4/8] Merge append implementation in batch accumulator --- .../org/apache/kafka/raft/RaftClient.java | 1 - .../raft/internals/BatchAccumulator.java | 55 ++++++++----------- .../kafka/raft/internals/BatchBuilder.java | 4 +- 3 files changed, 25 insertions(+), 35 deletions(-) diff --git a/raft/src/main/java/org/apache/kafka/raft/RaftClient.java b/raft/src/main/java/org/apache/kafka/raft/RaftClient.java index 5c60a185b3742..cf0fa1ec1d8c6 100644 --- a/raft/src/main/java/org/apache/kafka/raft/RaftClient.java +++ b/raft/src/main/java/org/apache/kafka/raft/RaftClient.java @@ -124,7 +124,6 @@ default void handleResign() {} */ Long scheduleAtomicAppend(int epoch, List records); - /** * Attempt a graceful shutdown of the client. This allows the leader to proactively * resign and help a new leader to get elected rather than forcing the remaining diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java index 817aa35ca8c2e..82c6e893e1585 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java @@ -97,34 +97,7 @@ public BatchAccumulator( * been committed */ public Long append(int epoch, List records) { - if (epoch != this.epoch) { - return Long.MAX_VALUE; - } - - ObjectSerializationCache serializationCache = new ObjectSerializationCache(); - - appendLock.lock(); - try { - maybeCompleteDrain(); - - for (T record : records) { - BatchBuilder batch = maybeAllocateBatch( - new int[] {serde.recordSize(record, serializationCache)} - ); - if (batch == null) { - return null; - } - - batch.appendRecord(record, serializationCache); - nextOffset += 1; - } - - maybeResetLinger(); - - return nextOffset - 1; - } finally { - appendLock.unlock(); - } + return append(epoch, records, false); } /** @@ -142,6 +115,10 @@ public Long append(int epoch, List records) { * committed */ public Long appendAtomic(int epoch, List records) { + return append(epoch, records, true); + } + + private Long append(int epoch, List records, boolean isAtomic) { if (epoch != this.epoch) { return Long.MAX_VALUE; } @@ -156,12 +133,22 @@ public Long appendAtomic(int epoch, List records) { try { maybeCompleteDrain(); - BatchBuilder batch = maybeAllocateBatch(recordSizes); - if (batch == null) { - return null; + BatchBuilder batch = null; + if (isAtomic) { + batch = maybeAllocateBatch(recordSizes); } for (T record : records) { + if (!isAtomic) { + batch = maybeAllocateBatch( + new int[] {serde.recordSize(record, serializationCache)} + ); + } + + if (batch == null) { + return null; + } + batch.appendRecord(record, serializationCache); nextOffset += 1; } @@ -181,7 +168,11 @@ private void maybeResetLinger() { } private BatchBuilder maybeAllocateBatch(int[] recordSizes) { - int bytesNeeded = BatchBuilder.batchSizeForRecordSizes(compressionType, nextOffset, recordSizes); + int bytesNeeded = BatchBuilder.batchSizeForRecordSizes( + compressionType, + nextOffset, + recordSizes + ); if (bytesNeeded > maxBatchSize) { throw new RecordBatchTooLargeException( String.format( diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java index 86f65121cbb9b..0315638e646cd 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java @@ -303,7 +303,7 @@ public int writeRecord( * @param compressionType the type of compression * @param baseOffset the base offset of the batch * @param recordSizes the size of the records - * @return the bytes required, {@link Integer.MAX_VALUE} when offset delta is too large. + * @return the bytes required, {@link Integer#MAX_VALUE} when offset delta is too large. */ public static int batchSizeForRecordSizes( CompressionType compressionType, @@ -337,7 +337,7 @@ static int batchHeaderSizeInBytes(CompressionType compressionType) { * @param baseOffset the base offset for the batch * @param startingOffset the offset of the first record * @param recordSizes the sizes of all of the records - * @return the bytes required to store records with the given sizes, {@link Integer.MAX_VALUE} + * @return the bytes required to store records with the given sizes, {@link Integer#MAX_VALUE} * when offset delta is too large. */ static int bytesNeededForRecordSizes( From 8ac0405a452aed0756382b286c579ed66a8cc89c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Armando=20Garc=C3=ADa=20Sancio?= Date: Wed, 17 Feb 2021 19:19:02 -0700 Subject: [PATCH 5/8] Calculate sizes once --- .../raft/internals/BatchAccumulator.java | 62 ++++++------- .../kafka/raft/internals/BatchBuilder.java | 90 +++++++------------ .../raft/internals/BatchBuilderTest.java | 27 +----- 3 files changed, 67 insertions(+), 112 deletions(-) diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java index 82c6e893e1585..3659aaa6d252b 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java @@ -27,8 +27,10 @@ import java.io.Closeable; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.Collection; import java.util.Collections; import java.util.List; +import java.util.OptionalInt; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.ReentrantLock; @@ -124,10 +126,6 @@ private Long append(int epoch, List records, boolean isAtomic) { } ObjectSerializationCache serializationCache = new ObjectSerializationCache(); - int[] recordSizes = records - .stream() - .mapToInt(record -> serde.recordSize(record, serializationCache)) - .toArray(); appendLock.lock(); try { @@ -135,14 +133,12 @@ private Long append(int epoch, List records, boolean isAtomic) { BatchBuilder batch = null; if (isAtomic) { - batch = maybeAllocateBatch(recordSizes); + batch = maybeAllocateBatch(records, serializationCache); } for (T record : records) { if (!isAtomic) { - batch = maybeAllocateBatch( - new int[] {serde.recordSize(record, serializationCache)} - ); + batch = maybeAllocateBatch(Collections.singleton(record), serializationCache); } if (batch == null) { @@ -167,26 +163,30 @@ private void maybeResetLinger() { } } - private BatchBuilder maybeAllocateBatch(int[] recordSizes) { - int bytesNeeded = BatchBuilder.batchSizeForRecordSizes( - compressionType, - nextOffset, - recordSizes - ); - if (bytesNeeded > maxBatchSize) { - throw new RecordBatchTooLargeException( - String.format( - "The total record(s) size of %s exceeds the maximum allowed batch size of %s", - bytesNeeded, - maxBatchSize - ) - ); - } else if (currentBatch == null) { - startNewBatch(); - } else if (!currentBatch.hasRoomFor(recordSizes)) { - completeCurrentBatch(); + private BatchBuilder maybeAllocateBatch( + Collection records, + ObjectSerializationCache serializationCache + ) { + if (currentBatch == null) { startNewBatch(); } + + if (currentBatch != null) { + OptionalInt bytesNeeded = currentBatch.roomFor(records, serializationCache); + if (bytesNeeded.isPresent() && bytesNeeded.getAsInt() > maxBatchSize) { + throw new RecordBatchTooLargeException( + String.format( + "The total record(s) size of %s exceeds the maximum allowed batch size of %s", + bytesNeeded.getAsInt(), + maxBatchSize + ) + ); + } else if (bytesNeeded.isPresent()) { + completeCurrentBatch(); + startNewBatch(); + } + } + return currentBatch; } @@ -342,20 +342,22 @@ public static class CompletedBatch { public final List records; public final MemoryRecords data; private final MemoryPool pool; - private final ByteBuffer buffer; + // Buffer that was allocated by the MemoryPool (pool). This may not be the buffer used in + // the MemoryRecords (data) object. + private final ByteBuffer pooledBuffer; private CompletedBatch( long baseOffset, List records, MemoryRecords data, MemoryPool pool, - ByteBuffer buffer + ByteBuffer pooledBuffer ) { this.baseOffset = baseOffset; this.records = records; this.data = data; this.pool = pool; - this.buffer = buffer; + this.pooledBuffer = pooledBuffer; } public int sizeInBytes() { @@ -363,7 +365,7 @@ public int sizeInBytes() { } public void release() { - pool.release(buffer); + pool.release(pooledBuffer); } } diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java index 0315638e646cd..6c8dd30797169 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java @@ -33,12 +33,14 @@ import java.io.DataOutputStream; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.Collection; import java.util.List; +import java.util.OptionalInt; /** * Collect a set of records into a single batch. New records are added * through {@link #appendRecord(Object, ObjectSerializationCache)}, but the caller must first - * check whether there is room using {@link #hasRoomFor(int)}. Once the + * check whether there is room using {@link #roomFor(Collection, ObjectSerializationCache)}. Once the * batch is ready, then {@link #build()} should be used to get the resulting * {@link MemoryRecords} instance. * @@ -85,7 +87,8 @@ public BatchBuilder( this.maxBytes = maxBytes; this.records = new ArrayList<>(); - int batchHeaderSizeInBytes = batchHeaderSizeInBytes(compressionType); + // field compressionType must be set before calculating the batch header size + int batchHeaderSizeInBytes = batchHeaderSizeInBytes(); batchOutput.position(initialPosition + batchHeaderSizeInBytes); this.recordOutput = new DataOutputStreamWritable(new DataOutputStream( @@ -94,7 +97,7 @@ public BatchBuilder( /** * Append a record to this patch. The caller must first verify there is room for the batch - * using {@link #hasRoomFor(int[])}. + * using {@link #roomFor(Collection, ObjectSerializationCache)}. * * @param record the record to append * @param serializationCache serialization cache for use in {@link RecordSerde#write(Object, ObjectSerializationCache, Writable)} @@ -124,27 +127,37 @@ public long appendRecord(T record, ObjectSerializationCache serializationCache) /** * Check whether the batch has enough room for all the record values. * - * @param recordSizes the sizes of the records to be appended - * @return true if there is room for the records to be appended, false otherwise + * Returns an empty {@link OptionalInt} if the batch builder has room for this list of records. + * Otherwise it returns the expected number of bytes needed for a batch to contain these records. + * + * @param records the records to use when checking for room + * @param serializationCache serialization cache for computing sizes + * @return empty {@link OptionalInt} if there is room for the records to be appended, otherwise + * returns the number of bytes needed */ - public boolean hasRoomFor(int[] recordSizes) { + public OptionalInt roomFor(Collection records, ObjectSerializationCache serializationCache) { + int bytesNeeded = bytesNeededForRecords( + records, + serializationCache + ); + if (!isOpenForAppends) { - return false; + return OptionalInt.of(batchHeaderSizeInBytes() + bytesNeeded); } - int bytesNeeded = bytesNeededForRecordSizes(baseOffset, nextOffset, recordSizes); - int approxUnusedSizeInBytes = maxBytes - approximateSizeInBytes(); if (approxUnusedSizeInBytes >= bytesNeeded) { - return true; + return OptionalInt.empty(); } else if (unflushedBytes > 0) { recordOutput.flush(); unflushedBytes = 0; int unusedSizeInBytes = maxBytes - flushedSizeInBytes(); - return unusedSizeInBytes >= bytesNeeded; - } else { - return false; + if (unusedSizeInBytes >= bytesNeeded) { + return OptionalInt.empty(); + } } + + return OptionalInt.of(batchHeaderSizeInBytes() + bytesNeeded); } private int flushedSizeInBytes() { @@ -297,57 +310,20 @@ public int writeRecord( return ByteUtils.sizeOfVarint(sizeInBytes) + sizeInBytes; } - /** - * Computes the bytes required to store records with the given sizes in a new batch. - * - * @param compressionType the type of compression - * @param baseOffset the base offset of the batch - * @param recordSizes the size of the records - * @return the bytes required, {@link Integer#MAX_VALUE} when offset delta is too large. - */ - public static int batchSizeForRecordSizes( - CompressionType compressionType, - long baseOffset, - int[] recordSizes - ) { - int bytesNeeded = bytesNeededForRecordSizes(baseOffset, 0, recordSizes); - if (bytesNeeded != Integer.MAX_VALUE) { - return batchHeaderSizeInBytes(compressionType) + bytesNeeded; - } else { - return Integer.MAX_VALUE; - } - } - - /** - * Computes the bytes overhead of an empty batch. - * - * @param compressionType the type of compression - * @return the minimal bytes size of a batch - */ - static int batchHeaderSizeInBytes(CompressionType compressionType) { + private int batchHeaderSizeInBytes() { return AbstractRecords.recordBatchHeaderSizeInBytes( RecordBatch.MAGIC_VALUE_V2, compressionType ); } - /** - * Compute the number of bytes required to store the give record sizes. - * - * @param baseOffset the base offset for the batch - * @param startingOffset the offset of the first record - * @param recordSizes the sizes of all of the records - * @return the bytes required to store records with the given sizes, {@link Integer#MAX_VALUE} - * when offset delta is too large. - */ - static int bytesNeededForRecordSizes( - long baseOffset, - long startingOffset, - int[] recordSizes + private int bytesNeededForRecords( + Collection records, + ObjectSerializationCache serializationCache ) { - long expectedNextOffset = startingOffset; + long expectedNextOffset = nextOffset; int bytesNeeded = 0; - for (int size: recordSizes) { + for (T record : records) { if (expectedNextOffset - baseOffset >= Integer.MAX_VALUE) { return Integer.MAX_VALUE; } @@ -356,7 +332,7 @@ static int bytesNeededForRecordSizes( (int) (expectedNextOffset - baseOffset), 0, -1, - size, + serde.recordSize(record, serializationCache), DefaultRecord.EMPTY_HEADERS ); diff --git a/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java b/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java index 42ba9882a57df..409f5b9f5e1ba 100644 --- a/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java +++ b/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java @@ -21,7 +21,6 @@ import org.apache.kafka.common.record.MutableRecordBatch; import org.apache.kafka.common.utils.MockTime; import org.apache.kafka.common.utils.Utils; -import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.junit.jupiter.params.provider.ValueSource; @@ -32,7 +31,6 @@ import java.util.stream.Collectors; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -70,7 +68,7 @@ void testBuildBatch(CompressionType compressionType) { records.forEach(record -> builder.appendRecord(record, null)); MemoryRecords builtRecordSet = builder.build(); - assertFalse(builder.hasRoomFor(new int[] {1})); + assertTrue(builder.roomFor(Arrays.asList("a"), null).isPresent()); assertThrows(IllegalStateException.class, () -> builder.appendRecord("a", null)); List builtBatches = Utils.toList(builtRecordSet.batchIterator()); @@ -113,9 +111,8 @@ public void testHasRoomForUncompressed(int batchSize) { ); String record = "i am a record"; - int recordSize = serde.recordSize(record); - while (builder.hasRoomFor(new int[] {recordSize})) { + while (!builder.roomFor(Arrays.asList(record), null).isPresent()) { builder.appendRecord(record, null); } @@ -126,24 +123,4 @@ public void testHasRoomForUncompressed(int batchSize) { assertTrue(sizeInBytes <= batchSize, "Built batch size " + sizeInBytes + " is larger than max batch size " + batchSize); } - - @Test - public void testSizeCalculations() { - CompressionType compressionType = CompressionType.NONE; - int batchHeaderSizeInBytes = BatchBuilder.batchHeaderSizeInBytes(compressionType); - long baseOffset = 57; - int recordSize = 1; - int bytesNeededForRecordSizes = BatchBuilder.bytesNeededForRecordSizes( - baseOffset, - baseOffset, - new int[] {recordSize} - ); - - assertEquals(61, batchHeaderSizeInBytes); - assertEquals(8, bytesNeededForRecordSizes); - assertEquals( - batchHeaderSizeInBytes + bytesNeededForRecordSizes, - BatchBuilder.batchSizeForRecordSizes(compressionType, baseOffset, new int[] {recordSize}) - ); - } } From a9e2bd9248b7de406f7f02fbfc4840a3d8f1a840 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Armando=20Garc=C3=ADa=20Sancio?= Date: Wed, 17 Feb 2021 19:27:48 -0700 Subject: [PATCH 6/8] Fix documentation --- raft/src/main/java/org/apache/kafka/raft/RaftClient.java | 5 +++-- .../org/apache/kafka/raft/internals/BatchAccumulator.java | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/raft/src/main/java/org/apache/kafka/raft/RaftClient.java b/raft/src/main/java/org/apache/kafka/raft/RaftClient.java index cf0fa1ec1d8c6..6231e09b69a4e 100644 --- a/raft/src/main/java/org/apache/kafka/raft/RaftClient.java +++ b/raft/src/main/java/org/apache/kafka/raft/RaftClient.java @@ -83,8 +83,9 @@ default void handleResign() {} * Append a list of records to the log. The write will be scheduled for some time * in the future. There is no guarantee that appended records will be written to * the log and eventually committed. While the order of the records is preserve, they can - * be appended to the log using one or more batches. This means that each record could - * be committed independently. + * be appended to the log using one or more batches. Each record may be committed independently. + * If a record is committed, then all records scheduled for append during this epoch + * and prior to this record are also committed. * * If the provided current leader epoch does not match the current epoch, which * is possible when the state machine has yet to observe the epoch change, then diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java index 3659aaa6d252b..745d85d0238f7 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java @@ -87,7 +87,7 @@ public BatchAccumulator( * The order of the elements in the records argument will match the order in the batches. * This method will use as many batches as necessary to serialize all of the records. Since * this method can split the records into multiple batches it is possible that some of the - * recors will get committed while other will not when the leader fails. + * records will get committed while other will not when the leader fails. * * @param epoch the expected leader epoch. If this does not match, then {@link Long#MAX_VALUE} * will be returned as an offset which cannot become committed From 2e884c4e42fde53520c4072ce63aad9015d85f2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Armando=20Garc=C3=ADa=20Sancio?= Date: Thu, 18 Feb 2021 13:16:05 -0700 Subject: [PATCH 7/8] Throw an exception if the number of records exceeds Integer.MAX_VALUE --- .../kafka/raft/internals/BatchAccumulator.java | 2 +- .../apache/kafka/raft/internals/BatchBuilder.java | 15 +++++++++++---- .../kafka/raft/internals/BatchBuilderTest.java | 4 ++-- 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java index 745d85d0238f7..7efa8d289f59c 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java @@ -172,7 +172,7 @@ private BatchBuilder maybeAllocateBatch( } if (currentBatch != null) { - OptionalInt bytesNeeded = currentBatch.roomFor(records, serializationCache); + OptionalInt bytesNeeded = currentBatch.bytesNeeded(records, serializationCache); if (bytesNeeded.isPresent() && bytesNeeded.getAsInt() > maxBatchSize) { throw new RecordBatchTooLargeException( String.format( diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java index 6c8dd30797169..4cc462e39ddb2 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java @@ -40,7 +40,7 @@ /** * Collect a set of records into a single batch. New records are added * through {@link #appendRecord(Object, ObjectSerializationCache)}, but the caller must first - * check whether there is room using {@link #roomFor(Collection, ObjectSerializationCache)}. Once the + * check whether there is room using {@link #bytesNeeded(Collection, ObjectSerializationCache)}. Once the * batch is ready, then {@link #build()} should be used to get the resulting * {@link MemoryRecords} instance. * @@ -97,7 +97,7 @@ public BatchBuilder( /** * Append a record to this patch. The caller must first verify there is room for the batch - * using {@link #roomFor(Collection, ObjectSerializationCache)}. + * using {@link #bytesNeeded(Collection, ObjectSerializationCache)}. * * @param record the record to append * @param serializationCache serialization cache for use in {@link RecordSerde#write(Object, ObjectSerializationCache, Writable)} @@ -135,7 +135,7 @@ public long appendRecord(T record, ObjectSerializationCache serializationCache) * @return empty {@link OptionalInt} if there is room for the records to be appended, otherwise * returns the number of bytes needed */ - public OptionalInt roomFor(Collection records, ObjectSerializationCache serializationCache) { + public OptionalInt bytesNeeded(Collection records, ObjectSerializationCache serializationCache) { int bytesNeeded = bytesNeededForRecords( records, serializationCache @@ -325,7 +325,14 @@ private int bytesNeededForRecords( int bytesNeeded = 0; for (T record : records) { if (expectedNextOffset - baseOffset >= Integer.MAX_VALUE) { - return Integer.MAX_VALUE; + throw new IllegalArgumentException( + String.format( + "Adding %s records to a batch with base offset of %s and next offset of %s", + records.size(), + baseOffset, + expectedNextOffset + ) + ); } int recordSizeInBytes = DefaultRecord.sizeOfBodyInBytes( diff --git a/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java b/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java index 409f5b9f5e1ba..e4611f1b8ca25 100644 --- a/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java +++ b/raft/src/test/java/org/apache/kafka/raft/internals/BatchBuilderTest.java @@ -68,7 +68,7 @@ void testBuildBatch(CompressionType compressionType) { records.forEach(record -> builder.appendRecord(record, null)); MemoryRecords builtRecordSet = builder.build(); - assertTrue(builder.roomFor(Arrays.asList("a"), null).isPresent()); + assertTrue(builder.bytesNeeded(Arrays.asList("a"), null).isPresent()); assertThrows(IllegalStateException.class, () -> builder.appendRecord("a", null)); List builtBatches = Utils.toList(builtRecordSet.batchIterator()); @@ -112,7 +112,7 @@ public void testHasRoomForUncompressed(int batchSize) { String record = "i am a record"; - while (!builder.roomFor(Arrays.asList(record), null).isPresent()) { + while (!builder.bytesNeeded(Arrays.asList(record), null).isPresent()) { builder.appendRecord(record, null); } From e61d63c06cb15713cd5fdc668d640244be86d189 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Armando=20Garc=C3=ADa=20Sancio?= Date: Thu, 18 Feb 2021 19:40:41 -0700 Subject: [PATCH 8/8] Check for overflow during addition --- .../org/apache/kafka/raft/internals/BatchAccumulator.java | 8 ++++---- .../org/apache/kafka/raft/internals/BatchBuilder.java | 3 ++- .../apache/kafka/raft/internals/BatchAccumulatorTest.java | 3 ++- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java index 7efa8d289f59c..07d1015c9da92 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java @@ -344,20 +344,20 @@ public static class CompletedBatch { private final MemoryPool pool; // Buffer that was allocated by the MemoryPool (pool). This may not be the buffer used in // the MemoryRecords (data) object. - private final ByteBuffer pooledBuffer; + private final ByteBuffer initialBuffer; private CompletedBatch( long baseOffset, List records, MemoryRecords data, MemoryPool pool, - ByteBuffer pooledBuffer + ByteBuffer initialBuffer ) { this.baseOffset = baseOffset; this.records = records; this.data = data; this.pool = pool; - this.pooledBuffer = pooledBuffer; + this.initialBuffer = initialBuffer; } public int sizeInBytes() { @@ -365,7 +365,7 @@ public int sizeInBytes() { } public void release() { - pool.release(pooledBuffer); + pool.release(initialBuffer); } } diff --git a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java index 4cc462e39ddb2..c953b6a6371c7 100644 --- a/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java +++ b/raft/src/main/java/org/apache/kafka/raft/internals/BatchBuilder.java @@ -343,7 +343,8 @@ private int bytesNeededForRecords( DefaultRecord.EMPTY_HEADERS ); - bytesNeeded += ByteUtils.sizeOfVarint(recordSizeInBytes) + recordSizeInBytes; + bytesNeeded = Math.addExact(bytesNeeded, ByteUtils.sizeOfVarint(recordSizeInBytes)); + bytesNeeded = Math.addExact(bytesNeeded, recordSizeInBytes); expectedNextOffset += 1; } diff --git a/raft/src/test/java/org/apache/kafka/raft/internals/BatchAccumulatorTest.java b/raft/src/test/java/org/apache/kafka/raft/internals/BatchAccumulatorTest.java index 0855bd80cedd3..b32168ec3101b 100644 --- a/raft/src/test/java/org/apache/kafka/raft/internals/BatchAccumulatorTest.java +++ b/raft/src/test/java/org/apache/kafka/raft/internals/BatchAccumulatorTest.java @@ -236,7 +236,8 @@ public void testMultipleBatchAccumulation() { }); } - @Test void testRecordsAreSplit() { + @Test + public void testRecordsAreSplit() { int leaderEpoch = 17; long baseOffset = 157; int lingerMs = 50;