[漏桶算法] 实现第一类漏桶算法 #python

leaky_bucket.py

#ref: https://en.wikipedia.org/wiki/Leaky_bucket#As_a_meter

class RateLimitExceededException(Exception):
    pass


class LeakyBucketRateLimiter:
    '''漏桶算法实现
    
    Attributes:
        capacity: 桶容量，多少份水
        leaky_rate: 消费速率，以s为单位(份/s)
    Raises:
        RateLimitExceededException: 当漏桶溢水抛出
    '''
    def __init__(self, capacity=10.0, leaky_rate=1.0, is_lock=True):
        self.lock = threading.RLock() if is_lock else None
        self._capacity = float(capacity)
        self._leak_rate = float(leaky_rate)
        # 上次漏水时间
        self._last_time = time.time()
        # 余水
        self._remained_water = float(0)
    
    def acquire(self, how_many=1):
        '''向漏桶加水

        Args:
            how_many: 加多少份水
        '''
        if self.lock:
            with self.lock:
                return self._acquire(how_many)
        return self._acquire(how_many)

    def _acquire(self, how_many):
        # 计算漏桶现在余下多少水
        now = time.time()
        delta = self._remained_water - self._leak_rate * (now - self._last_time)
        self._remained_water = float(min(0, delta))
        self._last_time = now
        # 向漏桶加水
        if self._remained_water + how_many >= self._capacity:
            # 溢水则抛出异常
            raise RateLimitExceededException
        self._remained_water += how_many