[Draft] [fix] [broker] Add a limitation of max memory usage of the replay queue

pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java

@@ -873,6 +873,13 @@ The delayed message index time step(in seconds) in per bucket snapshot segment,
             + " to `limit/2`. Using a value of 0, is disabling unackedMessage-limit check and dispatcher"
             + " can dispatch messages without any restriction")
     private int maxUnackedMessagesPerSubscription = 4 * 50000;
+    @FieldContext(
+            category = CATEGORY_POLICIES,
+            doc = "Max memory usage of the replay queue each subscription in bytes, Pulsar will stop dispatching"
+                    + " messages to client if memory usage is larger than expected. If this value is less than zero,"
+                    + " it means disabled this limitation. default: 5m.")
+    // TODO There needs to be a minimum limit.
+    private long maxMemoryUsageOfReplayQueueInBytesPerSubscription = 5 * 1024 * 1024;
     @FieldContext(
         category = CATEGORY_POLICIES,
         doc = "Max number of unacknowledged messages allowed per broker. \n\n"

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/MessageRedeliveryController.java

@@ -55,6 +55,12 @@ public MessageRedeliveryController(boolean allowOutOfOrderDelivery) {
         }
     }
 
+    public long getLongSizeInBytes() {
+        // TODO add method `getLongSizeInBytes` for `hashesRefCount`.
+        // TODO add test for method `getLongSizeInBytes`.
+        return messagesToRedeliver.getLongSizeInBytes() + hashesToBeBlocked.getLongSizeInBytes();
+    }
+
     public void add(long ledgerId, long entryId) {
         messagesToRedeliver.add(ledgerId, entryId);
     }

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java

@@ -60,6 +60,7 @@ public class PersistentStickyKeyDispatcherMultipleConsumers extends PersistentDi
 
     private final boolean allowOutOfOrderDelivery;
     private final StickyKeyConsumerSelector selector;
+    private final long maxMemoryUsageOfReplayQueueInBytes;
 
     private boolean isDispatcherStuckOnReplays = false;
     private final KeySharedMode keySharedMode;
@@ -83,6 +84,7 @@ public class PersistentStickyKeyDispatcherMultipleConsumers extends PersistentDi
         this.stuckConsumers = new HashSet<>();
         this.nextStuckConsumers = new HashSet<>();
         this.keySharedMode = ksm.getKeySharedMode();
+        this.maxMemoryUsageOfReplayQueueInBytes = conf.getMaxMemoryUsageOfReplayQueueInBytesPerSubscription();
         switch (this.keySharedMode) {
         case AUTO_SPLIT:
             if (conf.isSubscriptionKeySharedUseConsistentHashing()) {
@@ -102,6 +104,15 @@ public class PersistentStickyKeyDispatcherMultipleConsumers extends PersistentDi
         }
     }
 
+    @Override
+    public synchronized void readMoreEntries() {
+        if (isDispatcherStuckOnReplays && maxMemoryUsageOfReplayQueueInBytes > 0
+                && redeliveryMessages.getLongSizeInBytes() > maxMemoryUsageOfReplayQueueInBytes) {
+            return;
+        }
+        super.readMoreEntries();
+    }
+
     @VisibleForTesting
     public StickyKeyConsumerSelector getSelector() {
         return selector;

pulsar-broker/src/main/java/org/apache/pulsar/utils/ConcurrentBitmapSortedLongPairSet.java

@@ -69,6 +69,18 @@ public boolean contains(long item1, long item2) {
         }
     }
 
