@Test public void testTimeout() throws Exception { TokenBucket tokenBucket = new TokenBucket(100, 0); // If it cannot satisfy the request within the timeout, return false immediately Assert.assertFalse(tokenBucket.getTokens(100, 1, TimeUnit.MILLISECONDS)); Assert.assertFalse(tokenBucket.getTokens(100, 10, TimeUnit.MILLISECONDS)); Assert.assertFalse(tokenBucket.getTokens(100, 100, TimeUnit.MILLISECONDS)); Assert.assertTrue(tokenBucket.getTokens(10, 101, TimeUnit.MILLISECONDS)); // Can use stored tokens to satisfy request tokenBucket = new TokenBucket(100, 100); Thread.sleep(200); // fill up bucket Assert.assertTrue(tokenBucket.getTokens(20, 101, TimeUnit.MILLISECONDS)); }
long storedTokens = this.tokenBucket.getStoredTokens(); if (eagerTokens > requestedPermits && this.tokenBucket.getTokens(eagerTokens, 0, TimeUnit.MILLISECONDS)) { return eagerTokens; long millisToSatisfyMinPermits = (long) (minPermits / this.tokenBucket.getTokensPerMilli()); if (millisToSatisfyMinPermits > timeoutMillis) { return 0; if (this.tokenBucket.getTokens(requestedPermits, allowedTimeout, TimeUnit.MILLISECONDS)) { return requestedPermits; if (this.tokenBucket.getTokens(minPermits, allowedTimeout, TimeUnit.MILLISECONDS)) { return minPermits;
public TokenBucket(long qps, long maxBucketSizeInMillis) { this.nextTokenAvailableMillis = System.currentTimeMillis(); resetQPS(qps, maxBucketSizeInMillis); }
/** * @param qps the average qps desired. * @param fullRequestTimeoutMillis max time to fully satisfy a token request. This is generally a small timeout, on the * order of the network latency (e.g. ~100 ms). * @param maxBucketSizeMillis maximum number of unused tokens that can be stored during under-utilization time, in * milliseconds. The actual tokens stored will be 1000 * qps * maxBucketSizeMillis. */ DynamicTokenBucket(long qps, long fullRequestTimeoutMillis, long maxBucketSizeMillis) { this.tokenBucket = new TokenBucket(qps, maxBucketSizeMillis); this.baseTimeout = fullRequestTimeoutMillis; }
/** * Get the current number of stored tokens. Note this is a snapshot of the object, and there is no guarantee that those * tokens will be available at any point in the future. */ public long getStoredTokens() { synchronized (this) { updateTokensStored(System.currentTimeMillis()); } return (long) this.tokensStored; }
@Override public Boolean call() { try { return this.tokenBucket.getTokens(this.tokens, this.timeoutMillis, TimeUnit.MILLISECONDS); } catch (InterruptedException ie) { throw new RuntimeException(ie); } } }
@Override public void run() { DescriptiveStatistics stats = limiter.getRateStatsSinceLastReport(); if (stats != null) { log.info(String.format("Requests rate stats: count: %d, min: %f, max: %f, mean: %f, std: %f, sum: %f", stats.getN(), stats.getMin(), stats.getMax(), stats.getMean(), stats.getStandardDeviation(), stats.getSum())); } stats = limiter.getUnusedPermitsSinceLastReport(); if (stats != null) { log.info(String.format("Unused permits rate stats: count: %d, min: %f, max: %f, mean: %f, std: %f, sum: %f", stats.getN(), stats.getMin(), stats.getMax(), stats.getMean(), stats.getStandardDeviation(), stats.getSum())); } if (this.policy instanceof QPSPolicy) { QPSPolicy qpsPolicy = (QPSPolicy) this.policy; DynamicTokenBucket dynamicTokenBucket = qpsPolicy.getTokenBucket(); TokenBucket tokenBucket = dynamicTokenBucket.getTokenBucket(); log.info("Stored tokens: " + tokenBucket.getStoredTokens()); } } }
/** * Attempt to get the specified amount of tokens within the specified timeout. If the tokens cannot be retrieved in the * specified timeout, the call will return false immediately, otherwise, the call will block until the tokens are available. * * @return true if the tokens are granted. * @throws InterruptedException */ public boolean getTokens(long tokens, long timeout, TimeUnit timeoutUnit) throws InterruptedException { long timeoutMillis = timeoutUnit.toMillis(timeout); long wait; synchronized (this) { wait = tryReserveTokens(tokens, timeoutMillis); } if (wait < 0) { return false; } if (wait == 0) { return true; } Thread.sleep(wait); return true; }
private void testForQps(long qps) throws Exception { ExecutorService executorService = Executors.newFixedThreadPool(10); TokenBucket tokenBucket = new TokenBucket(qps, 1000); List<Future<Boolean>> futures = Lists.newArrayList(); long permitsPerRequest = qps / 10; long start = System.currentTimeMillis(); futures.add(executorService.submit(new MyRunnable(tokenBucket, permitsPerRequest, 1000))); futures.add(executorService.submit(new MyRunnable(tokenBucket, permitsPerRequest, 1000))); futures.add(executorService.submit(new MyRunnable(tokenBucket, permitsPerRequest, 1000))); futures.add(executorService.submit(new MyRunnable(tokenBucket, permitsPerRequest, 1000))); futures.add(executorService.submit(new MyRunnable(tokenBucket, permitsPerRequest, 1000))); futures.add(executorService.submit(new MyRunnable(tokenBucket, permitsPerRequest, 1000))); futures.add(executorService.submit(new MyRunnable(tokenBucket, permitsPerRequest, 1000))); futures.add(executorService.submit(new MyRunnable(tokenBucket, permitsPerRequest, 1000))); futures.add(executorService.submit(new MyRunnable(tokenBucket, permitsPerRequest, 1000))); futures.add(executorService.submit(new MyRunnable(tokenBucket, permitsPerRequest, 1000))); for (Future<Boolean> future : futures) { Assert.assertTrue(future.get()); } long end = System.currentTimeMillis(); double averageRate = 1000 * (double) (permitsPerRequest * futures.size()) / (end - start); Assert.assertTrue(Math.abs(averageRate - qps) / qps < 0.2, "Average rate: " + averageRate + " expected: 100"); }
public void resetQPS(long qps, long maxBucketSizeInMillis) { Preconditions.checkArgument(qps > 0, "QPS must be positive."); Preconditions.checkArgument(maxBucketSizeInMillis >= 0, "Max bucket size must be non-negative."); long now = System.currentTimeMillis(); synchronized (this) { updateTokensStored(now); if (this.nextTokenAvailableMillis > now) { this.tokensStored -= (this.nextTokenAvailableMillis - now) * this.tokensPerMilli; } this.tokensPerMilli = (double) qps / 1000; this.maxBucketSizeInTokens = this.tokensPerMilli * maxBucketSizeInMillis; } }
/** * Attempt to get the specified amount of tokens within the specified timeout. If the tokens cannot be retrieved in the * specified timeout, the call will return false immediately, otherwise, the call will block until the tokens are available. * * @return true if the tokens are granted. * @throws InterruptedException */ public boolean getTokens(long tokens, long timeout, TimeUnit timeoutUnit) throws InterruptedException { long timeoutMillis = timeoutUnit.toMillis(timeout); long wait; synchronized (this) { wait = tryReserveTokens(tokens, timeoutMillis); } if (wait < 0) { return false; } if (wait == 0) { return true; } Thread.sleep(wait); return true; }
long storedTokens = this.tokenBucket.getStoredTokens(); if (eagerTokens > requestedPermits && this.tokenBucket.getTokens(eagerTokens, 0, TimeUnit.MILLISECONDS)) { return eagerTokens; long millisToSatisfyMinPermits = (long) (minPermits / this.tokenBucket.getTokensPerMilli()); if (millisToSatisfyMinPermits > timeoutMillis) { return 0; if (this.tokenBucket.getTokens(requestedPermits, allowedTimeout, TimeUnit.MILLISECONDS)) { return requestedPermits; if (this.tokenBucket.getTokens(minPermits, allowedTimeout, TimeUnit.MILLISECONDS)) { return minPermits;
/** * @param qps the average qps desired. * @param fullRequestTimeoutMillis max time to fully satisfy a token request. This is generally a small timeout, on the * order of the network latency (e.g. ~100 ms). * @param maxBucketSizeMillis maximum number of unused tokens that can be stored during under-utilization time, in * milliseconds. The actual tokens stored will be 1000 * qps * maxBucketSizeMillis. */ DynamicTokenBucket(long qps, long fullRequestTimeoutMillis, long maxBucketSizeMillis) { this.tokenBucket = new TokenBucket(qps, maxBucketSizeMillis); this.baseTimeout = fullRequestTimeoutMillis; }
/** * Note: this method should only be called while holding the class lock. For performance, the lock is not explicitly * acquired. * * @return the wait until the tokens are available or negative if they can't be acquired in the give timeout. */ private long tryReserveTokens(long tokens, long maxWaitMillis) { long now = System.currentTimeMillis(); long waitUntilNextTokenAvailable = Math.max(0, this.nextTokenAvailableMillis - now); updateTokensStored(now); if (tokens <= this.tokensStored) { this.tokensStored -= tokens; return waitUntilNextTokenAvailable; } double additionalNeededTokens = tokens - this.tokensStored; // casting to long will round towards 0 long additionalWaitForEnoughTokens = (long) (additionalNeededTokens / this.tokensPerMilli) + 1; long totalWait = waitUntilNextTokenAvailable + additionalWaitForEnoughTokens; if (totalWait > maxWaitMillis) { return -1; } this.tokensStored = this.tokensPerMilli * additionalWaitForEnoughTokens - additionalNeededTokens; this.nextTokenAvailableMillis = this.nextTokenAvailableMillis + additionalWaitForEnoughTokens; return totalWait; }
public TokenBucket(long qps, long maxBucketSizeInMillis) { this.nextTokenAvailableMillis = System.currentTimeMillis(); resetQPS(qps, maxBucketSizeInMillis); }
/** * Get the current number of stored tokens. Note this is a snapshot of the object, and there is no guarantee that those * tokens will be available at any point in the future. */ public long getStoredTokens() { synchronized (this) { updateTokensStored(System.currentTimeMillis()); } return (long) this.tokensStored; }
public void resetQPS(long qps, long maxBucketSizeInMillis) { Preconditions.checkArgument(qps > 0, "QPS must be positive."); Preconditions.checkArgument(maxBucketSizeInMillis >= 0, "Max bucket size must be non-negative."); long now = System.currentTimeMillis(); synchronized (this) { updateTokensStored(now); if (this.nextTokenAvailableMillis > now) { this.tokensStored -= (this.nextTokenAvailableMillis - now) * this.tokensPerMilli; } this.tokensPerMilli = (double) qps / 1000; this.maxBucketSizeInTokens = this.tokensPerMilli * maxBucketSizeInMillis; } }
/** * Note: this method should only be called while holding the class lock. For performance, the lock is not explicitly * acquired. * * @return the wait until the tokens are available or negative if they can't be acquired in the give timeout. */ private long tryReserveTokens(long tokens, long maxWaitMillis) { long now = System.currentTimeMillis(); long waitUntilNextTokenAvailable = Math.max(0, this.nextTokenAvailableMillis - now); updateTokensStored(now); if (tokens <= this.tokensStored) { this.tokensStored -= tokens; return waitUntilNextTokenAvailable; } double additionalNeededTokens = tokens - this.tokensStored; // casting to long will round towards 0 long additionalWaitForEnoughTokens = (long) (additionalNeededTokens / this.tokensPerMilli) + 1; long totalWait = waitUntilNextTokenAvailable + additionalWaitForEnoughTokens; if (totalWait > maxWaitMillis) { return -1; } this.tokensStored = this.tokensPerMilli * additionalWaitForEnoughTokens - additionalNeededTokens; this.nextTokenAvailableMillis = this.nextTokenAvailableMillis + additionalWaitForEnoughTokens; return totalWait; }