Change the implementation of BucketRateLimiter and ItemExponentialFailureRateLimiter, add more unit tests

Author: cizezsy

extended/src/main/java/io/kubernetes/client/extended/workqueue/ratelimiter/BucketRateLimiter.java

@@ -1,15 +1,18 @@
 package io.kubernetes.client.extended.workqueue.ratelimiter;
 
-import io.github.bucket4j.*;
+import io.github.bucket4j.Bandwidth;
+import io.github.bucket4j.Bucket;
+import io.github.bucket4j.Bucket4j;
+import io.github.bucket4j.Refill;
+import io.github.bucket4j.local.SynchronizationStrategy;
 import java.time.Duration;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
 
 /** A light-weight token bucket implementation for RateLimiter. */
 public class BucketRateLimiter<T> implements RateLimiter<T> {
-
   private Bucket bucket;
-  private long tokensInQueue;
-  private long tokensGeneratedInPeriod;
-  private Duration period;
 
   /**
    * @param capacity Capacity is the maximum number of tokens can be consumed.
@@ -19,26 +22,18 @@ public class BucketRateLimiter<T> implements RateLimiter<T> {
   public BucketRateLimiter(long capacity, long tokensGeneratedInPeriod, Duration period) {
     Bandwidth bandwidth =
         Bandwidth.classic(capacity, Refill.greedy(tokensGeneratedInPeriod, period));
-
-    this.bucket = Bucket4j.builder().addLimit(bandwidth).build();
-    this.tokensInQueue = 0;
-    this.tokensGeneratedInPeriod = tokensGeneratedInPeriod;
-    this.period = period;
+    this.bucket =
+        Bucket4j.builder()
+            .addLimit(bandwidth)
+            .withSynchronizationStrategy(SynchronizationStrategy.SYNCHRONIZED)
+            .build();
   }
 
   @Override
   public synchronized Duration when(T item) {
-    tokensInQueue++;
-
-    long consumedTokens = bucket.tryConsumeAsMuchAsPossible(tokensInQueue);
-    if (tokensInQueue - consumedTokens == 0) {
-      tokensInQueue = 0;
-      return Duration.ZERO;
-    }
-
-    tokensInQueue = tokensInQueue - consumedTokens;
-
-    return durationFromTokens(tokensInQueue, tokensGeneratedInPeriod, period);
+    DelayGetter delayGetter = new DelayGetter();
+    bucket.asAsyncScheduler().consume(1, delayGetter).complete(null);
+    return delayGetter.getDelay();
   }
 
   @Override
@@ -49,8 +44,21 @@ public int numRequeues(T item) {
     return 0;
   }
 
-  private Duration durationFromTokens(
-      long tokensNeedToBeConsumed, long tokensGeneratedInPeriod, Duration period) {
-    return period.dividedBy(tokensGeneratedInPeriod).multipliedBy(tokensNeedToBeConsumed);
+  private class DelayGetter extends ScheduledThreadPoolExecutor {
+    private Duration delay = Duration.ZERO;
+
+    @Override
+    public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
+      this.delay = Duration.ofNanos(unit.toNanos(delay));
+      return null;
+    }
+
+    private DelayGetter() {
+      super(0);
+    }
+
+    private Duration getDelay() {
+      return delay;
+    }
   }
 }

extended/src/main/java/io/kubernetes/client/extended/workqueue/ratelimiter/ItemExponentialFailureRateLimiter.java

@@ -26,17 +26,8 @@ public ItemExponentialFailureRateLimiter(Duration baseDelay, Duration maxDelay)
   public synchronized Duration when(T item) {
     int exp = failures.getOrDefault(item, 0);
     failures.put(item, exp + 1);
-
-    double backOff = baseDelay.toNanos() * Math.pow(2, exp);
-    if (backOff > Long.MAX_VALUE) {
-      return maxDelay;
-    }
-
-    if (backOff > maxDelay.toNanos()) {
-      return maxDelay;
-    }
-
-    return Duration.ofNanos((long) backOff);
+    long d = maxDelay.toMillis() >> exp;
+    return d > baseDelay.toMillis() ? baseDelay.multipliedBy(1 << exp) : maxDelay;
   }
 
   @Override

extended/src/test/java/io/kubernetes/client/extended/workqueue/ratelimiter/BucketRateLimiterTest.java

@@ -1,21 +1,34 @@
 package io.kubernetes.client.extended.workqueue.ratelimiter;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import java.time.Duration;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 
 public class BucketRateLimiterTest {
+  @Rule public ExpectedException thrown = ExpectedException.none();
 
   @Test
   public void testBucketRateLimiterBasic() {
     RateLimiter<String> rateLimiter = new BucketRateLimiter<>(2, 1, Duration.ofMinutes(10));
     assertEquals(Duration.ZERO, rateLimiter.when("one"));
     assertEquals(Duration.ZERO, rateLimiter.when("one"));
 
-    assertEquals(Duration.ofMinutes(10), rateLimiter.when("one"));
-    assertEquals(Duration.ofMinutes(20), rateLimiter.when("one"));
-    assertEquals(Duration.ofMinutes(30), rateLimiter.when("one"));
+    Duration waitDuration = rateLimiter.when("one");
+    Duration expectDuration = Duration.ofMinutes(10);
+
+    Duration diff = waitDuration.minus(expectDuration);
+    // waitDuration might be smaller than expect duration because of time is elapsed.
+    assertTrue(diff.isZero() || (diff.isNegative() && !diff.plusSeconds(1).isNegative()));
+
+    waitDuration = rateLimiter.when("one");
+    expectDuration = Duration.ofMinutes(20);
+    diff = waitDuration.minus(expectDuration);
+
+    assertTrue(diff.isZero() || (diff.isNegative() && !diff.plusSeconds(1).isNegative()));
   }
 
   @Test
@@ -25,11 +38,44 @@ public void testBucketRateLimiterTokenAdded() throws InterruptedException {
     assertEquals(Duration.ZERO, rateLimiter.when("one"));
     assertEquals(Duration.ZERO, rateLimiter.when("one"));
 
-    assertEquals(Duration.ofSeconds(2), rateLimiter.when("one"));
+    Duration waitDuration = rateLimiter.when("one");
+    assertTrue(waitDuration.getSeconds() > 0);
 
     Thread.sleep(4000);
 
-    assertEquals(Duration.ZERO, rateLimiter.when("one"));
-    assertEquals(Duration.ofSeconds(2), rateLimiter.when("one"));
+    assertEquals(Duration.ZERO, rateLimiter.when("two"));
+
+    waitDuration = rateLimiter.when("two");
+    assertTrue(waitDuration.getSeconds() > 0);
+  }
+
+  @Test
+  public void testNegativeCapacity() {
+    thrown.expect(IllegalArgumentException.class);
+    thrown.expectMessage("-2 is wrong value for capacity, because capacity should be positive");
+    RateLimiter<String> rateLimiter = new BucketRateLimiter<>(-2, 1, Duration.ofSeconds(2));
+  }
+
+  @Test
+  public void testNegativeTokensGeneratedInPeriod() {
+    thrown.expect(IllegalArgumentException.class);
+    thrown.expectMessage("-1 is wrong value for period tokens, because tokens should be positive");
+    RateLimiter<String> rateLimiter = new BucketRateLimiter<>(2, -1, Duration.ofSeconds(2));
+  }
+
+  @Test
+  public void testNegativePeriod() {
+    thrown.expect(IllegalArgumentException.class);
+    thrown.expectMessage(
+        "-1 is wrong value for period of bandwidth, because period should be positive");
+    RateLimiter<String> rateLimiter = new BucketRateLimiter<>(2, 1, Duration.ofNanos(-1));
+  }
+
+  @Test
+  public void testTokensLargerThanNanos() {
+    thrown.expect(IllegalArgumentException.class);
+    thrown.expectMessage(
+        "100 token/nanosecond is not permitted refill rate, because highest supported rate is 1 token/nanosecond");
+    RateLimiter<String> rateLimiter = new BucketRateLimiter<>(2, 100, Duration.ofNanos(1));
   }
 }

extended/src/test/java/io/kubernetes/client/extended/workqueue/ratelimiter/ItemExponentialFailureRateLimiterTest.java

@@ -1,6 +1,7 @@
 package io.kubernetes.client.extended.workqueue.ratelimiter;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import java.time.Duration;
 import org.junit.Test;
@@ -56,4 +57,30 @@ public void testItemExponentialFailureRateLimiterOverFlow() {
     }
     assertEquals(Duration.ofHours(1000), rateLimiter.when("overflow2"));
   }
+
+  @Test
+  public void testNegativeBaseDelay() {
+    RateLimiter<String> rateLimiter =
+        new ItemExponentialFailureRateLimiter<>(Duration.ofMillis(-1), Duration.ofSeconds(1000));
+
+    for (int i = 0; i < 5; i++) {
+      rateLimiter.when("one");
+    }
+    assertEquals(Duration.ofMillis(-32), rateLimiter.when("one"));
+    for (int i = 0; i < 1000; i++) {
+      rateLimiter.when("overflow1");
+    }
+    assertTrue(rateLimiter.when("overflow1").isNegative());
+  }
+
+  @Test
+  public void testNegativeMaxDelay() {
+    RateLimiter<String> rateLimiter =
+        new ItemExponentialFailureRateLimiter<>(Duration.ofMillis(1), Duration.ofSeconds(-1000));
+
+    assertEquals(Duration.ofSeconds(-1000), rateLimiter.when("one"));
+    assertEquals(Duration.ofSeconds(-1000), rateLimiter.when("one"));
+    assertEquals(Duration.ofSeconds(-1000), rateLimiter.when("one"));
+    assertEquals(Duration.ofSeconds(-1000), rateLimiter.when("one"));
+  }
 }

extended/src/test/java/io/kubernetes/client/extended/workqueue/ratelimiter/ItemFastSlowRateLimiterTest.java

@@ -29,4 +29,19 @@ public void testFastSlowRateLimiter() {
     assertEquals(0, rateLimiter.numRequeues("one"));
     assertEquals(Duration.ofMillis(5), rateLimiter.when("one"));
   }
+
+  @Test
+  public void testNegativeOrZeroAttempts() {
+    RateLimiter<String> rateLimiter =
+        new ItemFastSlowRateLimiter<>(Duration.ofMillis(5), Duration.ofSeconds(10), -1);
+
+    assertEquals(Duration.ofSeconds(10), rateLimiter.when("one"));
+    assertEquals(Duration.ofSeconds(10), rateLimiter.when("one"));
+    assertEquals(Duration.ofSeconds(10), rateLimiter.when("one"));
+
+    rateLimiter = new ItemFastSlowRateLimiter<>(Duration.ofMillis(5), Duration.ofSeconds(10), 0);
+    assertEquals(Duration.ofSeconds(10), rateLimiter.when("two"));
+    assertEquals(Duration.ofSeconds(10), rateLimiter.when("two"));
+    assertEquals(Duration.ofSeconds(10), rateLimiter.when("two"));
+  }
 }

extended/src/test/java/io/kubernetes/client/extended/workqueue/ratelimiter/MaxOfRateLimiterTest.java

@@ -1,6 +1,7 @@
 package io.kubernetes.client.extended.workqueue.ratelimiter;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import java.time.Duration;
 import org.junit.Test;
@@ -31,4 +32,27 @@ public void testMaxOfRateLimiter() {
     assertEquals(0, rateLimiter.numRequeues("one"));
     assertEquals(Duration.ofMillis(5), rateLimiter.when("one"));
   }
+
+  @Test
+  public void testDefaultRateLimiter() {
+    RateLimiter<String> rateLimiter = new DefaultControllerRateLimiter<>();
+
+    assertEquals(Duration.ofMillis(5), rateLimiter.when("one"));
+    assertEquals(Duration.ofMillis(10), rateLimiter.when("one"));
+    assertEquals(Duration.ofMillis(20), rateLimiter.when("one"));
+
+    for (int i = 0; i < 20; i++) {
+      rateLimiter.when("one");
+    }
+
+    assertEquals(Duration.ofSeconds(1000), rateLimiter.when("one"));
+    assertEquals(Duration.ofSeconds(1000), rateLimiter.when("one"));
+
+    for (int i = 0; i < 75; i++) {
+      rateLimiter.when("one");
+    }
+
+    assertTrue(rateLimiter.when("one").getSeconds() > 0);
+    assertTrue(rateLimiter.when("two").getSeconds() > 0);
+  }
 }