本文目录导读：
我来为你介绍几种Python实现削峰填谷同步任务流量的方法：
使用队列和速率限制
生产者-消费者模式
import threading
import time
from queue import Queue
import random
class TrafficShapingSync:
def __init__(self, max_queue_size=100, rate_limit=10):
self.task_queue = Queue(maxsize=max_queue_size)
self.rate_limit = rate_limit # 每秒处理的任务数
self.last_time = time.time()
self.tokens = 0
self.lock = threading.Lock()
def token_bucket(self):
"""令牌桶算法控制速率"""
with self.lock:
now = time.time()
elapsed = now - self.last_time
self.tokens += elapsed * self.rate_limit
self.tokens = min(self.tokens, self.rate_limit * 2) # 最大积攒2秒的令牌
self.last_time = now
if self.tokens >= 1:
self.tokens -= 1
return True
return False
def producer(self, task_id):
"""生产者-添加任务"""
task = {
'id': task_id,
'data': f"Task_{task_id}",
'timestamp': time.time()
}
try:
self.task_queue.put(task, timeout=1)
print(f"生产者添加任务: {task_id}, 队列大小: {self.task_queue.qsize()}")
except:
print(f"生产者: 队列已满，丢弃任务 {task_id}")
def consumer(self, worker_id):
"""消费者-处理任务"""
while True:
if self.token_bucket():
try:
task = self.task_queue.get(timeout=0.1)
# 模拟任务处理
process_time = random.uniform(0.1, 0.5)
time.sleep(process_time)
print(f"消费者{worker_id}完成: 任务{task['id']}, 耗时:{process_time:.2f}s")
self.task_queue.task_done()
except:
time.sleep(0.1)
else:
time.sleep(0.05) # 等待令牌
# 使用示例
if __name__ == "__main__":
traffic = TrafficShapingSync(max_queue_size=50, rate_limit=5)
# 启动消费者
consumers = []
for i in range(3):
t = threading.Thread(target=traffic.consumer, args=(i,), daemon=True)
t.start()
consumers.append(t)
# 模拟突发流量
for i in range(100):
traffic.producer(i)
if i % 10 == 0:
# 模拟突发流量
for _ in range(5):
traffic.producer(i + 1000)
使用asyncio实现异步流量控制
import asyncio
import time
import random
class AsyncTrafficShaper:
def __init__(self, rate_limit=10, burst_size=20):
self.rate_limit = rate_limit
self.burst_size = burst_size
self.tokens = burst_size
self.last_refill = time.monotonic()
self.lock = asyncio.Lock()
async def _refill_tokens(self):
"""补充令牌"""
now = time.monotonic()
elapsed = now - self.last_refill
new_tokens = elapsed * self.rate_limit
self.tokens = min(self.burst_size, self.tokens + new_tokens)
self.last_refill = now
async def acquire(self):
"""获取处理权限"""
while True:
async with self.lock:
await self._refill_tokens()
if self.tokens >= 1:
self.tokens -= 1
return True
await asyncio.sleep(1.0 / self.rate_limit)
async def process_task(self, task_id, task_data):
"""处理单个任务"""
await self.acquire() # 等待令牌
# 模拟任务处理
process_time = random.uniform(0.1, 0.3)
await asyncio.sleep(process_time)
print(f"任务 {task_id} 完成，处理时间: {process_time:.2f}s")
return task_data
class AsyncScheduler:
def __init__(self, max_concurrent=5):
self.semaphore = asyncio.Semaphore(max_concurrent)
self.shaper = AsyncTrafficShaper(rate_limit=10)
async def safe_process(self, task_id, task_data):
"""安全处理任务（限流+并发控制）"""
async with self.semaphore: # 控制并发数
result = await self.shaper.process_task(task_id, task_data)
return result
async def process_batch(self, tasks):
"""批量处理任务"""
tasks_list = [self.safe_process(i, data) for i, data in enumerate(tasks)]
results = await asyncio.gather(*tasks_list, return_exceptions=True)
return results
# 使用示例
async def main():
scheduler = AsyncScheduler(max_concurrent=5)
# 模拟100个突发任务
tasks = [f"data_{i}" for i in range(100)]
# 分批处理，每批20个
batch_size = 20
for i in range(0, len(tasks), batch_size):
batch = tasks[i:i+batch_size]
print(f"\n处理批次 {i//batch_size + 1}, 任务数: {len(batch)}")
results = await scheduler.process_batch(batch)
await asyncio.sleep(1) # 批次间隔
if __name__ == "__main__":
asyncio.run(main())
使用第三方库实现高级控制
aio-limit 示例
pip install aio-limit
from aio_limit import Limiter
import asyncio
import time
class AdvancedTrafficController:
def __init__(self):
# 创建限流器：每秒10个请求，最大积攒20个
self.limiter = Limiter(rate=10, burst=20)
async def limited_task(self, task_id):
"""限流任务"""
async with self.limiter:
# 模拟任务处理
await asyncio.sleep(0.1)
print(f"任务 {task_id} 在 {time.time():.2f} 执行")
return task_id
async def main():
controller = AdvancedTrafficController()
# 模拟突发流量
tasks = [controller.limited_task(i) for i in range(50)]
results = await asyncio.gather(*tasks)
print(f"完成 {len(results)} 个任务")
if __name__ == "__main__":
asyncio.run(main())
使用Redis实现分布式限流
import redis
import time
import threading
class RedisTrafficShaper:
def __init__(self, redis_client, rate_limit=10, window_size=1):
self.redis = redis_client
self.rate_limit = rate_limit
self.window_size = window_size # 窗口大小（秒）
def acquire(self, key="traffic_shaper"):
"""获取处理权限"""
window_key = f"{key}:{int(time.time() / self.window_size)}"
# 使用Redis事务
pipe = self.redis.pipeline()
pipe.incr(window_key)
pipe.expire(window_key, self.window_size * 2)
current_count = pipe.execute()[0]
if current_count <= self.rate_limit:
return True
return False
# 使用示例
if __name__ == "__main__":
r = redis.Redis(host='localhost', port=6379, db=0)
shaper = RedisTrafficShaper(r, rate_limit=10)
def worker(worker_id):
for i in range(20):
if shaper.acquire():
print(f"Worker {worker_id}: 任务 {i} 执行")
time.sleep(0.1)
else:
print(f"Worker {worker_id}: 任务 {i} 限流")
time.sleep(0.2)
# 启动多个worker
threads = []
for i in range(3):
t = threading.Thread(target=worker, args=(i,))
t.start()
threads.append(t)
for t in threads:
t.join()
选择建议
| 场景 | 推荐方案 |
|---|---|
| 单机简单限流 | 令牌桶 + threading |
| 高并发异步任务 | asyncio + 信号量 |
| 精确速率控制 | aio-limit 第三方库 |
| 分布式系统 | Redis 实现 |
| 需要背压处理 | 队列 + 水位控制 |
选择合适的方法需要考虑：
- 并发量级
- 是否需要分布式支持
- 精度要求
- 系统复杂度
