|
59 | 59 | import org.apache.pulsar.broker.service.StickyKeyDispatcher; |
60 | 60 | import org.apache.pulsar.broker.service.Subscription; |
61 | 61 | import org.apache.pulsar.broker.service.Topic; |
| 62 | +import org.apache.pulsar.broker.service.persistent.AbstractPersistentDispatcherMultipleConsumers; |
62 | 63 | import org.apache.pulsar.broker.service.persistent.PersistentTopic; |
63 | 64 | import org.apache.pulsar.broker.service.plugin.EntryFilter; |
64 | 65 | import org.apache.pulsar.broker.service.plugin.EntryFilterProducerTest; |
@@ -106,6 +107,7 @@ protected void setup() throws Exception { |
106 | 107 | @Override |
107 | 108 | protected ServiceConfiguration getDefaultConf() { |
108 | 109 | ServiceConfiguration conf = super.getDefaultConf(); |
| 110 | + conf.setAcknowledgmentAtBatchIndexLevelEnabled(true); |
109 | 111 | conf.setMaxUnackedMessagesPerConsumer(0); |
110 | 112 | // wait for shutdown of the broker, this prevents flakiness which could be caused by metrics being |
111 | 113 | // unregistered asynchronously. This impacts the execution of the next test method if this would be happening. |
@@ -731,6 +733,170 @@ public void testKeySharedDrainingHashesConsumerStats() throws Exception { |
731 | 733 |
|
732 | 734 | } |
733 | 735 |
|
| 736 | + @DataProvider(name = "subscriptionTypes") |
| 737 | + public Object[][] subscriptionTypes() { |
| 738 | + return new Object[][]{ |
| 739 | + {SubscriptionType.Shared}, |
| 740 | + {SubscriptionType.Key_Shared} |
| 741 | + }; |
| 742 | + } |
| 743 | + |
| 744 | + /** |
| 745 | + * Verify unacked count is correctly decremented when removeAllUpTo removes non-batch |
| 746 | + * entries from pendingAcks after mark-delete advances via message expiry. |
| 747 | + */ |
| 748 | + @Test(dataProvider = "subscriptionTypes") |
| 749 | + public void testUnackedCountNonBatchAfterExpire(SubscriptionType subType) throws Exception { |
| 750 | + String topic = newTopicName(); |
| 751 | + String sub = "sub"; |
| 752 | + int numMessages = 10; |
| 753 | + |
| 754 | + @Cleanup Producer<byte[]> producer = pulsarClient.newProducer() |
| 755 | + .topic(topic).enableBatching(false).create(); |
| 756 | + @Cleanup Consumer<byte[]> consumer = pulsarClient.newConsumer() |
| 757 | + .topic(topic).subscriptionName(sub) |
| 758 | + .subscriptionType(subType) |
| 759 | + .subscribe(); |
| 760 | + |
| 761 | + for (int i = 0; i < numMessages; i++) { |
| 762 | + producer.send(("msg-" + i).getBytes()); |
| 763 | + } |
| 764 | + |
| 765 | + org.apache.pulsar.broker.service.Consumer svcConsumer = |
| 766 | + getTheUniqueServiceConsumer(topic, sub); |
| 767 | + for (int i = 0; i < numMessages; i++) { |
| 768 | + Message<byte[]> msg = consumer.receive(2, TimeUnit.SECONDS); |
| 769 | + Assert.assertNotNull(msg, "Expected to receive message " + i); |
| 770 | + } |
| 771 | + |
| 772 | + Awaitility.await().untilAsserted(() -> |
| 773 | + assertEquals(numMessages, svcConsumer.getUnackedMessages())); |
| 774 | + |
| 775 | + expireAndVerifyUnackedDrained(topic, sub, producer, consumer, svcConsumer); |
| 776 | + } |
| 777 | + |
| 778 | + /** |
| 779 | + * Verify unacked count is correctly decremented when removeAllUpTo removes batch |
| 780 | + * entries from pendingAcks after mark-delete advances via message expiry. |
| 781 | + */ |
| 782 | + @Test(dataProvider = "subscriptionTypes") |
| 783 | + public void testUnackedCountBatchAfterExpire(SubscriptionType subType) throws Exception { |
| 784 | + String topic = newTopicName(); |
| 785 | + String sub = "sub"; |
| 786 | + int numMessages = 10; |
| 787 | + |
| 788 | + @Cleanup Producer<byte[]> producer = pulsarClient.newProducer() |
| 789 | + .topic(topic) |
| 790 | + .batchingMaxMessages(20) |
| 791 | + .batchingMaxPublishDelay(1, TimeUnit.HOURS) |
| 792 | + .enableBatching(true) |
| 793 | + .create(); |
| 794 | + @Cleanup Consumer<byte[]> consumer = pulsarClient.newConsumer() |
| 795 | + .topic(topic).subscriptionName(sub) |
| 796 | + .subscriptionType(subType) |
| 797 | + .subscribe(); |
| 798 | + |
| 799 | + for (int i = 0; i < numMessages; i++) { |
| 800 | + producer.newMessage().value(("batch-" + i).getBytes()).sendAsync(); |
| 801 | + } |
| 802 | + producer.flush(); |
| 803 | + |
| 804 | + for (int i = 0; i < numMessages; i++) { |
| 805 | + Message<byte[]> msg = consumer.receive(2, TimeUnit.SECONDS); |
| 806 | + Assert.assertNotNull(msg, "Expected to receive message " + i); |
| 807 | + } |
| 808 | + |
| 809 | + org.apache.pulsar.broker.service.Consumer svcConsumer = |
| 810 | + getTheUniqueServiceConsumer(topic, sub); |
| 811 | + |
| 812 | + Awaitility.await().untilAsserted(() -> |
| 813 | + assertEquals(numMessages, svcConsumer.getUnackedMessages())); |
| 814 | + |
| 815 | + expireAndVerifyUnackedDrained(topic, sub, producer, consumer, svcConsumer); |
| 816 | + } |
| 817 | + |
| 818 | + /** |
| 819 | + * Verify unacked count is correctly decremented when removeAllUpTo removes a partially-acked |
| 820 | + * batch entry from pendingAcks after mark-delete advances via message expiry. |
| 821 | + * |
| 822 | + * <p>Flow: produce batch(batchSize=10) → consume all → ack 5 of 10 → expire → unacked should be 0. |
| 823 | + */ |
| 824 | + @Test(dataProvider = "subscriptionTypes") |
| 825 | + public void testUnackedCountBatchPartialAckAfterExpire(SubscriptionType subType) throws Exception { |
| 826 | + String topic = newTopicName(); |
| 827 | + String sub = "sub"; |
| 828 | + int numMessages = 10; |
| 829 | + int ackCount = 5; |
| 830 | + |
| 831 | + @Cleanup Producer<byte[]> producer = pulsarClient.newProducer() |
| 832 | + .topic(topic) |
| 833 | + .batchingMaxMessages(20) |
| 834 | + .batchingMaxPublishDelay(1, TimeUnit.HOURS) |
| 835 | + .enableBatching(true) |
| 836 | + .create(); |
| 837 | + @Cleanup Consumer<byte[]> consumer = pulsarClient.newConsumer() |
| 838 | + .topic(topic) |
| 839 | + .subscriptionName(sub) |
| 840 | + .enableBatchIndexAcknowledgment(true) |
| 841 | + .subscriptionType(subType) |
| 842 | + .subscribe(); |
| 843 | + |
| 844 | + for (int i = 0; i < numMessages; i++) { |
| 845 | + producer.newMessage().value(("batch-" + i).getBytes()).sendAsync(); |
| 846 | + } |
| 847 | + producer.flush(); |
| 848 | + |
| 849 | + List<Message<byte[]>> messages = new ArrayList<>(); |
| 850 | + for (int i = 0; i < numMessages; i++) { |
| 851 | + Message<byte[]> msg = consumer.receive(2, TimeUnit.SECONDS); |
| 852 | + Assert.assertNotNull(msg, "Expected to receive message " + i); |
| 853 | + messages.add(msg); |
| 854 | + } |
| 855 | + |
| 856 | + org.apache.pulsar.broker.service.Consumer svcConsumer = |
| 857 | + getTheUniqueServiceConsumer(topic, sub); |
| 858 | + |
| 859 | + Awaitility.await().untilAsserted(() -> |
| 860 | + assertEquals(numMessages, svcConsumer.getUnackedMessages())); |
| 861 | + |
| 862 | + // Partially ack — ack 5 of 10 batch indexes |
| 863 | + for (int i = 0; i < ackCount; i++) { |
| 864 | + consumer.acknowledge(messages.get(i)); |
| 865 | + } |
| 866 | + Awaitility.await().untilAsserted(() -> |
| 867 | + assertEquals(numMessages - ackCount, svcConsumer.getUnackedMessages())); |
| 868 | + |
| 869 | + expireAndVerifyUnackedDrained(topic, sub, producer, consumer, svcConsumer); |
| 870 | + } |
| 871 | + |
| 872 | + private void expireAndVerifyUnackedDrained(String topic, String sub, |
| 873 | + Producer<byte[]> producer, Consumer<byte[]> consumer, |
| 874 | + org.apache.pulsar.broker.service.Consumer svcConsumer) |
| 875 | + throws Exception { |
| 876 | + PersistentTopic pTopic = (PersistentTopic) pulsar.getBrokerService() |
| 877 | + .getTopicReference(topic).get(); |
| 878 | + |
| 879 | + Thread.sleep(1100); |
| 880 | + pTopic.getSubscription(sub).expireMessagesAsync(1).get(); |
| 881 | + |
| 882 | + // Trigger readMoreEntries to invoke removeAllUpTo |
| 883 | + producer.send("trigger".getBytes()); |
| 884 | + Message<byte[]> triggerMsg = consumer.receive(2, TimeUnit.SECONDS); |
| 885 | + Assert.assertNotNull(triggerMsg); |
| 886 | + consumer.acknowledge(triggerMsg); |
| 887 | + |
| 888 | + Awaitility.await().atMost(5, TimeUnit.SECONDS).untilAsserted(() -> |
| 889 | + assertEquals(0, svcConsumer.getUnackedMessages())); |
| 890 | + } |
| 891 | + |
| 892 | + private org.apache.pulsar.broker.service.Consumer getTheUniqueServiceConsumer(String topic, String sub) { |
| 893 | + PersistentTopic persistentTopic = |
| 894 | + (PersistentTopic) pulsar.getBrokerService().getTopic(topic, false).join().get(); |
| 895 | + AbstractPersistentDispatcherMultipleConsumers dispatcher = |
| 896 | + (AbstractPersistentDispatcherMultipleConsumers) persistentTopic.getSubscription(sub).getDispatcher(); |
| 897 | + return dispatcher.getConsumers().iterator().next(); |
| 898 | + } |
| 899 | + |
734 | 900 | private String findConsumerNameForHash(SubscriptionStats subscriptionStats, int hash) { |
735 | 901 | return findConsumerForHash(subscriptionStats, hash).map(ConsumerStats::getConsumerName).orElse(null); |
736 | 902 | } |
|
0 commit comments