/** * Enable the cache specifying size and expire time. * * @param bucketSize max number of jwks to deliver in the given rate. * @param refillRate amount of time to wait before a jwk can the jwk will be cached * @param unit unit of time for the expire of jwk * @return the builder */ public JwkProviderBuilder rateLimited(long bucketSize, long refillRate, TimeUnit unit) { bucket = new BucketImpl(bucketSize, refillRate, unit); return this; }
@Override public synchronized long willLeakIn(long count) { assertPositiveValue(count, size, String.format("Cannot consume %d tokens when the BucketImpl size is %d!", count, size)); updateAvailableTokens(); if (available >= count) { return 0; } long leakDelta = getTimeSinceLastTokenAddition(); if (leakDelta < getRatePerToken()) { leakDelta = getRatePerToken() - leakDelta; } final long remaining = count - available - 1; if (remaining > 0) { leakDelta += getRatePerToken() * remaining; } return leakDelta; }
@Override public synchronized boolean consume() { return consume(1); }
private void updateAvailableTokens() { final long ratePerToken = getRatePerToken(); final long elapsed = getTimeSinceLastTokenAddition(); if (elapsed < ratePerToken) { return; } accumDelta = elapsed % ratePerToken; long count = elapsed / ratePerToken; if (count > size - available) { count = size - available; } if (count > 0) { available += count; } restartStopWatch(); }
@Override public synchronized boolean consume(long count) { assertPositiveValue(count, size, String.format("Cannot consume %d tokens when the BucketImpl size is %d!", count, size)); updateAvailableTokens(); if (count <= available) { available -= count; return true; } return false; }
BucketImpl(long size, long rate, TimeUnit rateUnit) { assertPositiveValue(size, "Invalid bucket size."); assertPositiveValue(rate, "Invalid bucket refill rate."); this.stopwatch = Stopwatch.createStarted(); this.size = size; this.available = size; this.rate = rate; this.rateUnit = rateUnit; }
@Override public synchronized long willLeakIn() { return willLeakIn(1); }
private void assertPositiveValue(Number value, String exceptionMessage) { this.assertPositiveValue(value.intValue(), value.intValue(), exceptionMessage); }
/** * Creates a new Builder with the given URL where to load the jwks from. * * @param url to load the jwks * @throws IllegalStateException if url is null */ public JwkProviderBuilder(URL url) { if (url == null) { throw new IllegalStateException("Cannot build provider without url to jwks"); } this.url = url; this.cached = true; this.expiresIn = 10; this.expiresUnit = TimeUnit.HOURS; this.cacheSize = 5; this.rateLimited = true; this.bucket = new BucketImpl(10, 1, TimeUnit.MINUTES); }
@Test public void shouldConsumeAllBucketTokens() throws Exception { Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.MILLISECONDS); assertThat(bucket, notNullValue()); assertThat(bucket.consume(SIZE), equalTo(true)); assertThat(bucket.consume(), equalTo(false)); }
@Test public void shouldThrowOnCreateWithNegativeRate() throws Exception { exception.expect(IllegalArgumentException.class); exception.expectMessage("Invalid bucket refill rate."); new BucketImpl(10, -1, TimeUnit.SECONDS); }
@Test public void shouldCreateFullBucket() throws Exception { Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.MILLISECONDS); assertThat(bucket, notNullValue()); assertThat(bucket.willLeakIn(SIZE), equalTo(0L)); assertThat(bucket.willLeakIn(), equalTo(0L)); }
@Test public void shouldThrowOnCreateWithNegativeSize() throws Exception { exception.expect(IllegalArgumentException.class); exception.expectMessage("Invalid bucket size."); new BucketImpl(-1, 10, TimeUnit.SECONDS); }
@Test public void shouldThrowWhenLeakingMoreThanBucketSize() throws Exception { Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.SECONDS); exception.expect(IllegalArgumentException.class); exception.expectMessage(String.format("Cannot consume %d tokens when the BucketImpl size is %d!", SIZE + 1, SIZE)); bucket.willLeakIn(SIZE + 1); }
@Test public void shouldThrowWhenConsumingMoreThanBucketSize() throws Exception { Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.SECONDS); exception.expect(IllegalArgumentException.class); exception.expectMessage(String.format("Cannot consume %d tokens when the BucketImpl size is %d!", SIZE + 1, SIZE)); bucket.consume(SIZE + 1); }
@Test public void shouldNotAddMoreTokensThatTheBucketSize() throws Exception { Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.MILLISECONDS); assertThat(bucket, notNullValue()); assertThat(bucket.willLeakIn(SIZE), equalTo(0L)); //Give some time to fill the already full bucket pause(SIZE * RATE * 2); assertThat(bucket.consume(SIZE), equalTo(true)); assertThat(bucket.consume(), equalTo(false)); }
@Test public void shouldCalculateRemainingLeakTimeForOneToken() throws Exception { Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.MILLISECONDS); assertThat(bucket, notNullValue()); //Consume 5 tokens assertThat(bucket.consume(5), equalTo(true)); assertThat(bucket.willLeakIn(), allOf(greaterThan(0L), lessThanOrEqualTo(RATE))); // wait half rate time and check if the wait time is correct pause(RATE / 2); assertThat(bucket.willLeakIn(), allOf(greaterThan(0L), lessThanOrEqualTo(RATE / 2))); }
@Test public void shouldConsumeByOneToken() throws Exception { Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.MILLISECONDS); assertThat(bucket, notNullValue()); //Consume 5 tokens assertThat(bucket.consume(), equalTo(true)); assertThat(bucket.consume(), equalTo(true)); assertThat(bucket.consume(), equalTo(true)); assertThat(bucket.consume(), equalTo(true)); assertThat(bucket.consume(), equalTo(true)); //should not consume a 6th token assertThat(bucket.consume(), equalTo(false)); }
@Test public void shouldCarryDeltaWhenManyTokensAreRequested() throws Exception { Bucket bucket = new BucketImpl(5, 1000, TimeUnit.MILLISECONDS); assertThat(bucket, notNullValue()); //Consume all tokens. Expect to wait 5 seconds for refill assertThat(bucket.consume(5), equalTo(true)); assertThat(bucket.willLeakIn(5), allOf(greaterThanOrEqualTo(4900L), lessThanOrEqualTo(5000L))); //wait 1500ms to have 1 token. pause(1500); //Consume 1 and expect to wait 500 + 4000 ms if we want to consume 5 again. assertThat(bucket.consume(), equalTo(true)); assertThat(bucket.willLeakIn(5), allOf(greaterThanOrEqualTo(4400L), lessThanOrEqualTo(4500L))); }
@Test public void shouldCalculateRemainingLeakTimeForManyTokens() throws Exception { Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.MILLISECONDS); assertThat(bucket, notNullValue()); //Consume 3 tokens assertThat(bucket.consume(3), equalTo(true)); //Expected to wait 3 * RATE time at most to be able to consume 5 tokens assertThat(bucket.willLeakIn(5), allOf(greaterThanOrEqualTo(RATE * 2), lessThanOrEqualTo(RATE * 3))); pause(RATE * 3); assertThat(bucket.willLeakIn(5), allOf(greaterThanOrEqualTo(0L), lessThanOrEqualTo(RATE))); }