├── .gitignore ├── .travis.yml ├── LICENSE ├── README.md ├── pom.xml └── src ├── main └── java │ └── me │ └── sudohippie │ └── throttle │ ├── Throttle.java │ ├── strategy │ ├── ThrottleStrategy.java │ └── bucket │ │ ├── FixedTokenBucketStrategy.java │ │ ├── LeakyTokenBucketStrategy.java │ │ ├── StepDownLeakyTokenBucketStrategy.java │ │ ├── StepUpLeakyTokenBucketStrategy.java │ │ └── TokenBucketStrategy.java │ └── util │ └── Assert.java └── test └── java └── me └── sudohippie └── throttle └── strategy └── bucket ├── FixedTokenBucketTest.java ├── StepDownLeakyTokenBucketTest.java └── StepUpLeakyTokenBucketTest.java /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled source # 2 | ################### 3 | *.com 4 | *.class 5 | *.dll 6 | *.exe 7 | *.o 8 | *.so 9 | 10 | # Packages # 11 | ############ 12 | # it's better to unpack these files and commit the raw source 13 | # git has its own built in compression methods 14 | *.7z 15 | *.dmg 16 | *.gz 17 | *.iso 18 | *.jar 19 | *.rar 20 | *.tar 21 | *.zip 22 | 23 | # Logs and databases # 24 | ###################### 25 | *.log 26 | *.sql 27 | *.sqlite 28 | 29 | # OS generated files # 30 | ###################### 31 | .DS_Store 32 | .DS_Store? 33 | ._* 34 | .Spotlight-V100 35 | .Trashes 36 | ehthumbs.db 37 | Thumbs.db 38 | 39 | # IntelliJ files # 40 | ################## 41 | *.iml 42 | .idea/ 43 | 44 | # Apache Maven # 45 | ################ 46 | target/ 47 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: java 2 | jdk: 3 | - openjdk7 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2014 Raghav Sidhanti 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Throttling API [![Build Status](https://travis-ci.org/sudohippie/throttle.png)](https://travis-ci.org/sudohippie/throttle) 2 | 3 | ## Overview 4 | API to throttle/rate-limit requests 5 | 6 | This API implements two popular throttling strategies, namely: 7 | 8 | 1. Fixed token bucket 9 | 2. Leaky token bucket 10 | 11 | ### Fixed token bucket 12 | Details for this implementation can be found at: [Token Bucket](http://en.wikipedia.org/wiki/Token_bucket) 13 | 14 | ### Leaky token bucket 15 | Details for this implementation can be found at: [Leaky Bucket](http://en.wikipedia.org/wiki/Leaky_bucket) 16 | With in the API, Leaky buckets have been implemented as two types 17 | 18 | 1. StepDownLeakyTokenBucketStrategy 19 | 2. StepUpLeakyTokenBucketStrategy 20 | 21 | StepDownLeakyTokenBucketStrategy resembles a bucket which has been filled with tokens at the beginning but subsequently leaks tokens at a fixed interval. 22 | StepUpLeakyTokenBucketStrategy resemembles an empty bucket at the beginning but get filled will tokens over a fixed interval. 23 | 24 | ## Examples 25 | 26 | ### Fixed Bucket Example 27 | 28 | ```java 29 | // construct strategy 30 | ThrottleStrategy strategy = new FixedTokenBucketStrategy(100, 1, TimeUnit.MINUTES); 31 | 32 | // provide the strategy to the throttler 33 | Throttle throttle = new Throttle(strategy); 34 | 35 | // throttle :) 36 | boolean isThrottled = throttle.canProceed(); 37 | 38 | if(!isThrottled){ 39 | // your logic 40 | } 41 | ``` 42 | 43 | ### Step Up Leaky Bucket Example 44 | ```java 45 | // construct strategy 46 | ThrottleStrategy strategy = new StepUpLeakyTokenBucketStrategy(100, 1, TimeUnit.MINUTES, 25, 15, TimeUnit.SECONDS); 47 | 48 | // provide the strategy to the throttler 49 | Throttle throttle = new Throttle(strategy); 50 | 51 | // throttle :) 52 | boolean isThrottled = throttle.canProceed(); 53 | 54 | if(!isThrottled){ 55 | // your logic 56 | } 57 | ``` 58 | 59 | ### Step Down Leaky Bucket Example 60 | ```java 61 | // construct strategy 62 | ThrottleStrategy strategy = new StepDownLeakyTokenBucketStrategy(100, 1, TimeUnit.MINUTES, 25, 15, TimeUnit.SECONDS); 63 | 64 | // provide the strategy to the throttler 65 | Throttle throttle = new Throttle(strategy); 66 | 67 | // throttle :) 68 | boolean isThrottled = throttle.canProceed(); 69 | 70 | if(!isThrottled){ 71 | // your logic 72 | } 73 | ``` 74 | 75 | ### Wait For Token Availability Example 76 | ```java 77 | // construct strategy 78 | ThrottleStrategy strategy = new FixedTokenBucketStrategy(100, 1, TimeUnit.MINUTES); 79 | 80 | // provide the strategy to the throttler 81 | Throttle throttle = new Throttle(strategy); 82 | 83 | while(!throttle.canProceed()){ 84 | Thread.sleep(throttle.waitTime(TimeUnit.MILLISECONDS)); 85 | } 86 | 87 | // your logic 88 | ``` 89 | 90 | 91 | -------------------------------------------------------------------------------- /pom.xml: -------------------------------------------------------------------------------- 1 | 2 | 5 | 4.0.0 6 | 7 | me.sudohippie 8 | throttle 9 | 1.0-SNAPSHOT 10 | 11 | 12 | 13 | junit 14 | junit 15 | 4.10 16 | 17 | 18 | 19 | 20 | 21 | 22 | maven-compiler-plugin 23 | 3.0 24 | 25 | 1.7 26 | 1.7 27 | 28 | 29 | 30 | 31 | 32 | -------------------------------------------------------------------------------- /src/main/java/me/sudohippie/throttle/Throttle.java: -------------------------------------------------------------------------------- 1 | package me.sudohippie.throttle; 2 | 3 | import me.sudohippie.throttle.strategy.ThrottleStrategy; 4 | 5 | import java.util.concurrent.TimeUnit; 6 | 7 | /** 8 | * Bridge to enable throttling. 9 | * 10 | * User must specify his/her desired throttling strategy. 11 | * 12 | * Raghav Sidhanti 13 | * 9/25/13 14 | */ 15 | public class Throttle { 16 | 17 | private final ThrottleStrategy strategy; 18 | 19 | public Throttle(ThrottleStrategy strategy) { 20 | this.strategy = strategy; 21 | } 22 | 23 | /** 24 | * Based on the throttling strategy chosen, this method 25 | * will return true if the request can be serviced. 26 | * false is returned when the return can not be serviced. 27 | * 28 | * @return 29 | */ 30 | public boolean canProceed(){ 31 | return !strategy.isThrottled(); 32 | } 33 | 34 | /** 35 | * If the request can be serviced, the wait time will be 0 else a 36 | * positive value in the chosen {@code TimeUnit}. 37 | * @param timeUnit TimeUnit indicating the wait time. 38 | * @return 0 or more value in the time unit chosen 39 | */ 40 | public long waitTime(TimeUnit timeUnit){ 41 | return strategy.timeToRelease(1,timeUnit); 42 | } 43 | 44 | } 45 | -------------------------------------------------------------------------------- /src/main/java/me/sudohippie/throttle/strategy/ThrottleStrategy.java: -------------------------------------------------------------------------------- 1 | package me.sudohippie.throttle.strategy; 2 | 3 | import java.util.concurrent.TimeUnit; 4 | 5 | /** 6 | * Represents a throttling strategy. 7 | * 8 | * Raghav Sidhanti 9 | * 9/25/13 10 | */ 11 | public abstract class ThrottleStrategy { 12 | 13 | /** 14 | * This method will return true if a single request 15 | * has been throttled (rate-limited) or false if 16 | * the request is not throttled (not rate-limited). 17 | * 18 | * @return 19 | */ 20 | public abstract boolean isThrottled(); 21 | 22 | /** 23 | * This method will return true if {@code n} requests 24 | * have been throttled or false if {@code n} requests 25 | * have not been throttled. 26 | * 27 | * @param n 28 | * @return 29 | */ 30 | public abstract boolean isThrottled(long n); 31 | 32 | /** 33 | * Provides time for n tokens to be available. 34 | * 35 | * @param n 36 | * @param timeUnit 37 | * @return 38 | */ 39 | public abstract long timeToRelease(long n, TimeUnit timeUnit); 40 | 41 | /** 42 | * Defines the maximum tokens this strategy can handle. 43 | * 44 | * @return 45 | */ 46 | public abstract long getCapacity(); 47 | } 48 | -------------------------------------------------------------------------------- /src/main/java/me/sudohippie/throttle/strategy/bucket/FixedTokenBucketStrategy.java: -------------------------------------------------------------------------------- 1 | package me.sudohippie.throttle.strategy.bucket; 2 | 3 | import java.util.concurrent.TimeUnit; 4 | 5 | /** 6 | * FixedTokenBucketStrategy is a concrete implementation of {@link TokenBucketStrategy}. 7 | * 8 | * In this strategy, at the beginning of every refill interval the bucket 9 | * is filled to capacity with tokens. Every call to {@link FixedTokenBucketStrategy#isThrottled()} 10 | * reduces the number of tokens with in the bucket by a fixed amount. 11 | * 12 | * Details of this strategy can be found at {@see Token Bucket} 13 | * 14 | * Raghav Sidhanti 15 | * 9/25/13 16 | */ 17 | public class FixedTokenBucketStrategy extends TokenBucketStrategy { 18 | 19 | protected FixedTokenBucketStrategy(long bucketTokenCapacity, long refillInterval, TimeUnit refillIntervalTimeUnit) { 20 | super(bucketTokenCapacity, refillInterval, refillIntervalTimeUnit); 21 | } 22 | 23 | @Override 24 | protected void updateTokens() { 25 | // refill bucket if current time has exceed next refill time 26 | long currentTime = System.currentTimeMillis(); 27 | if(currentTime < nextRefillTime) return; 28 | 29 | tokens = bucketTokenCapacity; 30 | nextRefillTime = currentTime + refillInterval; 31 | } 32 | } 33 | -------------------------------------------------------------------------------- /src/main/java/me/sudohippie/throttle/strategy/bucket/LeakyTokenBucketStrategy.java: -------------------------------------------------------------------------------- 1 | package me.sudohippie.throttle.strategy.bucket; 2 | 3 | import me.sudohippie.throttle.util.Assert; 4 | 5 | import java.util.concurrent.TimeUnit; 6 | 7 | /** 8 | * LeakyTokenBucketStrategy is a representation of {@link TokenBucketStrategy}. 9 | * 10 | * Details of this strategy can be found at {@see Leaky Bucket} 11 | * 12 | * Raghav Sidhanti 13 | * 9/25/13 14 | */ 15 | public abstract class LeakyTokenBucketStrategy extends TokenBucketStrategy { 16 | 17 | // step time interval in millis 18 | protected final long stepInterval; 19 | 20 | // step token count 21 | protected final long stepTokens; 22 | 23 | protected LeakyTokenBucketStrategy(long maxTokens, long refillInterval, TimeUnit refillIntervalTimeUnit, long stepTokens, long stepInterval, TimeUnit stepIntervalTimeUnit) { 24 | super(maxTokens, refillInterval, refillIntervalTimeUnit); 25 | 26 | // preconditions 27 | Assert.isTrue(stepInterval >= 0, "Step interval can not be negative"); 28 | Assert.isTrue(stepTokens >= 0, "Step token can not be negative"); 29 | 30 | this.stepTokens = stepTokens; 31 | this.stepInterval = stepIntervalTimeUnit.toMillis(stepInterval); 32 | } 33 | } 34 | -------------------------------------------------------------------------------- /src/main/java/me/sudohippie/throttle/strategy/bucket/StepDownLeakyTokenBucketStrategy.java: -------------------------------------------------------------------------------- 1 | package me.sudohippie.throttle.strategy.bucket; 2 | 3 | import java.util.concurrent.TimeUnit; 4 | 5 | /** 6 | * StepDownLeakyTokenBucketStrategy is concrete representation of {@link LeakyTokenBucketStrategy}. 7 | * 8 | * This strategy is synonymous to a leaking bucket filled with some substance (in this case tokens). Here, 9 | * at the beginning of every refill interval, the bucket is filled to capacity with tokens. Over time, 10 | * the bucket leaks tokens at a constant rate as defined by input parameters stepTokens and stepInterval until 11 | * the start of the next interval. 12 | * 13 | * Raghav Sidhanti 14 | * 9/25/13 15 | */ 16 | public class StepDownLeakyTokenBucketStrategy extends LeakyTokenBucketStrategy { 17 | 18 | /** 19 | * Constructor to build a StepDownLeakyTokenBucketStrategy. 20 | * 21 | * @param maxTokens The maximum tokens this bucket can hold. 22 | * @param refillInterval The interval at which the bucket must be refilled to capacity with tokens. 23 | * @param refillIntervalTimeUnit {@link TimeUnit} class representing unit of time of refill interval 24 | * @param stepTokens The number of tokens this bucket leaks, at every step interval. 25 | * @param stepInterval The interval at which the token leaks tokens. 26 | * @param stepIntervalTimeUnit {@link TimeUnit} class representing unit of time of step interval 27 | */ 28 | public StepDownLeakyTokenBucketStrategy(long maxTokens, long refillInterval, TimeUnit refillIntervalTimeUnit, long stepTokens, long stepInterval, TimeUnit stepIntervalTimeUnit) { 29 | super(maxTokens, refillInterval, refillIntervalTimeUnit, stepTokens, stepInterval, stepIntervalTimeUnit); 30 | } 31 | 32 | @Override 33 | protected void updateTokens() { 34 | long currentTime = System.currentTimeMillis(); 35 | 36 | // if current time exceeds next refill time 37 | if(currentTime >= nextRefillTime){ 38 | // set tokens to max 39 | tokens = bucketTokenCapacity; 40 | // calculate next refill time 41 | nextRefillTime = currentTime + refillInterval; 42 | 43 | return; 44 | } 45 | // calculate max tokens possible till end 46 | long timeToNextRefill = nextRefillTime - currentTime; 47 | long stepsToNextRefill = timeToNextRefill / stepInterval; 48 | long maxPossibleTokens = stepsToNextRefill * stepTokens; 49 | // edge case, if current time not at edge of step 50 | if((timeToNextRefill % stepInterval) > 0) maxPossibleTokens += stepTokens; 51 | // tokens must be lesser of current and max possible tokens 52 | if(maxPossibleTokens < tokens) tokens = maxPossibleTokens; 53 | } 54 | } 55 | -------------------------------------------------------------------------------- /src/main/java/me/sudohippie/throttle/strategy/bucket/StepUpLeakyTokenBucketStrategy.java: -------------------------------------------------------------------------------- 1 | package me.sudohippie.throttle.strategy.bucket; 2 | 3 | import java.util.concurrent.TimeUnit; 4 | 5 | /** 6 | * StepUpLeakyTokenBucketStrategy is concrete representation of {@link LeakyTokenBucketStrategy}. 7 | * 8 | * This strategy is synonymous to an empty bucket being filled with some substance (here tokens) over time. 9 | * Here, at the beginning of every refill interval the bucket is emptied. The bucket is then gradually filled 10 | * with tokens and rate is defined by the input parameters stepTokens and stepInterval. 11 | * 12 | * Raghav Sidhanti 13 | * 9/25/13 14 | */ 15 | public class StepUpLeakyTokenBucketStrategy extends LeakyTokenBucketStrategy { 16 | 17 | private long lastActivityTime; 18 | 19 | /** 20 | * Constructor to build a StepUpLeakyTokenBucketStrategy. 21 | * 22 | * @param maxTokens The maximum tokens this bucket can hold. 23 | * @param refillInterval The interval at which the bucket must be emptied. 24 | * @param refillIntervalTimeUnit {@link java.util.concurrent.TimeUnit} class representing unit of time of refill interval 25 | * @param stepTokens The number of tokens added to the bucket at every step interval. 26 | * @param stepInterval The interval at which tokens are added. 27 | * @param stepIntervalTimeUnit {@link java.util.concurrent.TimeUnit} class representing unit of time of step interval 28 | */ 29 | public StepUpLeakyTokenBucketStrategy(long maxTokens, long refillInterval, TimeUnit refillIntervalTimeUnit, long stepTokens, long stepInterval, TimeUnit stepIntervalTimeUnit) { 30 | super(maxTokens, refillInterval, refillIntervalTimeUnit, stepTokens, stepInterval, stepIntervalTimeUnit); 31 | } 32 | 33 | @Override 34 | protected void updateTokens() { 35 | long currentTime = System.currentTimeMillis(); 36 | 37 | // if current time has exceeded next refill interval, 38 | if(currentTime >= nextRefillTime){ 39 | tokens = stepTokens; 40 | lastActivityTime = currentTime; 41 | nextRefillTime = currentTime + refillInterval; 42 | 43 | return; 44 | } 45 | 46 | // calculate tokens at current step 47 | long elapsedTimeSinceLastActivity = currentTime - lastActivityTime; 48 | long elapsedStepsSinceLastActivity = elapsedTimeSinceLastActivity / stepInterval; 49 | tokens += (elapsedStepsSinceLastActivity * stepTokens); 50 | 51 | // check for bucket overflow 52 | if(tokens > bucketTokenCapacity) tokens = bucketTokenCapacity; 53 | 54 | // update last activity time 55 | lastActivityTime = currentTime; 56 | } 57 | } 58 | -------------------------------------------------------------------------------- /src/main/java/me/sudohippie/throttle/strategy/bucket/TokenBucketStrategy.java: -------------------------------------------------------------------------------- 1 | package me.sudohippie.throttle.strategy.bucket; 2 | 3 | import me.sudohippie.throttle.strategy.ThrottleStrategy; 4 | import me.sudohippie.throttle.util.Assert; 5 | 6 | import java.util.concurrent.TimeUnit; 7 | 8 | /** 9 | * Abstract class representing a token bucket strategy. 10 | * 11 | * Using this strategy, throttling is enforced via the existence of tokens in a bucket. 12 | * 13 | * Raghav Sidhanti 14 | * 9/25/13 15 | */ 16 | public abstract class TokenBucketStrategy extends ThrottleStrategy { 17 | 18 | protected final long bucketTokenCapacity; 19 | protected final long refillInterval; 20 | 21 | // number of tokens in the bucket 22 | protected long tokens = 0; 23 | protected long nextRefillTime = 0; 24 | 25 | protected TokenBucketStrategy(long bucketTokenCapacity, long refillInterval, TimeUnit refillIntervalTimeUnit) { 26 | Assert.isTrue(bucketTokenCapacity >= 0, "Bucket token capacity can not be negative"); 27 | Assert.isTrue(refillInterval >= 0, "Bucket refill interval can not be negative"); 28 | 29 | this.bucketTokenCapacity = bucketTokenCapacity; 30 | this.refillInterval = refillIntervalTimeUnit.toMillis(refillInterval); 31 | } 32 | 33 | @Override 34 | public synchronized boolean isThrottled() { 35 | return isThrottled(1); 36 | } 37 | 38 | @Override 39 | public synchronized boolean isThrottled(long n) { 40 | // preconditions 41 | Assert.isTrue(n >= 0, "Invalid argument less than 0"); 42 | 43 | // check whether there exist at least n tokens in bucket 44 | if(getCurrentTokenCount() < n) return true; 45 | 46 | tokens -= n; 47 | return false; 48 | } 49 | 50 | @Override 51 | public long getCapacity() { 52 | return bucketTokenCapacity; 53 | } 54 | 55 | public synchronized long getCurrentTokenCount() { 56 | updateTokens(); 57 | return tokens; 58 | } 59 | 60 | @Override 61 | public synchronized long timeToRelease(long n, TimeUnit timeUnit){ 62 | // preconditions 63 | Assert.isTrue(n >= 0, "Invalid argument less than 0"); 64 | Assert.isTrue(timeUnit != null, "TimeUnit argument can not be null"); 65 | 66 | // check whether tokens exist 67 | if(getCurrentTokenCount() >= n){ 68 | return 0L; 69 | } else{ 70 | long timeToIntervalEnd = nextRefillTime - System.currentTimeMillis(); 71 | // edge case due to system slowness 72 | if(timeToIntervalEnd < 0) return timeToRelease(n, timeUnit); 73 | else return timeUnit.convert(timeToIntervalEnd, TimeUnit.MILLISECONDS); 74 | } 75 | } 76 | 77 | protected abstract void updateTokens(); 78 | } 79 | -------------------------------------------------------------------------------- /src/main/java/me/sudohippie/throttle/util/Assert.java: -------------------------------------------------------------------------------- 1 | package me.sudohippie.throttle.util; 2 | 3 | /** 4 | * Raghav Sidhanti 5 | * Date: 9/25/13 6 | */ 7 | public class Assert { 8 | public static void isTrue(boolean bool, String message){ 9 | if(!bool) throw new IllegalArgumentException(message); 10 | } 11 | } 12 | -------------------------------------------------------------------------------- /src/test/java/me/sudohippie/throttle/strategy/bucket/FixedTokenBucketTest.java: -------------------------------------------------------------------------------- 1 | package me.sudohippie.throttle.strategy.bucket; 2 | 3 | import org.junit.Before; 4 | import org.junit.Test; 5 | 6 | import java.util.concurrent.CountDownLatch; 7 | import java.util.concurrent.TimeUnit; 8 | 9 | import static org.junit.Assert.*; 10 | 11 | /** 12 | * Raghav Sidhanti 13 | * Date: 9/25/13 14 | */ 15 | public class FixedTokenBucketTest { 16 | 17 | /* test state */ 18 | // test max token when negative 19 | @Test(expected = IllegalArgumentException.class) 20 | public void testWhenMaxTokenIsNegative(){ 21 | FixedTokenBucketStrategy bucket = new FixedTokenBucketStrategy(-1, 0, TimeUnit.MILLISECONDS); 22 | } 23 | 24 | // test refill interval when negative 25 | @Test(expected = IllegalArgumentException.class) 26 | public void testWhenRefillIntervalIsNegative(){ 27 | FixedTokenBucketStrategy bucket = new FixedTokenBucketStrategy(0, -1, TimeUnit.MILLISECONDS); 28 | } 29 | 30 | /* test logic */ 31 | // test n throttle and fixed refill interval, fixed max tokens 32 | long MAX_TOKENS = 10; 33 | long REFILL_INTERVAL = 10; 34 | TimeUnit REFILL_INTERVAL_TIME_UNIT = TimeUnit.SECONDS; 35 | 36 | long N_LESS_THAN_MAX = 2; 37 | long N_GREATER_THAN_MAX = 12; 38 | int CUMULATIVE = 3; 39 | 40 | FixedTokenBucketStrategy bucket; 41 | 42 | @Before 43 | public void setUp() throws Exception { 44 | bucket = new FixedTokenBucketStrategy(MAX_TOKENS, REFILL_INTERVAL, REFILL_INTERVAL_TIME_UNIT); 45 | } 46 | 47 | // single threaded 48 | // when n is less than max token 49 | // token=max - n 50 | // isThrottle=false 51 | @Test 52 | public void testWhenNIsLessThanMaxTokens(){ 53 | boolean throttled = bucket.isThrottled(N_LESS_THAN_MAX); 54 | long tokens = bucket.getCurrentTokenCount(); 55 | 56 | assertFalse(throttled); 57 | assertEquals(MAX_TOKENS - N_LESS_THAN_MAX, tokens); 58 | } 59 | 60 | // when n is greater than max token 61 | // token=max 62 | // isThrottle=true 63 | @Test 64 | public void testWhenNIsGreaterThanMaxTokens() { 65 | boolean throttled = bucket.isThrottled(N_GREATER_THAN_MAX); 66 | long tokens = bucket.getCurrentTokenCount(); 67 | 68 | assertTrue(throttled); 69 | assertEquals(MAX_TOKENS, tokens); 70 | } 71 | 72 | // when cumulative n is less than max token 73 | // token=max - (n1+n2 ...) 74 | // isThrottle=false 75 | @Test 76 | public void testWhenCumulativeNIsLessThanMaxTokens() { 77 | // throttle 3 times 78 | for(int i = 0; i < CUMULATIVE; i++) assertFalse(bucket.isThrottled(N_LESS_THAN_MAX)); 79 | 80 | long tokens = bucket.getCurrentTokenCount(); 81 | assertEquals(MAX_TOKENS - (CUMULATIVE *N_LESS_THAN_MAX), tokens); 82 | } 83 | 84 | // when cumulative n is greater than max token 85 | // token=max 86 | // isThrottle=true 87 | 88 | @Test 89 | public void testWhenCumulativeNIsGreaterThanMaxTokens() { 90 | // throttle 3 times 91 | for(int i = 0; i < CUMULATIVE; i++) assertTrue(bucket.isThrottled(N_GREATER_THAN_MAX)); 92 | 93 | long tokens = bucket.getCurrentTokenCount(); 94 | assertEquals(MAX_TOKENS, tokens); 95 | } 96 | 97 | // when n is less than max token, sleep for refill interval, n is less than max token 98 | // tokens=max-n&&tokens=max-n 99 | // isThrottle=false && isThrottle=false 100 | 101 | @Test 102 | public void testWhenNLessThanMaxSleepNLessThanMax() throws InterruptedException { 103 | boolean before = bucket.isThrottled(N_LESS_THAN_MAX); 104 | long tokensBefore = bucket.getCurrentTokenCount(); 105 | 106 | assertFalse(before); 107 | assertEquals(MAX_TOKENS - N_LESS_THAN_MAX, tokensBefore); 108 | 109 | Thread.sleep(REFILL_INTERVAL_TIME_UNIT.toMillis(REFILL_INTERVAL)); 110 | 111 | boolean after = bucket.isThrottled(N_LESS_THAN_MAX); 112 | long afterTokens = bucket.getCurrentTokenCount(); 113 | 114 | assertFalse(after); 115 | assertEquals(MAX_TOKENS - N_LESS_THAN_MAX, afterTokens); 116 | } 117 | 118 | // when n is greater than max token, sleep for refill interval, n is greater than max token 119 | //tokens=max&&tokens=max 120 | // isThrottle=true && isThrottle=true 121 | public void testWhenNGreaterThanMaxSleepNGreaterThanMax() throws InterruptedException { 122 | boolean before = bucket.isThrottled(N_GREATER_THAN_MAX); 123 | long tokensBefore = bucket.getCurrentTokenCount(); 124 | 125 | assertTrue(before); 126 | assertEquals(MAX_TOKENS, tokensBefore); 127 | 128 | Thread.sleep(REFILL_INTERVAL_TIME_UNIT.toMillis(REFILL_INTERVAL)); 129 | 130 | boolean after = bucket.isThrottled(N_GREATER_THAN_MAX); 131 | long afterTokens = bucket.getCurrentTokenCount(); 132 | 133 | assertTrue(after); 134 | assertEquals(MAX_TOKENS, afterTokens); 135 | } 136 | 137 | // when cumulative n is less than max token, sleep for refill interval, cumulative n is less than max total 138 | // tokens=max-(n1+...)&&tokens=max-(n1+...) 139 | // isThrottle=false && isThrottle=false 140 | 141 | @Test 142 | public void testWhenCumulativeNLessThanMaxSleepCumulativeNLessThanMax() throws InterruptedException { 143 | // throttle 3 times 144 | int sum = 0; 145 | 146 | for(int i = 0; i < CUMULATIVE; i++){ 147 | assertFalse(bucket.isThrottled(N_LESS_THAN_MAX)); 148 | sum += N_LESS_THAN_MAX; 149 | } 150 | long beforeTokens = bucket.getCurrentTokenCount(); 151 | assertEquals(MAX_TOKENS - sum, beforeTokens); 152 | 153 | Thread.sleep(REFILL_INTERVAL_TIME_UNIT.toMillis(REFILL_INTERVAL)); 154 | 155 | // throttle 3 times 156 | for(int i = 0; i < CUMULATIVE; i++) assertFalse(bucket.isThrottled(N_LESS_THAN_MAX)); 157 | long afterTokens = bucket.getCurrentTokenCount(); 158 | assertEquals(MAX_TOKENS - sum, afterTokens); 159 | } 160 | 161 | // when cumulative n is less than max token, sleep for refill interval, cumulative n is greater than max token 162 | // tokens=max-(n1+...)&&tokens 0); 325 | assertTrue(nextRelease <= REFILL_INTERVAL); 326 | } 327 | 328 | // test when number of tokens requested are greater than existing in bucket 329 | @Test 330 | public void testNextReleaseWhenNumberOfTokensInBucketIsLessThanRequested(){ 331 | long limit = 3; 332 | bucket.isThrottled(MAX_TOKENS - limit); 333 | 334 | assertTrue(bucket.timeToRelease(limit + 1, TimeUnit.MILLISECONDS) > 0); 335 | } 336 | 337 | // test next release, when tokens don't exit in interval, sleep and check before start of next interval 338 | @Test 339 | public void testNextReleaseMultipleTimesAfterTokensAreExhaustedWithInInterval() throws InterruptedException { 340 | 341 | // overly glorified test. Realized all the tests are based on timing and system slowness can cause any of the 342 | // tests to fail. 343 | 344 | boolean hasExeededInterval = false; 345 | int repeats = 3; 346 | long releaseTime = 0L; 347 | 348 | do{ 349 | // max out tokens 350 | bucket.isThrottled(MAX_TOKENS); 351 | releaseTime = bucket.timeToRelease(1, TimeUnit.MILLISECONDS); 352 | 353 | assertTrue(releaseTime > 0); 354 | 355 | // sleep for a fraction 356 | long sleepInt = REFILL_INTERVAL_TIME_UNIT.toMillis(REFILL_INTERVAL) / 2; 357 | long start = System.currentTimeMillis(); 358 | Thread.sleep(sleepInt); 359 | long end = System.currentTimeMillis(); 360 | 361 | // check whether sleep time exceeded next release due to system slowness 362 | if(end - start >= releaseTime){ 363 | // if so, sleep until next release 364 | Thread.sleep(releaseTime); 365 | // repeat 366 | repeats --; 367 | hasExeededInterval = true; 368 | }else { 369 | hasExeededInterval = false; 370 | releaseTime = bucket.timeToRelease(1L, TimeUnit.MILLISECONDS); 371 | } 372 | }while(hasExeededInterval && repeats > 0); 373 | 374 | // test condition 375 | assertFalse(hasExeededInterval); 376 | assertTrue(releaseTime > 0); 377 | } 378 | 379 | // test next release, when tokens don't exit in interval, sleep for next release time and check after start of next interval 380 | @Test 381 | public void testNextReleaseWhenTokensHaveExhaustedButInNextInterval() throws InterruptedException { 382 | bucket.isThrottled(MAX_TOKENS); 383 | long nextRelease = bucket.timeToRelease(1L, TimeUnit.MILLISECONDS); 384 | 385 | assertTrue(nextRelease > 0L); 386 | 387 | Thread.sleep(nextRelease); 388 | 389 | assertEquals(0L, bucket.timeToRelease(1L, TimeUnit.MILLISECONDS)); 390 | } 391 | 392 | // test next release, unit conversion 393 | @Test 394 | public void testNextReleaseWhenUnitIsToMilliSeconds(){ 395 | bucket.isThrottled(MAX_TOKENS); 396 | long nextRelease = bucket.timeToRelease(1L, TimeUnit.MILLISECONDS); 397 | 398 | assertTrue(nextRelease > 0L); 399 | assertTrue(REFILL_INTERVAL_TIME_UNIT.convert(nextRelease, TimeUnit.MILLISECONDS) <= REFILL_INTERVAL); 400 | } 401 | 402 | @Test 403 | public void testNextReleaseWhenTimeUnitIsToSeconds(){ 404 | bucket.isThrottled(MAX_TOKENS); 405 | long nextRelease = bucket.timeToRelease(1L, TimeUnit.SECONDS); 406 | 407 | assertTrue(nextRelease > 0L); 408 | assertTrue(nextRelease <= REFILL_INTERVAL_TIME_UNIT.convert(nextRelease, TimeUnit.SECONDS)); 409 | } 410 | 411 | @Test 412 | public void testNextReleaseWhenTimeUnitIsHigherThanIntervalUnit(){ 413 | TimeUnit unit; 414 | 415 | switch (REFILL_INTERVAL_TIME_UNIT){ 416 | case NANOSECONDS: 417 | unit = TimeUnit.MICROSECONDS; 418 | break; 419 | case MICROSECONDS: 420 | unit = TimeUnit.MILLISECONDS; 421 | break; 422 | case MILLISECONDS: 423 | unit = TimeUnit.SECONDS; 424 | break; 425 | case SECONDS: 426 | unit = TimeUnit.MINUTES; 427 | break; 428 | case MINUTES: 429 | unit = TimeUnit.HOURS; 430 | break; 431 | case HOURS: 432 | unit = TimeUnit.DAYS; 433 | break; 434 | default: 435 | unit = TimeUnit.DAYS; 436 | break; 437 | } 438 | 439 | bucket.isThrottled(MAX_TOKENS); 440 | 441 | long nextRelease = bucket.timeToRelease(1L, unit); 442 | 443 | assertEquals(0L, nextRelease); 444 | 445 | } 446 | 447 | // test in multi threaded scenario, when both have tokens 448 | @Test 449 | public void testNextReleaseThreadedWhenTokensExist() throws InterruptedException { 450 | final long lessThanHalfTokens = MAX_TOKENS / 2 - 1; 451 | 452 | Thread t1 = new Thread(new Runnable() { 453 | @Override 454 | public void run() { 455 | bucket.isThrottled(lessThanHalfTokens); 456 | 457 | assertEquals(0L, bucket.timeToRelease(1L, TimeUnit.MILLISECONDS)); 458 | } 459 | }); 460 | 461 | Thread t2 = new Thread(new Runnable() { 462 | @Override 463 | public void run() { 464 | bucket.isThrottled(lessThanHalfTokens); 465 | 466 | assertEquals(0L, bucket.timeToRelease(1L, TimeUnit.MILLISECONDS)); 467 | } 468 | }); 469 | 470 | t1.start(); 471 | t2.start(); 472 | 473 | t1.join(); 474 | t2.join(); 475 | } 476 | 477 | // test in multi threaded scenario, when one thread gobbles all tokens and exits 478 | @Test 479 | public void testNextReleaseInAnIntervalWhenThread1GobblesAllTokens() throws InterruptedException { 480 | Thread t1 = new Thread(new Runnable() { 481 | @Override 482 | public void run() { 483 | bucket.isThrottled(MAX_TOKENS); 484 | 485 | assertTrue(bucket.timeToRelease(1L, TimeUnit.MILLISECONDS) > 0); 486 | } 487 | }); 488 | 489 | Thread t2 = new Thread(new Runnable() { 490 | @Override 491 | public void run() { 492 | try { 493 | assertTrue(bucket.timeToRelease(1L, TimeUnit.MILLISECONDS) > 0); 494 | 495 | Thread.sleep(REFILL_INTERVAL_TIME_UNIT.toMillis(REFILL_INTERVAL/2)); 496 | 497 | assertTrue(bucket.timeToRelease(1L, TimeUnit.MILLISECONDS) > 0); 498 | } catch (InterruptedException e) { 499 | e.printStackTrace(); 500 | } 501 | 502 | } 503 | }); 504 | 505 | t1.start(); 506 | t2.start(); 507 | 508 | t1.join(); 509 | t2.join(); 510 | } 511 | 512 | // test in multi threaded scenario, when one thread gobbles all tokens and exits other waits for next release 513 | @Test 514 | public void testNextReleaseWhenThread1GobblesAllTokensThread2WaitTillNextInterval() throws InterruptedException { 515 | 516 | final CountDownLatch latch = new CountDownLatch(1); 517 | 518 | Thread t1 = new Thread(new Runnable() { 519 | @Override 520 | public void run() { 521 | bucket.isThrottled(MAX_TOKENS); 522 | 523 | assertTrue(bucket.timeToRelease(1L, TimeUnit.MILLISECONDS) > 0); 524 | 525 | latch.countDown(); 526 | } 527 | }); 528 | 529 | Thread t2 = new Thread(new Runnable() { 530 | @Override 531 | public void run() { 532 | try { 533 | latch.await(); 534 | 535 | assertTrue(bucket.timeToRelease(1L, TimeUnit.MILLISECONDS) > 0); 536 | 537 | Thread.sleep(REFILL_INTERVAL_TIME_UNIT.toMillis(REFILL_INTERVAL)); 538 | 539 | assertEquals(0L, bucket.timeToRelease(1L, TimeUnit.MILLISECONDS)); 540 | } catch (InterruptedException e) { 541 | e.printStackTrace(); 542 | } 543 | 544 | } 545 | }); 546 | 547 | t1.start(); 548 | t2.start(); 549 | 550 | t1.join(); 551 | t2.join(); 552 | } 553 | } 554 | -------------------------------------------------------------------------------- /src/test/java/me/sudohippie/throttle/strategy/bucket/StepDownLeakyTokenBucketTest.java: -------------------------------------------------------------------------------- 1 | package me.sudohippie.throttle.strategy.bucket; 2 | 3 | import org.junit.Before; 4 | import org.junit.Test; 5 | 6 | import java.util.concurrent.TimeUnit; 7 | 8 | import static org.junit.Assert.*; 9 | 10 | /** 11 | * Raghav Sidhanti 12 | * 9/26/13 13 | */ 14 | public class StepDownLeakyTokenBucketTest { 15 | // test state only 16 | // when max tokens is negative 17 | @Test(expected = IllegalArgumentException.class) 18 | public void testWhenMaxTokensIsNegative() throws Exception { 19 | StepDownLeakyTokenBucketStrategy bucket = new StepDownLeakyTokenBucketStrategy(-1, REFILL_INTERVAL, TimeUnit.MILLISECONDS, STEP_TOKENS, STEP_INTERVAL, TimeUnit.MILLISECONDS); 20 | } 21 | 22 | // when refill interval is negative 23 | @Test(expected = IllegalArgumentException.class ) 24 | public void testWhenRefillIntervalIsNegative() throws Exception { 25 | StepDownLeakyTokenBucketStrategy bucket = new StepDownLeakyTokenBucketStrategy(MAX_TOKENS, -1, TimeUnit.MILLISECONDS, STEP_TOKENS, STEP_INTERVAL, TimeUnit.MILLISECONDS); 26 | } 27 | 28 | // when step token is negative 29 | @Test(expected = IllegalArgumentException.class) 30 | public void testWhenStepTokenIsNegative() throws Exception { 31 | StepDownLeakyTokenBucketStrategy bucket = new StepDownLeakyTokenBucketStrategy(MAX_TOKENS, REFILL_INTERVAL, TimeUnit.MILLISECONDS, -1, STEP_INTERVAL, TimeUnit.MILLISECONDS); 32 | } 33 | 34 | // when step interval is negative 35 | @Test(expected = IllegalArgumentException.class) 36 | public void testWhenStepIntervalIsNegative() throws Exception { 37 | StepDownLeakyTokenBucketStrategy bucket = new StepDownLeakyTokenBucketStrategy(MAX_TOKENS, REFILL_INTERVAL, TimeUnit.MILLISECONDS, STEP_TOKENS, -1, TimeUnit.MILLISECONDS); 38 | } 39 | 40 | @Before 41 | public void setUp() throws Exception { 42 | bucket = new StepDownLeakyTokenBucketStrategy(MAX_TOKENS, REFILL_INTERVAL, TimeUnit.MILLISECONDS, STEP_TOKENS, STEP_INTERVAL, TimeUnit.MILLISECONDS); 43 | } 44 | 45 | // setup state 46 | private final long MAX_TOKENS = 10; 47 | private final long REFILL_INTERVAL = 5000L; 48 | private final long STEP_TOKENS = 2; 49 | private final long STEP_INTERVAL = 1000L; 50 | 51 | private StepDownLeakyTokenBucketStrategy bucket; 52 | 53 | // test behaviour, single threaded 54 | // at the beginning, number of tokens should be equal to max 55 | 56 | @Test 57 | public void testTokensEqualsMaxAtFirstStep() throws Exception { 58 | assertEquals(MAX_TOKENS, bucket.getCurrentTokenCount()); 59 | } 60 | 61 | // after one whole interval, at the beginning number of tokens should be equal to max 62 | @Test 63 | public void testTokensEqualsMaxAtTheStartOfNextInterval() throws Exception { 64 | Thread.sleep(REFILL_INTERVAL); 65 | assertEquals(MAX_TOKENS, bucket.getCurrentTokenCount()); 66 | } 67 | 68 | // throttle max, full bucket at beginning of next interval 69 | @Test 70 | public void testTokensEqualsMaxAtTheStartOfNextIntervalAfterMaxThrottle() throws Exception { 71 | bucket.isThrottled(MAX_TOKENS); 72 | 73 | Thread.sleep(REFILL_INTERVAL); 74 | 75 | assertEquals(MAX_TOKENS, bucket.getCurrentTokenCount()); 76 | } 77 | 78 | // throttle max, not throttled at beginning of next interval 79 | 80 | @Test 81 | public void testNotThrottledWhenThrottledToMaxOnceWithInInterval() throws Exception { 82 | boolean throttled = bucket.isThrottled(MAX_TOKENS); 83 | 84 | assertFalse(throttled); 85 | } 86 | 87 | // throttle max, zero tokens till end of interval 88 | @Test 89 | public void testZeroTokensTillEndOfIntervalAfterThrottleMaxOnce() throws Exception { 90 | bucket.isThrottled(MAX_TOKENS); 91 | 92 | final long SLEEP_TIME = 1000L; 93 | for(int i = 1; i < (REFILL_INTERVAL / STEP_INTERVAL); i ++){ 94 | Thread.sleep(SLEEP_TIME); 95 | assertEquals(0, bucket.getCurrentTokenCount()); 96 | } 97 | } 98 | 99 | // throttle max, throttled till end of interval 100 | @Test 101 | public void testThrottledTillEndOfIntervalAfterThrottleMaxOnce() throws Exception { 102 | boolean throttled = bucket.isThrottled(MAX_TOKENS); 103 | assertFalse(throttled); 104 | 105 | for(int i = 1; i < (REFILL_INTERVAL / STEP_INTERVAL); i ++){ 106 | Thread.sleep(STEP_INTERVAL); 107 | assertTrue(bucket.isThrottled()); 108 | } 109 | } 110 | // sleep at every step interval, tokens must be equal to max - (i*stepTokens) at each interval 111 | 112 | @Test 113 | public void testTokensAtEachStepIntervalIsCorrect() throws Exception { 114 | for(int i = 0; i < (REFILL_INTERVAL / STEP_INTERVAL); i ++){ 115 | long expectedTokens = MAX_TOKENS - (i * STEP_TOKENS); 116 | 117 | assertEquals(expectedTokens, bucket.getCurrentTokenCount()); 118 | Thread.sleep(STEP_INTERVAL); 119 | } 120 | } 121 | 122 | // throttle n = 1 + max - (i * stepTokens), where i is a step. Should be throttled as we keep decreasing throttle n over time 123 | @Test 124 | public void testThrottledWhenNGreaterThanTokensAtEachStep() throws Exception { 125 | for(int i = 0; i < (REFILL_INTERVAL / STEP_INTERVAL); i ++){ 126 | long excessTokensForStep = MAX_TOKENS - (i * STEP_TOKENS) + 1; 127 | 128 | assertTrue(bucket.isThrottled(excessTokensForStep)); 129 | Thread.sleep(STEP_INTERVAL); 130 | } 131 | } 132 | 133 | // test behaviour, multi threaded 134 | // throttle n (n < max/2) each, number of tokens remaining should be max - 2*n 135 | 136 | @Test 137 | public void testTokenCountInMultiThreadedThrottleThread1AndThread2WithInAnInterval() throws Exception { 138 | 139 | Thread thread1 = new Thread(new Runnable() { 140 | @Override 141 | public void run() { 142 | bucket.isThrottled(); 143 | } 144 | }); 145 | 146 | Thread thread2 = new Thread(new Runnable() { 147 | @Override 148 | public void run() { 149 | bucket.isThrottled(); 150 | } 151 | }); 152 | 153 | thread1.start(); 154 | thread2.start(); 155 | 156 | thread1.join(); 157 | thread2.join(); 158 | 159 | assertEquals(MAX_TOKENS - 2, bucket.getCurrentTokenCount()); 160 | } 161 | 162 | // throttle t1 n > max, t2 n < max. t1 should be throttled, t2 should not be throttled. 163 | 164 | @Test 165 | public void testMultiThreadedWhenThread1ThrottledAboveMaxAndThread2BelowMax() throws Exception { 166 | Thread thread1 = new Thread(new Runnable() { 167 | @Override 168 | public void run() { 169 | boolean throttled = bucket.isThrottled(MAX_TOKENS + 1); 170 | assertTrue(throttled); 171 | } 172 | }); 173 | 174 | Thread thread2 = new Thread(new Runnable() { 175 | @Override 176 | public void run() { 177 | boolean throttled = bucket.isThrottled(MAX_TOKENS - 1); 178 | assertFalse(throttled); 179 | } 180 | }); 181 | 182 | thread1.start(); 183 | thread2.start(); 184 | 185 | thread1.join(); 186 | thread2.join(); 187 | } 188 | } 189 | -------------------------------------------------------------------------------- /src/test/java/me/sudohippie/throttle/strategy/bucket/StepUpLeakyTokenBucketTest.java: -------------------------------------------------------------------------------- 1 | package me.sudohippie.throttle.strategy.bucket; 2 | 3 | import org.junit.Before; 4 | import org.junit.Test; 5 | 6 | import java.util.concurrent.TimeUnit; 7 | 8 | import static org.junit.Assert.assertEquals; 9 | import static org.junit.Assert.assertFalse; 10 | 11 | /** 12 | * Raghav Sidhanti 13 | * 9/27/13 14 | */ 15 | public class StepUpLeakyTokenBucketTest { 16 | /* test state */ 17 | // max tokens can not be negative 18 | @Test(expected = IllegalArgumentException.class) 19 | public void testMaxTokensCanNotBeNegative() throws Exception { 20 | StepUpLeakyTokenBucketStrategy bucket = new StepUpLeakyTokenBucketStrategy(-1, REFILL_INTERVAL, TimeUnit.MILLISECONDS, STEP_TOKENS, STEP_INTERVAL, TimeUnit.MILLISECONDS); 21 | } 22 | 23 | // interval can not be negative 24 | @Test(expected = IllegalArgumentException.class) 25 | public void testRefillIntervalCanNotBeNegative() throws Exception { 26 | StepUpLeakyTokenBucketStrategy bucket = new StepUpLeakyTokenBucketStrategy(MAX_TOKENS, -1, TimeUnit.MILLISECONDS, STEP_TOKENS, STEP_INTERVAL, TimeUnit.MILLISECONDS); 27 | } 28 | 29 | // step tokens can not be negative 30 | @Test(expected = IllegalArgumentException.class) 31 | public void testStepTokensCanNotBeNegative() throws Exception { 32 | StepUpLeakyTokenBucketStrategy bucket = new StepUpLeakyTokenBucketStrategy(MAX_TOKENS, REFILL_INTERVAL, TimeUnit.MILLISECONDS, -1, STEP_INTERVAL, TimeUnit.MILLISECONDS); 33 | } 34 | 35 | // step interval can not be negative 36 | @Test(expected = IllegalArgumentException.class) 37 | public void testStepIntervalCanNotBeNegative() throws Exception { 38 | StepUpLeakyTokenBucketStrategy bucket = new StepUpLeakyTokenBucketStrategy(MAX_TOKENS, REFILL_INTERVAL, TimeUnit.MILLISECONDS, STEP_TOKENS, -1, TimeUnit.MILLISECONDS); 39 | } 40 | 41 | /* test behaviour for given state */ 42 | // setup state 43 | private final long MAX_TOKENS = 10; 44 | private final long REFILL_INTERVAL = 5000L; 45 | private final long STEP_TOKENS = 2; 46 | private final long STEP_INTERVAL = 1000L; 47 | 48 | private StepUpLeakyTokenBucketStrategy bucket; 49 | @Before 50 | public void setUp() throws Exception { 51 | bucket = new StepUpLeakyTokenBucketStrategy(MAX_TOKENS, REFILL_INTERVAL, TimeUnit.MILLISECONDS, STEP_TOKENS, STEP_INTERVAL, TimeUnit.MILLISECONDS); 52 | } 53 | 54 | // single threaded 55 | // for each step, token count must increase as (i*stepTokens) where i=1 <= n 56 | @Test 57 | public void testTokenCountIsCorrectAtEachStepWithInInterval() throws Exception { 58 | for (int i = 1; i <= (REFILL_INTERVAL/STEP_INTERVAL); i ++){ 59 | long tokenCount = bucket.getCurrentTokenCount(); 60 | long expectedTokensAtStep = STEP_TOKENS * i; 61 | 62 | assertEquals(expectedTokensAtStep, tokenCount); 63 | 64 | Thread.sleep(STEP_INTERVAL); 65 | } 66 | } 67 | 68 | // for each step, throttle should be false when n <= stepTokens 69 | @Test 70 | public void testThrottleIsFalseAtEachStepWithInInterval() throws Exception { 71 | for (int i = 1; i <= (REFILL_INTERVAL/STEP_INTERVAL); i ++){ 72 | boolean throttled = bucket.isThrottled(STEP_TOKENS); 73 | 74 | assertFalse(throttled); 75 | 76 | Thread.sleep(STEP_INTERVAL); 77 | } 78 | } 79 | 80 | // for each step, token should equal 0 after throttled at n=stepToken 81 | @Test 82 | public void testTokenEqualsZeroAfterThrottlingAtStepTokensAtEachInterval() throws Exception { 83 | for (int i = 1; i <= (REFILL_INTERVAL/STEP_INTERVAL); i ++){ 84 | bucket.isThrottled(STEP_TOKENS); 85 | 86 | assertEquals(0, bucket.getCurrentTokenCount()); 87 | } 88 | } 89 | 90 | // At the start of next interval, tokens should equal stepTokens 91 | @Test 92 | public void testTokensEqualsMaxAtStartOfNextInterval() throws Exception { 93 | bucket.isThrottled(); 94 | 95 | Thread.sleep(REFILL_INTERVAL); 96 | 97 | assertEquals(STEP_TOKENS, bucket.getCurrentTokenCount()); 98 | } 99 | 100 | // multi threaded 101 | // thread1 throttle n= steptoken, thread2 sleep for stepInterval. token must be 0 102 | @Test 103 | public void testMultiThreadedWhenThread1NEqualsToStepTokenAndThread2SleepForStepIntervalFollowedByNStepTokenThrottle() throws Exception { 104 | Thread t1 = new Thread(new Runnable() { 105 | @Override 106 | public void run() { 107 | boolean throttled = bucket.isThrottled(STEP_TOKENS); 108 | 109 | assertFalse(throttled); 110 | assertEquals(0, bucket.getCurrentTokenCount()); 111 | } 112 | }); 113 | 114 | Thread t2 = new Thread(new Runnable() { 115 | @Override 116 | public void run() { 117 | try { 118 | Thread.sleep(STEP_INTERVAL); 119 | boolean throttled = bucket.isThrottled(STEP_TOKENS); 120 | 121 | assertFalse(throttled); 122 | assertEquals(0, bucket.getCurrentTokenCount()); 123 | } catch (InterruptedException e) { 124 | e.printStackTrace(); //To change body of catch statement use File | Settings | File Templates. 125 | } 126 | } 127 | }); 128 | 129 | t1.start(); 130 | t2.start(); 131 | 132 | t1.join(); 133 | t2.join(); 134 | } 135 | } 136 | --------------------------------------------------------------------------------