[do not merge] [improve] [broker] Make the dispatch rate limit more precise when subscription is created

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractBaseDispatcher.java

@@ -18,6 +18,7 @@
  */
 package org.apache.pulsar.broker.service;
 
+import static org.apache.pulsar.broker.service.persistent.PersistentTopic.MESSAGE_RATE_BACKOFF_MS;
 import io.netty.buffer.ByteBuf;
 import io.prometheus.client.Gauge;
 import java.util.ArrayList;
@@ -27,6 +28,7 @@
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.LongAdder;
+import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.bookkeeper.mledger.Entry;
 import org.apache.bookkeeper.mledger.ManagedCursor;
@@ -46,6 +48,7 @@
 import org.apache.pulsar.common.api.proto.ReplicatedSubscriptionsSnapshot;
 import org.apache.pulsar.common.protocol.Commands;
 import org.apache.pulsar.common.protocol.Markers;
+import org.apache.pulsar.common.util.Codec;
 
 @Slf4j
 public abstract class AbstractBaseDispatcher extends EntryFilterSupport implements Dispatcher {
@@ -381,4 +384,96 @@ public long getFilterRescheduledMsgCount() {
     protected final void updatePendingBytesToDispatch(long size) {
         PENDING_BYTES_TO_DISPATCH.inc(size);
     }
+
+    /**
+     * Calculate messages count & bytes size to read by rate-limiters.
+     * @return left pair is messagesToRead, right pair is bytesToRead
+     */
+    protected Pair<Integer, Long> calculateToReadByRateLimiter(int messageCountToReadDefault,
+                                                  @Nullable ManagedCursor cursor,
+                                                  Optional<DispatchRateLimiter>...rateLimiters){
+        int messageCountToRead = messageCountToReadDefault;
+        long bytesToReadDefault = serviceConfig.getDispatcherMaxReadSizeBytes();
+
+        // throttle only if: (1) cursor is not active (or flag for throttle-nonBacklogConsumer is enabled) bcz
+        // active-cursor reads message from cache rather from bookkeeper (2) if topic has reached message-rate
+        // threshold: then schedule the read after MESSAGE_RATE_BACKOFF_MS
+        boolean cursorActive = cursor == null /* NonDurableCursor has no backlog */ || cursor.isActive();
+        if (!serviceConfig.isDispatchThrottlingOnNonBacklogConsumerEnabled() && cursorActive) {
+            return Pair.of(messageCountToRead, bytesToReadDefault);
+        }
+
+        for (Optional<DispatchRateLimiter> rateLimiterOptional : rateLimiters){
+            if (!rateLimiterOptional.isPresent()){
+                continue;
+            }
+            DispatchRateLimiter rateLimiter = rateLimiterOptional.get();
+            if (reachDispatchRateLimit(rateLimiter)) {
+                if (log.isDebugEnabled()) {
+                    log.debug("[{}] message-read exceeded broker message-rate {}/{}, schedule after a {}",
+                            buildName(cursor), rateLimiter.getDispatchRateOnMsg(), rateLimiter.getDispatchRateOnByte(),
+                            MESSAGE_RATE_BACKOFF_MS);
+                }
+                return Pair.of(-1, -1L);
+            } else {
+                if (rateLimiter.getAvailableDispatchRateLimitOnMsg() > 0) {
+                    messageCountToRead =
+                            Math.min(messageCountToRead, (int) rateLimiter.getAvailableDispatchRateLimitOnMsg());
+                }
+                if (rateLimiter.getAvailableDispatchRateLimitOnByte() > 0){
+                    bytesToReadDefault =
+                            Math.min(bytesToReadDefault, rateLimiter.getAvailableDispatchRateLimitOnByte());
+                }
+            }
+        }
+        // If messagesToRead is 0 or less, correct it to 1 to prevent IllegalArgumentException
+        messageCountToRead = Math.max(messageCountToRead, 1);
+        bytesToReadDefault = Math.max(bytesToReadDefault, 1);
+        return Pair.of(messageCountToRead, bytesToReadDefault);
+    }
+
+    protected String buildName(@Nullable ManagedCursor cursor){
+        // NonDurableCursor doesn't have cursor.
+        String cursorName = cursor == null || cursor.getName() == null ? "" : Codec.decode(cursor.getName());
+        return subscription.getTopic().getName() + " / " + cursorName;
+    }
+
+    protected int calculateEntryCountToReadIfEnabledBatch(Topic topic, int messageCountToRead){
+        // if turn of precise dispatcher flow control, adjust the records to read.
+        if (!serviceConfig.isPreciseDispatcherFlowControl()) {
+            return messageCountToRead;
+        }
+        // When calculating the "avgMessagesPerEntry,"
+        // 1. look up the attribute avgMessagesPerEntry from the consumers under subscription.
+        // 2. if all consumers avgMessagesPerEntry is zero, look up from other subscriptions under this topic.
+        // 3. if still is zero, directly call "topic.getAvgMessagesPerEntryAccumulator()".
+        // 3. if still is zero too, just read one entry.
+        int avgMessagesPerEntry = calculateAvgMessagesPerEntryInTopic(topic);
+        if (avgMessagesPerEntry < 1){
+            return 1;
+        }
+        return Math.min((int) Math.ceil(messageCountToRead * 1.0 / avgMessagesPerEntry),
+                serviceConfig.getDispatcherMaxReadBatchSize());
+    }
+
+    protected int calculateAvgMessagesPerEntry(){
+        return 0;
+    }
+
+    protected int calculateAvgMessagesPerEntryInTopic(Topic topic){
+        int avgMessagesPerEntry = calculateAvgMessagesPerEntry();
+        if (avgMessagesPerEntry > 0){
+            return avgMessagesPerEntry;
+        }
+        for (Subscription sub : topic.getSubscriptions().values()){
+            if (sub.getDispatcher() instanceof AbstractBaseDispatcher otherDispatcher){
+                int avgMessagesPerEntryInOtherSub = otherDispatcher.calculateAvgMessagesPerEntry();
+                if (avgMessagesPerEntryInOtherSub > 0){
+                    return avgMessagesPerEntryInOtherSub;
+                }
+            }
+        }
+        return Math.max(0,
+                Double.valueOf(topic.getAvgMessagesPerEntryAccumulator().getAvgMessagesPerEntry()).intValue());
+    }
 }

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractDispatcherMultipleConsumers.java

@@ -237,7 +237,23 @@ private int getFirstConsumerIndexOfPriority(int targetPriority) {
         return -1;
     }
 
-    private static final Logger log = LoggerFactory.getLogger(PersistentStickyKeyDispatcherMultipleConsumers.class);
-
-
+    /**
+     * @return If the consumer knows, the correct value is returned; otherwise it returns 0.
+     */
+    protected int calculateAvgMessagesPerEntry(){
+        if (consumerList.isEmpty() || IS_CLOSED_UPDATER.get(this) == TRUE) {
+            return 0;
+        }
+        Consumer randomConsumer = null;
+        int nextConsumerIndex = random.nextInt(consumerList.size());
+        for (int i = 0; i < consumerList.size(); i++){
+            randomConsumer = consumerList.get(nextConsumerIndex);
+            if (randomConsumer.getAvgMessagesPerEntry() > 0){
+                return randomConsumer.getAvgMessagesPerEntry();
+            }
+            nextConsumerIndex++;
+            nextConsumerIndex = nextConsumerIndex % consumerList.size();
+        }
+        return 0;
+    }
 }

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractDispatcherSingleActiveConsumer.java

