Rate Limiter #1 –Token Bucket Algorithm

Advantages / Disadvantages

Advantages

Flexibility
- If the request is not over the bucket capacity, it will be able to handle a burst of traffic
- It is possible to adjust the bucket capacity dynamically
Ease of Implementation

Disadvantages

Hard to find an adequate bucket capacity size and refill rate
The time frame is fixed.
- If a request is concentrated in the back of the previous time frame and in the front time of the next time frame, the sum of them can exceed the server’s capacity

The requests in the red box can exceed capacity

Overhead of Atomic Operations
- In high-concurrency environments, managing atomic operations can introduce latency and reduce throughput.

Implementation

import ratelimiter.TokenBucketRateLimiter;

public class Main {
    public static void main(String args[]) throws Exception {
        TokenBucketRateLimiter rateLimiter = new TokenBucketRateLimiter(5, 5);

        for (int i=0;i<100;i++) {
            boolean canConsume = rateLimiter.tryConsume();
            if (canConsume) {
                System.out.println("Request " + i + " can be consumed!");
            } else {
                System.out.println("Request " + i + " is blocked by rate limiter");
            }
            Thread.sleep(100);
        }
    }

}

package ratelimiter;

public class TokenBucketRateLimiter {
    private final int capacity;
    private final int refillRatePerSecond;
    private int availableTokens;
    private long lastRefillTimeMillis;

    /**
     * Initialize token bucket with specified capacity and refill rate
     *
     * @param capacity
     * @param refillRatePerSecond
     */
    public TokenBucketRateLimiter(int capacity, int refillRatePerSecond) {
        this.capacity = capacity;
        this.refillRatePerSecond = refillRatePerSecond;

        this.availableTokens = capacity;
        this.lastRefillTimeMillis = System.currentTimeMillis();
    }

    /**
     * 1. Refill first. This allows we don't block unnecessarily
     * 2. Check available token count.
     *      - If there are tokens available, pass the request.
     *      - Otherwise, block.
     */
    public boolean tryConsume() {
        refill();

        if (availableTokens > 0) {
            availableTokens--;
            return true;
        } else {
            return false;
        }
    }

    /**
     * 1. Calculate how long it's been after the last refill
     * 2. Calculate how many tokens need to be added based on the result of #1.
     * 3. If tokens are > 0
     *      - Update available tokens count. The count shouldn't exceed the capacity
     */
    private void refill() {
        long currentTimeMillis = System.currentTimeMillis();
        long secondsSinceRefill = (currentTimeMillis - lastRefillTimeMillis) / 1_000;
        long tokensToAdd = secondsSinceRefill * refillRatePerSecond;

        System.out.println("  [" + currentTimeMillis + "] - Current availableTokens: " + availableTokens + ", tokensToAdd: " + tokensToAdd);
        if (tokensToAdd > 0) {
            availableTokens += tokensToAdd;
            if (availableTokens > capacity) {
                availableTokens = capacity;
            }

            lastRefillTimeMillis = currentTimeMillis;
        }
    }
}

Implementation – Output

  [1736670196156] - Current availableTokens: 5, tokensToAdd: 0
Request 0 can be consumed!
  [1736670196261] - Current availableTokens: 4, tokensToAdd: 0
Request 1 can be consumed!
  [1736670196367] - Current availableTokens: 3, tokensToAdd: 0
Request 2 can be consumed!
  [1736670196472] - Current availableTokens: 2, tokensToAdd: 0
Request 3 can be consumed!
  [1736670196577] - Current availableTokens: 1, tokensToAdd: 0
Request 4 can be consumed!
  [1736670196682] - Current availableTokens: 0, tokensToAdd: 0
Request 5 is blocked by rate limiter
  [1736670196787] - Current availableTokens: 0, tokensToAdd: 0
Request 6 is blocked by rate limiter
  [1736670196892] - Current availableTokens: 0, tokensToAdd: 0
Request 7 is blocked by rate limiter
  [1736670196997] - Current availableTokens: 0, tokensToAdd: 0
Request 8 is blocked by rate limiter
  [1736670197103] - Current availableTokens: 0, tokensToAdd: 0
