Implementing Rate Limiting in Java Spring Boot using Token Bucket Algorithm

Introduction

Rate limiting is an essential feature for controlling the number of requests a client can make to an API over a given period. It helps protect APIs from abuse and overuse by ensuring that clients adhere to predefined limits. In this blog, we’ll implement a rate limiter using the token bucket algorithm in Java Spring Boot.

What is the Token Bucket Algorithm?

The token bucket algorithm is a simple mechanism used to control data flow. Tokens are added to a bucket at a fixed rate, and each request consumes one token. If there are no tokens available, the request is denied. Over time, the bucket refills with tokens at the configured rate.

Key Concepts:

Bucket Size: Maximum number of tokens available at any time.
Refill Rate: How often tokens are added to the bucket.

Step 1: Create the TokenBucket Class

We'll start by defining the TokenBucket class that implements the token bucket algorithm. The class will maintain the number of tokens available and handle the logic for refilling the bucket and checking if a request can be allowed.

package com.sid121212.rate_limiter.algo;

public class TokenBucket {
    private final long maxTokenLimit;
    private final long refillRate;
    private long availableTokens;
    private long lastFilled;

    public TokenBucket() {
        this.maxTokenLimit = 5;
        this.refillRate = 10000; // Refill 1 token every 10 seconds
        this.availableTokens = 5;
        this.lastFilled = System.currentTimeMillis();
    }

    public synchronized boolean allowRequest() {
        refill();
        if (availableTokens > 0) {
            availableTokens--;
            return true;
        }
        return false;
    }

    public void refill() {
        long now = System.currentTimeMillis();
        if (now - lastFilled >= refillRate) {
            availableTokens = Math.min(maxTokenLimit, availableTokens + 1);
            lastFilled = now;
        }
    }
}

Step 2: Rate Limiting Service Implementation

Now, let’s implement the RateLimiterService, which will handle the rate-limiting logic for different clients by creating and managing token buckets.

package com.sid121212.rate_limiter.service;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.stereotype.Service;
import com.sid121212.rate_limiter.algo.TokenBucket;

@Service
public class RateLimiterService {
    private final Map<String, TokenBucket> buckets = new ConcurrentHashMap<>();

    public boolean isAllowed(String clientId) {
        TokenBucket tokenBucket = buckets.get(clientId);
        if (tokenBucket == null) {
            tokenBucket = new TokenBucket();
            buckets.put(clientId, tokenBucket);
        }
        return tokenBucket.allowRequest();
    }
}

Step 3: Adding the Service to the Controller

In this step, we will integrate the RateLimiterService with a RESTful controller to handle incoming HTTP requests. This controller will use the service to apply rate limiting based on client identifiers.

To create a REST endpoint that leverages our rate limiter, we will build a RateLimiterController. This controller will be responsible for processing client requests and checking if they comply with the rate limits defined in the service.

package com.sid121212.rate_limiter.controller;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestHeader;
import org.springframework.web.bind.annotation.RestController;
import com.sid121212.rate_limiter.service.RateLimiterService;

@RestController
public class RateLimiterController {

    @Autowired
    private RateLimiterService rateLimiterService;

    @GetMapping("/rateLimiter")
    public String limitedEndpoint(@RequestHeader("Client-Id") String clientId) {
        if (rateLimiterService.isAllowed(clientId)) {
            return "Request successful!";
        } else {
            return "Rate limit exceeded. Please try again later.";
        }
    }
}

Step 4: Configuring `pom.xml` and Running the Application

To get your rate limiter application up and running, you need to configure your Maven dependencies and follow the steps to start the application.

4.1 Configuring `pom.xml`

Ensure your pom.xml file includes the necessary dependencies for Spring Boot and any additional libraries you might need. Below is an example configuration for a Spring Boot project using Maven:

<dependencies>
    <!-- Spring Boot Starter Web for building web applications -->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-web</artifactId>
    </dependency>
</dependencies>

Now Run your Java-Spring Boot Project

And there you go !! Thanks for reading !