分布式多层次限流概述

分布式限流

分布式限流常见于对外提供服务的API，例如阿里云的语音服务接口等。对于这类的服务一般常见于资源消耗比较大，服务相对比较敏感，而且某些提供的限流功能要满足多层次的限流功能，比如一分钟一次，一小时十五次，一天三十次这种功能，相对于当前常见的分布式限流方案有采用zk、etcd等组件来实现的，优点就是分布式限流过程中对于使用方来说保证了高可用，也有使用redis来实现的，在redis的实现方案中如果需要保证高可用可考虑redis的哨兵功能来实现（如果集群模式下要保持访问的客户端的时钟一致），本文就基于redis来实现。

redis分布式限流

对于单次的限流，比如一分钟一次，我们可以通过单次的redis操作就可以进行了，如果既有分钟、小时和天的就需要再redis中一起执行，综合判断最终的执行结果。

import time
from functools import total_ordering

import redis


@total_ordering
class Limiter(object):
    ALLOW_FAIL_FLAG = None
    TIME_LIMIT = None
    BUCKET = 0

    def __init__(self, false_flag: int, time_limite: int, bucket=1):
        if not isinstance(false_flag, int):
            raise Exception("base limiter fail flag must be int")
        self.ALLOW_FAIL_FLAG = false_flag
        self.TIME_LIMIT = time_limite
        self.BUCKET = bucket
        if self.TIME_LIMIT <= 0:
            raise Exception("time limit should gt 0")
        if self.BUCKET <= 0:
            raise Exception("bucket should gt 0")
        self.rate = float(self.BUCKET / self.TIME_LIMIT)

    def __ge__(self, other):
        return self.rate > other.rate

    def __le__(self, other):
        return self.rate < other.rate

    def __eq__(self, other):
        return self.rate == other.rate

    def generate(self):
        raise NotImplementedError


class BaseLimiter(Limiter):

    def generate(self):
        return """
        local min_limit = {0}
        if key_len >= min_limit then
            local litem = redis.call('lindex', key, -min_limit)
            if not(timestamp - litem > {1}) then
                is_limit = true
            end
        end
        if is_limit == true then
            return {2}
        end
        """.format(self.BUCKET, self.TIME_LIMIT, self.ALLOW_FAIL_FLAG)


class RateLimter(object):
    ALLOW_SUC_FLAG = 1
    LUA_COMMAND_HEAD = """
    local key = KEYS[1]
    local timestamp = ARGV[1]

    local key_len = redis.call('llen', key)
    if key_len == 0 then
        redis.call('rpush', key, timestamp)
        return {0}
    end
    local is_limit = false
    
    """
    LUA_REMOVE = """
    
    while true do
        local item = redis.call('lindex', key, 0) or timestamp
        if timestamp - item > {0} then
            redis.call('lpop', key)
        else
            break
        end
    end

    local key_len = redis.call('llen', key)
    """
    LUA_END = """
    redis.call('rpush', key, timestamp)
    return {0}
    """

    def __init__(self):
        self.limters = []
        self.client = None
        self.lua_command = ""

    def init_redis_client(self, host: str, port=6379):
        self.client = redis.Redis(host=host, port=port)
        self.lua_command = self.generate()

    def init_redis_sential(self, hosts, service_name):
        from redis.sentinel import Sentinel
        try:
            client = Sentinel(hosts)
            master_msg = client.discover_master(service_name)
            print(master_msg)
            self.client = client.master_for(service_name)
        except Exception as e:
            raise e

        self.lua_command = self.generate()

    def allow(self, key):
        if not self.lua_command:
            raise Exception("lua command must init")
        if self.client is None:
            raise Exception("client must init")
        try:
            now_time = int(time.time())
            r = self.client.eval(self.lua_command, 1, key, now_time)
            if r == self.ALLOW_SUC_FLAG:
                return True
            return r
        except Exception as e:
            print(e)
            return

    def register(self, limiter):
        self.limters.append(limiter)
        self.limters = sorted(self.limters, reverse=True)
        for i in range(len(self.limters) - 1):
            if self.limters[i + 1].rate == self.limters[i].rate:
                raise Exception(" limiter have same rate ")
            if self.limters[i + 1].TIME_LIMIT <= self.limters[i].TIME_LIMIT:
                raise Exception(" limiter Time limit is error")
            if self.limters[i + 1].BUCKET == self.limters[i].BUCKET:
                raise Exception(" limiter Bucket different has same bucket")

    def generate(self):
        lua_command = self.LUA_COMMAND_HEAD.format(self.ALLOW_SUC_FLAG)
        lua_command += self.LUA_REMOVE.format(self.limters[-1].TIME_LIMIT)
        for v in self.limters:
            lua_command += v.generate()
        lua_command += self.LUA_END.format(self.ALLOW_SUC_FLAG)
        return lua_command


def main():
    m = BaseLimiter(20, 30, 1)
    h = BaseLimiter(30, 500, 15)
    d = BaseLimiter(40, 7000, 16)
    r = RateLimter()
    r.register(m)
    r.register(h)
    r.register(d)
    r.init_redis_sential([("192.168.10.204", 27000), ("192.168.10.204", 27001), ("192.168.10.204", 27002)], "mymaster")
    while True:
        print(r.allow("13247199655"))
        time.sleep(10)


if __name__ == '__main__':
    main()

示例代码主要展示了保持了一个队列，这个队列最长的长度就是16，在7000秒里面最多保持16个时间戳，通过保存的时间戳列表来判断是否达到限流条件。而且这个实现可以基于毫秒级也可以是秒级的队列。通过lua脚本的原子操作，保证了再进行每个限流条件判断的时候都能够是原子的操作，从而保证了分布式的限流。

总结

本文只是简单的封装了lua脚本来实现分布式限流的功能，通过redis的原子操作来实现分层分级的限流功能。相对于当下许多就行一次频率限制的方案，该方案适合多层级限流，权当记录。由于本人才疏学浅，如有错误请批评指正。