[ISSUE #9080] Fix tranfer logic when get large messages from cache in…

… tiered storage (#9079)
apache · Dec 26, 2024 · 2089abd · 2089abd
1 parent 1c35adb
commit 2089abd
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/tieredstore/src/main/java/org/apache/rocketmq/tieredstore/core/MessageStoreFetcherImpl.java b/tieredstore/src/main/java/org/apache/rocketmq/tieredstore/core/MessageStoreFetcherImpl.java
@@ -56,6 +56,7 @@ public class MessageStoreFetcherImpl implements MessageStoreFetcher {
     private final String brokerName;
     private final MetadataStore metadataStore;
     private final MessageStoreConfig storeConfig;
+    private final org.apache.rocketmq.store.config.MessageStoreConfig messageStoreConfig;
     private final TieredMessageStore messageStore;
     private final IndexService indexService;
     private final FlatFileStore flatFileStore;
@@ -71,6 +72,7 @@ public MessageStoreFetcherImpl(TieredMessageStore messageStore, MessageStoreConf
         FlatFileStore flatFileStore, IndexService indexService) {
 
         this.storeConfig = storeConfig;
+        this.messageStoreConfig = messageStore.getMessageStoreConfig();
         this.brokerName = storeConfig.getBrokerName();
         this.flatFileStore = flatFileStore;
         this.messageStore = messageStore;
@@ -148,6 +150,9 @@ protected GetMessageResultExt getMessageFromCache(
             if (result.getMessageCount() == maxCount) {
                 break;
             }
+            if (result.getBufferTotalSize() >= messageStoreConfig.getMaxTransferBytesOnMessageInMemory()) {
+                break;
+            }
         }
         result.setStatus(result.getMessageCount() > 0 ?
             GetMessageStatus.FOUND : GetMessageStatus.NO_MATCHED_MESSAGE);

diff --git a/...re/src/test/java/org/apache/rocketmq/tieredstore/core/MessageStoreDispatcherImplTest.java b/...re/src/test/java/org/apache/rocketmq/tieredstore/core/MessageStoreDispatcherImplTest.java
@@ -106,6 +106,7 @@ public void dispatchFromCommitLogTest() throws Exception {
         Mockito.when(messageStore.getStoreExecutor()).thenReturn(executor);
         Mockito.when(messageStore.getFlatFileStore()).thenReturn(fileStore);
         Mockito.when(messageStore.getIndexService()).thenReturn(indexService);
+        Mockito.when(messageStore.getMessageStoreConfig()).thenReturn(new org.apache.rocketmq.store.config.MessageStoreConfig());
 
         // mock message
         ByteBuffer buffer = MessageFormatUtilTest.buildMockedMessageBuffer();