/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.common.record;

import org.apache.kafka.common.InvalidRecordException;
import org.apache.kafka.common.compress.Compression;
import org.apache.kafka.common.compress.ZstdCompression;
import org.apache.kafka.common.errors.CorruptRecordException;
import org.apache.kafka.common.header.Header;
import org.apache.kafka.common.header.internals.RecordHeader;
import org.apache.kafka.common.utils.BufferSupplier;
import org.apache.kafka.common.utils.ChunkedBytesStream;
import org.apache.kafka.common.utils.CloseableIterator;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.test.TestUtils;
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 java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.List;
import java.util.Random;
import java.util.stream.Stream;

import static org.apache.kafka.common.record.DefaultRecordBatch.RECORDS_COUNT_OFFSET;
import static org.apache.kafka.common.record.DefaultRecordBatch.RECORDS_OFFSET;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertIterableEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;

public class DefaultRecordBatchTest {
    // We avoid SecureRandom.getInstanceStrong() here because it reads from /dev/random and blocks on Linux. Since these
    // tests don't require cryptographically strong random data, we avoid a CSPRNG (SecureRandom) altogether.
    private static final Random RANDOM = new Random(20231025);

    @Test
    public void testWriteEmptyHeader() {
        long producerId = 23423L;
        short producerEpoch = 145;
        int baseSequence = 983;
        long baseOffset = 15L;
        long lastOffset = 37;
        int partitionLeaderEpoch = 15;
        long timestamp = System.currentTimeMillis();

        for (TimestampType timestampType : Arrays.asList(TimestampType.CREATE_TIME, TimestampType.LOG_APPEND_TIME)) {
            for (boolean isTransactional : Arrays.asList(true, false)) {
                for (boolean isControlBatch : Arrays.asList(true, false)) {
                    ByteBuffer buffer = ByteBuffer.allocate(2048);
                    DefaultRecordBatch.writeEmptyHeader(buffer, RecordBatch.CURRENT_MAGIC_VALUE, producerId,
                            producerEpoch, baseSequence, baseOffset, lastOffset, partitionLeaderEpoch, timestampType,
                            timestamp, isTransactional, isControlBatch);
                    buffer.flip();
                    DefaultRecordBatch batch = new DefaultRecordBatch(buffer);
                    assertEquals(producerId, batch.producerId());
                    assertEquals(producerEpoch, batch.producerEpoch());
                    assertEquals(baseSequence, batch.baseSequence());
                    assertEquals(baseSequence + ((int) (lastOffset - baseOffset)), batch.lastSequence());
                    assertEquals(baseOffset, batch.baseOffset());
                    assertEquals(lastOffset, batch.lastOffset());
                    assertEquals(partitionLeaderEpoch, batch.partitionLeaderEpoch());
                    assertEquals(isTransactional, batch.isTransactional());
                    assertEquals(timestampType, batch.timestampType());
                    assertEquals(timestamp, batch.maxTimestamp());
                    assertEquals(RecordBatch.NO_TIMESTAMP, batch.baseTimestamp());
                    assertEquals(isControlBatch, batch.isControlBatch());
                }
            }
        }
    }

    @Test
    public void buildDefaultRecordBatch() {
        ByteBuffer buffer = ByteBuffer.allocate(2048);

        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, RecordBatch.MAGIC_VALUE_V2, Compression.NONE,
                TimestampType.CREATE_TIME, 1234567L);
        builder.appendWithOffset(1234567, 1L, "a".getBytes(), "v".getBytes());
        builder.appendWithOffset(1234568, 2L, "b".getBytes(), "v".getBytes());

