RateLimiter { public boolean checkCredit(double itemCost) { long cost = (long) (itemCost / creditsPerNanosecond); long credit; long currentDebit; long balance; do { currentDebit = debit.get(); credit = clock.currentNanoTicks(); balance = credit - currentDebit; if (balance > maxBalance) { balance = maxBalance; } balance -= cost; if (balance < 0) { return false; } } while (!debit.compareAndSet(currentDebit, credit - balance)); return true; } RateLimiter(double creditsPerSecond, double maxBalance);  RateLimiter(double creditsPerSecond, double maxBalance, Clock clock); boolean checkCredit(double itemCost);  }
@Test public void testRateLimiterWholeNumber() { MockClock clock = new MockClock(); RateLimiter limiter = new RateLimiter(2.0, 2.0, clock); long currentTime = TimeUnit.MICROSECONDS.toNanos(100); clock.timeNanos = currentTime; assertTrue(limiter.checkCredit(1.0)); assertTrue(limiter.checkCredit(1.0)); assertFalse(limiter.checkCredit(1.0)); currentTime += TimeUnit.MILLISECONDS.toNanos(250); clock.timeNanos = currentTime; assertFalse(limiter.checkCredit(1.0)); currentTime += TimeUnit.MILLISECONDS.toNanos(500); clock.timeNanos = currentTime; assertTrue(limiter.checkCredit(1.0)); assertFalse(limiter.checkCredit(1.0)); currentTime += TimeUnit.MILLISECONDS.toNanos(5000); clock.timeNanos = currentTime; assertTrue(limiter.checkCredit(1.0)); assertTrue(limiter.checkCredit(1.0)); assertFalse(limiter.checkCredit(1.0)); assertFalse(limiter.checkCredit(1.0)); assertFalse(limiter.checkCredit(1.0)); }
@Test public void testRateLimiterLessThanOne() { MockClock clock = new MockClock(); RateLimiter limiter = new RateLimiter(0.5, 0.5, clock); long currentTime = TimeUnit.MICROSECONDS.toNanos(100); clock.timeNanos = currentTime; assertTrue(limiter.checkCredit(0.25)); assertTrue(limiter.checkCredit(0.25)); assertFalse(limiter.checkCredit(0.25)); currentTime += TimeUnit.MILLISECONDS.toNanos(250); clock.timeNanos = currentTime; assertFalse(limiter.checkCredit(0.25)); currentTime += TimeUnit.MILLISECONDS.toNanos(500); clock.timeNanos = currentTime; assertTrue(limiter.checkCredit(0.25)); assertFalse(limiter.checkCredit(0.25)); currentTime += TimeUnit.MILLISECONDS.toNanos(5000); clock.timeNanos = currentTime; assertTrue(limiter.checkCredit(0.25)); assertTrue(limiter.checkCredit(0.25)); assertFalse(limiter.checkCredit(0.25)); assertFalse(limiter.checkCredit(0.25)); assertFalse(limiter.checkCredit(0.25)); }
@Test public void testRateLimiterMaxBalance() { MockClock clock = new MockClock(); RateLimiter limiter = new RateLimiter(0.1, 1.0, clock); long currentTime = TimeUnit.MICROSECONDS.toNanos(100); clock.timeNanos = currentTime; assertTrue(limiter.checkCredit(1.0)); assertFalse(limiter.checkCredit(1.0)); currentTime += TimeUnit.MILLISECONDS.toNanos(20000); clock.timeNanos = currentTime; assertTrue(limiter.checkCredit(1.0)); assertFalse(limiter.checkCredit(1.0)); }
@Test public void testRateLimiterInitial() { MockClock clock = new MockClock(); clock.timeNanos = TimeUnit.MILLISECONDS.toNanos(-1_000_000); RateLimiter limiter = new RateLimiter(1000, 100, clock); assertTrue(limiter.checkCredit(100)); assertFalse(limiter.checkCredit(1)); clock.timeNanos += TimeUnit.MILLISECONDS.toNanos(49); assertFalse(limiter.checkCredit(50)); clock.timeNanos += TimeUnit.MILLISECONDS.toNanos(1); assertTrue(limiter.checkCredit(50)); assertFalse(limiter.checkCredit(1)); clock.timeNanos += TimeUnit.MILLISECONDS.toNanos(1_000_000); assertTrue(limiter.checkCredit(1)); clock.timeNanos += TimeUnit.MILLISECONDS.toNanos(1_000_000); assertFalse(limiter.checkCredit(101)); assertTrue(limiter.checkCredit(100)); assertFalse(limiter.checkCredit(1)); }
@Test public void testRateLimiterConcurrency() { int numWorkers = ForkJoinPool.getCommonPoolParallelism(); int creditsPerWorker = 1000; MockClock clock = new MockClock(); RateLimiter limiter = new RateLimiter(1, numWorkers * creditsPerWorker, clock); AtomicInteger count = new AtomicInteger(); for (int w = 0; w < numWorkers; ++w) { ForkJoinPool.commonPool().execute(() -> { for (int i = 0; i < creditsPerWorker * 2; ++i) { if (limiter.checkCredit(1)) { count.getAndIncrement(); } } }); } ForkJoinPool.commonPool().awaitQuiescence(1, TimeUnit.SECONDS); assertEquals("Exactly the allocated number of credits must be consumed", numWorkers * creditsPerWorker,count.get()); assertFalse(limiter.checkCredit(1)); }