令牌桶算法是一种常用的限流算法，特点是允许一定程度的突发流量。Redis 提供了令牌桶限流的功能，可以通过 INCR 和 LTRIM 命令组合实现。
以下是一个简单的 Redis 令牌桶限流的实现示例：
import redis
import time
# 连接到Redis
r = redis.Redis(host='localhost', port=6379, db=0)
# 令牌桶的大小和填充速率
bucket_size = 10
fill_rate = 1.0 / 5 # 每5秒填充1个令牌
def try_consume(token_key, count):
# 获取当前时间戳
now = time.time()
# 使用Lua脚本原子性地填充令牌桶并尝试消费令牌
lua_script = """
local tokens_key = KEYS[1]
local bucket_size = tonumber(ARGV[1])
local fill_rate = tonumber(ARGV[2])
local count = tonumber(ARGV[3])
local now = tonumber(ARGV[4])
local last_fill_time = redis.call('GET', tokens_key..':last_fill_time')
if last_fill_time then
last_fill_time = tonumber(last_fill_time)
else
last_fill_time = 0
end
-- 计算应该填充的令牌数
local tokens_added = math.floor((now - last_fill_time) * fill_rate)
local bucket_count = redis.call('GET', tokens_key)
if bucket_count then
bucket_count = tonumber(bucket_count)
else
bucket_count = 0
end
-- 限流，如果需要的令牌数大于当前桶中的令牌数，返回0
if count > (bucket_size - bucket_count + tokens_added) then
return 0
else
-- 尝试消费令牌
local consumed = redis.call('DECRBY', tokens_key, count)
-- 如果桶中的令牌数小于0，则设置为0
if consumed < 0 then
redis.call('SET', tokens_key, 0)
end
return consumed
end
"""
# 执行Lua脚本
consumed = r.eval(lua_script, 1, token_key, bucket_size, fill_rate, count, now)
if consumed == 0:
return False
else:
return True
# 使用令牌桶限流的关键在于确保令牌桶的大小和填充速率的合理设置
# 令牌桶的大小是10，填充速率是每5秒填充1个令牌
token_key = 'my_rate_limited_resource'
# 尝试消费2个令牌
if try_consume(token_key, 2):
print("请求被允许")
else:
print("请求超出限制")
在这个示例中，我们定义了一个 try_consume 函数，它尝试从令牌桶中消费指定数量的令牌。如果令牌足够，消费成功，并允许访问资源；如果令牌不足，则拒绝访问资源。
Lua 脚本用于原子性地处理令牌桶的填充和令牌的消费，确保在高并发情况下的正确性。