Skip to content

Commit 82dee61

Browse files
authored
Add TokenBucket algorithm (#5840)
1 parent 9dbdaf6 commit 82dee61

File tree

3 files changed

+153
-0
lines changed

3 files changed

+153
-0
lines changed

DIRECTORY.md

+2
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,7 @@
201201
* [PriorityQueues](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/queues/PriorityQueues.java)
202202
* [Queue](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/queues/Queue.java)
203203
* [QueueByTwoStacks](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/queues/QueueByTwoStacks.java)
204+
* [TokenBucket](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/queues/TokenBucket.java)
204205
* stacks
205206
* [NodeStack](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/stacks/NodeStack.java)
206207
* [ReverseStack](https://github.com/TheAlgorithms/Java/blob/master/src/main/java/com/thealgorithms/datastructures/stacks/ReverseStack.java)
@@ -865,6 +866,7 @@
865866
* [PriorityQueuesTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/queues/PriorityQueuesTest.java)
866867
* [QueueByTwoStacksTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/queues/QueueByTwoStacksTest.java)
867868
* [QueueTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/queues/QueueTest.java)
869+
* [TokenBucketTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/queues/TokenBucketTest.java)
868870
* stacks
869871
* [NodeStackTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/stacks/NodeStackTest.java)
870872
* [StackArrayListTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/stacks/StackArrayListTest.java)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
package com.thealgorithms.datastructures.queues;
2+
3+
import java.util.concurrent.TimeUnit;
4+
5+
/**
6+
* TokenBucket implements a token bucket rate limiter algorithm.
7+
* This class is used to control the rate of requests in a distributed system.
8+
* It allows a certain number of requests (tokens) to be processed in a time frame,
9+
* based on the defined refill rate.
10+
*
11+
* Applications: Computer networks, API rate limiting, distributed systems, etc.
12+
*
13+
* @author Hardvan
14+
*/
15+
public final class TokenBucket {
16+
private final int maxTokens;
17+
private final int refillRate; // tokens per second
18+
private int tokens;
19+
private long lastRefill; // Timestamp in nanoseconds
20+
21+
/**
22+
* Constructs a TokenBucket instance.
23+
*
24+
* @param maxTokens Maximum number of tokens the bucket can hold.
25+
* @param refillRate The rate at which tokens are refilled (tokens per second).
26+
*/
27+
public TokenBucket(int maxTokens, int refillRate) {
28+
this.maxTokens = maxTokens;
29+
this.refillRate = refillRate;
30+
this.tokens = maxTokens;
31+
this.lastRefill = System.nanoTime();
32+
}
33+
34+
/**
35+
* Attempts to allow a request based on the available tokens.
36+
* If a token is available, it decrements the token count and allows the request.
37+
* Otherwise, the request is denied.
38+
*
39+
* @return true if the request is allowed, false if the request is denied.
40+
*/
41+
public synchronized boolean allowRequest() {
42+
refillTokens();
43+
if (tokens > 0) {
44+
tokens--;
45+
return true;
46+
}
47+
return false;
48+
}
49+
50+
/**
51+
* Refills the tokens based on the time elapsed since the last refill.
52+
* The number of tokens to be added is calculated based on the elapsed time
53+
* and the refill rate, ensuring the total does not exceed maxTokens.
54+
*/
55+
private void refillTokens() {
56+
long now = System.nanoTime();
57+
long tokensToAdd = (now - lastRefill) / TimeUnit.SECONDS.toNanos(1) * refillRate;
58+
tokens = Math.min(maxTokens, tokens + (int) tokensToAdd);
59+
lastRefill = now;
60+
}
61+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
package com.thealgorithms.datastructures.queues;
2+
3+
import static org.junit.jupiter.api.Assertions.assertFalse;
4+
import static org.junit.jupiter.api.Assertions.assertTrue;
5+
6+
import org.junit.jupiter.api.Test;
7+
8+
public class TokenBucketTest {
9+
10+
@Test
11+
public void testRateLimiter() throws InterruptedException {
12+
TokenBucket bucket = new TokenBucket(5, 1);
13+
for (int i = 0; i < 5; i++) {
14+
assertTrue(bucket.allowRequest());
15+
}
16+
assertFalse(bucket.allowRequest());
17+
Thread.sleep(1000);
18+
assertTrue(bucket.allowRequest());
19+
}
20+
21+
@Test
22+
public void testRateLimiterWithExceedingRequests() throws InterruptedException {
23+
TokenBucket bucket = new TokenBucket(3, 1);
24+
25+
for (int i = 0; i < 3; i++) {
26+
assertTrue(bucket.allowRequest());
27+
}
28+
assertFalse(bucket.allowRequest());
29+
30+
Thread.sleep(1000);
31+
assertTrue(bucket.allowRequest());
32+
assertFalse(bucket.allowRequest());
33+
}
34+
35+
@Test
36+
public void testRateLimiterMultipleRefills() throws InterruptedException {
37+
TokenBucket bucket = new TokenBucket(2, 1);
38+
39+
assertTrue(bucket.allowRequest());
40+
assertTrue(bucket.allowRequest());
41+
assertFalse(bucket.allowRequest());
42+
43+
Thread.sleep(1000);
44+
assertTrue(bucket.allowRequest());
45+
46+
Thread.sleep(1000);
47+
assertTrue(bucket.allowRequest());
48+
assertFalse(bucket.allowRequest());
49+
}
50+
51+
@Test
52+
public void testRateLimiterEmptyBucket() {
53+
TokenBucket bucket = new TokenBucket(0, 1);
54+
55+
assertFalse(bucket.allowRequest());
56+
}
57+
58+
@Test
59+
public void testRateLimiterWithHighRefillRate() throws InterruptedException {
60+
TokenBucket bucket = new TokenBucket(5, 10);
61+
62+
for (int i = 0; i < 5; i++) {
63+
assertTrue(bucket.allowRequest());
64+
}
65+
66+
assertFalse(bucket.allowRequest());
67+
68+
Thread.sleep(1000);
69+
70+
for (int i = 0; i < 5; i++) {
71+
assertTrue(bucket.allowRequest());
72+
}
73+
}
74+
75+
@Test
76+
public void testRateLimiterWithSlowRequests() throws InterruptedException {
77+
TokenBucket bucket = new TokenBucket(5, 1);
78+
79+
for (int i = 0; i < 5; i++) {
80+
assertTrue(bucket.allowRequest());
81+
}
82+
83+
Thread.sleep(1000);
84+
assertTrue(bucket.allowRequest());
85+
86+
Thread.sleep(2000);
87+
assertTrue(bucket.allowRequest());
88+
assertTrue(bucket.allowRequest());
89+
}
90+
}

0 commit comments

Comments
 (0)