Repository: AutoMQ/automq (2025-07-25T08:35:02Z)

Issue:
Cold Read on Compaction Topic Fails to Fully Reuse StreamReader Due to Offset Gaps
During cold reads of a Compaction Topic, the StreamReader fails to be reused when encountering offset gaps (a natural result of log compaction). The current logic miscalculates nextFetchOffset after a gap, causing new StreamReader instances to be unnecessarily created instead of reusing the existing reader. 


Discussion:

Patch:
diff --git a/build.gradle b/build.gradle
index 85ff8f9aac..0b4b803398 100644
--- a/build.gradle
+++ b/build.gradle
@@ -1060,6 +1060,7 @@ project(':core') {
     testImplementation project(':storage:storage-api').sourceSets.test.output
     testImplementation project(':server').sourceSets.test.output
     testImplementation libs.bcpkix
+    testImplementation libs.mockitoJunitJupiter // supports MockitoExtension
     testImplementation libs.mockitoCore
     testImplementation libs.guava
     testImplementation(libs.apacheda) {
@@ -2218,9 +2219,9 @@ project(':s3stream') {
     implementation 'com.ibm.async:asyncutil:0.1.0'
 
     testImplementation 'org.slf4j:slf4j-simple:2.0.9'
-    testImplementation 'org.junit.jupiter:junit-jupiter:5.10.0'
-    testImplementation 'org.mockito:mockito-core:5.5.0'
-    testImplementation 'org.mockito:mockito-junit-jupiter:5.5.0'
+    testImplementation libs.junitJupiter
+    testImplementation libs.mockitoCore
+    testImplementation libs.mockitoJunitJupiter // supports MockitoExtension
     testImplementation 'org.awaitility:awaitility:4.2.1'
   }
 
diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLogFileRecords.java b/core/src/main/scala/kafka/log/streamaspect/ElasticLogFileRecords.java
index ba9832d4de..8652079687 100644
--- a/core/src/main/scala/kafka/log/streamaspect/ElasticLogFileRecords.java
+++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLogFileRecords.java
@@ -46,6 +46,7 @@
 import com.automq.stream.s3.context.FetchContext;
 import com.automq.stream.s3.trace.TraceUtils;
 import com.automq.stream.utils.FutureUtil;
+import com.google.common.annotations.VisibleForTesting;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -406,7 +407,8 @@ public long lastOffset() {
     }
 
     static class StreamSegmentInputStream implements LogInputStream<RecordBatch> {
-        private static final int FETCH_BATCH_SIZE = 64 * 1024;
+        @VisibleForTesting
+        protected static final int FETCH_BATCH_SIZE = 64 * 1024;
         private final ElasticLogFileRecords elasticLogFileRecords;
         private final Queue<RecordBatch> remaining = new LinkedList<>();
         private final int maxSize;
@@ -446,9 +448,9 @@ public RecordBatch nextBatch() throws IOException {
                                 buf = heapBuf;
                             }
                             readSize += buf.remaining();
+                            nextFetchOffset = Math.max(streamRecord.lastOffset(), nextFetchOffset);
                             for (RecordBatch r : MemoryRecords.readableRecords(buf).batches()) {
                                 remaining.offer(r);
-                                nextFetchOffset = r.lastOffset() - elasticLogFileRecords.baseOffset + 1;
                             }
                         } catch (Throwable e) {
                             ElasticStreamSlice slice = elasticLogFileRecords.streamSlice;
diff --git a/core/src/main/scala/kafka/log/streamaspect/MemoryClient.java b/core/src/main/scala/kafka/log/streamaspect/MemoryClient.java
index 75074f5b77..8182801834 100644
--- a/core/src/main/scala/kafka/log/streamaspect/MemoryClient.java
+++ b/core/src/main/scala/kafka/log/streamaspect/MemoryClient.java
@@ -80,7 +80,7 @@ public CompletableFuture<FailoverResponse> failover(FailoverRequest request) {
         return FutureUtil.failedFuture(new UnsupportedOperationException());
     }
 
-    static class StreamImpl implements Stream {
+    public static class StreamImpl implements Stream {
         private final AtomicLong nextOffsetAlloc = new AtomicLong();
         private NavigableMap<Long, RecordBatchWithContext> recordMap = new ConcurrentSkipListMap<>();
         private final long streamId;
@@ -132,10 +132,26 @@ public CompletableFuture<FetchResult> fetch(FetchContext context, long startOffs
             if (floorKey == null) {
                 return CompletableFuture.completedFuture(ArrayList::new);
             }
-            List<RecordBatchWithContext> records = new ArrayList<>(recordMap.subMap(floorKey, endOffset).values());
+            NavigableMap<Long, RecordBatchWithContext> subMap = recordMap.subMap(floorKey, true, endOffset, false);
+            List<RecordBatchWithContext> records = new ArrayList<>();
+            int accumulatedSize = 0;
+            for (Map.Entry<Long, RecordBatchWithContext> entry : subMap.entrySet()) {
+                RecordBatchWithContext batch = entry.getValue();
+                int batchSize = batch.rawPayload().remaining();
+                if (accumulatedSize + batchSize > maxSizeHint && !records.isEmpty()) {
+                    break;
+                }
+                records.add(batch);
+                accumulatedSize += batchSize;
+
+                if (accumulatedSize > maxSizeHint) {
+                    break;
+                }
+            }
             return CompletableFuture.completedFuture(() -> records);
         }
 
+
         @Override
         public CompletableFuture<Void> trim(long newStartOffset) {
             recordMap = new ConcurrentSkipListMap<>(recordMap.tailMap(newStartOffset));


Test patch:
diff --git a/core/src/test/java/kafka/log/streamaspect/ElasticLogFileRecordsTest.java b/core/src/test/java/kafka/log/streamaspect/ElasticLogFileRecordsTest.java
new file mode 100644
index 0000000000..5ea93a3dbd
--- /dev/null
+++ b/core/src/test/java/kafka/log/streamaspect/ElasticLogFileRecordsTest.java
@@ -0,0 +1,332 @@
+/*
+ * Copyright 2025, AutoMQ HK Limited.
+ *
+ * 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 kafka.log.streamaspect;
+
+import org.apache.kafka.common.compress.NoCompression;
+import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.common.record.MemoryRecordsBuilder;
+import org.apache.kafka.common.record.Record;
+import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.record.SimpleRecord;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.utils.ByteBufferOutputStream;
+import org.apache.kafka.common.utils.Time;
+
+import com.automq.stream.api.Stream;
+import com.automq.stream.s3.context.FetchContext;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.concurrent.ExecutionException;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+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.ArgumentMatchers.eq;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+
+@Tag("S3Unit")
+@ExtendWith(MockitoExtension.class)
+class ElasticLogFileRecordsTest {
+
+    private ElasticStreamSlice streamSlice;
+    private Stream stream;
+
+    private ElasticLogFileRecords elasticLogFileRecords;
+
+    private final Random random = new Random();
+
+    @BeforeEach
+    void setUp() {
+        stream = spy(new MemoryClient.StreamImpl(1));
+        streamSlice = spy(new DefaultElasticStreamSlice(stream, SliceRange.of(0, Offsets.NOOP_OFFSET)));
+        elasticLogFileRecords = new ElasticLogFileRecords(streamSlice, 0, 0);
+    }
+
+    /**
+     * Test reading a single, complete batch of records.
+     */
+    @Test
+    void testReadSingleBatch() throws ExecutionException, InterruptedException, IOException {
+        // Arrange
+        long startOffset = 0;
+        int recordCount = 100;
+        long maxOffset = startOffset + recordCount;
+        final int maxReadBytes = 4096;
+
+        Map<Long, SimpleRecord> expectedRecords = prepareRecords(startOffset, recordCount);
+        MemoryRecords memoryRecords = createMemoryRecords(expectedRecords);
+        elasticLogFileRecords.append(memoryRecords, maxOffset);
+
+        // Act
+        ReadHint.setReadAll(false);
+        org.apache.kafka.common.record.Records readRecords = elasticLogFileRecords.read(startOffset, maxOffset, maxReadBytes).get();
+
+        // Assert
+        assertNotNull(readRecords);
+        assertTrue(readRecords instanceof ElasticLogFileRecords.BatchIteratorRecordsAdaptor);
+        assertRecords(expectedRecords, readRecords);
+        verify(streamSlice).fetch(any(FetchContext.class), eq(0L), eq(maxOffset), eq(maxReadBytes));
+    }
+
+    /**
+     * Test reading data that spans across multiple record batches.
+     */
+    @Test
+    void testReadAcrossMultipleBatches() throws ExecutionException, InterruptedException, IOException {
+        // Arrange
+        long startOffset = 0;
+        int recordsPerBatch = 50;
+        int batchCount = 3;
+        long totalRecords = (long) recordsPerBatch * batchCount;
+        long maxOffset = startOffset + totalRecords;
+        final int maxReadBytes = Integer.MAX_VALUE;
+
+        Map<Long, SimpleRecord> allExpectedRecords = new HashMap<>();
+        long currentStartOffset = startOffset;
+        List<Long> nextOffsetList = new ArrayList<>();
+        for (int i = 0; i < batchCount; i++) {
+            Map<Long, SimpleRecord> batchRecords = prepareRecords(currentStartOffset, recordsPerBatch);
+            allExpectedRecords.putAll(batchRecords);
+            MemoryRecords memoryRecords = createMemoryRecords(batchRecords);
+            elasticLogFileRecords.append(memoryRecords, currentStartOffset + recordsPerBatch);
+
+            nextOffsetList.add(currentStartOffset);
+            currentStartOffset += recordsPerBatch;
+        }
+
+        // Act
+        ReadHint.setReadAll(false);
+        org.apache.kafka.common.record.Records readRecords = elasticLogFileRecords.read(startOffset, maxOffset, maxReadBytes).get();
+
+        // Assert
+        assertNotNull(readRecords);
+        assertTrue(readRecords instanceof ElasticLogFileRecords.BatchIteratorRecordsAdaptor);
+        assertRecords(allExpectedRecords, readRecords);
+        nextOffsetList.forEach(nextOffset -> {
+            verify(streamSlice).fetch(any(FetchContext.class), eq(nextOffset), eq(maxOffset), anyInt());
+        });
+    }
+
+    /**
+     * Test reading a batch that has been compacted, resulting in offset gaps.
+     * The reader should correctly iterate over the existing records and skip the gaps.
+     */
+    @Test
+    void testReadCompactedBatchWithGaps() throws ExecutionException, InterruptedException, IOException {
+        // Arrange
+        long batchStartOffset = 0;
+        long lastOffsetInBatch = batchStartOffset + 5; // Batch spans from 0 to 5
+        final int maxReadBytes = 4096;
+
+        // Create a MemoryRecords buffer with offset gaps to simulate compaction.
+        // We will only include records for offsets 0, 2, 4.
+        Map<Long, SimpleRecord> expectedRecords = new HashMap<>();
+        expectedRecords.put(batchStartOffset, createSimpleRecord("key" + batchStartOffset, "value" + batchStartOffset));
+        // Skip 1
+        expectedRecords.put(batchStartOffset + 2, createSimpleRecord("key" + (batchStartOffset + 2), "value" + (batchStartOffset + 2)));
+        // Skip 3
+        expectedRecords.put(batchStartOffset + 4, createSimpleRecord("key" + (batchStartOffset + 4), "value" + (batchStartOffset + 4)));
+        // Skip 5
+        MemoryRecords memoryRecords = createMemoryRecords(expectedRecords);
+        elasticLogFileRecords.append(memoryRecords, lastOffsetInBatch + 1);
+
+        // Act: Read the entire range that contains the gappy batch.
+        ReadHint.setReadAll(false);
+        org.apache.kafka.common.record.Records readRecords = elasticLogFileRecords.read(batchStartOffset, lastOffsetInBatch + 1, maxReadBytes).get();
+
+        // Assert: Verify that only the existing records are returned.
+        assertNotNull(readRecords);
+        assertTrue(readRecords instanceof ElasticLogFileRecords.BatchIteratorRecordsAdaptor);
+        assertRecords(expectedRecords, readRecords);
+        verify(streamSlice).fetch(any(FetchContext.class), eq(batchStartOffset), eq(lastOffsetInBatch + 1), eq(maxReadBytes));
+    }
+
+    /**
+     * Test reading across multiple batches where at least one batch is compacted and has offset gaps.
+     */
+    @Test
+    void testReadAcrossMultipleBatchesWithGaps() throws ExecutionException, InterruptedException, IOException {
+        // Arrange
+        long startOffset = 0;
+        int recordsPerBatch = 50;
+        int batchCount = 3;
+        long totalRecords = (long) recordsPerBatch * batchCount;
+        long maxOffset = startOffset + totalRecords;
+        final int maxReadBytes = Integer.MAX_VALUE;
+
+        Map<Long, SimpleRecord> allExpectedRecords = new HashMap<>();
+        long currentStartOffset = startOffset;
+
+        List<Long> nextOffsetList = new ArrayList<>();
+        for (int i = 0; i < batchCount; i++) {
+            int skippedRecords = random.nextInt(recordsPerBatch);
+            Map<Long, SimpleRecord> batchRecords = prepareRecords(currentStartOffset, recordsPerBatch - skippedRecords);
+            allExpectedRecords.putAll(batchRecords);
+            MemoryRecords memoryRecords = createMemoryRecords(batchRecords);
+            elasticLogFileRecords.append(memoryRecords, currentStartOffset + recordsPerBatch);
+
+            nextOffsetList.add(currentStartOffset);
+            currentStartOffset += recordsPerBatch;
+        }
+
+        // Act
+        ReadHint.setReadAll(false);
+        org.apache.kafka.common.record.Records readRecords = elasticLogFileRecords.read(startOffset, maxOffset, maxReadBytes).get();
+
+        // Assert
+        assertNotNull(readRecords);
+        assertTrue(readRecords instanceof ElasticLogFileRecords.BatchIteratorRecordsAdaptor);
+        assertRecords(allExpectedRecords, readRecords);
+        nextOffsetList.forEach(nextOffset -> {
+            verify(streamSlice).fetch(any(FetchContext.class), eq(nextOffset), eq(maxOffset), anyInt());
+        });
+    }
+
+
+    // Helper for preparing records to avoid code duplication
+    private Map<Long, SimpleRecord> prepareRecords(long startOffset, int count) {
+        final int fetchBatchSize = ElasticLogFileRecords.StreamSegmentInputStream.FETCH_BATCH_SIZE;
+        Map<Long, SimpleRecord> records = new HashMap<>();
+
+        // Calculate approximate size per record to ensure total size > FETCH_BATCH_SIZE
+        int estimatedRecordOverhead = 50; // Approximate overhead per record (headers, etc.)
+        int targetValueSize = Math.max(100, (fetchBatchSize / count) + estimatedRecordOverhead);
+
+        // Create a large value string to ensure we exceed FETCH_BATCH_SIZE
+        StringBuilder largeValue = new StringBuilder();
+        for (int j = 0; j < targetValueSize; j++) {
+            largeValue.append('a');
+        }
+        String valueTemplate = largeValue.toString();
+
+        for (int i = 0; i < count; i++) {
+            long currentOffset = startOffset + i;
+            String value = valueTemplate + "_" + currentOffset; // Make each value unique
+            records.put(currentOffset, createSimpleRecord("key" + currentOffset, value));
+        }
+        return records;
+    }
+
+    // Helper for asserting records to avoid code duplication
+    private void assertRecords(Map<Long, SimpleRecord> expectedRecords, org.apache.kafka.common.record.Records actualRecords) {
+        List<Record> records = new ArrayList<>();
+        actualRecords.records().forEach(records::add);
+
+        assertEquals(expectedRecords.size(), records.size());
+
+        for (Record record : records) {
+            SimpleRecord expectedRecord = expectedRecords.get(record.offset());
+            assertNotNull(expectedRecord, "Unexpected record with offset " + record.offset());
+            assertEquals(new String(expectedRecord.key().array()), new String(readBytes(record.key())));
+            assertEquals(new String(expectedRecord.value().array()), new String(readBytes(record.value())));
+        }
+    }
+
+    /**
+     * Test reading an empty range (startOffset >= maxOffset).
+     * Expects an empty Records object and no interaction with the underlying streamSlice.
+     */
+    @Test
+    void testReadEmptyRange() throws ExecutionException, InterruptedException, IOException {
+        // Arrange
+        long startOffset = 100;
+        long maxOffset = 100; // startOffset >= maxOffset
+
+        // Act
+        ReadHint.setReadAll(false);
+        org.apache.kafka.common.record.Records readRecords = elasticLogFileRecords.read(startOffset, maxOffset, 4096).get();
+
+        // Assert
+        assertNotNull(readRecords);
+        assertEquals(0, readRecords.sizeInBytes());
+        // Verify that streamSlice.fetch was NOT called
+        verify(streamSlice, never()).fetch(any(FetchContext.class), anyLong(), anyLong(), anyInt());
+    }
+
+    // Helper methods
+    private SimpleRecord createSimpleRecord(String key, String value) {
+        return new SimpleRecord(System.currentTimeMillis(), key.getBytes(), value.getBytes());
+    }
+
+    private MemoryRecords createMemoryRecords(Map<Long, SimpleRecord> records) {
+        ByteBuffer buffer = ByteBuffer.allocate(128 * 1024); // Increased buffer size to accommodate larger records
+        long baseOffset = records.keySet().stream().min(Long::compare).get();
+        ByteBufferOutputStream stream = new ByteBufferOutputStream(buffer);
+        try (MemoryRecordsBuilder builder = new MemoryRecordsBuilder(stream,
+            RecordBatch.CURRENT_MAGIC_VALUE,
+            NoCompression.NONE,
+            TimestampType.CREATE_TIME,
+            baseOffset,
+            Time.SYSTEM.milliseconds(),
+            RecordBatch.NO_PRODUCER_ID,
+            RecordBatch.NO_PRODUCER_EPOCH,
+            RecordBatch.NO_SEQUENCE,
+            false,
+            false,
+            RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.limit(), 0L)) {
+
+            records.keySet().stream().sorted().forEach(
+                offset -> builder.appendWithOffset(offset, records.get(offset))
+            );
+            return builder.build();
+        }
+    }
+
+    private byte[] readBytes(ByteBuffer buffer) {
+        if (buffer == null) return new byte[0];
+        byte[] bytes = new byte[buffer.remaining()];
+        buffer.get(bytes);
+        return bytes;
+    }
+
+    static class ReadHint {
+        private static final ThreadLocal<Boolean> READ_ALL = ThreadLocal.withInitial(() -> false);
+        private static final ThreadLocal<Boolean> FAST_READ = ThreadLocal.withInitial(() -> false);
+
+        public static boolean isReadAll() {
+            return READ_ALL.get();
+        }
+
+        public static void setReadAll(boolean readAll) {
+            READ_ALL.set(readAll);
+        }
+
+        public static boolean isFastRead() {
+            return FAST_READ.get();
+        }
+    }
+}