[fix][broker] Fix NonDurable Subscription msgBackLog incorrect after topic unload

summeriiii · summeriiii · commit 6df5fa535c56 · 2024-09-23T15:19:32.000+08:00
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java
@@ -1232,6 +1232,7 @@ protected void handleSubscribe(final CommandSubscribe subscribe) {
         final long consumerEpoch = subscribe.hasConsumerEpoch() ? subscribe.getConsumerEpoch() : DEFAULT_CONSUMER_EPOCH;
         final Optional<Map<String, String>> subscriptionProperties = SubscriptionOption.getPropertiesMap(
                 subscribe.getSubscriptionPropertiesList());
+        final boolean resetIncludeHead = subscribe.isResetIncludeHead();
 
         if (log.isDebugEnabled()) {
             log.debug("Topic name = {}, subscription name = {}, schema is {}", topicName, subscriptionName,
@@ -1356,6 +1357,7 @@ protected void handleSubscribe(final CommandSubscribe subscribe) {
                                                 .subscriptionProperties(subscriptionProperties)
                                                 .consumerEpoch(consumerEpoch)
                                                 .schemaType(schema == null ? null : schema.getType())
+                                                .resetIncludeHead(resetIncludeHead)
                                                 .build();
                                         if (schema != null && schema.getType() != SchemaType.AUTO_CONSUME) {
                                             return ignoreUnrecoverableBKException
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SubscriptionOption.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SubscriptionOption.java
@@ -51,6 +51,7 @@ public class SubscriptionOption {
     private Optional<Map<String, String>> subscriptionProperties;
     private long consumerEpoch;
     private SchemaType schemaType;
+    private boolean resetIncludeHead;
 
     public static Optional<Map<String, String>> getPropertiesMap(List<KeyValue> list) {
         if (list == null) {
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentTopic.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentTopic.java
@@ -887,7 +887,7 @@ public CompletableFuture<Consumer> subscribe(SubscriptionOption option) {
                 option.getInitialPosition(), option.getStartMessageRollbackDurationSec(),
                 option.isReplicatedSubscriptionStateArg(), option.getKeySharedMeta(),
                 option.getSubscriptionProperties().orElse(Collections.emptyMap()),
-                option.getConsumerEpoch(), option.getSchemaType());
+                option.getConsumerEpoch(), option.getSchemaType(), option.isResetIncludeHead());
     }
 
     private CompletableFuture<Consumer> internalSubscribe(final TransportCnx cnx, String subscriptionName,
@@ -901,7 +901,8 @@ private CompletableFuture<Consumer> internalSubscribe(final TransportCnx cnx, St
                                                           KeySharedMeta keySharedMeta,
                                                           Map<String, String> subscriptionProperties,
                                                           long consumerEpoch,
-                                                          SchemaType schemaType) {
+                                                          SchemaType schemaType,
+                                                          boolean resetIncludeHead) {
         if (readCompacted && !(subType == SubType.Failover || subType == SubType.Exclusive)) {
             return FutureUtil.failedFuture(new NotAllowedException(
                     "readCompacted only allowed on failover or exclusive subscriptions"));
@@ -984,7 +985,7 @@ private CompletableFuture<Consumer> internalSubscribe(final TransportCnx cnx, St
                     ? getDurableSubscription(subscriptionName, initialPosition, startMessageRollbackDurationSec,
                             replicatedSubscriptionState, subscriptionProperties)
                     : getNonDurableSubscription(subscriptionName, startMessageId, initialPosition,
-                    startMessageRollbackDurationSec, readCompacted, subscriptionProperties);
+                    startMessageRollbackDurationSec, readCompacted, subscriptionProperties, resetIncludeHead);
 
             CompletableFuture<Consumer> future = subscriptionFuture.thenCompose(subscription -> {
                 Consumer consumer = new Consumer(subscription, subType, topic, consumerId, priorityLevel,
@@ -1068,7 +1069,7 @@ public CompletableFuture<Consumer> subscribe(final TransportCnx cnx, String subs
                                                  KeySharedMeta keySharedMeta) {
         return internalSubscribe(cnx, subscriptionName, consumerId, subType, priorityLevel, consumerName,
                 isDurable, startMessageId, metadata, readCompacted, initialPosition, startMessageRollbackDurationSec,
-                replicatedSubscriptionStateArg, keySharedMeta, null, DEFAULT_CONSUMER_EPOCH, null);
+                replicatedSubscriptionStateArg, keySharedMeta, null, DEFAULT_CONSUMER_EPOCH, null, false);
     }
 
     private CompletableFuture<Subscription> getDurableSubscription(String subscriptionName,
@@ -1134,7 +1135,7 @@ public void openCursorFailed(ManagedLedgerException exception, Object ctx) {
 
     private CompletableFuture<? extends Subscription> getNonDurableSubscription(String subscriptionName,
             MessageId startMessageId, InitialPosition initialPosition, long startMessageRollbackDurationSec,
-            boolean isReadCompacted, Map<String, String> subscriptionProperties) {
+            boolean isReadCompacted, Map<String, String> subscriptionProperties, boolean resetIncludeHead) {
         log.info("[{}][{}] Creating non-durable subscription at msg id {} - {}",
                 topic, subscriptionName, startMessageId, subscriptionProperties);
 
@@ -1157,7 +1158,8 @@ private CompletableFuture<? extends Subscription> getNonDurableSubscription(Stri
                 long entryId = msgId.getEntryId();
                 // Ensure that the start message id starts from a valid entry.
                 if (ledgerId >= 0 && entryId >= 0
-                        && msgId instanceof BatchMessageIdImpl) {
+                        && msgId instanceof BatchMessageIdImpl
+                        && (msgId.getBatchIndex() >= 0 || resetIncludeHead)) {
                     // When the start message is relative to a batch, we need to take one step back on the previous
                     // message,
                     // because the "batch" might not have been consumed in its entirety.
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/NonDurableSubscriptionTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/NonDurableSubscriptionTest.java
@@ -256,6 +256,57 @@ public void testNonDurableSubscriptionRecovery(SubscriptionType subscriptionType
 
     }
 
+    @Test
+    public void testNonDurableSubscriptionBackLogAfterTopicUnload() throws Exception {
+        String topicName = "persistent://my-property/my-ns/nonDurable-sub-test";
+        String subName = "test-sub";
+
+        admin.topics().createNonPartitionedTopic(topicName);
+        Producer<byte[]> producer = pulsarClient.newProducer().topic(topicName).create();
+
+        Consumer<byte[]> consumer = pulsarClient.newConsumer()
+                .topic(topicName)
+                .subscriptionName(subName)
+                .subscriptionMode(SubscriptionMode.NonDurable).subscribe();
+
+        // 1. send message
+        for (int i = 0; i < 10; i++) {
+            String message = "my-message-" + i;
+            producer.send(message.getBytes());
+        }
+        producer.close();
+
+        assertEquals(admin.topics().getStats(topicName).getSubscriptions().get(subName).getMsgBacklog(), 10);
+
+        // 2. receive the message
+        Thread t = new Thread(() -> {
+            while (true) {
+                Message<byte[]> msg;
+                try {
+                    msg = consumer.receive();
+                    consumer.acknowledge(msg);
+                } catch (PulsarClientException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        });
+        t.start();
+
+        // 3. consumed all messages and the msgBacklog is 0
+        Awaitility.await().untilAsserted(() ->
+                assertEquals(admin.topics().getStats(topicName).getSubscriptions().get(subName).getMsgBacklog(), 0));
+
+        // 4. unload the topic
+        admin.topics().unload(topicName);
+
+        // 5. wait the consumer reconnect
+        Awaitility.await().until(() -> admin.topics().getStats(topicName).getSubscriptions() != null);
+
+        // 6. the backlog should still be 0
+        Awaitility.await().untilAsserted(() -> assertEquals(admin.topics().getStats(topicName).
+                getSubscriptions().get(subName).getMsgBacklog(), 0));
+    }
+
     @Test
     public void testFlowCountForMultiTopics() throws Exception {
         String topicName = "persistent://my-property/my-ns/test-flow-count";
@@ -464,7 +515,7 @@ public void testInitReaderAtSpecifiedPosition() throws Exception {
         // A middle ledger id, and entry id is "-1".
         log.info("start test s8");
         String s8 = "s8";
-        MessageIdImpl startMessageId8 = new MessageIdImpl(ledgers.get(2), 0, -1);
+        MessageIdImpl startMessageId8 = new MessageIdImpl(ledgers.get(2), -1, -1);
         Reader<String> reader8 = pulsarClient.newReader(Schema.STRING).topic(topicName).subscriptionName(s8)
                 .receiverQueueSize(0).startMessageId(startMessageId8).create();
         ManagedLedgerInternalStats.CursorStats cursor8 = admin.topics().getInternalStats(topicName).cursors.get(s8);
@@ -497,7 +548,7 @@ public void testInitReaderAtSpecifiedPosition() throws Exception {
         ManagedLedgerInternalStats.CursorStats cursor10 = admin.topics().getInternalStats(topicName).cursors.get(s10);
         log.info("cursor10 readPosition: {}, markDeletedPosition: {}", cursor10.readPosition, cursor10.markDeletePosition);
         Position p10 = parseReadPosition(cursor10);
-        assertEquals(p10.getLedgerId(), ledgers.get(2));
+        assertEquals(p10.getLedgerId(), ledgers.get(3));
         assertEquals(p10.getEntryId(), 0);
         reader10.close();
 
diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java
@@ -895,7 +895,7 @@ public CompletableFuture<Void> connectionOpened(final ClientCnx cnx) {
                     InitialPosition.valueOf(subscriptionInitialPosition.getValue()),
                     startMessageRollbackDuration, si, createTopicIfDoesNotExist, conf.getKeySharedPolicy(),
                     // Use the current epoch to subscribe.
-                    conf.getSubscriptionProperties(), CONSUMER_EPOCH.get(this));
+                    conf.getSubscriptionProperties(), CONSUMER_EPOCH.get(this), resetIncludeHead);
 
             cnx.sendRequestWithId(request, requestId).thenRun(() -> {
                 synchronized (ConsumerImpl.this) {
diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java b/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java
@@ -588,15 +588,15 @@ public static ByteBuf newSubscribe(String topic, String subscription, long consu
         return newSubscribe(topic, subscription, consumerId, requestId, subType, priorityLevel, consumerName,
                 isDurable, startMessageId, metadata, readCompacted, isReplicated, subscriptionInitialPosition,
                 startMessageRollbackDurationInSec, schemaInfo, createTopicIfDoesNotExist, null,
-                Collections.emptyMap(), DEFAULT_CONSUMER_EPOCH);
+                Collections.emptyMap(), DEFAULT_CONSUMER_EPOCH, false);
     }
 
     public static ByteBuf newSubscribe(String topic, String subscription, long consumerId, long requestId,
                SubType subType, int priorityLevel, String consumerName, boolean isDurable, MessageIdData startMessageId,
                Map<String, String> metadata, boolean readCompacted, boolean isReplicated,
                InitialPosition subscriptionInitialPosition, long startMessageRollbackDurationInSec,
                SchemaInfo schemaInfo, boolean createTopicIfDoesNotExist, KeySharedPolicy keySharedPolicy,
-               Map<String, String> subscriptionProperties, long consumerEpoch) {
+               Map<String, String> subscriptionProperties, long consumerEpoch, boolean resetIncludeHead) {
         BaseCommand cmd = localCmd(Type.SUBSCRIBE);
         CommandSubscribe subscribe = cmd.setSubscribe()
                 .setTopic(topic)
@@ -611,7 +611,8 @@ public static ByteBuf newSubscribe(String topic, String subscription, long consu
                 .setInitialPosition(subscriptionInitialPosition)
                 .setReplicateSubscriptionState(isReplicated)
                 .setForceTopicCreation(createTopicIfDoesNotExist)
-                .setConsumerEpoch(consumerEpoch);
+                .setConsumerEpoch(consumerEpoch)
+                .setResetIncludeHead(resetIncludeHead);
 
         if (subscriptionProperties != null && !subscriptionProperties.isEmpty()) {
             List<KeyValue> keyValues = new ArrayList<>();
diff --git a/pulsar-common/src/main/proto/PulsarApi.proto b/pulsar-common/src/main/proto/PulsarApi.proto
@@ -401,6 +401,8 @@ message CommandSubscribe {
 
     // The consumer epoch, when exclusive and failover consumer redeliver unack message will increase the epoch
     optional uint64 consumer_epoch = 19;
+
+    optional bool reset_include_head = 20 [default = false];
 }
 
 message CommandPartitionedTopicMetadata {

Original file line number	Diff line number	Diff line change
`@@ -401,6 +401,8 @@ message CommandSubscribe {`
`401`	`401`
`402`	`402`	`// The consumer epoch, when exclusive and failover consumer redeliver unack message will increase the epoch`
`403`	`403`	`optional uint64 consumer_epoch = 19;`
	`404`	`+`
	`405`	`+ optional bool reset_include_head = 20 [default = false];`
`404`	`406`	`}`
`405`	`407`
`406`	`408`	`message CommandPartitionedTopicMetadata {`