[Bugfix] Preserve retained subscription progress after the last consumer closes

integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/IoTDBConsensusSubscriptionBasicTableIT.java

@@ -19,8 +19,12 @@
 
 package org.apache.iotdb.subscription.it.consensus.local.tablemodel;
 
+import org.apache.iotdb.it.env.EnvFactory;
 import org.apache.iotdb.it.framework.IoTDBTestRunner;
 import org.apache.iotdb.itbase.category.LocalStandaloneIT;
+import org.apache.iotdb.rpc.subscription.config.TopicConstant;
+import org.apache.iotdb.session.subscription.ISubscriptionTableSession;
+import org.apache.iotdb.session.subscription.SubscriptionTableSessionBuilder;
 import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumer;
 import org.apache.iotdb.subscription.it.consensus.local.AbstractSubscriptionConsensusLocalIT;
 
@@ -31,12 +35,68 @@
 
 import java.time.Duration;
 import java.util.LinkedHashSet;
+import java.util.Properties;
 import java.util.Set;
 
 @RunWith(IoTDBTestRunner.class)
 @Category({LocalStandaloneIT.class})
 public class IoTDBConsensusSubscriptionBasicTableIT extends AbstractSubscriptionConsensusLocalIT {
 
+  @Test
+  public void testLastConsumerCloseRetainsProgress() throws Exception {
+    for (final boolean unsubscribeFirst : new boolean[] {false, true}) {
+      final ConsensusSubscriptionTableITSupport.TestIdentifiers ids =
+          ConsensusSubscriptionTableITSupport.newIdentifiers("retained_last_close");
+      SubscriptionTablePullConsumer consumer = null;
+      try {
+        ConsensusSubscriptionTableITSupport.bootstrapDatabaseAndTable(
+            ids.getDatabase(), "t1", ConsensusSubscriptionTableITSupport.DEFAULT_TABLE_SCHEMA);
+        ConsensusSubscriptionTableITSupport.createConsensusTopic(
+            ids.getTopic(), ids.getDatabase(), "t1");
+        try (final ISubscriptionTableSession session =
+            new SubscriptionTableSessionBuilder()
+                .host(EnvFactory.getEnv().getIP())
+                .port(Integer.parseInt(EnvFactory.getEnv().getPort()))
+                .build()) {
+          session.open();
+          final Properties config = new Properties();
+          config.setProperty(TopicConstant.RETAIN_PROGRESS_AFTER_UNSUBSCRIBE_KEY, "true");
+          session.alterTopic(ids.getTopic(), config);
+        }
+        consumer =
+            ConsensusSubscriptionTableITSupport.createConsumer(
+                ids.getConsumerId(), ids.getConsumerGroupId());
+        consumer.subscribe(ids.getTopic());
+        final Set<String> firstRows =
+            ConsensusSubscriptionTableITSupport.insertRows(ids.getDatabase(), "t1", 100L, 10, true);
+        ConsensusSubscriptionTableITSupport.assertExactRowKeys(
+            firstRows,
+            ConsensusSubscriptionTableITSupport.pollAndCommitUntilAtLeast(
+                consumer, firstRows.size(), 40));
+        if (unsubscribeFirst) {
+          consumer.unsubscribe(ids.getTopic());
+        }
+        consumer.close();
+        consumer = null;
+
+        final Set<String> offlineRows =
+            ConsensusSubscriptionTableITSupport.insertRows(ids.getDatabase(), "t1", 200L, 10, true);
+        consumer =
+            ConsensusSubscriptionTableITSupport.createConsumer(
+                ids.consumer("replacement"), ids.getConsumerGroupId());
+        consumer.subscribe(ids.getTopic());
+        ConsensusSubscriptionTableITSupport.assertExactRowKeys(
+            offlineRows,
+            ConsensusSubscriptionTableITSupport.pollAndCommitUntilAtLeast(
+                consumer, offlineRows.size(), 40));
+        ConsensusSubscriptionTableITSupport.assertNoMoreMessages(
+            consumer, 3, Duration.ofMillis(500));
+      } finally {
+        ConsensusSubscriptionTableITSupport.cleanup(consumer, ids.getTopic(), ids.getDatabase());
+      }
+    }
+  }
+
   @Test
   public void testRealtimeOnlyAfterSubscribe() throws Exception {
     final ConsensusSubscriptionTableITSupport.TestIdentifiers ids =

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/subscription/SubscriptionInfo.java

@@ -929,6 +929,7 @@ private TSStatus alterTopicInternal(final AlterTopicPlan plan) {
                 plan.getTopicMeta().getTopicName())) {
           commitProgressKeeper.removeTopicProgress(
               consumerGroupMeta.getConsumerGroupId(), plan.getTopicMeta().getTopicName());
+          removeEmptyConsumerGroupWithoutProgress(consumerGroupMeta);
         }
       }
     }
