feat(log): enhance reading logic to handle offset gaps and add unit tests (#2699)

Author: Gezi-lzq

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'
   }
 

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;

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));