Request 9 is blocked by rate limiter
  [1736670197208] - Current availableTokens: 0, tokensToAdd: 5
Request 10 can be consumed!
  [1736670197309] - Current availableTokens: 4, tokensToAdd: 0
Request 11 can be consumed!
  [1736670197414] - Current availableTokens: 3, tokensToAdd: 0

FollowUp Question

#1. Is it thread-safe? How does it handle concurrent requests?

It is not thread-safe. The current implementation can cause race conditions
- To improve that, we can use synchronized keywords or Reentrant Locks.

public synchronized boolean tryConsume() {
    refill();

    if (availableTokens > 0) {
        availableTokens--;
        return true;
    } else {
        return false;
    }
}

...

 private final ReentrantLock lock = new ReentrantLock();

    public boolean tryConsume() {
        lock.lock();
        try {
            refill();

            if (availableTokens > 0) {
                availableTokens--;
                return true;
            } else {
                return false;
            }
        } finally {
            lock.unlock();
        }
    }

#2. Can you implement a lock-free approach for Rate Limiter?

AtomicReference + CAS(Compare-And-Set) + Immutable State
- We can use AtomicReference which allows to update reference to an object atomically, ensuring that updates are thread-safe.
- Also, using an immutable objects for Rate Limiter’s state ensures that updates produce a new state instances without altering the existing ones, facilitating safe concurrent access.

public class TokenBucketState {
    private final double availableTokens;
    private final long lastRefillTimestamp; // in nanoseconds

    public TokenBucketState(double availableTokens, long lastRefillTimestamp) {
        this.availableTokens = availableTokens;
        this.lastRefillTimestamp = lastRefillTimestamp;
    }

    public double getAvailableTokens() {
        return availableTokens;
    }

    public long getLastRefillTimestamp() {
        return lastRefillTimestamp;
    }
}

public class TokenBucketRateLimiter {
    private final int capacity;
    private final double refillRatePerSecond;
    private final AtomicReference<TokenBucketState> state;

    /**
     * Initialize token bucket with specified capacity and refill rate
     *
     * @param capacity
     * @param refillRatePerSecond
     */
    public TokenBucketRateLimiter(int capacity, double refillRatePerSecond) {
        if (capacity <= 0 || refillRatePerSecond <= 0) {
            throw new IllegalArgumentException("Capacity and refill rate must be positive.");
        }
        this.capacity = capacity;
        this.refillRatePerSecond = refillRatePerSecond;
        this.state = new AtomicReference<>(new TokenBucketState(capacity, System.nanoTime()));
    }

    /**
     * Attempts to consume a token from the bucket.
     *
     * @return true if a token was consumed successfully, false otherwise.
     */
    public boolean tryConsume() {
        while (true) {
            TokenBucketState currentState = state.get();
            TokenBucketState newState = refill(currentState);

            if (newState.getAvailableTokens() < 1) {
                // Not enough tokens to consume
                return false;
            }

            TokenBucketState updatedState = new TokenBucketState(
                newState.getAvailableTokens() - 1,
                newState.getLastRefillTimestamp()
            );

            // Attempt to set the new state atomically
            if (state.compareAndSet(currentState, updatedState)) {
                // Successfully consumed a token
                return true;
            }
            // Else, retry due to concurrent modification
        }
    }

    /**
     * Refills tokens based on the elapsed time since the last refill.
     *
     * @param currentState The current state of the token bucket.
     * @return The updated state after refilling.
     */
    private TokenBucketState refill(TokenBucketState currentState) {
        long now = System.nanoTime();
        double secondsSinceLastRefill = (now - currentState.getLastRefillTimestamp()) / 1_000_000_000.0;
        double tokensToAdd = secondsSinceLastRefill * refillRatePerSecond;
        double newTokens = Math.min(capacity, currentState.getAvailableTokens() + tokensToAdd);

        // Update the last refill timestamp only if tokens were added
        if (tokensToAdd > 0) {
            return new TokenBucketState(newTokens, now);
        } else {
            return currentState;
        }
    }