@@ -282,6 +282,14 @@ public boolean isConsumerConnected() {
         return ACTIVE_CONSUMER_UPDATER.get(this) != null;
     }
 
+    protected int calculateAvgMessagesPerEntry(){
+        Consumer activeConsumer = getActiveConsumer();
+        if (activeConsumer == null){
+            return 0;
+        }
+        return Math.max(0, activeConsumer.getAvgMessagesPerEntry());
+    }
+
     private static final Logger log = LoggerFactory.getLogger(AbstractDispatcherSingleActiveConsumer.class);
 
 }

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractTopic.java

@@ -150,6 +150,10 @@ public abstract class AbstractTopic implements Topic, TopicPolicyListener<TopicP
     protected final LongAdder bytesOutFromRemovedSubscriptions = new LongAdder();
     protected Map<String, EntryFilterWithClassLoader> entryFilters;
 
+    @Getter
+    protected final AvgMessagesPerEntryAccumulator avgMessagesPerEntryAccumulator =
+            new AvgMessagesPerEntryAccumulator();
+
     public AbstractTopic(String topic, BrokerService brokerService) {
         this.topic = topic;
         this.brokerService = brokerService;

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AvgMessagesPerEntryAccumulator.java

@@ -0,0 +1,51 @@
+/*
+ * 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.broker.service;
+
+import com.google.common.util.concurrent.AtomicDouble;
+
+/**
+ * It starts keep tracking the average messages per entry.
+ * The initial value is 0, when new value comes, it will update with
+ * avgMessagesPerEntry = avgMessagePerEntry * avgPercent + (1 - avgPercent) * new Value.
+ */
+public class AvgMessagesPerEntryAccumulator {
+
+    private static final double avgPercent = 0.9;
+
+    private final AtomicDouble avgMessagesPerEntry = new AtomicDouble(0);
+
+    public double getAvgMessagesPerEntry(){
+        return avgMessagesPerEntry.get();
+    }
+
+    public void setAvgMessagesPerEntry(double avgMessagesPerEntry){
+        this.avgMessagesPerEntry.set(avgMessagesPerEntry);
+    }
+
+    public void accumulate(int totalMessages, int totalEntries) {
+        if (avgMessagesPerEntry.get() < 1) { //valid avgMessagesPerEntry should always >= 1
+            // set init value.
+            avgMessagesPerEntry.set(1.0 * totalMessages / totalEntries);
+        } else {
+            avgMessagesPerEntry.set(avgMessagesPerEntry.get() * avgPercent
+                    + (1 - avgPercent) * totalMessages / totalEntries);
+        }
+    }
+}

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java

@@ -121,12 +121,7 @@ public class Consumer {
 
     private final KeySharedMeta keySharedMeta;
 
-    /**
-     * It starts keep tracking the average messages per entry.
-     * The initial value is 0, when new value comes, it will update with
-     * avgMessagesPerEntry = avgMessagePerEntry * avgPercent + (1 - avgPercent) * new Value.
-     */
-    private final AtomicDouble avgMessagesPerEntry = new AtomicDouble(0);
+    private AvgMessagesPerEntryAccumulator avgMessagesPerEntryAccumulator = new AvgMessagesPerEntryAccumulator();
     private static final long [] EMPTY_ACK_SET = new long[0];
 
     private static final double avgPercent = 0.9;
@@ -319,22 +314,17 @@ public Future<Void> sendMessages(final List<? extends Entry> entries, EntryBatch
             }
         }
 
-        // calculate avg message per entry
-        if (avgMessagesPerEntry.get() < 1) { //valid avgMessagesPerEntry should always >= 1
-            // set init value.
-            avgMessagesPerEntry.set(1.0 * totalMessages / totalEntries);
-        } else {
-            avgMessagesPerEntry.set(avgMessagesPerEntry.get() * avgPercent
-                    + (1 - avgPercent) * totalMessages / totalEntries);
-        }
+        avgMessagesPerEntryAccumulator.accumulate(totalMessages, totalEntries);
+        getSubscription().getTopic().getAvgMessagesPerEntryAccumulator().accumulate(totalMessages, totalEntries);
 
         // reduce permit and increment unackedMsg count with total number of messages in batch-msgs
         int ackedCount = batchIndexesAcks == null ? 0 : batchIndexesAcks.getTotalAckedIndexCount();
         MESSAGE_PERMITS_UPDATER.addAndGet(this, ackedCount - totalMessages);
         if (log.isDebugEnabled()){
             log.debug("[{}-{}] Added {} minus {} messages to MESSAGE_PERMITS_UPDATER in broker.service.Consumer"
                             + " for consumerId: {}; avgMessagesPerEntry is {}",
-                   topicName, subscription, ackedCount, totalMessages, consumerId, avgMessagesPerEntry.get());
+                   topicName, subscription, ackedCount, totalMessages, consumerId,
+                    avgMessagesPerEntryAccumulator.getAvgMessagesPerEntry());
         }
         incrementUnackedMessages(unackedMessages);
         Future<Void> writeAndFlushPromise =
