[fix] [broker] Fix problem at RateLimiter#tryAcquire (apache#15306)

datastax · Apr 25, 2022 · 84b6559 · 84b6559
1 parent 35c4b68
commit 84b6559
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/util/RateLimiter.java b/pulsar-common/src/main/java/org/apache/pulsar/common/util/RateLimiter.java
@@ -189,8 +189,7 @@ public synchronized boolean tryAcquire(long acquirePermit) {
             canAcquire = acquirePermit < 0 || acquiredPermits < this.permits;
         } else {
             // acquired-permits can't be larger than the rate
-            if (acquirePermit > this.permits) {
-                acquiredPermits = this.permits;
+            if (acquirePermit + acquiredPermits > this.permits) {
                 return false;
             }
 

diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/util/RateLimiterTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/util/RateLimiterTest.java
@@ -133,6 +133,24 @@ public void testTryAcquire() {
         rate.close();
     }
 
+    @Test
+    public void testTryAcquireMoreThanPermits() {
+        final long rateTimeMSec = 1000;
+        RateLimiter rate = RateLimiter.builder().permits(3).rateTime(rateTimeMSec).timeUnit(TimeUnit.MILLISECONDS)
+                .build();
+        assertTrue(rate.tryAcquire(2));
+        assertEquals(rate.getAvailablePermits(), 1);
+
+        //try to acquire failed, not decrease availablePermits.
+        assertFalse(rate.tryAcquire(2));
+        assertEquals(rate.getAvailablePermits(), 1);
+
+        assertTrue(rate.tryAcquire(1));
+        assertEquals(rate.getAvailablePermits(), 0);
+
+        rate.close();
+    }
+
     @Test
     public void testMultipleTryAcquire() {
         final long rateTimeMSec = 1000;
@@ -189,7 +207,7 @@ public void testDispatchRate() throws Exception {
 
         Thread.sleep(rateTimeMSec);
         // check after three rate-time: acquiredPermits is 0
-        assertEquals(rate.getAvailablePermits() > 0, true);
+        assertTrue(rate.getAvailablePermits() > 0);
 
         rate.close();
     }