[fix][broker] Fix rate limiter token bucket and clock consistency issues causing excessive throttling and connection timeouts

microbench/src/main/java/org/apache/pulsar/broker/qos/AsyncTokenBucketBenchmark.java

@@ -33,6 +33,7 @@
 import org.openjdk.jmh.annotations.TearDown;
 import org.openjdk.jmh.annotations.Threads;
 import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
 
 @Fork(3)
 @BenchmarkMode(Mode.Throughput)
@@ -59,23 +60,26 @@ public void teardown() {
     @Benchmark
     @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
     @Warmup(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
-    public void consumeTokensBenchmark001Threads() {
+    public void consumeTokensBenchmark001Threads(Blackhole blackhole) {
         asyncTokenBucket.consumeTokens(1);
+        blackhole.consume(asyncTokenBucket.getTokens());
     }
 
     @Threads(10)
     @Benchmark
     @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
     @Warmup(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
-    public void consumeTokensBenchmark010Threads() {
+    public void consumeTokensBenchmark010Threads(Blackhole blackhole) {
         asyncTokenBucket.consumeTokens(1);
+        blackhole.consume(asyncTokenBucket.getTokens());
     }
 
     @Threads(100)
     @Benchmark
     @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
     @Warmup(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
-    public void consumeTokensBenchmark100Threads() {
+    public void consumeTokensBenchmark100Threads(Blackhole blackhole) {
         asyncTokenBucket.consumeTokens(1);
+        blackhole.consume(asyncTokenBucket.getTokens());
     }
 }

pulsar-broker/src/main/java/org/apache/pulsar/broker/qos/AsyncTokenBucket.java

@@ -159,16 +159,15 @@ private long consumeTokensAndMaybeUpdateTokensBalance(long consumeTokens, boolea
         if (shouldUpdateTokensImmediately(currentNanos, forceUpdateTokens)) {
             // calculate the number of new tokens since the last update
             long newTokens = calculateNewTokensSinceLastUpdate(currentNanos);
-            // calculate the total amount of tokens to consume in this update
             // flush the pendingConsumedTokens by calling "sumThenReset"
-            long totalConsumedTokens = consumeTokens + pendingConsumedTokens.sumThenReset();
+            long currentPendingConsumedTokens = pendingConsumedTokens.sumThenReset();
             // update the tokens and return the current token value
             return TOKENS_UPDATER.updateAndGet(this,
-                    currentTokens ->
-                            // after adding new tokens, limit the tokens to the capacity
-                            Math.min(currentTokens + newTokens, getCapacity())
-                                    // subtract the consumed tokens
-                                    - totalConsumedTokens);
+                    // after adding new tokens subtract the pending consumed tokens and
+                    // limit the tokens to the capacity of the bucket
+                    currentTokens -> Math.min(currentTokens + newTokens - currentPendingConsumedTokens, getCapacity())
+                            // subtract the consumed tokens
+                            - consumeTokens);
         } else {
             // eventual consistent fast path, tokens are not updated immediately
 
@@ -211,8 +210,10 @@ private boolean shouldUpdateTokensImmediately(long currentNanos, boolean forceUp
      */
     private long calculateNewTokensSinceLastUpdate(long currentNanos) {
         long newTokens;
-        long previousLastNanos = LAST_NANOS_UPDATER.getAndSet(this, currentNanos);
-        if (previousLastNanos == 0) {
+        // update lastNanos if currentNanos is greater than the current lastNanos
+        long previousLastNanos = LAST_NANOS_UPDATER.getAndUpdate(this,
+                currentLastNanos -> Math.max(currentNanos, currentLastNanos));
+        if (previousLastNanos == 0 || previousLastNanos >= currentNanos) {
             newTokens = 0;
         } else {
             long durationNanos = currentNanos - previousLastNanos + REMAINDER_NANOS_UPDATER.getAndSet(this, 0);

pulsar-broker/src/main/java/org/apache/pulsar/broker/qos/DefaultMonotonicSnapshotClock.java

@@ -29,13 +29,17 @@
  *
  * Starts a daemon thread that updates the snapshot value periodically with a configured interval. The close method
  * should be called to stop the thread.
+ * A single thread is used to update the monotonic clock value so that the snapshot value is always increasing,
+ * even if the clock source is not strictly monotonic across all CPUs. This might be the case in some virtualized
+ * environments.
  */
 public class DefaultMonotonicSnapshotClock implements MonotonicSnapshotClock, AutoCloseable {
     private static final Logger LOG = LoggerFactory.getLogger(DefaultMonotonicSnapshotClock.class);
     private final long sleepMillis;
     private final int sleepNanos;
     private final LongSupplier clockSource;
-    private final Thread thread;
+    private final TickUpdaterThread tickUpdaterThread;
+    private final long snapshotIntervalNanos;
     private volatile long snapshotTickNanos;
 
     public DefaultMonotonicSnapshotClock(long snapshotIntervalNanos, LongSupplier clockSource) {
@@ -45,45 +49,170 @@ public DefaultMonotonicSnapshotClock(long snapshotIntervalNanos, LongSupplier cl
         this.sleepMillis = TimeUnit.NANOSECONDS.toMillis(snapshotIntervalNanos);
         this.sleepNanos = (int) (snapshotIntervalNanos - TimeUnit.MILLISECONDS.toNanos(sleepMillis));
         this.clockSource = clockSource;
-        updateSnapshotTickNanos();
-        thread = new Thread(this::snapshotLoop, getClass().getSimpleName() + "-update-loop");
-        thread.setDaemon(true);
-        thread.start();
+        this.snapshotIntervalNanos = snapshotIntervalNanos;
+        tickUpdaterThread = new TickUpdaterThread();
+        tickUpdaterThread.start();
     }
 
     /** {@inheritDoc} */
     @Override
     public long getTickNanos(boolean requestSnapshot) {
         if (requestSnapshot) {
-            updateSnapshotTickNanos();
+            tickUpdaterThread.requestUpdate();
         }
         return snapshotTickNanos;
     }
 
-    private void updateSnapshotTickNanos() {
-        snapshotTickNanos = clockSource.getAsLong();
-    }
+    /**
+     * A thread that updates snapshotTickNanos value periodically with a configured interval.
+     * The thread is started when the DefaultMonotonicSnapshotClock is created and runs until the close method is
+     * called.
+     * A single thread is used to read the clock source value since on some hardware of virtualized platforms,
+     * System.nanoTime() isn't strictly monotonic across all CPUs. Reading by a single thread will improve the
+     * stability of the read value since a single thread is scheduled on a single CPU. If the thread is migrated
+     * to another CPU, the clock source value might leap backward or forward, but logic in this class will handle it.
+     */
+    private class TickUpdaterThread extends Thread {
+        private final Object tickUpdateDelayMonitor = new Object();
+        private final Object tickUpdatedMonitor = new Object();
+        private final long maxDelta;
+        private long referenceClockSourceValue;
+        private long baseSnapshotTickNanos;
+        private long previousSnapshotTickNanos;
+        private volatile boolean running;
+        private boolean tickUpdateDelayMonitorNotified;
+        private long requestCount;
+
+        TickUpdaterThread() {
+            super(DefaultMonotonicSnapshotClock.class.getSimpleName() + "-update-loop");
+            // set as daemon thread so that it doesn't prevent the JVM from exiting
+            setDaemon(true);
+            // set the highest priority
+            setPriority(MAX_PRIORITY);
+            this.maxDelta = 2 * snapshotIntervalNanos;
+        }
+
+        @Override
+        public void run() {
+            try {
+                running = true;
+                long updatedForRequestCount = -1;
+                while (!isInterrupted()) {
+                    try {
+                        boolean snapshotRequested = false;
+                        // sleep for the configured interval on a monitor that can be notified to stop the sleep
+                        // and update the tick value immediately. This is used in requestUpdate method.
+                        synchronized (tickUpdateDelayMonitor) {
+                            tickUpdateDelayMonitorNotified = false;
+                            // only wait if no explicit request has been made since the last update
+                            if (requestCount == updatedForRequestCount) {
+                                // if no request has been made, sleep for the configured interval
+                                tickUpdateDelayMonitor.wait(sleepMillis, sleepNanos);
+                                snapshotRequested = tickUpdateDelayMonitorNotified;
+                            }
+                            updatedForRequestCount = requestCount;
+                        }
+                        updateSnapshotTickNanos(snapshotRequested);
+                        notifyAllTickUpdated();
+                    } catch (InterruptedException e) {
+                        interrupt();
+                        break;
+                    }
+                }
+            } catch (Throwable t) {
+                // report unexpected error since this would be a fatal error when the clock doesn't progress anymore
+                // this is very unlikely to happen, but it's better to log it in any case
+                LOG.error("Unexpected fatal error that stopped the clock.", t);
+            } finally {
+                LOG.info("DefaultMonotonicSnapshotClock's TickUpdaterThread stopped. {},tid={}", this, getId());
+                running = false;
+                notifyAllTickUpdated();
+            }
+        }
+
+        private void updateSnapshotTickNanos(boolean snapshotRequested) {
+            long clockValue = clockSource.getAsLong();
+
+            // Initialization
+            if (referenceClockSourceValue == 0) {
+                referenceClockSourceValue = clockValue;
+                baseSnapshotTickNanos = clockValue;
+                snapshotTickNanos = clockValue;
+                previousSnapshotTickNanos = clockValue;
+                return;
+            }
+
+            // calculate the duration since the reference clock source value
+            // so that the snapshot value is always increasing and tolerates it when the clock source is not strictly
+            // monotonic across all CPUs and leaps backward or forward
+            long durationSinceReference = clockValue - referenceClockSourceValue;
+            long newSnapshotTickNanos = baseSnapshotTickNanos + durationSinceReference;
+
+            // reset the reference clock source value if the clock source value leaps backward or forward
+            if (newSnapshotTickNanos < previousSnapshotTickNanos - maxDelta
+                    || newSnapshotTickNanos > previousSnapshotTickNanos + maxDelta) {
+                referenceClockSourceValue = clockValue;
+                baseSnapshotTickNanos = previousSnapshotTickNanos;
+                if (!snapshotRequested) {
+                    // if the snapshot value is not requested, increment by the snapshot interval
+                    baseSnapshotTickNanos += snapshotIntervalNanos;
+                }
+                newSnapshotTickNanos = baseSnapshotTickNanos;
+            }
+
+            // update snapshotTickNanos value if the new value is greater than the previous value
+            if (newSnapshotTickNanos > previousSnapshotTickNanos) {
+                snapshotTickNanos = newSnapshotTickNanos;
+                // store into a field so that we don't need to do a volatile read to find out the previous value
+                previousSnapshotTickNanos = newSnapshotTickNanos;
+            }
+        }
+
+        private void notifyAllTickUpdated() {
+            synchronized (tickUpdatedMonitor) {
+                // notify all threads that are waiting for the tick value to be updated
+                tickUpdatedMonitor.notifyAll();
+            }
+        }
+
+        public void requestUpdate() {
+            if (!running) {
+                // thread has stopped running, fallback to update the value directly without any optimizations
+                snapshotTickNanos = clockSource.getAsLong();
+                return;
+            }
+            synchronized (tickUpdatedMonitor) {
+                // notify the thread to stop waiting and update the tick value
+                synchronized (tickUpdateDelayMonitor) {
+                    tickUpdateDelayMonitorNotified = true;
+                    requestCount++;
+                    tickUpdateDelayMonitor.notify();
+                }
+                // wait until the tick value has been updated
+                try {
+                    tickUpdatedMonitor.wait();
+                } catch (InterruptedException e) {
+                    currentThread().interrupt();
+                }
+            }
+        }
 
-    private void snapshotLoop() {
-        try {
-            while (!Thread.currentThread().isInterrupted()) {
-                updateSnapshotTickNanos();
+        @Override
+        public synchronized void start() {
+            super.start();
+            // wait until the thread is started and the tick value has been updated
+            synchronized (tickUpdatedMonitor) {
                 try {
-                    Thread.sleep(sleepMillis, sleepNanos);
+                    tickUpdatedMonitor.wait();
                 } catch (InterruptedException e) {
-                    Thread.currentThread().interrupt();
-                    break;
+                    currentThread().interrupt();
                 }
             }
-        } catch (Throwable t) {
-            // report unexpected error since this would be a fatal error when the clock doesn't progress anymore
-            // this is very unlikely to happen, but it's better to log it in any case
-            LOG.error("Unexpected fatal error that stopped the clock.", t);
         }
     }
 
     @Override
     public void close() {
-        thread.interrupt();
+        tickUpdaterThread.interrupt();
     }
 }

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

@@ -20,18 +20,19 @@
 package org.apache.pulsar.broker.service;
 
 import com.google.common.annotations.VisibleForTesting;
-import io.netty.channel.EventLoopGroup;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.pulsar.broker.qos.AsyncTokenBucket;
 import org.apache.pulsar.broker.qos.MonotonicSnapshotClock;
 import org.apache.pulsar.common.policies.data.Policies;
 import org.apache.pulsar.common.policies.data.PublishRate;
 import org.jctools.queues.MessagePassingQueue;
 import org.jctools.queues.MpscUnboundedArrayQueue;
 
+@Slf4j
 public class PublishRateLimiterImpl implements PublishRateLimiter {
     private volatile AsyncTokenBucket tokenBucketOnMessage;
     private volatile AsyncTokenBucket tokenBucketOnByte;
@@ -80,7 +81,7 @@ private void scheduleDecrementThrottleCount(Producer producer) {
         // schedule unthrottling when the throttling count is incremented to 1
         // this is to avoid scheduling unthrottling multiple times for concurrent producers
         if (throttledProducersCount.incrementAndGet() == 1) {
-            EventLoopGroup executor = producer.getCnx().getBrokerService().executor();
+            ScheduledExecutorService executor = producer.getCnx().getBrokerService().executor().next();
             scheduleUnthrottling(executor, calculateThrottlingDurationNanos());
         }
     }
@@ -134,7 +135,11 @@ private void unthrottleQueuedProducers(ScheduledExecutorService executor) {
             // unthrottle as many producers as possible while there are token available
             while ((throttlingDuration = calculateThrottlingDurationNanos()) == 0L
                     && (producer = unthrottlingQueue.poll()) != null) {
-                producer.decrementThrottleCount();
+                try {
+                    producer.decrementThrottleCount();
+                } catch (Exception e) {
+                    log.error("Failed to unthrottle producer {}", producer, e);
+                }
                 throttledProducersCount.decrementAndGet();
             }
             // if there are still producers to be unthrottled, schedule unthrottling again

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

@@ -23,6 +23,7 @@
 import java.util.concurrent.TimeUnit;
 import org.apache.pulsar.broker.ServiceConfiguration;
 import org.apache.pulsar.broker.qos.AsyncTokenBucket;
+import org.apache.pulsar.broker.qos.MonotonicSnapshotClock;
 import org.apache.pulsar.broker.service.BrokerService;
 import org.apache.pulsar.common.naming.NamespaceName;
 import org.apache.pulsar.common.naming.TopicName;
@@ -216,18 +217,22 @@ public synchronized void updateDispatchRate(DispatchRate dispatchRate) {
         long msgRate = dispatchRate.getDispatchThrottlingRateInMsg();
         long byteRate = dispatchRate.getDispatchThrottlingRateInByte();
         long ratePeriodNanos = TimeUnit.SECONDS.toNanos(Math.max(dispatchRate.getRatePeriodInSecond(), 1));
+        MonotonicSnapshotClock clock = brokerService.getPulsar().getMonotonicSnapshotClock();
 
         // update msg-rateLimiter
         if (msgRate > 0) {
             if (dispatchRate.isRelativeToPublishRate()) {
                 this.dispatchRateLimiterOnMessage =
                         AsyncTokenBucket.builderForDynamicRate()
+                                .clock(clock)
                                 .rateFunction(() -> getRelativeDispatchRateInMsg(dispatchRate))
                                 .ratePeriodNanosFunction(() -> ratePeriodNanos)
                                 .build();
             } else {
                 this.dispatchRateLimiterOnMessage =
-                        AsyncTokenBucket.builder().rate(msgRate).ratePeriodNanos(ratePeriodNanos)
+                        AsyncTokenBucket.builder()
+                                .clock(clock)
+                                .rate(msgRate).ratePeriodNanos(ratePeriodNanos)
                                 .build();
             }
         } else {
@@ -239,12 +244,15 @@ public synchronized void updateDispatchRate(DispatchRate dispatchRate) {
             if (dispatchRate.isRelativeToPublishRate()) {
                 this.dispatchRateLimiterOnByte =
                         AsyncTokenBucket.builderForDynamicRate()
+                                .clock(clock)
                                 .rateFunction(() -> getRelativeDispatchRateInByte(dispatchRate))
                                 .ratePeriodNanosFunction(() -> ratePeriodNanos)
                                 .build();
             } else {
                 this.dispatchRateLimiterOnByte =
-                        AsyncTokenBucket.builder().rate(byteRate).ratePeriodNanos(ratePeriodNanos)
+                        AsyncTokenBucket.builder()
+                                .clock(clock)
+                                .rate(byteRate).ratePeriodNanos(ratePeriodNanos)
                                 .build();
             }
         } else {

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

@@ -117,7 +117,9 @@ private synchronized void updateSubscribeRate(ConsumerIdentifier consumerIdentif
         // update subscribe-rateLimiter
         if (ratePerConsumer > 0) {
             AsyncTokenBucket tokenBucket =
-                    AsyncTokenBucket.builder().rate(ratePerConsumer).ratePeriodNanos(ratePeriodNanos).build();
+                    AsyncTokenBucket.builder()
+                            .clock(brokerService.getPulsar().getMonotonicSnapshotClock())
+                            .rate(ratePerConsumer).ratePeriodNanos(ratePeriodNanos).build();
             this.subscribeRateLimiter.put(consumerIdentifier, tokenBucket);
         } else {
             // subscribe-rate should be disable and close