fix(store): fix end offset of fetch result

Signed-off-by: SSpirits <admin@lv5.moe>
AutoMQ · Nov 24, 2023 · 4f60039 · 4f60039
1 parent e7f310d
commit 4f60039
Showing 1 changed file with 13 additions and 6 deletions.
diff --git a/store/src/main/java/com/automq/rocketmq/store/queue/StreamLogicQueue.java b/store/src/main/java/com/automq/rocketmq/store/queue/StreamLogicQueue.java
@@ -385,8 +385,14 @@ public CompletableFuture<PopResult> popRetry(StoreContext context, long consumer
         return retryStreamIdFuture.thenCompose(retryStreamId -> pop(context, consumerGroupId, retryStreamId, offset, PopOperation.PopOperationType.POP_RETRY, filter, batchSize, invisibleDuration));
     }
 
+    record FetchResult(List<FlatMessageExt> messageList, long endOffset) {
+        public int size() {
+            return messageList.size();
+        }
+    }
+
     @WithSpan(kind = SpanKind.SERVER)
-    private CompletableFuture<List<FlatMessageExt>> fetchMessages(StoreContext context, @SpanAttribute long streamId,
+    private CompletableFuture<FetchResult> fetchMessages(StoreContext context, @SpanAttribute long streamId,
         @SpanAttribute long offset, @SpanAttribute int batchSize) {
         long startOffset = streamStore.startOffset(streamId);
         if (offset < startOffset) {
@@ -395,13 +401,14 @@ private CompletableFuture<List<FlatMessageExt>> fetchMessages(StoreContext conte
 
         long confirmOffset = streamStore.confirmOffset(streamId);
         if (offset >= confirmOffset) {
-            return CompletableFuture.completedFuture(Collections.emptyList());
+            return CompletableFuture.completedFuture(new FetchResult(Collections.emptyList(), confirmOffset));
         }
 
         if (offset + batchSize > confirmOffset) {
             batchSize = (int) (confirmOffset - offset);
         }
 
+        long finalOffset = offset;
         return streamStore.fetch(context, streamId, offset, batchSize)
             .thenApply(fetchResult -> {
                 AtomicLong fetchBytes = new AtomicLong();
@@ -422,7 +429,7 @@ private CompletableFuture<List<FlatMessageExt>> fetchMessages(StoreContext conte
 
                 context.span().ifPresent(span -> span.setAttribute("fetchBytes", fetchBytes.get()));
 
-                return resultList;
+                return new FetchResult(resultList, finalOffset + resultList.size());
             });
     }
 
@@ -436,23 +443,23 @@ private CompletableFuture<FilterFetchResult> fetchAndFilterMessages(StoreContext
         return fetchMessages(context, streamId, offset, fetchBatchSize)
             .thenCompose(fetchResult -> {
                 // Add filter result to message list.
-                List<FlatMessageExt> matchedMessageList = filter.doFilter(fetchResult);
+                List<FlatMessageExt> matchedMessageList = filter.doFilter(fetchResult.messageList());
                 // Update end offset
                 int index = batchSize - result.size();
                 if (matchedMessageList.size() > index) {
                     FlatMessageExt messageExt = matchedMessageList.get(index);
                     result.setEndOffset(messageExt.offset());
                     result.addMessageList(matchedMessageList.subList(0, index));
                 } else {
-                    result.setEndOffset(offset + fetchResult.size());
+                    result.setEndOffset(fetchResult.endOffset());
                     result.addMessageList(matchedMessageList);
                 }
                 // If not enough messages after applying filter, fetch more messages.
                 boolean needToFetch = result.size() < batchSize;
                 boolean hasMoreMessages = fetchResult.size() >= fetchBatchSize;
 
                 int newFetchCount = fetchCount + fetchResult.size();
-                long newFetchBytes = fetchBytes + fetchResult.stream()
+                long newFetchBytes = fetchBytes + fetchResult.messageList().stream()
                     .map(messageExt -> (long) messageExt.message().getByteBuffer().limit())
                     .reduce(0L, Long::sum);
                 boolean notExceedLimit = newFetchCount < config.maxFetchCount() &&