@@ -971,6 +972,7 @@ public TSStatus dropTopic(DropTopicPlan plan) {
         if (!plan.isTableModelSet() || consumerGroupMeta.visibleUnder(plan.isTableModel())) {
           commitProgressKeeper.removeTopicProgress(
               consumerGroupMeta.getConsumerGroupId(), plan.getTopicName());
+          removeEmptyConsumerGroupWithoutProgress(consumerGroupMeta);
         }
       }
       if (plan.isTableModelSet()) {
@@ -1146,24 +1148,32 @@ public TSStatus alterConsumerGroup(AlterConsumerGroupPlan plan) {
                 ? consumerGroupMetaKeeper.getConsumerGroupMeta(consumerGroupId)
                 : null;
         if (Objects.nonNull(currentConsumerGroupMeta)) {
+          ConsumerGroupMeta.getTopicsUnsubByGroup(currentConsumerGroupMeta, consumerGroupMeta)
+              .forEach(
+                  topicName -> {
+                    final TopicMeta topicMeta =
+                        topicMetaKeeper.getTopicMeta(
+                            topicName, currentConsumerGroupMeta.visibleUnder(true));
+                    if (Objects.isNull(topicMeta)
+                        || !topicMeta.getConfig().isProgressRetainedAfterUnsubscribe()) {
+                      commitProgressKeeper.removeTopicProgress(consumerGroupId, topicName);
+                    }
+                  });
           if (consumerGroupMeta.isEmpty()) {
-            commitProgressKeeper.removeConsumerGroupProgress(consumerGroupId);
-          } else {
-            ConsumerGroupMeta.getTopicsUnsubByGroup(currentConsumerGroupMeta, consumerGroupMeta)
-                .forEach(
-                    topicName -> {
-                      final TopicMeta topicMeta =
-                          topicMetaKeeper.getTopicMeta(
-                              topicName, currentConsumerGroupMeta.visibleUnder(true));
-                      if (Objects.isNull(topicMeta)
-                          || !topicMeta.getConfig().isProgressRetainedAfterUnsubscribe()) {
-                        commitProgressKeeper.removeTopicProgress(consumerGroupId, topicName);
-                      }
-                    });
+            final Set<String> retainedTopicNames = new HashSet<>();
+            for (final TopicMeta topicMeta : topicMetaKeeper.getAllTopicMeta()) {
+              if (topicMeta.visibleUnderTableModel() == currentConsumerGroupMeta.visibleUnder(true)
+                  && topicMeta.getConfig().isProgressRetainedAfterUnsubscribe()) {
+                retainedTopicNames.add(topicMeta.getTopicName());
+              }
+            }
+            commitProgressKeeper.removeConsumerGroupProgressExceptTopics(
+                consumerGroupId, retainedTopicNames);
           }
         }
         consumerGroupMetaKeeper.removeConsumerGroupMeta(consumerGroupId);
-        if (!consumerGroupMeta.isEmpty()) {
+        if (!consumerGroupMeta.isEmpty()
+            || commitProgressKeeper.hasConsumerGroupProgress(consumerGroupId)) {
           consumerGroupMetaKeeper.addConsumerGroupMeta(consumerGroupId, consumerGroupMeta);
         }
       }
@@ -1173,6 +1183,13 @@ public TSStatus alterConsumerGroup(AlterConsumerGroupPlan plan) {
     }
   }
 
+  private void removeEmptyConsumerGroupWithoutProgress(final ConsumerGroupMeta consumerGroupMeta) {
+    if (consumerGroupMeta.isEmpty()
+        && !commitProgressKeeper.hasConsumerGroupProgress(consumerGroupMeta.getConsumerGroupId())) {
+      consumerGroupMetaKeeper.removeConsumerGroupMeta(consumerGroupMeta.getConsumerGroupId());
+    }
+  }
+
   public TSStatus handleConsumerGroupMetaChanges(ConsumerGroupHandleMetaChangePlan plan) {
     acquireWriteLock();
     try {

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/subscription/SubscriptionInfoTest.java

@@ -34,10 +34,13 @@
 import org.apache.iotdb.confignode.rpc.thrift.TShowTopicResp;
 import org.apache.iotdb.mpp.rpc.thrift.TTopicOwnerLeaseEntry;
 import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.rpc.subscription.config.ConsumerConstant;
 import org.apache.iotdb.rpc.subscription.config.TopicConstant;
 
 import org.junit.Assert;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
 
 import java.nio.ByteBuffer;
 import java.util.Collections;
@@ -49,6 +52,102 @@
 
 public class SubscriptionInfoTest {
 
+  @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+  @Test
+  public void testLastConsumerClosePreservesConfiguredProgressAcrossSnapshot() throws Exception {
+    for (final boolean tableModel : new boolean[] {false, true}) {
+      for (final boolean unsubscribeFirst : new boolean[] {false, true}) {
+        assertLastConsumerClosePreservesProgress(tableModel, unsubscribeFirst);
+      }
+    }
+  }
+
+  private void assertLastConsumerClosePreservesProgress(
+      final boolean tableModel, final boolean unsubscribeFirst) throws Exception {
+    final SubscriptionInfo info = new SubscriptionInfo();
+    final Map<String, String> attributes = new HashMap<>();
+    attributes.put(TopicConstant.MODE_KEY, TopicConstant.MODE_INCREMENTAL_VALUE);
+    attributes.put(TopicConstant.RETAIN_PROGRESS_AFTER_UNSUBSCRIBE_KEY, "true");
+    attributes.put(SystemConstant.SQL_DIALECT_KEY, tableModel ? "table" : "tree");
+    info.createTopic(new CreateTopicPlan(new TopicMeta("kept", 1L, attributes)));
+    final Map<String, String> discardedAttributes = new HashMap<>(attributes);
+    discardedAttributes.put(TopicConstant.RETAIN_PROGRESS_AFTER_UNSUBSCRIBE_KEY, "false");
+    info.createTopic(new CreateTopicPlan(new TopicMeta("discarded", 1L, discardedAttributes)));
+    final ConsumerGroupMeta group =
+        new ConsumerGroupMeta(
+            "group",
+            1L,
+            new ConsumerMeta(
+                "consumer",
+                1L,
+                Collections.singletonMap(
+                    ConsumerConstant.SQL_DIALECT_KEY, tableModel ? "table" : "tree")));
+    group.addSubscription(
+        "consumer", new java.util.HashSet<>(java.util.Arrays.asList("kept", "discarded")));
+    info.alterConsumerGroup(new AlterConsumerGroupPlan(group));
+    final String keptKey = CommitProgressKeeper.generateKey("group", "kept", "DataRegion[1]", 1);
+    final String legacyKey =
+        CommitProgressKeeper.generateLegacyKey("group", "kept", "DataRegion[1]", 2);
+    final String discardedKey =
+        CommitProgressKeeper.generateKey("group", "discarded", "DataRegion[1]", 1);
+    for (final String key : new String[] {keptKey, legacyKey, discardedKey}) {
+      info.getCommitProgressKeeper().updateRegionProgress(key, ByteBuffer.wrap(new byte[] {1}));
+    }
+    final ConsumerGroupMeta updated = group.deepCopy();
+    if (unsubscribeFirst) {
+      updated.removeSubscription(
+          "consumer", new java.util.HashSet<>(java.util.Arrays.asList("kept", "discarded")));
+      info.alterConsumerGroup(new AlterConsumerGroupPlan(updated.deepCopy()));
+    }
+    updated.removeConsumer("consumer");
+    info.alterConsumerGroup(new AlterConsumerGroupPlan(updated));
+    Assert.assertNotNull(
+        "tableModel=" + tableModel + ", unsubscribeFirst=" + unsubscribeFirst,
+        info.getCommitProgressKeeper().getRegionProgress(keptKey));
+    Assert.assertNotNull(info.getCommitProgressKeeper().getRegionProgress(legacyKey));
+    Assert.assertNull(info.getCommitProgressKeeper().getRegionProgress(discardedKey));
+    Assert.assertTrue(info.getConsumerGroupMeta("group").isEmpty());
+    Assert.assertTrue(info.getConsumerGroupMeta("group").visibleUnder(tableModel));
+
+    final java.io.File snapshotDir = temporaryFolder.newFolder();
+    Assert.assertTrue(info.processTakeSnapshot(snapshotDir));
+    final SubscriptionInfo restored = new SubscriptionInfo();
+    restored.processLoadSnapshot(snapshotDir);
+    Assert.assertNotNull(restored.getCommitProgressKeeper().getRegionProgress(keptKey));
+    final ConsumerGroupMeta restoredGroup = restored.deepCopyConsumerGroupMeta("group");
+    Assert.assertTrue(restoredGroup.isEmpty());
+    Assert.assertTrue(restoredGroup.visibleUnder(tableModel));
+    final ConsumerMeta replacement =
+        new ConsumerMeta(
+            "replacement",
+            2L,
+            Collections.singletonMap(
+                ConsumerConstant.SQL_DIALECT_KEY, tableModel ? "table" : "tree"));
+    restoredGroup.checkAuthorityBeforeJoinConsumerGroup(replacement);
+    restoredGroup.addConsumer(replacement);
+    restored.alterConsumerGroup(new AlterConsumerGroupPlan(restoredGroup));
+    Assert.assertNotNull(restored.getCommitProgressKeeper().getRegionProgress(keptKey));
+    restoredGroup.removeConsumer("replacement");
+    restored.alterConsumerGroup(new AlterConsumerGroupPlan(restoredGroup));
+
+    // Same-name topics in the other model must not remove this retained group's progress.
+    final Map<String, String> otherAttributes = new HashMap<>(attributes);
+    otherAttributes.put(SystemConstant.SQL_DIALECT_KEY, tableModel ? "tree" : "table");
+    restored.createTopic(new CreateTopicPlan(new TopicMeta("kept", 2L, otherAttributes)));
+    restored.dropTopic(new DropTopicPlan("kept", !tableModel));
+    Assert.assertNotNull(restored.getCommitProgressKeeper().getRegionProgress(keptKey));
+    if (unsubscribeFirst) {
+      attributes.put(TopicConstant.RETAIN_PROGRESS_AFTER_UNSUBSCRIBE_KEY, "false");
+      restored.alterTopic(new AlterTopicPlan(new TopicMeta("kept", 2L, attributes)));
+    } else {
+      restored.dropTopic(new DropTopicPlan("kept", tableModel));
+    }
+    Assert.assertNull(restored.getCommitProgressKeeper().getRegionProgress(keptKey));
+    Assert.assertNull(restored.getCommitProgressKeeper().getRegionProgress(legacyKey));
+    Assert.assertFalse(restored.isConsumerGroupExisted("group"));
+  }
+
   @Test
   public void testSameNameTreeAndTableTopicsCanBeManagedIndependently() {
     final String topicName = "same-name-topic";

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/IoTConsensusServerImpl.java

@@ -1267,7 +1267,20 @@ public void registerSubscriptionQueue(
       final BlockingQueue<IndexedConsensusRequest> queue,
       final SubscriptionWalRetentionPolicy retentionPolicy,
       final LongSupplier retainedMinVersionIdSupplier) {
-    subscriptionQueueRegistry.register(queue, retentionPolicy, retainedMinVersionIdSupplier);
+    registerSubscriptionQueue(queue, retentionPolicy, retainedMinVersionIdSupplier, null);
+  }
+
+  public void registerSubscriptionQueue(
+      final BlockingQueue<IndexedConsensusRequest> queue,
+      final SubscriptionWalRetentionPolicy retentionPolicy,
+      final LongSupplier retainedMinVersionIdSupplier,
+      final String retentionId) {
+    if (retentionId == null) {
+      subscriptionQueueRegistry.register(queue, retentionPolicy, retainedMinVersionIdSupplier);
+    } else {
+      subscriptionQueueRegistry.register(
+          queue, retentionPolicy, retainedMinVersionIdSupplier, retentionId);
+    }
     // Immediately re-evaluate the safe delete index with new subscription awareness
     checkAndUpdateSafeDeletedSearchIndex();
     logger.info(