[Rate Limit - step 3] Token Bucket 알고리즘 구현 (rate limiting)

필자는 현재 API서버를 열심히 개발 하고 있다. 개발한 서비스가 트래픽을 많이 받다보니, 트래픽을 적절히 제한하지 않았을 때 장애가 발생했고, 그에 따라 API의 Rate Limit을 구현해야 했다.

 

공부하다보니 흥미가 생겨 나중에 내가 다시 챙겨보기 위해 글을 정리하려고 한다.

 

Rate Limit시리즈의 목차는 다음과 같다.

 

[Rate Limit - step 1] Rate Limit이란? (개념, Throttling, 분산환경에서의 구현)

[Rate Limit - step 2] Leaky Bucket 알고리즘 구현 (rate limiting) (p.s. memory estimation 하는 법)

[Rate Limit - step 3] Token Bucket 알고리즘 구현 (rate limiting) <- 이번 글

[Rate Limit - step 4] Fixed Window 알고리즘 구현 (rate limiting) 

[Rate Limit - step 5] Sliding Window 알고리즘 구현 (rate limiting)

 

본 글은 step 3로 Rate Limit의 대표적인 알고리즘 중 하나인 Token Bucket 알고리즘에 대해 알아본다. 개인적으로 이전 시리즈를 읽고 보는 것을 추천한다. 물론 이것만 보고도 알 수 있지만, 더 잘 이해할 수 있을거라 생각한다. 

 

Token Bucket 알고리즘

Leaky Bucket과 유사한 구조지만, Bucket을 FIFO 큐로 사용하지 않고, 트래픽을 제어하기 위한 제어용 토큰을 관리하는 용도로 사용한다. Leaky Bucket과의 가장 큰 차이는 요청을 처리할 때 leaky bucket은 constant한 속도로 처리할 수 있는 반면 token bucket은 버스트 요청도 허용한다.

구현할 때 주의사항은 token을 fixed window처럼 특정 기간 단위로 추가하는게 아니라 일정 rate(속도)에 맞추어 추가해주어야 한다. (1분에 10개면 6초당 1개씩 추가해주는식)

• 트래픽은 토큰의 유무에 따라 흐름을 제어받는다.
  • 버킷에 토큰이 있으면 요청을 허용하고, 토큰이 고갈되면 에러를 리턴한다.
• 버킷에 정해진 rate에 맞추어 토큰을 다시 채워준다(refill).
  • 특정 시간 구간동안 모든 토큰을 다 사용한 유저에게는 request를 허용하지 않는다.
• 토큰 버킷은 토큰이 배치되는 속도를 기반으로 액세스 속도를 제어한다.

 

그림을 보면서 생각해보자.

1. bucket에는 d개만큼의 토큰이 채워져있다. bucket에는 최대 n개만큼의 토큰이 채워질 수 있다.

2. 앞으로는 1/r 초마다 의 속도만큼 토큰이 버킷에 채워진다. (n개를 넘지 않음)

3. 요청이 들어오게 되면 버킷에서 토큰을 꺼낸다. 토큰이 부족하면 요청을 거절한다.

 

버킷을 조금 더 시각화 해보면

 

 

그림에서 보이듯 10:10에는 3개 토큰이 있고, 하나씩 줄다가 10:10:45에는 0개라 요청을 받지 않는다.

• 이 때부터 10:11까지는 요청을 받을 수 없음

 

그리고 10:11이 되고나니 토큰이 3개가 리필된다.

• 이제 요청을 받을 수 있음

 

Token Bucket의 장점

1. Max concurrency관리에 용이하다

2. burst of request도 일정수준은 유연하게 처리 가능하다.

 

Token Bucket의 단점

1. distributed환경에서 그냥 그대로 사용할 경우 race condition이 발생할 수 있다. (step 1 글 참고 바람)

 

Token Bucket의 샘플 구현

코드에 대한 line by line설명보다는 주석을 좋아하는 편이라 코드에 주석으로 설명해두었으니, 코드 주석을 잘 읽어보면서 참고하길 바란다. 

import time

class TokenBucket:

    def __init__(self, tokens, time_unit):
        self.tokens = tokens
        self.time_unit = time_unit
        self.bucket = tokens
        self.last_check = time.time()

    def handle(self, packet):
        #The current timestamp is put into current.
        current = time.time()
        #The time passed between now and the previous handle call is put into time_passed.
        time_passed = current - self.last_check
        #We update the last_check to the current timestamp.
        self.last_check = current

        # By multiplying the time_passed and our rate (which is tokens / time_unit) we find out how many token needs to be added to the bucket.
        self.bucket = self.bucket + \
            time_passed * (self.tokens / self.time_unit)

        # Reset the bucket if it has more tokens than the default value.
        if (self.bucket > self.tokens):
            self.bucket = self.tokens

        # If the bucket doesn't have enough token, drop the packet.
        if (self.bucket < 1):
            return False
        # If the bucket has enough token remove one token and forward the packet.
        else:
            return True

throttle = TokenBucket(1, 1)

packet = 0

while True:
    time.sleep(0.2)
    throttle.handle(packet)
    packet += 1

참고로 잘 구현된 go의 ratelimit패키지 사용해서 테스트 해보면 잘 refill되는 거 확인할 수 있음

 

---

자, 이렇게 Token Bucket 알고리즘을 알아보면서 절반까지 왔다. 다음으로 두 개의 알고리즘이 남았다. 실제로 구현을 어떻게 했는지까지 적어두었으니, 각 코드별로 직접 구현해보고 API설계나 고객 패턴에 따라 성능이 다를 수 있으니 테스트 해보는 것을 추천한다.