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

microbench/README.md

@@ -41,3 +41,29 @@ For fast recompiling of the benchmarks (without compiling Pulsar modules) and cr
 mvn -Pmicrobench -pl microbench clean package
 ```
 
+### Running specific benchmarks
+
+Display help:
+
+```shell
+java -jar microbench/target/microbenchmarks.jar -h
+```
+
+Listing all benchmarks:
+
+```shell
+java -jar microbench/target/microbenchmarks.jar -l
+```
+
+Running specific benchmarks:
+
+```shell
+java -jar microbench/target/microbenchmarks.jar ".*BenchmarkName.*"
+```
+
+Checking what benchmarks match the pattern:
+
+```shell
+java -jar microbench/target/microbenchmarks.jar ".*BenchmarkName.*" -lp
+```
+

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,29 @@ public void teardown() {
     @Benchmark
     @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
     @Warmup(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
-    public void consumeTokensBenchmark001Threads() {
-        asyncTokenBucket.consumeTokens(1);
+    public void consumeTokensBenchmark001Threads(Blackhole blackhole) {
+        consumeTokenAndGetTokens(blackhole);
     }
 
     @Threads(10)
     @Benchmark
     @Measurement(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
     @Warmup(time = 10, timeUnit = TimeUnit.SECONDS, iterations = 1)
-    public void consumeTokensBenchmark010Threads() {
-        asyncTokenBucket.consumeTokens(1);
+    public void consumeTokensBenchmark010Threads(Blackhole blackhole) {
+        consumeTokenAndGetTokens(blackhole);
     }
 
     @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) {
+        consumeTokenAndGetTokens(blackhole);
+    }
+
+    private void consumeTokenAndGetTokens(Blackhole blackhole) {
         asyncTokenBucket.consumeTokens(1);
+        // blackhole is used to ensure that the compiler doesn't do dead code elimination
+        blackhole.consume(asyncTokenBucket.getTokens());
     }
 }

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

@@ -0,0 +1,102 @@
+/*
+ * 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.qos;
+
+import java.util.concurrent.TimeUnit;
+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;
+
+@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();
+    }
+
+    @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);
+    }
+
+    @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);
+    }
+
+    @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));
+    }
+}

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

@@ -122,6 +122,7 @@ public static void resetToDefaultEventualConsistentTokensView() {
     protected AsyncTokenBucket(MonotonicSnapshotClock clockSource, long resolutionNanos) {
         this.clockSource = clockSource;
         this.resolutionNanos = resolutionNanos;
+        this.lastNanos = Long.MIN_VALUE;
     }
 
     public static FinalRateAsyncTokenBucketBuilder builder() {
@@ -155,20 +156,23 @@ private long consumeTokensAndMaybeUpdateTokensBalance(long consumeTokens, boolea
             throw new IllegalArgumentException("consumeTokens must be >= 0");
         }
         long currentNanos = clockSource.getTickNanos(forceUpdateTokens);
+        long newTokens = 0;
         // check if the tokens should be updated immediately
         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
+            newTokens = calculateNewTokensSinceLastUpdate(currentNanos);
+        }
+        if (newTokens > 0) {
             // flush the pendingConsumedTokens by calling "sumThenReset"
-            long totalConsumedTokens = consumeTokens + pendingConsumedTokens.sumThenReset();
+            long currentPendingConsumedTokens = pendingConsumedTokens.sumThenReset();
+            // calculate the token delta by subtracting the consumed tokens from the new tokens
+            long tokenDelta = newTokens - currentPendingConsumedTokens;
             // 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);
+                    // limit the tokens to the capacity of the bucket
+                    currentTokens -> Math.min(currentTokens + tokenDelta, getCapacity())
+                            // subtract the consumed tokens from the capped tokens
+                            - consumeTokens);
         } else {
             // eventual consistent fast path, tokens are not updated immediately
 
@@ -177,8 +181,13 @@ private long consumeTokensAndMaybeUpdateTokensBalance(long consumeTokens, boolea
                 pendingConsumedTokens.add(consumeTokens);
             }
 
-            // return Long.MIN_VALUE if the current value of tokens is unknown due to the eventual consistency
-            return Long.MIN_VALUE;
+            if (forceUpdateTokens) {
+                // return the current tokens balance without updating the tokens and resetting the pendingConsumedTokens
+                return tokens - pendingConsumedTokens.sum();
+            } else {
+                // return Long.MIN_VALUE if the current value of tokens is unknown due to the eventual consistency
+                return Long.MIN_VALUE;
+            }
         }
     }
 
@@ -211,8 +220,16 @@ 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) {
+        long previousLastNanos = LAST_NANOS_UPDATER.getAndUpdate(this,
+                currentLastNanos -> {
+                    // update lastNanos only if at least resolutionNanos/2 nanoseconds has passed since the last update
+                    if (currentNanos >= currentLastNanos + resolutionNanos / 2) {
+                        return currentNanos;
+                    } else {
+                        return currentLastNanos;
+                    }
+                });
+        if (previousLastNanos == Long.MIN_VALUE || previousLastNanos >= currentNanos) {
             newTokens = 0;
         } else {
             long durationNanos = currentNanos - previousLastNanos + REMAINDER_NANOS_UPDATER.getAndSet(this, 0);