[fix][broker] PIP-322 Fix issue with rate limiters where rates can exceed limits initially and consumption pauses until token balance is positive (apache#24012)

Co-authored-by: Matteo Merli <[email protected]>
(cherry picked from commit e547bea)

Author: lhotari

conf/broker.conf

@@ -243,7 +243,7 @@ messageExpiryCheckIntervalInMinutes=5
 # How long to delay rewinding cursor and dispatching messages when active consumer is changed
 activeConsumerFailoverDelayTimeMillis=1000
 
-# Enable consistent hashing for selecting the active consumer in partitioned topics with Failover subscription type. 
+# Enable consistent hashing for selecting the active consumer in partitioned topics with Failover subscription type.
 # For non-partitioned topics, consistent hashing is used by default.
 activeConsumerFailoverConsistentHashing=false
 
@@ -419,6 +419,11 @@ subscribeThrottlingRatePerConsumer=0
 # Rate period for {subscribeThrottlingRatePerConsumer}. Default is 30s.
 subscribeRatePeriodPerConsumerInSecond=30
 
+# The class name of the factory that creates DispatchRateLimiter implementations. Current options are
+# org.apache.pulsar.broker.service.persistent.DispatchRateLimiterFactoryAsyncTokenBucket (default, PIP-322 implementation)
+# and org.apache.pulsar.broker.service.persistent.DispatchRateLimiterFactoryClassic (legacy implementation)
+dispatchRateLimiterFactoryClassName=org.apache.pulsar.broker.service.persistent.DispatchRateLimiterFactoryAsyncTokenBucket
+
 # Default messages per second dispatch throttling-limit for whole broker. Using a value of 0, is disabling default
 # message dispatch-throttling
 dispatchThrottlingRateInMsg=0

conf/standalone.conf

@@ -152,7 +152,7 @@ maxMessageSizeCheckIntervalInSeconds=60
 # How long to delay rewinding cursor and dispatching messages when active consumer is changed
 activeConsumerFailoverDelayTimeMillis=1000
 
-# Enable consistent hashing for selecting the active consumer in partitioned topics with Failover subscription type. 
+# Enable consistent hashing for selecting the active consumer in partitioned topics with Failover subscription type.
 # For non-partitioned topics, consistent hashing is used by default.
 activeConsumerFailoverConsistentHashing=false
 
@@ -276,6 +276,11 @@ brokerPublisherThrottlingMaxMessageRate=0
 # (Disable byte rate limit with value 0)
 brokerPublisherThrottlingMaxByteRate=0
 
+# The class name of the factory that creates DispatchRateLimiter implementations. Current options are
+# org.apache.pulsar.broker.service.persistent.DispatchRateLimiterFactoryAsyncTokenBucket (default, PIP-322 implementation)
+# and org.apache.pulsar.broker.service.persistent.DispatchRateLimiterFactoryClassic (legacy implementation)
+dispatchRateLimiterFactoryClassName=org.apache.pulsar.broker.service.persistent.DispatchRateLimiterFactoryAsyncTokenBucket
+
 # Default messages per second dispatch throttling-limit for every topic. Using a value of 0, is disabling default
 # message dispatch-throttling
 dispatchThrottlingRatePerTopicInMsg=0

microbench/README.md

@@ -67,3 +67,12 @@ Checking what benchmarks match the pattern:
 java -jar microbench/target/microbenchmarks.jar ".*BenchmarkName.*" -lp
 ```
 
+Profiling benchmarks with [async-profiler](https://github.com/async-profiler/async-profiler):
+
+```shell
+# example of profiling with async-profiler
+# download async-profiler from https://github.com/async-profiler/async-profiler/releases
+LIBASYNCPROFILER_PATH=$HOME/async-profiler/lib/libasyncProfiler.dylib
+java -jar microbench/target/microbenchmarks.jar -prof async:libPath=$LIBASYNCPROFILER_PATH\;output=flamegraph\;dir=profile-results ".*BenchmarkName.*"
+```
+

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

@@ -30,19 +30,22 @@
 import org.openjdk.jmh.annotations.Scope;
 import org.openjdk.jmh.annotations.Setup;
 import org.openjdk.jmh.annotations.State;
-import org.openjdk.jmh.annotations.TearDown;
 import org.openjdk.jmh.annotations.Threads;
 import org.openjdk.jmh.annotations.Warmup;
 import org.openjdk.jmh.infra.Blackhole;
 
+/**
+ * On MacOS, the performance of System.nanoTime() is not great. Running benchmarks on Linux is recommended due
+ * to the bottleneck of System.nanoTime() implementation on MacOS.
+ */
 @Fork(3)
 @BenchmarkMode(Mode.Throughput)
 @OutputTimeUnit(TimeUnit.SECONDS)
 @State(Scope.Thread)
 public class AsyncTokenBucketBenchmark {
     private AsyncTokenBucket asyncTokenBucket;
-    private DefaultMonotonicSnapshotClock monotonicSnapshotClock =
-            new DefaultMonotonicSnapshotClock(TimeUnit.MILLISECONDS.toNanos(8), System::nanoTime);
+    private DefaultMonotonicClock monotonicSnapshotClock =
+            new DefaultMonotonicClock();
 
     @Setup(Level.Iteration)
     public void setup() {
@@ -51,11 +54,6 @@ public void setup() {
                 .initialTokens(2 * ratePerSecond).capacity(2 * ratePerSecond).build();
     }
 
-    @TearDown(Level.Iteration)
-    public void teardown() {
-        monotonicSnapshotClock.close();
-    }
-
     @Threads(1)
     @Benchmark
     @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)

microbench/src/main/java/org/apache/pulsar/broker/qos/DefaultMonotonicSnapshotClockBenchmark.java renamed to microbench/src/main/java/org/apache/pulsar/broker/qos/DefaultMonotonicClockBenchmark.java

@@ -23,80 +23,48 @@
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
 import org.openjdk.jmh.annotations.Fork;
-import org.openjdk.jmh.annotations.Level;
 import org.openjdk.jmh.annotations.Measurement;
 import org.openjdk.jmh.annotations.Mode;
 import org.openjdk.jmh.annotations.OutputTimeUnit;
 import org.openjdk.jmh.annotations.Scope;
 import org.openjdk.jmh.annotations.State;
-import org.openjdk.jmh.annotations.TearDown;
 import org.openjdk.jmh.annotations.Threads;
 import org.openjdk.jmh.annotations.Warmup;
 import org.openjdk.jmh.infra.Blackhole;
 
+/**
+ * On MacOS, the performance of System.nanoTime() is not great. Running benchmarks on Linux is recommended due
+ * to the bottleneck of System.nanoTime() implementation on MacOS.
+ */
 @Fork(3)
 @BenchmarkMode(Mode.Throughput)
 @OutputTimeUnit(TimeUnit.SECONDS)
 @State(Scope.Thread)
-public class DefaultMonotonicSnapshotClockBenchmark {
-    private DefaultMonotonicSnapshotClock monotonicSnapshotClock =
-            new DefaultMonotonicSnapshotClock(TimeUnit.MILLISECONDS.toNanos(1), System::nanoTime);
-
-    @TearDown(Level.Iteration)
-    public void teardown() {
-        monotonicSnapshotClock.close();
-    }
+public class DefaultMonotonicClockBenchmark {
+    private DefaultMonotonicClock monotonicSnapshotClock =
+            new DefaultMonotonicClock();
 
     @Threads(1)
     @Benchmark
     @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
     @Warmup(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
     public void getTickNanos001Threads(Blackhole blackhole) {
-        consumeTokenAndGetTokens(blackhole, false);
+        blackhole.consume(monotonicSnapshotClock.getTickNanos());
     }
 
     @Threads(10)
     @Benchmark
     @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
     @Warmup(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
     public void getTickNanos010Threads(Blackhole blackhole) {
-        consumeTokenAndGetTokens(blackhole, false);
+        blackhole.consume(monotonicSnapshotClock.getTickNanos());
     }
 
     @Threads(100)
     @Benchmark
     @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
     @Warmup(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
     public void getTickNanos100Threads(Blackhole blackhole) {
-        consumeTokenAndGetTokens(blackhole, false);
-    }
-
-    @Threads(1)
-    @Benchmark
-    @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
-    @Warmup(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
-    public void getTickNanosRequestSnapshot001Threads(Blackhole blackhole) {
-        consumeTokenAndGetTokens(blackhole, true);
-    }
-
-    @Threads(10)
-    @Benchmark
-    @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
-    @Warmup(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
-    public void getTickNanosRequestSnapshot010Threads(Blackhole blackhole) {
-        consumeTokenAndGetTokens(blackhole, true);
-    }
-
-    @Threads(100)
-    @Benchmark
-    @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
-    @Warmup(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
-    public void getTickNanosRequestSnapshot100Threads(Blackhole blackhole) {
-        consumeTokenAndGetTokens(blackhole, true);
-    }
-
-    private void consumeTokenAndGetTokens(Blackhole blackhole, boolean requestSnapshot) {
-        // blackhole is used to ensure that the compiler doesn't do dead code elimination
-        blackhole.consume(monotonicSnapshotClock.getTickNanos(requestSnapshot));
+        blackhole.consume(monotonicSnapshotClock.getTickNanos());
     }
 }

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

@@ -1042,6 +1042,7 @@ The max allowed delay for delayed delivery (in milliseconds). If the broker rece
             doc = "Enable precise rate limit for topic publish"
     )
     private boolean preciseTopicPublishRateLimiterEnable = false;
+
     @FieldContext(
         category = CATEGORY_SERVER,
         dynamic = true,
@@ -1064,6 +1065,16 @@ The max allowed delay for delayed delivery (in milliseconds). If the broker rece
             + "when broker publish rate limiting enabled. (Disable byte rate limit with value 0)"
     )
     private long brokerPublisherThrottlingMaxByteRate = 0;
+
+    @FieldContext(category = CATEGORY_SERVER, doc =
+            "The class name of the factory that creates DispatchRateLimiter implementations. Current options are "
+                    + "org.apache.pulsar.broker.service.persistent.DispatchRateLimiterFactoryAsyncTokenBucket "
+                    + "(default, PIP-322 implementation) "
+                    + "org.apache.pulsar.broker.service.persistent.DispatchRateLimiterFactoryClassic (legacy "
+                    + "implementation)")
+    private String dispatchRateLimiterFactoryClassName =
+            "org.apache.pulsar.broker.service.persistent.DispatchRateLimiterFactoryAsyncTokenBucket";
+
     @FieldContext(
             category = CATEGORY_SERVER,
             dynamic = true,

pulsar-broker/src/main/java/org/apache/pulsar/broker/PulsarService.java

@@ -96,8 +96,8 @@
 import org.apache.pulsar.broker.lookup.v1.TopicLookup;
 import org.apache.pulsar.broker.namespace.NamespaceService;
 import org.apache.pulsar.broker.protocol.ProtocolHandlers;
-import org.apache.pulsar.broker.qos.DefaultMonotonicSnapshotClock;
-import org.apache.pulsar.broker.qos.MonotonicSnapshotClock;
+import org.apache.pulsar.broker.qos.DefaultMonotonicClock;
+import org.apache.pulsar.broker.qos.MonotonicClock;
 import org.apache.pulsar.broker.resourcegroup.ResourceGroupService;
 import org.apache.pulsar.broker.resourcegroup.ResourceUsageTopicTransportManager;
 import org.apache.pulsar.broker.resourcegroup.ResourceUsageTransportManager;
@@ -300,7 +300,7 @@ public class PulsarService implements AutoCloseable, ShutdownService {
 
     private TransactionPendingAckStoreProvider transactionPendingAckStoreProvider;
     private final ExecutorProvider transactionExecutorProvider;
-    private final DefaultMonotonicSnapshotClock monotonicSnapshotClock;
+    private final MonotonicClock monotonicClock;
     private String brokerId;
     private final CompletableFuture<Void> readyForIncomingRequestsFuture = new CompletableFuture<>();
     private final List<Runnable> pendingTasksBeforeReadyForIncomingRequests = new ArrayList<>();
@@ -395,8 +395,11 @@ public PulsarService(ServiceConfiguration config,
         // here in the constructor we don't have the offloader scheduler yet
         this.offloaderStats = LedgerOffloaderStats.create(false, false, null, 0);
 
-        this.monotonicSnapshotClock = new DefaultMonotonicSnapshotClock(TimeUnit.MILLISECONDS.toNanos(
-                DEFAULT_MONOTONIC_CLOCK_GRANULARITY_MILLIS), System::nanoTime);
+        this.monotonicClock = createMonotonicClock(config);
+    }
+
+    protected MonotonicClock createMonotonicClock(ServiceConfiguration config) {
+        return new DefaultMonotonicClock();
     }
 
     public MetadataStore createConfigurationMetadataStore(PulsarMetadataEventSynchronizer synchronizer,
@@ -701,7 +704,9 @@ public CompletableFuture<Void> closeAsync() {
             brokerClientSharedScheduledExecutorProvider.shutdownNow();
             brokerClientSharedLookupExecutorProvider.shutdownNow();
             brokerClientSharedTimer.stop();
-            monotonicSnapshotClock.close();
+            if (monotonicClock instanceof AutoCloseable c) {
+                c.close();
+            }
 
             if (openTelemetryTransactionPendingAckStoreStats != null) {
                 openTelemetryTransactionPendingAckStoreStats.close();
@@ -1982,8 +1987,8 @@ public Optional<Integer> getBrokerListenPortTls() {
         return brokerService.getListenPortTls();
     }
 
-    public MonotonicSnapshotClock getMonotonicSnapshotClock() {
-        return monotonicSnapshotClock;
+    public MonotonicClock getMonotonicClock() {
+        return monotonicClock;
     }
 
     public static WorkerConfig initializeWorkerConfigFromBrokerConfig(ServiceConfiguration brokerConfig,