@@ -777,7 +767,7 @@ public int getAvailablePermits() {
      * return 0 if there is no entry dispatched yet.
      */
     public int getAvgMessagesPerEntry() {
-        return (int) Math.round(avgMessagesPerEntry.get());
+        return (int) Math.round(avgMessagesPerEntryAccumulator.getAvgMessagesPerEntry());
     }
 
     public boolean isBlocked() {
@@ -836,7 +826,7 @@ public void updateStats(ConsumerStatsImpl consumerStats) {
         }
         unackedMessages = consumerStats.unackedMessages;
         blockedConsumerOnUnackedMsgs = consumerStats.blockedConsumerOnUnackedMsgs;
-        avgMessagesPerEntry.set(consumerStats.avgMessagesPerEntry);
+        avgMessagesPerEntryAccumulator.setAvgMessagesPerEntry(consumerStats.avgMessagesPerEntry);
     }
 
     public ConsumerStatsImpl getStats() {

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Topic.java

@@ -366,4 +366,5 @@ default boolean isSystemTopic() {
      */
     HierarchyTopicPolicies getHierarchyTopicPolicies();
 
+    AvgMessagesPerEntryAccumulator getAvgMessagesPerEntryAccumulator();
 }

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentTopic.java

@@ -37,13 +37,15 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLongFieldUpdater;
+import lombok.Getter;
 import org.apache.bookkeeper.mledger.Entry;
 import org.apache.bookkeeper.mledger.Position;
 import org.apache.pulsar.broker.PulsarServerException;
 import org.apache.pulsar.broker.namespace.NamespaceService;
 import org.apache.pulsar.broker.resources.NamespaceResources;
 import org.apache.pulsar.broker.service.AbstractReplicator;
 import org.apache.pulsar.broker.service.AbstractTopic;
+import org.apache.pulsar.broker.service.AvgMessagesPerEntryAccumulator;
 import org.apache.pulsar.broker.service.BrokerService;
 import org.apache.pulsar.broker.service.BrokerServiceException;
 import org.apache.pulsar.broker.service.BrokerServiceException.ConsumerBusyException;
@@ -107,6 +109,10 @@ public class NonPersistentTopic extends AbstractTopic implements Topic, TopicPol
             AtomicLongFieldUpdater.newUpdater(NonPersistentTopic.class, "entriesAddedCounter");
     private volatile long entriesAddedCounter = 0;
 
+    @Getter
+    protected final AvgMessagesPerEntryAccumulator avgMessagesPerEntryAccumulator =
+            new AvgMessagesPerEntryAccumulator();
+
     private volatile boolean migrated = false;
     private static final FastThreadLocal<TopicStats> threadLocalTopicStats = new FastThreadLocal<TopicStats>() {
         @Override