[ISSUE #10949] Improve pop consumer offset committing

broker/src/main/java/org/apache/rocketmq/broker/pop/PopConsumerCache.java

@@ -29,6 +29,7 @@
 import org.apache.rocketmq.broker.BrokerController;
 import org.apache.rocketmq.broker.offset.ConsumerOffsetManager;
 import org.apache.rocketmq.common.BrokerConfig;
+import org.apache.rocketmq.common.KeyBuilder;
 import org.apache.rocketmq.common.ServiceThread;
 import org.apache.rocketmq.common.constant.LoggerName;
 import org.apache.rocketmq.common.utils.ConcurrentHashMapUtils;
@@ -84,6 +85,13 @@ public long getMinOffsetInCache(String groupId, String topicId, int queueId) {
         return consumerRecords != null ? consumerRecords.getMinOffsetInBuffer() : OFFSET_NOT_EXIST;
     }
 
+    public void setPendingCommitOffset(String groupId, String topicId, int queueId, long offset) {
+        ConsumerRecords consumerRecords = consumerRecordTable.get(this.getKey(groupId, topicId, queueId));
+        if (consumerRecords != null) {
+            consumerRecords.setPendingCommitOffset(offset);
+        }
+    }
+
     public long getPopInFlightMessageCount(String groupId, String topicId, int queueId) {
         ConsumerRecords consumerRecords = consumerRecordTable.get(this.getKey(groupId, topicId, queueId));
         return consumerRecords != null ? consumerRecords.getInFlightRecordCount() : 0L;
@@ -126,16 +134,28 @@ public int cleanupRecords(Consumer<PopConsumerRecord> consumer) {
                 records.getGroupId(), records.getTopicId());
 
             if (timeout) {
-                records.stageExpiredRecords(Long.MAX_VALUE);
-                List<PopConsumerRecord> writeConsumerRecords =
-                    new ArrayList<>(records.getRemoveTreeMap().values());
-                if (!writeConsumerRecords.isEmpty()) {
-                    consumerRecordStore.writeRecords(writeConsumerRecords);
+                // hold the same lock as PopConsumerService#popAsync to prevent evicting records that a
+                // concurrent pop is still writing.
+                String lockTopicId = KeyBuilder.parseNormalTopic(records.getTopicId(), records.getGroupId());
+                if (!consumerLockService.tryLock(records.getGroupId(), lockTopicId)) {
+                    remain += records.getInFlightRecordCount();
+                    continue;
+                }
+                try {
+                    records.stageExpiredRecords(Long.MAX_VALUE);
+                    List<PopConsumerRecord> writeConsumerRecords =
+                        new ArrayList<>(records.getRemoveTreeMap().values());
+                    if (!writeConsumerRecords.isEmpty()) {
+                        consumerRecordStore.writeRecords(writeConsumerRecords);
+                    }
+                    records.clearStagedRecords();
+                    log.info("PopConsumerOffline, so clean expire records, groupId={}, topic={}, queueId={}, records={}",
+                        records.getGroupId(), records.getTopicId(), records.getQueueId(), writeConsumerRecords.size());
+                    iterator.remove();
+                } finally {
+                    consumerLockService.unlock(records.getGroupId(), lockTopicId);
                 }
-                records.clearStagedRecords();
-                log.info("PopConsumerOffline, so clean expire records, groupId={}, topic={}, queueId={}, records={}",
-                    records.getGroupId(), records.getTopicId(), records.getQueueId(), records.getInFlightRecordCount());
-                iterator.remove();
+                commitPendingOffset(records);
                 continue;
             }
 
@@ -159,13 +179,42 @@ public int cleanupRecords(Consumer<PopConsumerRecord> consumer) {
             if (offset > OFFSET_NOT_EXIST) {
                 this.commitOffset("PopConsumerCache",
                     records.getGroupId(), records.getTopicId(), records.getQueueId(), offset);
+            } else {
+                this.commitPendingOffset(records);
             }
 
             remain += records.getInFlightRecordCount();
         }
         return remain;
     }
 
+    /**
+     * Commit the offset left over by pop once no record of the queue remains in the cache.
+     */
+    private void commitPendingOffset(ConsumerRecords records) {
+        // Read the pending offset before checking the in-flight records, never the other way
+        // around. A concurrent pop writes its records into the cache before it updates the pending
+        // offset, so this order guarantees that either the check below sees those records, or the
+        // pending offset read here is the older one.
+        long pendingCommitOffset = records.getPendingCommitOffset();
+        if (pendingCommitOffset == OFFSET_NOT_EXIST || records.getInFlightRecordCount() != 0) {
+            return;
+        }
+
+        String groupId = records.getGroupId();
+        String topicId = records.getTopicId();
+        int queueId = records.getQueueId();
+
+        ConsumerOffsetManager consumerOffsetManager = brokerController.getConsumerOffsetManager();
+        if (consumerOffsetManager.hasOffsetReset(topicId, groupId, queueId)) {
+            return;
+        }
+
+        if (pendingCommitOffset > consumerOffsetManager.queryOffset(groupId, topicId, queueId)) {
+            this.commitOffset("PopConsumerCache", groupId, topicId, queueId, pendingCommitOffset);
+        }
+    }
+
     public void commitOffset(String clientHost, String groupId, String topicId, int queueId, long offset) {
         if (!consumerLockService.tryLock(groupId, topicId)) {
             return;
@@ -214,6 +263,9 @@ protected static class ConsumerRecords {
         private final ConcurrentSkipListMap<Long /* offset */, PopConsumerRecord> removeTreeMap;
         private final ConcurrentSkipListMap<Long /* offset */, PopConsumerRecord> recordTreeMap;
 
+        // The consumer offset to commit once no record of this queue is in cache
+        private volatile long pendingCommitOffset = OFFSET_NOT_EXIST;
+
         public ConsumerRecords(BrokerConfig brokerConfig, String groupId, String topicId, int queueId) {
             this.groupId = groupId;
             this.topicId = topicId;
@@ -244,6 +296,14 @@ public int getInFlightRecordCount() {
             return removeTreeMap.size() + recordTreeMap.size();
         }
 
+        public void setPendingCommitOffset(long pendingCommitOffset) {
+            this.pendingCommitOffset = pendingCommitOffset;
+        }
+
+        public long getPendingCommitOffset() {
+            return pendingCommitOffset;
+        }
+
         public void stageExpiredRecords(long currentTime) {
             Iterator<Map.Entry<Long, PopConsumerRecord>>
                 iterator = recordTreeMap.entrySet().iterator();

broker/src/main/java/org/apache/rocketmq/broker/pop/PopConsumerContext.java

@@ -51,6 +51,8 @@ public class PopConsumerContext {
 
     private List<PopConsumerRecord> popConsumerRecordList;
 
+    private List<PendingCommit> pendingCommitList;
+
     public PopConsumerContext(String clientHost,
         long popTime, long invisibleTime, String groupId, boolean fifo, int initMode, String attemptId) {
 
@@ -164,6 +166,44 @@ public List<PopConsumerRecord> getPopConsumerRecordList() {
         return popConsumerRecordList;
     }
 
+    public void addPendingCommit(String topicId, int queueId, long commitOffset) {
+        if (this.pendingCommitList == null) {
+            this.pendingCommitList = new ArrayList<>();
+        }
+        this.pendingCommitList.add(new PendingCommit(topicId, queueId, commitOffset));
+    }
+
+    public List<PendingCommit> getPendingCommitList() {
+        return pendingCommitList;
+    }
+
+    public static class PendingCommit {
+
+        private final String topicId;
+
+        private final int queueId;
+
+        private final long commitOffset;
+
+        public PendingCommit(String topicId, int queueId, long commitOffset) {
+            this.topicId = topicId;
+            this.queueId = queueId;
+            this.commitOffset = commitOffset;
+        }
+
+        public String getTopicId() {
+            return topicId;
+        }
+
+        public int getQueueId() {
+            return queueId;
+        }
+
+        public long getCommitOffset() {
+            return commitOffset;
+        }
+    }
+
     @Override
     public String toString() {
         return "PopConsumerContext{" +

broker/src/main/java/org/apache/rocketmq/broker/pop/PopConsumerService.java

@@ -39,6 +39,7 @@
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Triple;
 import org.apache.rocketmq.broker.BrokerController;
+import org.apache.rocketmq.broker.offset.ConsumerOffsetManager;
 import org.apache.rocketmq.common.BrokerConfig;
 import org.apache.rocketmq.common.KeyBuilder;
 import org.apache.rocketmq.common.MixAll;
@@ -185,23 +186,18 @@ public PopConsumerContext handleGetMessageResult(PopConsumerContext context, Get
             }
         }
 
-        long commitOffset = offset;
         if (context.isFifo()) {
+            long commitOffset = offset;
             if (!GetMessageStatus.FOUND.equals(result.getStatus())) {
                 commitOffset = result.getNextBeginOffset();
             }
+            this.brokerController.getConsumerOffsetManager().commitOffset(
+                context.getClientHost(), context.getGroupId(), topicId, queueId, commitOffset);
         } else {
             this.brokerController.getConsumerOffsetManager().commitPullOffset(
                 context.getClientHost(), context.getGroupId(), topicId, queueId, result.getNextBeginOffset());
-            if (brokerConfig.isEnablePopBufferMerge() && popConsumerCache != null) {
-                long minOffset = popConsumerCache.getMinOffsetInCache(context.getGroupId(), topicId, queueId);
-                if (minOffset != OFFSET_NOT_EXIST) {
-                    commitOffset = minOffset;
-                }
-            }
+            context.addPendingCommit(topicId, queueId, result.getNextBeginOffset());
         }
-        this.brokerController.getConsumerOffsetManager().commitOffset(
-            context.getClientHost(), context.getGroupId(), topicId, queueId, commitOffset);
         return context;
     }
 
@@ -447,6 +443,7 @@ public CompletableFuture<PopConsumerContext> popAsync(String clientHost, long po
                         }
                     }
                 }
+                this.commitPendingOffset(result);
                 return CompletableFuture.completedFuture(result);
             }).whenComplete((result, throwable) -> {
                 try {
@@ -470,6 +467,48 @@ public CompletableFuture<PopConsumerContext> popAsync(String clientHost, long po
         return getMessageFuture;
     }
 
+    /**
+     * Commit the consumer offset up to nextBeginOffset of this pop, instead of the batch start
+     * offset which always lags one round behind. Acknowledgements only delete records, so the offset would
+     * otherwise never catch up once a consumer acknowledged everything and went offline.
+     * <p>
+     * Records written to the kv store are durable and unacknowledged ones are redelivered by revive, so
+     * the offset may pass them. Records held in the cache are not persisted yet, so they cap the
+     * offset at minOffsetInCache here, and nextBeginOffset is kept on the cache entry for
+     * PopConsumerCache#cleanupRecords to commit once they are acknowledged or persisted.
+     */
+    protected void commitPendingOffset(PopConsumerContext context) {
+        List<PopConsumerContext.PendingCommit> pendingCommitList = context.getPendingCommitList();
+        if (pendingCommitList == null) {
+            return;
+        }
+
+        String groupId = context.getGroupId();
+        ConsumerOffsetManager consumerOffsetManager = this.brokerController.getConsumerOffsetManager();
+
+        for (PopConsumerContext.PendingCommit pendingCommit : pendingCommitList) {
+            String topicId = pendingCommit.getTopicId();
+            int queueId = pendingCommit.getQueueId();
+
+            if (consumerOffsetManager.hasOffsetReset(topicId, groupId, queueId)) {
+                continue;
+            }
+
+            long commitOffset = pendingCommit.getCommitOffset();
+            if (popConsumerCache != null) {
+                long minOffset = popConsumerCache.getMinOffsetInCache(groupId, topicId, queueId);
+                if (minOffset != OFFSET_NOT_EXIST) {
+                    popConsumerCache.setPendingCommitOffset(groupId, topicId, queueId, commitOffset);
+                    commitOffset = Math.min(commitOffset, minOffset);
+                }
+            }
+
+            if (commitOffset > consumerOffsetManager.queryOffset(groupId, topicId, queueId)) {
+                consumerOffsetManager.commitOffset(context.getClientHost(), groupId, topicId, queueId, commitOffset);
+            }
+        }
+    }
+
     /**
      * Fifo pops carrying an attemptId already registered in OrderInfo are in-flight retries
      * of the same receive attempt. Instead of failing fast on lock contention (which leaves

broker/src/test/java/org/apache/rocketmq/broker/pop/PopConsumerCacheTest.java

@@ -30,6 +30,9 @@
 import org.mockito.Mockito;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
 
 public class PopConsumerCacheTest {
 
@@ -142,5 +145,81 @@ record = new PopConsumerRecord(2L, groupId, topicId, queueId,
         consumerRecordList.clear();
         consumerCache.cleanupRecords(consumerRecordList::add);
         Assert.assertEquals(0, consumerRecordList.size());
+
+        // timeout cleanup is skipped when tryLock fails: concurrent pop holds the lock
+        record = new PopConsumerRecord(System.currentTimeMillis(),
+            groupId, topicId, queueId, 0, 20000, 105, attemptId);
+        consumerCache.writeRecords(Collections.singletonList(record));
+        Mockito.when(consumerLockService.tryLock(eq(groupId), eq(topicId))).thenReturn(false);
+        int remain = consumerCache.cleanupRecords(consumerRecordList::add);
+        Assert.assertEquals(1, remain);
+        Assert.assertEquals(1, consumerCache.getCacheKeySize());
+        Assert.assertEquals(1, consumerCache.getPopInFlightMessageCount(groupId, topicId, queueId));
+
+        // timeout cleanup proceeds and persists records to KV store when the lock is acquired
+        Mockito.when(consumerLockService.tryLock(eq(groupId), eq(topicId))).thenReturn(true);
+        consumerCache.cleanupRecords(consumerRecordList::add);
+        Assert.assertEquals(0, consumerCache.getCacheKeySize());
+        Mockito.verify(consumerKVStore).writeRecords(Collections.singletonList(record));
+    }
+
+    @Test
+    public void commitPendingOffsetTest() {
+        BrokerController brokerController = Mockito.mock(BrokerController.class);
+        PopConsumerKVStore consumerKVStore = Mockito.mock(PopConsumerRocksdbStore.class);
+        PopConsumerLockService consumerLockService = Mockito.mock(PopConsumerLockService.class);
+        ConsumerOffsetManager consumerOffsetManager = Mockito.mock(ConsumerOffsetManager.class);
+        Mockito.when(brokerController.getBrokerConfig()).thenReturn(new BrokerConfig());
+        Mockito.when(brokerController.getConsumerOffsetManager()).thenReturn(consumerOffsetManager);
+        Mockito.when(consumerLockService.tryLock(groupId, topicId)).thenReturn(true);
+
+        PopConsumerCache consumerCache =
+            new PopConsumerCache(brokerController, consumerKVStore, consumerLockService, null);
+
+        // the record is still in cache, so it bounds the commit
+        PopConsumerRecord record = new PopConsumerRecord(System.currentTimeMillis(),
+            groupId, topicId, queueId, 0, 20000, 100, attemptId);
+        consumerCache.writeRecords(Collections.singletonList(record));
+        consumerCache.setPendingCommitOffset(groupId, topicId, queueId, 110L);
+        consumerCache.cleanupRecords(consumerRecord -> {
+        });
+        Mockito.verify(consumerOffsetManager).commitOffset(
+            anyString(), eq(groupId), eq(topicId), eq(queueId), eq(100L));
+
+        // nothing remains in cache after the acknowledgement, so the pending offset is committed
+        Assert.assertTrue(consumerCache.deleteRecords(Collections.singletonList(record)).isEmpty());
+        consumerCache.cleanupRecords(consumerRecord -> {
+        });
+        Mockito.verify(consumerOffsetManager).commitOffset(
+            anyString(), eq(groupId), eq(topicId), eq(queueId), eq(110L));
+
+        // the offset store has caught up, so the pending offset is not committed again
+        Mockito.when(consumerOffsetManager.queryOffset(groupId, topicId, queueId)).thenReturn(110L);
+        consumerCache.cleanupRecords(consumerRecord -> {
+        });
+        Mockito.verify(consumerOffsetManager, Mockito.times(1)).commitOffset(
+            anyString(), eq(groupId), eq(topicId), eq(queueId), eq(110L));
+
+        // reset offset wins over the pending offset
+        Mockito.when(consumerOffsetManager.hasOffsetReset(topicId, groupId, queueId)).thenReturn(true);
+        consumerCache.setPendingCommitOffset(groupId, topicId, queueId, 120L);
+        consumerCache.cleanupRecords(consumerRecord -> {
+        });
+        Mockito.verify(consumerOffsetManager, Mockito.never()).commitOffset(
+            anyString(), anyString(), anyString(), anyInt(), eq(120L));
+
+        // records of an offline consumer are persisted before the entry is dropped
+        Mockito.when(consumerOffsetManager.hasOffsetReset(topicId, groupId, queueId)).thenReturn(false);
+        Mockito.when(consumerLockService.isLockTimeout(any(), any())).thenReturn(true);
+        consumerCache.writeRecords(Collections.singletonList(new PopConsumerRecord(
+            System.currentTimeMillis(), groupId, topicId, queueId, 0, 20000, 200, attemptId)));
+        consumerCache.setPendingCommitOffset(groupId, topicId, queueId, 210L);
+        consumerCache.cleanupRecords(consumerRecord -> {
+        });
+        Mockito.verify(consumerKVStore).writeRecords(Mockito.argThat(records ->
+            records.size() == 1 && records.get(0).getOffset() == 200L));
+        Mockito.verify(consumerOffsetManager).commitOffset(
+            anyString(), eq(groupId), eq(topicId), eq(queueId), eq(210L));
+        Assert.assertEquals(0, consumerCache.getCacheKeySize());
     }
 }