基於Redis實現令牌桶限流

常用限流算法有漏桶算法和令牌桶算法，本文借助Redis的redis_cell模塊來實現令牌桶算法限流。

構建鏡像並啟動容器

FROM redis:latest

ARG cell_dir=/lib/redis_modules/redis_cell
RUN mkdir -p ${cell_dir}
WORKDIR ${cell_dir}

RUN apt-get update \
     && apt-get -y install wget \
     && wget https://github.com/brandur/redis-cell/releases/download/v0.3.0/redis-cell-v0.3.0-x86_64-unknown-linux-gnu.tar.gz \
     && tar -zxvf redis-cell-v0.3.0-x86_64-unknown-linux-gnu.tar.gz \
     && rm redis-cell-v0.3.0-x86_64-unknown-linux-gnu.tar.gz \
     && apt-get -y remove wget \
     && apt -y autoremove \
     && ls -alh

ENTRYPOINT ["redis-server","--loadmodule","/lib/redis_modules/redis_cell/libredis_cell.so"]

  

docker build -t redis_limit .
docker run -d -p 6379:6379 --rm --name redis_limit redis_limit

 

模擬有波動的請求

redis_cell模塊提供了原子性命令來實現限流，我們只需要根據命令執行結果來做響應處理即可，不用自己再處理令牌發放相關邏輯：

import random
import time
import redis

r: redis.Redis = redis.Redis(host='localhost', port=6379, db=0)

exec_result = r.execute_command('cl.throttle rate_limit 100 50 10 1')
while True:
    if exec_result[0] == 1:
        seconds = exec_result[3]
        if seconds <= 0:
            seconds = 1
        print(f'等待{seconds}秒后重試')
        time.sleep(seconds)
        print('重試')
    token_num = int(random.random() * 30)
    print(f'剩余{exec_result[2]}個令牌，即將獲取{token_num}個令牌')
    if token_num <= 0:
        token_num = 1
    exec_result = r.execute_command(f'cl.throttle rate_limit 100 50 10 {token_num}')

r.close()

# 輸出如下：
# 剩余81個令牌，即將獲取3個令牌
# 剩余78個令牌，即將獲取28個令牌
# 剩余50個令牌，即將獲取24個令牌
# 剩余26個令牌，即將獲取23個令牌
# 剩余3個令牌，即將獲取26個令牌
# 等待4秒后重試
# 重試
# 剩余3個令牌，即將獲取3個令牌
# 剩余20個令牌，即將獲取25個令牌
# 等待1秒后重試
# 重試
# 剩余20個令牌，即將獲取3個令牌
# 剩余22個令牌，即將獲取28個令牌
# 等待1秒后重試

推薦閱讀

012-redis應用-05-限流【簡單限流、漏斗限流】