API token bucket: API 요청수 관리

API 요청을 관리하는 데에는 여러 이유가 있습니다. 그 중에서도 토큰 버킷을 사용하는 이유는 다음과 같습니다:

1. 트래픽 제어: 토큰 버킷을 사용하면 API로 들어오는 트래픽을 제어할 수 있습니다. 이를 통해 서버에 과도한 부하가 걸리는 것을 방지하고, 서비스의 안정성을 유지할 수 있습니다.
2. 사용량 제한: 토큰 버킷을 이용하면 API를 사용하는 클라이언트의 요청 수를 제한할 수 있습니다. 이는 과도한 사용량으로 인한 서버 과부하를 방지하고, 공정한 서비스 이용을 보장합니다.

 

TokenBucket 만들기

토큰 버킷 패턴은 고정된 속도로 토큰을 생성하고, 요청이 들어올 때마다 토큰을 소비하여 일정량의 토큰을 가지고 있는지 확인하여 처리하는 방식입니다. 알고리즘은 아래와 같습니다.

1. TokenBucket은 버켓(바구니)가 있다고 가정합니다.
2. 요청이 들어오면, 버켓에서 1개를 소비해서 그 다음의 로직을 진행시킵니다. 요청이 들어왔을 때 버킷에 토큰이 없으면 drop합니다.
3. 고정된 속도(r)로 버킷에 토큰채웁니다. 

https://www.researchgate.net/figure/Diagram-representing-token-bucket-algorithm_fig1_351452059

 

TokenBucket 을 middle ware 에 추가

- FastAPI의 starlette을 베이스로 작성되어있어서, 이 미들웨어를 서브클레싱해서 손쉽게 token bucket을 사용할 수 있습니다.

- starlette.middleware.base.BaseHTTPMiddleware을 초기화시 app을 받게되어있고 이 베이스클레스는 dispatch라는 함수를 구현하게어 있습니다.

class BaseHTTPMiddleware:
    ...
    
    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        raise NotImplementedError()  # pragma: no cover

 

dispatch을 함수를 아래와 같이 구현해서 bucket.consume_token()이라는 함수가 bool으로 false or True을 구현하게 함에따라, 예외를 일으킬것인지 다음 로직을 실행시킬것인지를 결정하게 작성합니다.

 

위의 구현을 위해 아래와 같이 tokenbucket을 정의합니다.

초기화시, 버킷의 토큰 최대 저장량(max_len)과 토큰을 채우는 시간(r, refill_rate)을 지정하여 아래와 같이 작성합니다.

 

실행

- 다음과 같이 python3으로 fastAPI을 구동시킵니다.

(misc) (base) heon@gpusvr03:~/repositories/misc/tokenbucket_ex$ python3 app.py 
INFO:     Started server process [4121818]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:30000 (Press CTRL+C to quit)

 

그 후 요청 15개를 보내봅니다. 아래와 같이 10개를 토큰을 다 소비한경우 5개의 요청이 예외처리됩니다.

>>> import requests
>>> requests.get("http://0.0.0.0:30000/")
<Response [200]>
>>> for _ in range(15):
...     requests.get("http://0.0.0.0:30000/")
... 
<Response [200]>
<Response [200]>
<Response [200]>
<Response [200]>
<Response [200]>
<Response [200]>
<Response [200]>
<Response [200]>
<Response [200]>
<Response [200]>
<Response [500]> # 토큰 10개 소비되어 예외처리
<Response [500]>
<Response [500]>
<Response [500]>
<Response [500]>