+    public long getLongSizeInBytes() {
+        lock.readLock().lock();
+        try {
+            long valuesSize = map.values().stream().map(RoaringBitmap::getLongSizeInBytes)
+                    .reduce((l1, l2) -> l1 + l2).orElse(0L);
+            long keysSize = map.size() * 8L;
+            return valuesSize + keysSize;
+        } finally {
+            lock.readLock().unlock();
+        }
+    }
+
     public void removeUpTo(long item1, long item2) {
         lock.writeLock().lock();
         try {

pulsar-broker/src/test/java/org/apache/pulsar/client/api/MaxUnAckMessagesTest.java

@@ -0,0 +1,195 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pulsar.client.api;
+
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import java.util.Collections;
+import java.util.LinkedHashSet;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import lombok.Cleanup;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.pulsar.broker.BrokerTestUtil;
+import org.apache.pulsar.broker.service.persistent.MessageRedeliveryController;
+import org.apache.pulsar.broker.service.persistent.PersistentStickyKeyDispatcherMultipleConsumers;
+import org.apache.pulsar.broker.service.persistent.PersistentTopic;
+import org.apache.pulsar.client.impl.ConsumerImpl;
+import org.apache.pulsar.common.util.collections.ConcurrentLongLongPairHashMap;
+import org.awaitility.Awaitility;
+import org.awaitility.reflect.WhiteboxImpl;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+@Slf4j
+@Test(groups = "broker-api")
+public class MaxUnAckMessagesTest  extends ProducerConsumerBase {
+
+    private int maxUnackedMessagesPerSubscription = 20;
+
+    @BeforeClass(alwaysRun = true)
+    @Override
+    protected void setup() throws Exception {
+        super.internalSetup();
+        super.producerBaseSetup();
+    }
+
+    @AfterClass(alwaysRun = true)
+    @Override
+    protected void cleanup() throws Exception {
+        super.internalCleanup();
+    }
+
+    @Override
+    protected void doInitConf() throws Exception {
+        conf.setMaxUnackedMessagesPerSubscription(maxUnackedMessagesPerSubscription);
+    }
+
+    private String uniquePersistentTopicName() {
+        return BrokerTestUtil.newUniqueName("persistent://my-property/my-ns/tp_");
+    }
+
+    @Test
+    public void testLimitationTurnOnAndTurnOff() throws Exception {
+        final String topicName = uniquePersistentTopicName();
+        final String subName = "sub1";
+        final int incomingQueueCapacity = maxUnackedMessagesPerSubscription * 2;
+        admin.topics().createNonPartitionedTopic(topicName);
+        admin.topics().createSubscription(topicName, subName, MessageId.earliest);
+
+        @Cleanup
+        ConsumerImpl<String> consumer = (ConsumerImpl<String>) pulsarClient.newConsumer(Schema.STRING)
+                .topic(topicName)
+                .subscriptionName(subName)
+                .subscriptionType(SubscriptionType.Key_Shared)
+                .receiverQueueSize(incomingQueueCapacity)
+                .subscribe();
+        @Cleanup
+        Producer<String> producer = pulsarClient.newProducer(Schema.STRING)
+                .enableBatching(false)
+                .topic(topicName)
+                .create();
+
+        // Send messages;
+        int messageCount = incomingQueueCapacity * 2;
+        for (int i = 0; i < messageCount; i++) {
+            producer.newMessage().key("" + i).value("" + i).send();
+        }
+
+        // Wait the limitation of maxUnackedMessagesPerSubscription is turned on.
+        Awaitility.await().untilAsserted(() -> {
+            log.info("consumer.numMessagesInQueue: {}", consumer.numMessagesInQueue());
+            assertTrue(consumer.numMessagesInQueue() >= maxUnackedMessagesPerSubscription);
+        });
+
+        // Verify consumer can receive all the messages.
+        Set<String> receivedMessages = Collections.synchronizedSet(new LinkedHashSet<>());
+        Awaitility.await().atMost(30, TimeUnit.SECONDS).untilAsserted(() -> {
+            Message<String> message = consumer.receive(2, TimeUnit.SECONDS);
+            assertNotNull(message);
+            log.info("received message: {}", message.getValue());
+            receivedMessages.add(message.getValue());
+            consumer.acknowledge(message);
+            assertEquals(receivedMessages.size(), messageCount);
+        });
+    }
+
+    @Test
+    public void testReplayQueueInfiniteExpand() throws Exception {
+        final String topicName = BrokerTestUtil.newUniqueName("persistent://my-property/my-ns/tp_");
+        final String subName = "sub1";
+        final long maxMemoryUsageOfReplayQueueInBytes = 32 * 1024;
+        conf.setMaxMemoryUsageOfReplayQueueInBytesPerSubscription(maxMemoryUsageOfReplayQueueInBytes);
+
+        Consumer<String> consumer1 = pulsarClient.newConsumer(Schema.STRING)
+                .topic(topicName).subscriptionName(subName)
+                .subscriptionType(SubscriptionType.Key_Shared)
+                .receiverQueueSize(5)
+                .acknowledgmentGroupTime(0, TimeUnit.SECONDS)
+                .subscribe();
+        Consumer<String> consumer2 = pulsarClient.newConsumer(Schema.STRING)
+                .topic(topicName).subscriptionName(subName)
+                .subscriptionType(SubscriptionType.Key_Shared)
+                .receiverQueueSize(5)
+                .acknowledgmentGroupTime(0, TimeUnit.SECONDS)
+                .subscribe();
+        Producer<String> producer = pulsarClient.newProducer(Schema.STRING)
+                .enableBatching(false)
+                .topic(topicName)
+                .create();
+
+        // Send 10,000 messages and only consumer1 works.
+        AtomicInteger sendMessageCounter = new AtomicInteger();
+        Thread sendTask = new Thread(() -> {
+            while(sendMessageCounter.get() < 5000) {
+                int i = sendMessageCounter.incrementAndGet();
+                producer.newMessage().key("" + i).value("" + i).sendAsync().join();
+                try {
+                    // sleep 5 ms to avoid high cpu loads.
+                    if (sendMessageCounter.get() % 500 == 0) {
+                        Thread.sleep(5);
+                    }
+                } catch (InterruptedException e) {
+                }
+            }
+        });
+        AtomicBoolean consumeTaskShouldStop = new AtomicBoolean(true);
+        Thread consumeTask = new Thread(() -> {
+            while(true) {
+                try {
+                    Message<String> msg = consumer1.receive(2, TimeUnit.SECONDS);
+                    if (msg != null) {
+                        consumer1.acknowledge(msg);
+                    } else if (consumeTaskShouldStop.get()) {
+                        break;
+                    }
+                } catch (Exception e){
+                    throw new RuntimeException(e);
+                }
+            }
+        });
+        consumeTask.start();
+        sendTask.start();
+        sendTask.join();
+        consumeTaskShouldStop.set(true);
+
+        assertTrue(getMemoryUsageOfReplayQueue(topicName, subName) <= maxMemoryUsageOfReplayQueueInBytes * 2);
+
+        consumer1.close();
+        consumer2.close();
+        producer.close();
+        admin.topics().delete(topicName, false);
+    }
+
+    private long getMemoryUsageOfReplayQueue(String topicName, String subName) {
+        PersistentTopic persistentTopic =
+                (PersistentTopic) pulsar.getBrokerService().getTopic(topicName, false).join().get();
+        PersistentStickyKeyDispatcherMultipleConsumers dispatcher =
+                (PersistentStickyKeyDispatcherMultipleConsumers) persistentTopic
+                        .getSubscription(subName).getDispatcher();
+        MessageRedeliveryController redeliveryMessages =
+                WhiteboxImpl.getInternalState(dispatcher, "redeliveryMessages");
+        return redeliveryMessages.getLongSizeInBytes();
+    }
+}

pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/ConcurrentLongLongPairHashMap.java

@@ -159,6 +159,14 @@ private ConcurrentLongLongPairHashMap(int expectedItems, int concurrencyLevel,
         }
     }
 
+    public long getLongSizeInBytes() {
+        long result = 0;
+        for (Section section : sections) {
+            result += section.getLongSizeInBytes();
+        }
+        return result;
+    }
+
     public long size() {
         long size = 0;
         for (Section s : sections) {
@@ -319,6 +327,11 @@ private static final class Section extends StampedLock {
             Arrays.fill(table, EmptyKey);
         }
 
+        long getLongSizeInBytes() {
+            int baseSize = 4 * 6 + 2 * 5;
+            return (capacity << 6) + baseSize;
+        }
+
         LongPair get(long key1, long key2, int keyHash) {
             long stamp = tryOptimisticRead();
             boolean acquiredLock = false;