＜异步令牌桶＞ 解决爬虫访问限制,解决aiomysql异步访问时多个task抢占句柄bug

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

---

什么是令牌桶？

用来限制访问频率以及多线程情况下对某一资源访问资格的问题的解决方案：想象一个桶，每次你访问网络时候都要看下桶内是否有剩余的令牌，如果有你可以拿走令牌去访问，没有你就等着令牌投进来再访问。

一、最简单的：单令牌

class OneToken():
    def __init__(self):
        self.current_token = 1

    async def greenlight(self):
        while self.consume() == False:
            # print(".")
            await asyncio.sleep(0.01)
    def consume(self):
        if self.current_token:
            self.current_token=0
            return True
        else:
            return False
    def giveback(self):
        self.current_token=1

在要用令牌的异步task/coroutine 的语句前面添加 await self.mytoken.greenlight()
这样就可以限制后面部分的代码不会被两个协程同时调用。
这个思路帮我解决了aiomysql 里，当执行一个插入语句和commit这两个动作都需要await时候，会出现其他的协程调用数据库，在两个动作期间插入其他动作造成bug的情况。
请注意： 这里单令牌的交回也是由自己的程序做出的 建议用 try except finally 这种语法来释放token
与线程锁lock差不多

二、总量限制桶

理解了简单的单个令牌后，我们设想一下，在一个小时内只能访问50次，但是不限制你这50的间隔。

class LongWaitToken():

    def __init__(self,split,mytime):

        self.split = split
        self.currenttoken = split
        self.waittime = mytime
        self.lastconsumetime = time.time()

    def consume(self):
        if self.currenttoken>0:
            self.currenttoken-=1
            self.lastconsumetime = time.time()
            return True
        else:
            if time.time() -self.lastconsumetime>self.waittime:
                self.currenttoken = self.split
                self.lastconsumetime = time.time()
                return  True
            else:
                return False

    async def greenlight(self, number=1):
        """
        Block and yeild the contral from the subprocess.
        :param number:
        :return:
        """
        while self.consume() == False:
            # print(".")
            await asyncio.sleep(1)

逻辑与上面的单令牌类似，只是在查询还剩多少个令牌这一步（consume）时候，当所有令牌都没了以后，等待一个固定的时间，然后把桶装满。

三、水流桶（完整桶）

代码如下（示例）：

class TokenBucket(object):

    def __init__(self, name, rate, capacity,longwait = {}):
        """
        :param name: 单例模式下的 key
        :param rate:  这个实际上是多少秒增加一个token
        :param capacity: 桶内最大token数量（并发数量）
        """
        if name in created_bucket.keys():
            self = created_bucket[name]
            # 更新实例内参数 方便调用者 继承或者覆盖使用
           

        else:
            self._name = name
            self._rate = rate
            self._capacity = capacity
            self._current_token = capacity
            self._last_consum_time = time.time()
            if longwait!={}:
                self.longwait_init(longwait)
            else:
                self.use_longwait = False

            created_bucket[name] = self

    def longwait_init(self,longwait):
        self._longwait_split = longwait["split"]
        self._longwait_time = longwait["time"]
        self._longwait_current_left = longwait["split"]
        self.use_longwait = True


    def update(self,rate,capacity,longwait = {}):
        self._rate = rate
        self._capacity = capacity
        self._last_consum_time = time.time()
        if longwait != {}:
            self.longwait_init(longwait)

        created_bucket[self._name] = self

    def giveback(self, giveback: int = 1):
        """
        通知桶，有人将令牌拿回来了
        :return:
        """
        self._current_token = min(self._capacity, self._current_token + giveback)
        if self.use_longwait:
            self._longwait_current_left = min(self._longwait_split,self._longwait_current_left+1)

    def consume(self, token_amount):
        """
        根据时间计算桶内现有token量，待桶内量大于拿出量，程序返回True并在b桶内减少相应的token更取token时间 否则返回False 继续计时
        :param token_amount: 单次消耗几个token
        :return:
        """
        
        if token_amount > self._capacity:
            print("错误: 需求token大于桶内容量！")
            raise TokenExceedMax("希望获取的token数量大于桶内最大容量，逻辑死锁！！！")
        
        timefly = (time.time() - self._last_consum_time)
        increasement = int(timefly/ self._rate)
        # print(increasement)
        # 水流不能超过水桶容量
        self._current_token = min(self._capacity, self._current_token + increasement)
        # print("token token")

        if token_amount > self._current_token:
            return False
        else:
            if self.use_longwait:
                if self._longwait_current_left > 0:
                    self._longwait_current_left -= 1

                    self._current_token -= token_amount
                    self._last_consum_time = time.time()
                    return True
                else:
                    if time.time()  - self._last_consum_time>self._longwait_time :
                        self._longwait_current_left = self._longwait_split

                        self._current_token -= token_amount
                        self._last_consum_time = time.time()
                        return True
                    else:
                        return False
            else:
                self._current_token -= token_amount
                self._last_consum_time = time.time()
                return True

    async def greenlight(self, number=1):
        """
        Block and yeild the contral from the subprocess.
        :param number:
        :return:
        """
        while self.consume(number) == False:
            # print(".")
            await asyncio.sleep(0.01)

与之前两个桶不同的在于 这个桶的填装方式是基于时间的，而不是基于动作的，每隔一段时间就会往桶里放相应多的令牌，令牌的数量可以跟你锁并行的访问数量相同，这样就相当于并发了，这种桶往往更容易调整到服务器访问限制相同的情况
注意： 不要使用 await asyncio.sleep(0) 会让CPU跑满

总结

本文由简到繁介绍了三种程度的令牌桶，以及其正确的异步调用方法。
在拿到令牌之前，实际的代码是在不停的await asyncio.sleep(0.01) 因此进程可以去执行其他协程，这里的0.01只是一个代表 我实际使用时候用的是0.2（根据你自己的情况调整，间隔小CPU占用稍高，间隔大，桶不精确程度增加）