Java Token Bucket Rate Limiter Algorithm


Token Bucket Algorithm

The Token Bucket Algorithm is a more sophisticated rate limiting algorithm. It maintains a bucket of tokens, where each token represents a single request. Tokens are added to the bucket at a fixed rate, and requests are allowed only if there are enough tokens in the bucket. If the bucket is empty, then any subsequent requests are blocked until a token becomes available.

Example Java code for the Token Bucket Algorithm used in rate limiting:

import java.time.Instant;

public class TokenBucketRateLimiter {
    private final int capacity; // Maximum number of tokens the bucket can hold
    private final int refillRate; // Tokens to be added per second
    private int tokens; // Current number of tokens in the bucket
    private Instant lastRefillTime; // Last time the bucket was refilled

    public TokenBucketRateLimiter(int capacity, int refillRate) {
        this.capacity = capacity;
        this.refillRate = refillRate;
        this.tokens = capacity; // Initially, the bucket is full
        this.lastRefillTime = Instant.now(); // Initialize with current time
    }

    public synchronized boolean allowRequest() {
        refillTokens(); // Refill the bucket with tokens

        if (tokens > 0) {
            tokens--; // Consume a token
            return true; // Request is allowed
        } else {
            return false; // Request is not allowed
        }
    }

    private void refillTokens() {
        Instant now = Instant.now();
        long timeElapsed = now.getEpochSecond() - lastRefillTime.getEpochSecond();
        int tokensToAdd = (int) (timeElapsed * refillRate); 
       // Calculate tokens to add based on refill rate

        if (tokensToAdd > 0) {
            tokens = Math.min(capacity, tokens + tokensToAdd); 
           // Refill the bucket up to its capacity
            lastRefillTime = now; // Update the last refill time
        }
    }
}

In the above code, we have a TokenBucketRateLimiter class that implements the Token Bucket Algorithm for rate limiting. It takes the capacity (maximum number of tokens the bucket can hold) and the refill rate (tokens to be added per second) as parameters in the constructor.

The allowRequest() method is used to check if a request is allowed or not based on the availability of tokens in the bucket. It synchronizes the method to ensure thread safety.

The refillTokens() method is called before checking for token availability. It calculates the number of tokens to add based on the refill rate and the time elapsed since the last refill. It then refills the bucket up to its capacity and updates the last refill time.

You can create an instance of TokenBucketRateLimiter and use the allowRequest() method to control the rate at which requests are allowed.

Note that this is a simplified example, and in a real-world scenario, you may need to consider additional factors like token expiration, concurrency, and integration with your application’s framework or middleware.