        MemoryRecords records = builder.build();
        for (MutableRecordBatch batch : records.batches()) {
            assertTrue(batch.isValid());
            assertEquals(1234567, batch.baseOffset());
            assertEquals(1234568, batch.lastOffset());
            assertEquals(2L, batch.maxTimestamp());
            assertEquals(RecordBatch.NO_PRODUCER_ID, batch.producerId());
            assertEquals(RecordBatch.NO_PRODUCER_EPOCH, batch.producerEpoch());
            assertEquals(RecordBatch.NO_SEQUENCE, batch.baseSequence());
            assertEquals(RecordBatch.NO_SEQUENCE, batch.lastSequence());

            for (Record record : batch) record.ensureValid();
        }
    }

    @Test
    public void buildDefaultRecordBatchWithProducerId() {
        long pid = 23423L;
        short epoch = 145;
        int baseSequence = 983;

        ByteBuffer buffer = ByteBuffer.allocate(2048);

        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, RecordBatch.MAGIC_VALUE_V2, Compression.NONE,
                TimestampType.CREATE_TIME, 1234567L, RecordBatch.NO_TIMESTAMP, pid, epoch, baseSequence);
        builder.appendWithOffset(1234567, 1L, "a".getBytes(), "v".getBytes());
        builder.appendWithOffset(1234568, 2L, "b".getBytes(), "v".getBytes());

        MemoryRecords records = builder.build();
        for (MutableRecordBatch batch : records.batches()) {
            assertTrue(batch.isValid());
            assertEquals(1234567, batch.baseOffset());
            assertEquals(1234568, batch.lastOffset());
            assertEquals(2L, batch.maxTimestamp());
            assertEquals(pid, batch.producerId());
            assertEquals(epoch, batch.producerEpoch());
            assertEquals(baseSequence, batch.baseSequence());
            assertEquals(baseSequence + 1, batch.lastSequence());

            for (Record record : batch) record.ensureValid();
        }
    }

    @Test
    public void buildDefaultRecordBatchWithSequenceWrapAround() {
        long pid = 23423L;
        short epoch = 145;
        int baseSequence = Integer.MAX_VALUE - 1;
        ByteBuffer buffer = ByteBuffer.allocate(2048);

        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, RecordBatch.MAGIC_VALUE_V2, Compression.NONE,
                TimestampType.CREATE_TIME, 1234567L, RecordBatch.NO_TIMESTAMP, pid, epoch, baseSequence);
        builder.appendWithOffset(1234567, 1L, "a".getBytes(), "v".getBytes());
        builder.appendWithOffset(1234568, 2L, "b".getBytes(), "v".getBytes());
        builder.appendWithOffset(1234569, 3L, "c".getBytes(), "v".getBytes());

        MemoryRecords records = builder.build();
        List<MutableRecordBatch> batches = TestUtils.toList(records.batches());
        assertEquals(1, batches.size());
        RecordBatch batch = batches.get(0);

        assertEquals(pid, batch.producerId());
        assertEquals(epoch, batch.producerEpoch());
        assertEquals(baseSequence, batch.baseSequence());
        assertEquals(0, batch.lastSequence());
        List<Record> allRecords = TestUtils.toList(batch);
        assertEquals(3, allRecords.size());
        assertEquals(Integer.MAX_VALUE - 1, allRecords.get(0).sequence());
        assertEquals(Integer.MAX_VALUE, allRecords.get(1).sequence());
        assertEquals(0, allRecords.get(2).sequence());
    }

    @Test
    public void testSizeInBytes() {
        Header[] headers = new Header[] {
            new RecordHeader("foo", "value".getBytes()),
            new RecordHeader("bar", null)
        };

        long timestamp = System.currentTimeMillis();
        SimpleRecord[] records = new SimpleRecord[] {
            new SimpleRecord(timestamp, "key".getBytes(), "value".getBytes()),
            new SimpleRecord(timestamp + 30000, null, "value".getBytes()),
            new SimpleRecord(timestamp + 60000, "key".getBytes(), null),
            new SimpleRecord(timestamp + 60000, "key".getBytes(), "value".getBytes(), headers)
        };
        int actualSize = MemoryRecords.withRecords(Compression.NONE, records).sizeInBytes();
        assertEquals(actualSize, DefaultRecordBatch.sizeInBytes(Arrays.asList(records)));
    }

    @Test
    public void testInvalidRecordSize() {
        MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
                Compression.NONE, TimestampType.CREATE_TIME,
                new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
                new SimpleRecord(2L, "b".getBytes(), "2".getBytes()),
                new SimpleRecord(3L, "c".getBytes(), "3".getBytes()));

        ByteBuffer buffer = records.buffer();
        buffer.putInt(DefaultRecordBatch.LENGTH_OFFSET, 10);

        DefaultRecordBatch batch = new DefaultRecordBatch(buffer);
        assertFalse(batch.isValid());
        assertThrows(CorruptRecordException.class, batch::ensureValid);
    }

    @Test
    public void testInvalidRecordCountTooManyNonCompressedV2() {
        long now = System.currentTimeMillis();
        DefaultRecordBatch batch = recordsWithInvalidRecordCount(RecordBatch.MAGIC_VALUE_V2, now, CompressionType.NONE, 5);
        // force iteration through the batch to execute validation
        // batch validation is a part of normal workflow for LogValidator.validateMessagesAndAssignOffsets
        assertThrows(InvalidRecordException.class, () -> batch.forEach(Record::ensureValid));
    }

    @Test
    public void testInvalidRecordCountTooLittleNonCompressedV2() {
        long now = System.currentTimeMillis();
        DefaultRecordBatch batch = recordsWithInvalidRecordCount(RecordBatch.MAGIC_VALUE_V2, now, CompressionType.NONE, 2);
        // force iteration through the batch to execute validation
        // batch validation is a part of normal workflow for LogValidator.validateMessagesAndAssignOffsets
        assertThrows(InvalidRecordException.class, () -> batch.forEach(Record::ensureValid));
    }

    @Test
    public void testInvalidRecordCountTooManyCompressedV2() {
        long now = System.currentTimeMillis();
        DefaultRecordBatch batch = recordsWithInvalidRecordCount(RecordBatch.MAGIC_VALUE_V2, now, CompressionType.GZIP, 5);
        // force iteration through the batch to execute validation
        // batch validation is a part of normal workflow for LogValidator.validateMessagesAndAssignOffsets
        assertThrows(InvalidRecordException.class, () -> batch.forEach(Record::ensureValid));
    }

    @Test
    public void testInvalidRecordCountTooLittleCompressedV2() {
        long now = System.currentTimeMillis();
        DefaultRecordBatch batch = recordsWithInvalidRecordCount(RecordBatch.MAGIC_VALUE_V2, now, CompressionType.GZIP, 2);
        // force iteration through the batch to execute validation
        // batch validation is a part of normal workflow for LogValidator.validateMessagesAndAssignOffsets
        assertThrows(InvalidRecordException.class, () -> batch.forEach(Record::ensureValid));
    }

    @Test
    public void testInvalidCrc() {
        MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
                Compression.NONE, TimestampType.CREATE_TIME,
                new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
                new SimpleRecord(2L, "b".getBytes(), "2".getBytes()),
                new SimpleRecord(3L, "c".getBytes(), "3".getBytes()));

        ByteBuffer buffer = records.buffer();
        buffer.putInt(DefaultRecordBatch.LAST_OFFSET_DELTA_OFFSET, 23);

        DefaultRecordBatch batch = new DefaultRecordBatch(buffer);
        assertFalse(batch.isValid());
        assertThrows(CorruptRecordException.class, batch::ensureValid);
    }

    @Test
    public void testSetLastOffset() {
        SimpleRecord[] simpleRecords = new SimpleRecord[] {
            new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
            new SimpleRecord(2L, "b".getBytes(), "2".getBytes()),
            new SimpleRecord(3L, "c".getBytes(), "3".getBytes())
        };
        MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
                Compression.NONE, TimestampType.CREATE_TIME, simpleRecords);

        long lastOffset = 500L;
        long firstOffset = lastOffset - simpleRecords.length + 1;

        DefaultRecordBatch batch = new DefaultRecordBatch(records.buffer());
        batch.setLastOffset(lastOffset);
        assertEquals(lastOffset, batch.lastOffset());
        assertEquals(firstOffset, batch.baseOffset());
        assertTrue(batch.isValid());

        List<MutableRecordBatch> recordBatches = Utils.toList(records.batches().iterator());
        assertEquals(1, recordBatches.size());
        assertEquals(lastOffset, recordBatches.get(0).lastOffset());

        long offset = firstOffset;
        for (Record record : records.records())
            assertEquals(offset++, record.offset());
    }

    @Test
    public void testSetPartitionLeaderEpoch() {
        MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
                Compression.NONE, TimestampType.CREATE_TIME,
                new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
                new SimpleRecord(2L, "b".getBytes(), "2".getBytes()),
                new SimpleRecord(3L, "c".getBytes(), "3".getBytes()));

        int leaderEpoch = 500;

        DefaultRecordBatch batch = new DefaultRecordBatch(records.buffer());
        batch.setPartitionLeaderEpoch(leaderEpoch);
        assertEquals(leaderEpoch, batch.partitionLeaderEpoch());
        assertTrue(batch.isValid());

        List<MutableRecordBatch> recordBatches = Utils.toList(records.batches().iterator());
        assertEquals(1, recordBatches.size());
        assertEquals(leaderEpoch, recordBatches.get(0).partitionLeaderEpoch());
    }

    @Test
    public void testSetLogAppendTime() {
        MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
                Compression.NONE, TimestampType.CREATE_TIME,
                new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
                new SimpleRecord(2L, "b".getBytes(), "2".getBytes()),
                new SimpleRecord(3L, "c".getBytes(), "3".getBytes()));

        long logAppendTime = 15L;

        DefaultRecordBatch batch = new DefaultRecordBatch(records.buffer());
        batch.setMaxTimestamp(TimestampType.LOG_APPEND_TIME, logAppendTime);
        assertEquals(TimestampType.LOG_APPEND_TIME, batch.timestampType());
        assertEquals(logAppendTime, batch.maxTimestamp());
        assertTrue(batch.isValid());

        List<MutableRecordBatch> recordBatches = Utils.toList(records.batches().iterator());
        assertEquals(1, recordBatches.size());
        assertEquals(logAppendTime, recordBatches.get(0).maxTimestamp());
        assertEquals(TimestampType.LOG_APPEND_TIME, recordBatches.get(0).timestampType());

        for (Record record : records.records())
            assertEquals(logAppendTime, record.timestamp());
    }

    @Test
    public void testSetNoTimestampTypeNotAllowed() {
        MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
                Compression.NONE, TimestampType.CREATE_TIME,
                new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
                new SimpleRecord(2L, "b".getBytes(), "2".getBytes()),
                new SimpleRecord(3L, "c".getBytes(), "3".getBytes()));
        DefaultRecordBatch batch = new DefaultRecordBatch(records.buffer());
        assertThrows(IllegalArgumentException.class, () -> batch.setMaxTimestamp(TimestampType.NO_TIMESTAMP_TYPE, RecordBatch.NO_TIMESTAMP));
    }

    @Test
    public void testReadAndWriteControlBatch() {
        long producerId = 1L;
        short producerEpoch = 0;
        int coordinatorEpoch = 15;

        ByteBuffer buffer = ByteBuffer.allocate(128);
        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.CURRENT_MAGIC_VALUE,
                Compression.NONE, TimestampType.CREATE_TIME, 0L, RecordBatch.NO_TIMESTAMP, producerId,
                producerEpoch, RecordBatch.NO_SEQUENCE, true, true, RecordBatch.NO_PARTITION_LEADER_EPOCH,
                buffer.remaining());

        EndTransactionMarker marker = new EndTransactionMarker(ControlRecordType.COMMIT, coordinatorEpoch);
        builder.appendEndTxnMarker(System.currentTimeMillis(), marker);
        MemoryRecords records = builder.build();

        List<MutableRecordBatch> batches = TestUtils.toList(records.batches());
        assertEquals(1, batches.size());

        MutableRecordBatch batch = batches.get(0);
        assertTrue(batch.isControlBatch());

        List<Record> logRecords = TestUtils.toList(records.records());
        assertEquals(1, logRecords.size());

        Record commitRecord = logRecords.get(0);
        assertEquals(marker, EndTransactionMarker.deserialize(commitRecord));
    }

    @ParameterizedTest
    @EnumSource(value = CompressionType.class)
    public void testStreamingIteratorConsistency(CompressionType compressionType) {
        Compression compression = Compression.of(compressionType).build();
        MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
                compression, TimestampType.CREATE_TIME,
                new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
                new SimpleRecord(2L, "b".getBytes(), "2".getBytes()),
                new SimpleRecord(3L, "c".getBytes(), "3".getBytes()));
        DefaultRecordBatch batch = new DefaultRecordBatch(records.buffer());
        try (CloseableIterator<Record> streamingIterator = batch.streamingIterator(BufferSupplier.create())) {
            TestUtils.checkEquals(streamingIterator, batch.iterator());
        }
    }

    @ParameterizedTest
    @EnumSource(value = CompressionType.class)
    public void testSkipKeyValueIteratorCorrectness(CompressionType compressionType) {
        Compression compression = Compression.of(compressionType).build();
        Header[] headers = {new RecordHeader("k1", "v1".getBytes()), new RecordHeader("k2", null)};
        byte[] largeRecordValue = new byte[200 * 1024]; // 200KB
        RANDOM.nextBytes(largeRecordValue);

        MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
            compression, TimestampType.CREATE_TIME,
            // one sample with small value size
            new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
            // one sample with null value
            new SimpleRecord(2L, "b".getBytes(), null),
            // one sample with null key
            new SimpleRecord(3L, null, "3".getBytes()),
            // one sample with null key and null value
            new SimpleRecord(4L, null, (byte[]) null),
            // one sample with large value size
            new SimpleRecord(1000L, "abc".getBytes(), largeRecordValue),
            // one sample with headers, one of the header has null value
            new SimpleRecord(9999L, "abc".getBytes(), "0".getBytes(), headers)
            );

        DefaultRecordBatch batch = new DefaultRecordBatch(records.buffer());

        try (BufferSupplier bufferSupplier = BufferSupplier.create();
             CloseableIterator<Record> skipKeyValueIterator = batch.skipKeyValueIterator(bufferSupplier)) {

            if (CompressionType.NONE == compressionType) {
                // assert that for uncompressed data stream record iterator is not used
                assertInstanceOf(DefaultRecordBatch.RecordIterator.class, skipKeyValueIterator);
                // superficial validation for correctness. Deep validation is already performed in other tests
                assertEquals(Utils.toList(records.records()).size(), Utils.toList(skipKeyValueIterator).size());
            } else {
                // assert that a streaming iterator is used for compressed records
                assertInstanceOf(DefaultRecordBatch.StreamRecordIterator.class, skipKeyValueIterator);
                // assert correctness for compressed records
                assertIterableEquals(Arrays.asList(
                        new PartialDefaultRecord(9, (byte) 0, 0L, 1L, -1, 1, 1),
                        new PartialDefaultRecord(8, (byte) 0, 1L, 2L, -1, 1, -1),
                        new PartialDefaultRecord(8, (byte) 0, 2L, 3L, -1, -1, 1),
                        new PartialDefaultRecord(7, (byte) 0, 3L, 4L, -1, -1, -1),
                        new PartialDefaultRecord(15 + largeRecordValue.length, (byte) 0, 4L, 1000L, -1, 3, largeRecordValue.length),
                        new PartialDefaultRecord(23, (byte) 0, 5L, 9999L, -1, 3, 1)
                    ), Utils.toList(skipKeyValueIterator));
            }
        }
    }

    @ParameterizedTest
    @MethodSource
    public void testBufferReuseInSkipKeyValueIterator(CompressionType compressionType, int expectedNumBufferAllocations, byte[] recordValue) {
        Compression compression = Compression.of(compressionType).build();
        MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
            compression, TimestampType.CREATE_TIME,
            new SimpleRecord(1000L, "a".getBytes(), "0".getBytes()),
            new SimpleRecord(9999L, "b".getBytes(), recordValue)
        );

        DefaultRecordBatch batch = new DefaultRecordBatch(records.buffer());

        try (BufferSupplier bufferSupplier = spy(BufferSupplier.create());
             CloseableIterator<Record> streamingIterator = batch.skipKeyValueIterator(bufferSupplier)) {

            // Consume through the iterator
            Utils.toList(streamingIterator);

            // Close the iterator to release any buffers
            streamingIterator.close();

            // assert number of buffer allocations
            verify(bufferSupplier, times(expectedNumBufferAllocations)).get(anyInt());
            verify(bufferSupplier, times(expectedNumBufferAllocations)).release(any(ByteBuffer.class));
        }
    }
    private static Stream<Arguments> testBufferReuseInSkipKeyValueIterator() {
        byte[] smallRecordValue = "1".getBytes();
        byte[] largeRecordValue = new byte[512 * 1024]; // 512KB
        RANDOM.nextBytes(largeRecordValue);

        return Stream.of(
            /*
             * 1 allocation per batch (i.e. per iterator instance) for buffer holding uncompressed data
             * = 1 buffer allocations
             */
            Arguments.of(CompressionType.GZIP, 1, smallRecordValue),
            Arguments.of(CompressionType.GZIP, 1, largeRecordValue),
            Arguments.of(CompressionType.SNAPPY, 1, smallRecordValue),
            Arguments.of(CompressionType.SNAPPY, 1, largeRecordValue),
            /*
             * 1 allocation per batch (i.e. per iterator instance) for buffer holding compressed data
             * 1 allocation per batch (i.e. per iterator instance) for buffer holding uncompressed data
             * = 2 buffer allocations
             */
            Arguments.of(CompressionType.LZ4, 2, smallRecordValue),
            Arguments.of(CompressionType.LZ4, 2, largeRecordValue),
            Arguments.of(CompressionType.ZSTD, 2, smallRecordValue),
            Arguments.of(CompressionType.ZSTD, 2, largeRecordValue)
        );
    }

    @ParameterizedTest
    @MethodSource
    public void testZstdJniForSkipKeyValueIterator(int expectedJniCalls, byte[] recordValue) throws IOException {
        MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
            Compression.zstd().build(), TimestampType.CREATE_TIME,
            new SimpleRecord(9L, "hakuna-matata".getBytes(), recordValue)
        );

        // Buffer containing compressed data
        final ByteBuffer compressedBuf = records.buffer();
        // Create a RecordBatch object
        final DefaultRecordBatch batch = spy(new DefaultRecordBatch(compressedBuf.duplicate()));
        // Buffer containing compressed records to be used for creating zstd-jni stream
        ByteBuffer recordsBuffer = compressedBuf.duplicate();
        recordsBuffer.position(RECORDS_OFFSET);

        try (final BufferSupplier bufferSupplier = BufferSupplier.create();
             final InputStream zstdStream = spy(ZstdCompression.wrapForZstdInput(recordsBuffer, bufferSupplier));
             final InputStream chunkedStream = new ChunkedBytesStream(zstdStream, bufferSupplier, 16 * 1024, false)
        ) {
            doReturn(chunkedStream).when(batch).recordInputStream(any());
            try (CloseableIterator<Record> streamingIterator = batch.skipKeyValueIterator(bufferSupplier)) {
                assertNotNull(streamingIterator);
                Utils.toList(streamingIterator);
                // verify the number of read() calls to zstd JNI stream. Each read() call is a JNI call.
                verify(zstdStream, times(expectedJniCalls)).read(any(byte[].class), anyInt(), anyInt());
                // verify that we don't use the underlying skip() functionality. The underlying skip() allocates
                // 1 buffer per skip call from he buffer pool whereas our implementation does not perform any allocation
                // during skip.
                verify(zstdStream, never()).skip(anyLong());
            }
        }
    }

    private static Stream<Arguments> testZstdJniForSkipKeyValueIterator() {
        byte[] smallRecordValue = "1".getBytes();
        byte[] largeRecordValue = new byte[40 * 1024]; // 40KB
        RANDOM.nextBytes(largeRecordValue);

        return Stream.of(
            /*
             * We expect exactly 2 read call to the JNI:
             * 1 for fetching the full data (size < 16KB)
             * 1 for detecting end of stream by trying to read more data
             */
            Arguments.of(2, smallRecordValue),
            /*
             * We expect exactly 4 read call to the JNI:
             * 3 for fetching the full data (Math.ceil(40/16))
             * 1 for detecting end of stream by trying to read more data
             */
            Arguments.of(4, largeRecordValue)
        );
    }

    @Test
    public void testIncrementSequence() {
        assertEquals(10, DefaultRecordBatch.incrementSequence(5, 5));
        assertEquals(0, DefaultRecordBatch.incrementSequence(Integer.MAX_VALUE, 1));
        assertEquals(4, DefaultRecordBatch.incrementSequence(Integer.MAX_VALUE - 5, 10));
    }

    @Test
    public void testDecrementSequence() {
        assertEquals(0, DefaultRecordBatch.decrementSequence(5, 5));
        assertEquals(Integer.MAX_VALUE, DefaultRecordBatch.decrementSequence(0, 1));
    }

    private static DefaultRecordBatch recordsWithInvalidRecordCount(Byte magicValue, long timestamp,
                                              CompressionType codec, int invalidCount) {
        ByteBuffer buf = ByteBuffer.allocate(512);
        MemoryRecordsBuilder builder = MemoryRecords.builder(buf, magicValue, Compression.of(codec).build(), TimestampType.CREATE_TIME, 0L);
        builder.appendWithOffset(0, timestamp, null, "hello".getBytes());
        builder.appendWithOffset(1, timestamp, null, "there".getBytes());
        builder.appendWithOffset(2, timestamp, null, "beautiful".getBytes());
        MemoryRecords records = builder.build();
        ByteBuffer buffer = records.buffer();
        buffer.position(0);
        buffer.putInt(RECORDS_COUNT_OFFSET, invalidCount);
        buffer.position(0);
        return new DefaultRecordBatch(buffer);
    }
}
