Implementing different Rate limiter algorithms in Springboot

Sun Mar 17

Implementing different Rate limiter algorithms in Springboot
- 1. Tokenbucket Algo:
  - The models:

Implementing different Rate limiter algorithms in Springboot

In this article I will talk about my solution to the challenge: Build Your Own Rate Limiter in codingchallenges.

As I’ve been lately learning Springboot, I decided to do this challenge to practice the bean and dependency injection concept.

1. Tokenbucket Algo:

The models:

The Token model:

import java.time.LocalDateTime;
import java.util.UUID;
public class Token {
    private String id;
    private LocalDateTime created;
    public Token() {
        this.id = UUID.randomUUID().toString();
        this.created = LocalDateTime.now();
    }
}

The Bucket Model:

public class Bucket {
    private String id;
    private String ip;
    private List<Token> tokens = new ArrayList<Token>();

    public Bucket(String ip) {
        this.id = UUID.randomUUID().toString();
        this.ip = ip;
    }

    public void addToken(Token token) {
        tokens.add(token);
    }

    public void removeToken(Token token) {
        tokens.remove(token);
    }

    public String getId() {
        return id;
    }

    public String getIp() {
        return ip;
    }

    public List<Token> getTokens() {
        return tokens;
    }
}

The Token bucket Algorithm implementation:

@Component
public class TokenBucket {
    private final ConcurrentHashMap<String, Bucket> buckets;
    int threshold;

    public TokenBucket() {
        this.threshold = 3;
        this.buckets = new ConcurrentHashMap<>();

    }

    public void registerIp(String ip) {
        if (!buckets.containsKey(ip)) {
            Bucket bucket = new Bucket(ip);
            for (int i = 0; i < threshold; i++) {
                bucket.addToken(new Token());
            }
            buckets.put(ip, bucket);

        }
    }

    public synchronized boolean consumeToken(String ip) {
        if (buckets.containsKey(ip)) {
            List<Token> tokens = buckets.get(ip).getTokens();
            if (tokens.size() > threshold || tokens.size() == 0) {
                return false;
            } else {
                buckets.get(ip).removeToken(tokens.get(0));
                return true;
            }
        } else {
            return false;
        }
    }
    @Scheduled(cron = "*/5 * * * * ?")
    public void addToken() {
        for (Bucket bucket : buckets.values()) {
            if (bucket.getTokens().size() < threshold) {
                bucket.addToken(new Token());
                System.out.println("Token added to bucket: " + bucket.getId());
            }
        }
    }
}

Implementing the algorithm in the controller:

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import com.ratelimiter.algorithms.TokenBucket;

import jakarta.servlet.http.HttpServletRequest;

@RestController
@RequestMapping("/")
@EnableScheduling
public class controller {
    private TokenBucket tokenBucket;

    @Autowired
    public void Controller(TokenBucket tokenBucket) {
        this.tokenBucket = tokenBucket;

    }

    @GetMapping("limited")
    public String limitedApi(HttpServletRequest request) {
        String ipAddress = request.getRemoteAddr();
        tokenBucket.registerIp(ipAddress);
        if (tokenBucket.consumeToken(ipAddress)) {
            return "Success: Request allowed";
        } else {
            return "Error: Rate limit exceeded";
        }

    }

    @GetMapping("unlimited")
    public String unlimitedApi() {

        return "Unlimited! Let's Go!";
    }

}