[improve][broker] Part-2 of PIP-434: Use ServerCnxThrottleTracker, instead of modifying channel.readable directly

pulsar-broker/src/main/java/org/apache/pulsar/broker/resourcegroup/ResourceGroupPublishLimiter.java

@@ -21,6 +21,7 @@
 import org.apache.pulsar.broker.qos.MonotonicClock;
 import org.apache.pulsar.broker.resourcegroup.ResourceGroup.BytesAndMessagesCount;
 import org.apache.pulsar.broker.service.PublishRateLimiterImpl;
+import org.apache.pulsar.broker.service.ServerCnxThrottleTracker;
 import org.apache.pulsar.common.policies.data.Policies;
 import org.apache.pulsar.common.policies.data.PublishRate;
 import org.apache.pulsar.common.policies.data.ResourceGroup;
@@ -30,7 +31,13 @@ public class ResourceGroupPublishLimiter extends PublishRateLimiterImpl  {
     private volatile long publishMaxByteRate;
 
     public ResourceGroupPublishLimiter(ResourceGroup resourceGroup, MonotonicClock monotonicClock) {
-        super(monotonicClock);
+        super(monotonicClock, producer -> {
+            producer.getCnx().getThrottleTracker().markThrottled(
+                ServerCnxThrottleTracker.ThrottleType.ResourceGroupPublishRate);
+        }, producer -> {
+            producer.getCnx().getThrottleTracker().unmarkThrottled(
+                ServerCnxThrottleTracker.ThrottleType.ResourceGroupPublishRate);
+        });
         update(resourceGroup);
     }
 

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

@@ -21,6 +21,7 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
 import static org.apache.bookkeeper.mledger.impl.ManagedLedgerMBeanImpl.ENTRY_LATENCY_BUCKETS_USEC;
+import static org.apache.pulsar.broker.service.ServerCnxThrottleTracker.ThrottleType;
 import static org.apache.pulsar.compaction.Compactor.COMPACTION_SUBSCRIPTION;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
@@ -193,7 +194,12 @@ public AbstractTopic(String topic, BrokerService brokerService) {
         updateTopicPolicyByBrokerConfig();
 
         this.lastActive = System.nanoTime();
-        topicPublishRateLimiter = new PublishRateLimiterImpl(brokerService.getPulsar().getMonotonicClock());
+        topicPublishRateLimiter = new PublishRateLimiterImpl(brokerService.getPulsar().getMonotonicClock(),
+            producer -> {
+                producer.getCnx().getThrottleTracker().markThrottled(ThrottleType.TopicPublishRate);
+            }, producer -> {
+                producer.getCnx().getThrottleTracker().unmarkThrottled(ThrottleType.TopicPublishRate);
+            });
         updateActiveRateLimiters();
 
         additionalSystemCursorNames = brokerService.pulsar().getConfiguration().getAdditionalSystemCursorNames();

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

@@ -339,7 +339,13 @@ public BrokerService(PulsarService pulsar, EventLoopGroup eventLoopGroup) throws
         this.pulsar = pulsar;
         this.clock = pulsar.getClock();
         this.dynamicConfigurationMap = prepareDynamicConfigurationMap();
-        this.brokerPublishRateLimiter = new PublishRateLimiterImpl(pulsar.getMonotonicClock());
+        this.brokerPublishRateLimiter = new PublishRateLimiterImpl(pulsar.getMonotonicClock(), producer -> {
+            producer.getCnx().getThrottleTracker().markThrottled(
+                    ServerCnxThrottleTracker.ThrottleType.BrokerPublishRate);
+        }, producer -> {
+            producer.getCnx().getThrottleTracker().unmarkThrottled(
+                    ServerCnxThrottleTracker.ThrottleType.BrokerPublishRate);
+        });
         this.dispatchRateLimiterFactory = createDispatchRateLimiterFactory(pulsar.getConfig());
         this.managedLedgerStorage = pulsar.getManagedLedgerStorage();
         this.keepAliveIntervalSeconds = pulsar.getConfiguration().getKeepAliveIntervalSeconds();

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

@@ -899,30 +899,6 @@ public boolean isDisconnecting() {
 
     private static final Logger log = LoggerFactory.getLogger(Producer.class);
 
-    /**
-     * This method increments a counter that is used to control the throttling of a connection.
-     * The connection's read operations are paused when the counter's value is greater than 0, indicating that
-     * throttling is in effect.
-     * It's important to note that after calling this method, it is the caller's responsibility to ensure that the
-     * counter is decremented by calling the {@link #decrementThrottleCount()} method when throttling is no longer
-     * needed on the connection.
-     */
-    public void incrementThrottleCount() {
-        cnx.incrementThrottleCount();
-    }
-
-    /**
-     * This method decrements a counter that is used to control the throttling of a connection.
-     * The connection's read operations are resumed when the counter's value is 0, indicating that
-     * throttling is no longer in effect.
-     * It's important to note that before calling this method, the caller should have previously
-     * incremented the counter by calling the {@link #incrementThrottleCount()} method when throttling
-     * was needed on the connection.
-     */
-    public void decrementThrottleCount() {
-        cnx.decrementThrottleCount();
-    }
-
     public Attributes getOpenTelemetryAttributes() {
         if (attributes != null) {
             return attributes;

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

@@ -24,6 +24,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.pulsar.broker.qos.AsyncTokenBucket;
 import org.apache.pulsar.broker.qos.MonotonicClock;
@@ -42,9 +43,14 @@ public class PublishRateLimiterImpl implements PublishRateLimiter {
 
     private final AtomicInteger throttledProducersCount = new AtomicInteger(0);
     private final AtomicBoolean processingQueuedProducers = new AtomicBoolean(false);
+    private final Consumer<Producer> throttleAction;
+    private final Consumer<Producer> unthrottleAction;
 
-    public PublishRateLimiterImpl(MonotonicClock monotonicClock) {
+    public PublishRateLimiterImpl(MonotonicClock monotonicClock, Consumer<Producer> throttleAction,
+                                  Consumer<Producer> unthrottleAction) {
         this.monotonicClock = monotonicClock;
+        this.throttleAction = throttleAction;
+        this.unthrottleAction = unthrottleAction;
     }
 
     /**
@@ -68,7 +74,7 @@ public void handlePublishThrottling(Producer producer, int numOfMessages,
         }
         if (shouldThrottle) {
             // throttle the producer by incrementing the throttle count
-            producer.incrementThrottleCount();
+            throttleAction.accept(producer);
             // schedule decrementing the throttle count to possibly unthrottle the producer after the
             // throttling period
             scheduleDecrementThrottleCount(producer);
@@ -136,7 +142,8 @@ private void unthrottleQueuedProducers(ScheduledExecutorService executor) {
             while ((throttlingDuration = calculateThrottlingDurationNanos()) == 0L
                     && (producer = unthrottlingQueue.poll()) != null) {
                 try {
-                    producer.decrementThrottleCount();
+                    final Producer producerFinal = producer;
+                    producer.getCnx().execute(() -> unthrottleAction.accept(producerFinal));
                 } catch (Exception e) {
                     log.error("Failed to unthrottle producer {}", producer, e);
                 }

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

@@ -25,6 +25,7 @@
 import static org.apache.commons.lang3.StringUtils.isNotBlank;
 import static org.apache.pulsar.broker.admin.impl.PersistentTopicsBase.unsafeGetPartitionedTopicMetadataAsync;
 import static org.apache.pulsar.broker.lookup.TopicLookupBase.lookupTopicAsync;
+import static org.apache.pulsar.broker.service.ServerCnxThrottleTracker.ThrottleType;
 import static org.apache.pulsar.broker.service.persistent.PersistentTopic.getMigratedClusterUrl;
 import static org.apache.pulsar.broker.service.schema.BookkeeperSchemaStorage.ignoreUnrecoverableBKException;
 import static org.apache.pulsar.common.api.proto.ProtocolVersion.v5;
@@ -283,7 +284,8 @@ public void incrementPublishBytes(long bytes, long maxPendingBytesPerThread) {
             if (maxPendingBytesPerThread > 0 && pendingBytes > maxPendingBytesPerThread
                     && !limitExceeded) {
                 limitExceeded = true;
-                cnxsPerThread.get().forEach(cnx -> cnx.throttleTracker.setPublishBufferLimiting(true));
+                cnxsPerThread.get().forEach(cnx -> cnx.throttleTracker.markThrottled(
+                        ThrottleType.IOThreadMaxPendingPublishBytesExceeded));
             }
         }
 
@@ -293,7 +295,8 @@ public void decrementPublishBytes(long bytes, long resumeThresholdPendingBytesPe
             // we resume all connections sharing the same thread
             if (limitExceeded && pendingBytes <= resumeThresholdPendingBytesPerThread) {
                 limitExceeded = false;
-                cnxsPerThread.get().forEach(cnx -> cnx.throttleTracker.setPublishBufferLimiting(false));
+                cnxsPerThread.get().forEach(cnx -> cnx.throttleTracker.unmarkThrottled(
+                        ThrottleType.IOThreadMaxPendingPublishBytesExceeded));
             }
         }
     }
@@ -311,6 +314,7 @@ enum State {
         Start, Connected, Failed, Connecting
     }
 
+    @Getter
     private final ServerCnxThrottleTracker throttleTracker;
 
     public ServerCnx(PulsarService pulsar) {
@@ -481,12 +485,12 @@ private void checkPauseReceivingRequestsAfterResumeRateLimit(BaseCommand cmd) {
             log.warn("[{}] Reached rate limitation", this);
             // Stop receiving requests.
             pausedDueToRateLimitation = true;
-            ctx.channel().config().setAutoRead(false);
+            getThrottleTracker().markThrottled(ThrottleType.ConnectionPauseReceivingCooldownRateLimit);
             // Resume after 1 second.
             ctx.channel().eventLoop().schedule(() -> {
                 if (pausedDueToRateLimitation) {
                     log.info("[{}] Resuming connection after rate limitation", this);
-                    ctx.channel().config().setAutoRead(true);
+                    getThrottleTracker().unmarkThrottled(ThrottleType.ConnectionPauseReceivingCooldownRateLimit);
                     pausedDueToRateLimitation = false;
                 }
             }, requestRateLimiter.getPeriodAtMs(), TimeUnit.MILLISECONDS);
@@ -497,7 +501,7 @@ private void checkPauseReceivingRequestsAfterResumeRateLimit(BaseCommand cmd) {
     public void channelWritabilityChanged(ChannelHandlerContext ctx) throws Exception {
         if (pauseReceivingRequestsIfUnwritable && ctx.channel().isWritable()) {
             log.info("[{}] is writable, turn on channel auto-read", this);
-            ctx.channel().config().setAutoRead(true);
+            getThrottleTracker().unmarkThrottled(ThrottleType.ConnectionOutboundBufferFull);
             requestRateLimiter.timingOpen(pauseReceivingCooldownMilliSeconds, TimeUnit.MILLISECONDS);
         } else if (pauseReceivingRequestsIfUnwritable && !ctx.channel().isWritable()) {
             final ChannelOutboundBuffer outboundBuffer = ctx.channel().unsafe().outboundBuffer();
@@ -511,7 +515,7 @@ public void channelWritabilityChanged(ChannelHandlerContext ctx) throws Exceptio
                     PAUSE_RECEIVING_LOG.debug("[{}] is not writable, turn off channel auto-read", this);
                 }
             }
-            ctx.channel().config().setAutoRead(false);
+            getThrottleTracker().markThrottled(ThrottleType.ConnectionOutboundBufferFull);
         }
         ctx.fireChannelWritabilityChanged();
     }
@@ -3399,7 +3403,7 @@ public boolean isWritable() {
     // or the pending publish bytes
     private void increasePendingSendRequestsAndPublishBytes(int msgSize) {
         if (++pendingSendRequest == maxPendingSendRequests) {
-            throttleTracker.setPendingSendRequestsExceeded(true);
+            throttleTracker.markThrottled(ThrottleType.ConnectionMaxPendingPublishRequestsExceeded);
         }
         PendingBytesPerThreadTracker.getInstance().incrementPublishBytes(msgSize, maxPendingBytesPerThread);
     }
@@ -3424,7 +3428,7 @@ public void completedSendOperation(boolean isNonPersistentTopic, int msgSize) {
         PendingBytesPerThreadTracker.getInstance().decrementPublishBytes(msgSize, resumeThresholdPendingBytesPerThread);
 
         if (--pendingSendRequest == resumeReadsThreshold) {
-            throttleTracker.setPendingSendRequestsExceeded(false);
+            throttleTracker.unmarkThrottled(ThrottleType.ConnectionMaxPendingPublishRequestsExceeded);
         }
 
         if (isNonPersistentTopic) {
@@ -3803,22 +3807,6 @@ protected void setAuthRole(String authRole) {
         this.authRole = authRole;
     }
 
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public void incrementThrottleCount() {
-        throttleTracker.incrementThrottleCount();
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public void decrementThrottleCount() {
-        throttleTracker.decrementThrottleCount();
-    }
-
     @VisibleForTesting
     void setAuthState(AuthenticationState authState) {
         this.authState = authState;