说明
业务中需要对一些行为进行数量限制，例如：直播间内的横幅数量和公屏的礼物特效数量，用户每天的请求次数等。
固定窗口限流
思路
固定窗口的思路是根据时间去计算窗口，然后在窗口上累积已执行次数，根据历史执行次数，窗口的大小，和本次要执行的次数，生成用户最终可执行的次数。
实现
分布式的限流器要将数据存储在分布式数据库中，因为数据具有时间属性，redis的自动过期策略可以很好实现，redis的数据结构选择string，直接利用incrby方法去增加一个key的value。
代码实现
# 定义lua脚本
# KEYS参数为[string_key]
# ARGV参数为[acqure_num, window_ttl]
# 返回incrby后的值
INC_WINDOW_NUM = """
local new_count = redis.call("incrby",KEYS[1], ARGV[1])
local old_count = new_count - ARGV[1]
if old_count == 0 then
redis.call('expire',KEYS[1],ARGV[2])
end
return new_count
"""
class FixWindowLimiter(object):
"""固定窗口限流器"""
def __init__(self, limiter_name, limit_num, window_ttl, **kwargs):
"""
:param limiter_name: 限制器的名称，例如subcid，gametype等
:param limit_num: 限制数量
:param window_ttl: 窗口有效期，也即窗口的大小（单位秒）
"""
# 从context中拿到执行的任务信息
self.task = TaskContext.current_task() # type: BaseContextTask
self.appprefix = self.task.service_appprefix # 服务名前缀
self.now_ts = self.task.now_ts # 请求的时间
self.rdb = self.task.rdb # redis的client
self.limiter_name = limiter_name
self.limit_num = limit_num
self.window_ttl = window_ttl
self.kwargs = kwargs
self._inc_window_num = None
@property
def inc_window_num(self):
if not self._inc_window_num:
self._inc_window_num = self.rdb.register_script(INC_WINDOW_NUM)
return self._inc_window_num
def acquire(self, limit_value, acquire_num):
"""
:param limit_value: 限制的值
:param acquire_num: 申请数量
:return: 实际申请到的数量
"""
# 根据时间计算窗口
window = interval_start_ts(self.window_ttl, self.now_ts)
key = "%s_%s_%s_%s_fwl" % (self.appprefix, self.limiter_name, limit_value, window)
new_num = self.inc_window_num(keys=[key], args=[acquire_num, self.window_ttl])
# 窗口累积的值不超过阈值，直接返还请求的acquire_num，否则返回实际能申请到的数量
return acquire_num if new_num <= self.limit_num else max(self.limit_num + acquire_num - new_num, 0)
基于限流器实现的装饰器
被装饰的方法，如果执行超过阈值，不会被执行
from functools import wraps
from fix_window_limiter import FixWindowLimiter
def fix_window_limiter(limit_num, window_ttl, value_params=None, limiter_name=None):
"""
固定窗口限制期
:param limit_num: 限制数量
:param window_ttl: 窗口有效期，也即窗口的大小（单位秒），如：一分钟
:param value_params: 限制的值对应的参数，数组类型
:param limiter_name: 限制名称
:return:
"""
ACQUIRE_NUM = 1
def _get_limit_value(func, *args, **kwargs):
# 不传value_params即不根据方法参数进行限制
if not value_params:
return ""
param_dict = locals()
values = list()
for param in value_params:
values.append(param_dict.get(param) or kwargs.get(param, ""))
return ":".join(values)
def decorate(func):
@wraps(func)
def wrapper(self, *args, **kwargs):
limit_value = _get_limit_value(func, *args, **kwargs)
limiter = FixWindowLimiter(limiter_name or func.__name__, limit_num, window_ttl)
if limiter.acquire(limit_value, ACQUIRE_NUM):
func(self, *args, **kwargs)
return wrapper
return decorate
滑动窗口限流器
思路
固定窗口的流量限制不均匀，可以通过滑动窗口进行动态限制
实现
滑动窗口可以根据窗口数目window_num生成多个window，每个window会分别累计自己窗口内的执行次数，所有有效的window进行流量限制
代码实现
# lua脚本执行执行申请操作
# KEYS是window_num个数量的window，从后往前
# ARGV是[acquire_num, limit_num, window_ttl]
DO_ACQUIRE = """
# 计算所有窗口的累积数量
local results = redis.call('mget', unpack(KEYS))
local total_count = 0
for i, v in ipairs(results) do
if results[i] then
total_count = total_count + tonumber(results[i])
end
end
# 计算可增加数量，不满足条件直接返回
local add_num = math.min(ARGV[2] - total_count, ARGV[1])
if add_num <= 0 then
return add_num
end
# 增加最后一个窗口的数值
local new_count = redis.call('incrby', KEYS[1], add_num)
local old_count = new_count - add_num
if old_count == 0 then
redis.call('expire',KEYS[1],ARGV[3])
end
return add_num
class SlidingWindowLimiter(object):
def __init__(self, limiter_name, limit_num, window_num, window_ttl, **kwargs):
"""
:param limiter_name: 限制器的名称（如：subcid，gametype等）
:param limit_num: 总共限制的数量
:param window_num: 滑动窗口数量
:param window_ttl: 窗口有效期，也即窗口的大小（单位秒）
"""
# 从context中拿到执行的任务信息
self.task = TaskContext.current_task() # type: BaseContextTask
self.appprefix = self.task.service_appprefix # 服务名前缀
self.now_ts = self.task.now_ts # 请求的时间
self.rdb = self.task.rdb # redis的client
self.window_num = window_num
self.window_ttl = window_ttl
self.limiter_name = limiter_name
self.limit_num = limit_num
self.kwargs = kwargs
self._do_acquire = None
@property
def do_acquire(self):
if not self._do_acquire:
self._do_acquire = self.rdb.register_script(DO_ACQUIRE)
return self._do_acquire
def _get_key(self, limit_value, end_window, idx):
"""key的前缀用${}包住，保证滑动窗口的所有key都落在redis集群的同一个slot中"""
return "${%s_%s_%s}_%s_swl" % (
self.appprefix, self.limiter_name, limit_value, end_window - self.window_ttl * idx
)
def acquire(self, limit_value, acquire_num=1):
"""
:param limit_value: 限制的值
:param acquire_num: 申请的数量
:return: 实际申请到的数量
"""
end_window = interval_start_ts(self.window_ttl, self.now_ts)
keys = [self._get_key(limit_value, end_window, idx) for idx in range(0, self.window_num)]
args = [acquire_num, self.limit_num, self.window_ttl]
return self.do_acquire(keys=keys, args=args)
基于限流器实现的装饰器
from functools import wraps
from module.limiter.sliding_window_limiter import SlidingWindowLimiter
def sliding_window_limiter(limit_num, window_num, window_ttl, value_params=None, limiter_name=None):
"""
固定窗口限制期
:param limit_num: 限制数量
:param window_num: 滑动窗口数量
:param window_ttl: 窗口有效期
:param value_params: 限制的值对应的参数，数组类型
:param limiter_name: 限制名称
:return:
"""
ACQUIRE_NUM = 1
def _get_limit_value(func, *args, **kwargs):
# 不传value_params即不根据方法参数进行限制
if not value_params:
return ""
param_dict = locals()
values = list()
for param in value_params:
values.append(param_dict.get(param) or kwargs.get(param, ""))
return ":".join(values)
def decorate(func):
@wraps(func)
def wrapper(self, *args, **kwargs):
limit_value = _get_limit_value(func, *args, **kwargs)
limiter = SlidingWindowLimiter(limiter_name or func.__name__, limit_num, window_num, window_ttl)
if limiter.acquire(limit_value, ACQUIRE_NUM):
func(self, *args, **kwargs)
return wrapper
return decorate