    /**
     * Retrieves the current number of available tokens.
     *
     * @return The number of tokens currently available.
     */
    public double getAvailableTokens() {
        TokenBucketState currentState = state.get();
        TokenBucketState newState = refill(currentState);
        // Update state if refilled
        if (newState != currentState) {
            state.compareAndSet(currentState, newState);
        }
        return newState.getAvailableTokens();
    }
}

Immutable State Object (TokenBucketState):
- Encapsulates the availableTokens and lastRefillTimestamp as immutable fields.
- Ensures that once a state object is created, it cannot be altered, preventing unintended side effects during concurrent access.
AtomicReference (state):
- Holds the current TokenBucketState object.
- Allows atomic updates to the state without explicit locks, ensuring thread safety.
Atomic Update with CAS:
- Uses compareAndSet to atomically update the state.
- If the state has not been modified by another thread, the update succeeds. Otherwise, the loop retries with the new current state. (a.k.a – Compare And Set)

#3. How would you adapt your rate limiter to function correctly in distributed system?

Challeges
- Consistent token state across multiple instances
- Latency and performance
- Fault tolerance and availability
Solution
- Use Redis as a centralized token storage
  - Redis is in-memory storage which provides low-latency and high-throughput rate limiting
  - Supports atomic operations using its Lua script

#4. How should you update your Rate Limiter to support multiple users with different configurations?

Create a separate class that manages a map that contains the configuration

public class RateLimiterManager {
    private final Map<UserTier, RateLimitConfig> tierConfigMap = new HashMap<>();
    private final Map<UserTier, TokenBucketRateLimiter> rateLimiterMap = new HashMap<>();

    public RateLimiterManager() {
        // Initialize configurations for each tier
        tierConfigMap.put(UserTier.FREE, new RateLimitConfig(5, 5));     // Example: 5 tokens capacity, 5 tokens/sec
        tierConfigMap.put(UserTier.PREMIUM, new RateLimitConfig(10, 10)); // Example: 10 tokens capacity, 10 tokens/sec

        // Initialize rate limiters for each tier
        for (Map.Entry<UserTier, RateLimitConfig> entry : tierConfigMap.entrySet()) {
            rateLimiterMap.put(entry.getKey(), new TokenBucketRateLimiter(
                entry.getValue().getCapacity(),
                entry.getValue().getRefillRatePerSecond()
            ));
        }
    }

    public boolean allowRequest(UserTier tier) {
        TokenBucketRateLimiter rateLimiter = rateLimiterMap.get(tier);
        if (rateLimiter == null) {
            throw new IllegalArgumentException("Unsupported user tier: " + tier);
        }
        return rateLimiter.tryConsume();
    }
}

#5. How can your Rate Limiter accommodate dynamic configuration updates like adjusting refill rates without any re-deployment or downtime?

Solution: Utilize external configuration service
- ex) Zookeeper, Consul, Spring Cloud Config,.. etc
- Make the application server to watch configuration service to check if there’s any updates
  - Polling, Push

Reference

Compare And Set
- https://velog.io/@joosing/high-performance-multithreaded-sync-techniques-compare-and-swap

2 responses to “Rate Limiter #1 –Token Bucket Algorithm”

Rate Limiter #2 – Leaky Bucket Algorithm – Will the Guru

January 12, 2025

[…] The same FollowUp Question #1~#5 in https://willtheguru.me/2025/01/12/rate-limiter-token-bucket-algorithm/ is […]

LikeLike

Rate Limiter #3 – Sliding Window Algorithm – Will the Guru

January 12, 2025

[…] The same FollowUp Question #1~#5 in https://willtheguru.me/2025/01/12/rate-limiter-token-bucket-algorithm/ is […]

LikeLike

Rate Limiter #1 –Token Bucket Algorithm

Advantages / Disadvantages

Implementation

Implementation – Output

FollowUp Question

#1. Is it thread-safe? How does it handle concurrent requests?

#2. Can you implement a lock-free approach for Rate Limiter?

#3. How would you adapt your rate limiter to function correctly in distributed system?

#4. How should you update your Rate Limiter to support multiple users with different configurations?

#5. How can your Rate Limiter accommodate dynamic configuration updates like adjusting refill rates without any re-deployment or downtime?

Reference

Share this:

2 responses to “Rate Limiter #1 –Token Bucket Algorithm”

Leave a comment Cancel reply