토큰 버킷 알고리즘 RateLimiter

자바로 구현하는 RateLimiter

API 서버를 운영하다 보면 요청이 몰려 서버가 감당하기 힘든 상황이 발생한다. 이를 방지하기 위해 Rate Limiting 기법을 사용하는데, 그 중 하나가 바로 토큰 버킷이다.

토큰 버킷 알고리즘?

• 버킷에는 일정 개수의 토큰이 들어있음
• 요청이 들어올 때마다 토큰을 1개씩 사용
• 이때 토큰이 없다면 요청은 거부됨
• 시간이 흐르면 일정 비율로 토큰이 채워짐
• 이때 최대 버킷 용량을 넘지 않음
• 이를 통해 초당 요청 수 제한 가능

버킷 = 허용 가능한 요청 수 저장하는 공간
토큰 = 허용할 수 있는 요청 단위

구현 코드

public class RateLimit {
    public static void main(String[] args) throws InterruptedException {
        RateLimiter limiter = new RateLimiter(3, 1);

        System.out.println("초기 토큰: " + limiter.getCurrentTokens()); // 예상: 3

        System.out.println("1차 요청: " + limiter.tryAcquire()); // true
        System.out.println("2차 요청: " + limiter.tryAcquire()); // true
        System.out.println("3차 요청: " + limiter.tryAcquire()); // true
        System.out.println("4차 요청: " + limiter.tryAcquire()); // false, 토큰 소진

        System.out.println("소진 후 토큰: " + limiter.getCurrentTokens()); // 예상: 0

        System.out.println("\n2초 대기...\n");
        Thread.sleep(2000); // 2초 동안 대기

        System.out.println("2초 후 토큰: " + limiter.getCurrentTokens()); // 예상: 2 (2초 동안 2개 리필)
        
        System.out.println("5차 요청: " + limiter.tryAcquire()); // true
        System.out.println("6차 요청: " + limiter.tryAcquire()); // true
        System.out.println("7차 요청: " + limiter.tryAcquire()); // false

        System.out.println("최종 토큰: " + limiter.getCurrentTokens()); // 예상: 0
    }
}

class RateLimiter{
    private int capacity, refillRatePerSecond;
    private long currentTokens;
    private long lastRefillTimestamp;

    public RateLimiter(int capacity, int refillRatePerSecond){
        this.capacity = capacity;
        this.refillRatePerSecond = refillRatePerSecond;
        this.currentTokens = capacity;
        this.lastRefillTimestamp = System.currentTimeMillis(); // 밀리초
    }

    // 메소드 전체를 동기화하기 위해 synchronized 사용
    public synchronized boolean tryAcquire(){
        refillTokens();

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

    private void refillTokens(){
        long now = System.currentTimeMillis(); //밀리초
        long time = now - lastRefillTimestamp;

        if(time>0){
            long tokensToRefill = (time*refillRatePerSecond)/1000; // 초로 바꾸기 위해 /1000 연산 수행

            if(tokensToRefill>0){
                currentTokens = Math.min(currentTokens+tokensToRefill, capacity); // 토큰 최대 용량을 넘지 않도록
                this.lastRefillTimestamp = now;
            }
        }

    }

    public synchronized long getCurrentTokens(){
        refillTokens();
        return currentTokens;
    }
}

• 초기에는 토큰이 3개이며, 1초당 1개씩 채워짐
• 첫 3번 요청 -> 각각 토큰 하나씩 모두 소모해 true 반환
• 4번째 요청 -> 토큰이 없기에 false 반환
• 2초 대기 후 토큰 2개 리필
• 다시 요청 -> 4, 5번째는 성공
• 6번째부터 또 2초 대기 ... (반복)

이때 synchronized 키워드를 사용해 동시성 문제를 방지했으며 밀리초 단위로 계산해 정밀한 제어가 가능하도록 구현했다.
비록 간단한 예제로 구현했지만 실제 API 서버에서 이를 이용해 초당 요청을 제한한다면 보안, 비용적으로도 큰 도움이 될 것이다.