Skip to content

Instantly share code, notes, and snippets.

@mywaiting
Created June 3, 2024 02:09
Show Gist options
  • Save mywaiting/8b58d242c5624591610cfbd265a156af to your computer and use it in GitHub Desktop.
Save mywaiting/8b58d242c5624591610cfbd265a156af to your computer and use it in GitHub Desktop.
基于时间计算的令牌桶实现 - 注意:令牌桶算法实现上另一个细节就是进程/线程锁,如果不使用后端组件比如 Redis 则实现上必须加入锁实现(此时后端组件相当于另一种形式上锁实现), 尽管单进程/线程执行代码的无所谓,但多进程/线程执行代码时,令牌桶的数据就会混乱,此处还是带锁执行比较稳妥 - 注意:令牌桶原始意思应该有个定时器,定期往桶里塞入固定数量的令牌除非溢出, 但多数代码实现上都是取令牌时计算上次消耗到本次消耗时间差中生成令牌数量,没有原始意义上塞令牌的代码
#!/usr/bin/env python
#
# 精简实现的令牌桶算法
# 注意:此处实现上都是取令牌时计算上次消耗到本次消耗时间差中生成令牌数量,没有原始意义上往桶塞令牌的实现
# 注意:此处实现上有简单的进程锁,防止多线程/进程读取数据产生混乱
#
import time
import threading
class Bucket(object):
"""基于时间计算的令牌桶实现
- 注意:令牌桶算法实现上另一个细节就是进程/线程锁,如果不使用后端组件比如 Redis 则实现上必须加入锁实现(此时后端组件相当于另一种形式上锁实现),
尽管单进程/线程执行代码的无所谓,但多进程/线程执行代码时,令牌桶的数据就会混乱,此处还是带锁执行比较稳妥
- 注意:令牌桶原始意思应该有个定时器,定期往桶里塞入固定数量的令牌除非溢出,
但多数代码实现上都是取令牌时计算上次消耗到本次消耗时间差中生成令牌数量,没有原始意义上塞令牌的代码
"""
def __init__(self, rate=1, burst=None):
"""参数说明
- rate 代表每秒往当前桶放入多少个令牌,相当于令牌生成速率,单位秒
- burst 代表当前每次并发读取多少个令牌,相当于每次取用令牌数量
"""
self.rate = float(rate)
self.burst = float(burst) if burst is not None else float(rate) * 10
self.mutex = threading.Lock() # 进程/线程锁,强制保证令牌数据原子化读取/写入
self.bucket = self.burst # 初始化令牌桶中令牌的数量,默认就是 burst
self.last_update = time.time() # 上次更新令牌时间(取令牌时计算上次消耗到本次消耗时间差中生成令牌数量)
def get(self):
"""而是返回当前全部令牌的数量,当前令牌的数量不会超过 burst 数量的限制
- 注意:此处是“偷窥”一下当前有多少可用的令牌,不是直接取出并消灭其令牌数量
"""
now = time.time()
# 如果当前桶中令牌数量大于并发读取数量则直接返回所有令牌
if self.bucket >= self.burst:
self.last_update = now
return self.bucket
# 计算上次消耗到本次消耗时间差中生成令牌数量(该数量至少大于一才有意义)
bucket = self.rate * (now - self.last_update)
self.mutex.acquire()
# 如果时间差中确实生成令牌数量大于一,那么更新到令牌桶(进程锁保证原子化写入)
if bucket > 1:
self.bucket += bucket
# 确保存入桶的令牌数量不会超过当前并发读取数
if self.bucket > self.burst:
self.bucket = self.burst
self.last_update = now
self.mutex.release()
return self.bucket
def incr(self, value=0):
"""设置当前令牌桶中令牌的数量,该方法很少直接使用,为了方便还是加上吧
"""
self.bucket += int(value) # 强制传入的必须是整数
def desc(self, value=1):
"""取出对应的令牌并直接消灭令牌桶的令牌数量
"""
self.bucket -= int(value)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment