0%

流量控制算法

对外的服务一般需要对某些接口或者对某些业务方进行限流,否则突然流量增大可能会造成服务崩溃。常见的用来控制流量大小的两种算法,令牌桶和漏桶算法。当超过限定的数值就进行缓存或者直接拒绝等其他操作。

令牌桶算法

令牌桶算法就是按照一定的速度将令牌放到桶中,每次请求来的时候需要从桶中取出令牌。假设你想把速度限制到 r,那么每隔 1/r 就放入到桶中。如果桶满了,那么新的令牌就被丢弃。如果可用的令牌过少,那么新来的请求也会丢弃或者缓存住。但是我们用程序实现的时候可以记录上次发令牌的时间,当前时间减去上次时间,用时间间隔乘上发送的令牌的速率,就知道这段时间内应该发了多少令牌。那么当前的令牌数量就是桶的容量和本可以发送令牌数量的最小值。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
from time import time
from threading import RLock

__all__ = ("TokenBucket", )


class TokenBucket(object):

def __init__(self, capacity, fill_rate, is_lock=False):
"""
:param capacity: The total tokens in the bucket.
:param fill_rate: The rate in tokens/second that the bucket will be refilled
"""
self._capacity = float(capacity)
self._tokens = float(capacity)
self._fill_rate = float(fill_rate)
self._last_time = time()
self._is_lock = is_lock
self._lock = RLock() if is_lock else None

def _get_cur_tokens(self):
if self._tokens < self._capacity:
now = time()
delta = self._fill_rate * (now - self._last_time)
self._tokens = min(self._capacity, self._tokens + delta)
self._last_time = now
return self._tokens

def get_cur_tokens(self):
if self._is_lock:
with self._lock:
return self._get_cur_tokens()
else:
return self._get_cur_tokens()

def _consume(self, tokens):
if tokens <= self.get_cur_tokens():
self._tokens -= tokens
return True
return False

def consume(self, tokens):
if self._is_lock:
with self._lock:
return self._consume(tokens)
else:
return self._consume(tokens)


if __name__ == "__main__":
tb = TokenBucket(5, 1, True)
for _ in range(10):
print(tb.consume(1))

漏桶算法

漏桶算法则是把数据先放到桶中,然后从桶中控制消费的速度。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
from time import time
from threading import RLock

__all__ = ("LeakyBucket", )


class LeakyBucket(object):

def __init__(self, capacity, leak_rate, is_lock=False):
"""
:param capacity: The total tokens in the bucket.
:param leak_rate: The rate in tokens/second that the bucket leaks
"""
self._capacity = float(capacity)
self._used_tokens = 0
self._leak_rate = float(leak_rate)
self._last_time = time()
self._lock = RLock() if is_lock else None

def get_used_tokens(self):
if self._lock:
with self._lock:
return self._get_used_tokens()
else:
return self._get_used_tokens()

def _get_used_tokens(self):
now = time()
delta = self._leak_rate * (now - self._last_time)
self._used_tokens = max(0, self._used_tokens - delta)
return self._used_tokens

def _consume(self, tokens):
if tokens + self._get_used_tokens() <= self._capacity:
self._used_tokens += tokens
self._last_time = time()
return True
return False

def consume(self, tokens):
if self._lock:
with self._lock:
return self._consume(tokens)
else:
return self._consume(tokens)


if __name__ == "__main__":
lb = LeakyBucket(5, 1, True)
for _ in range(10):
print(lb.consume(1))

总结

两者的程序可以看到不同,前者是控制还剩下的令牌,后者是控制已经用了多少流量。漏桶算法的漏出速度是固定的,而令牌桶算法则是控制平均速度,并且允许突发性的提供速度,反正只要平均速度是稳定的就好了。

参考链接