hello-dd-code
77 lines
2.0 KiB
Python
77 lines
2.0 KiB
Python
"""
|
|
全局请求速率限制器
|
|
确保代理每秒不超过5次请求
|
|
"""
|
|
import time
|
|
import threading
|
|
from collections import deque
|
|
|
|
|
|
class RateLimiter:
|
|
"""
|
|
令牌桶算法实现的速率限制器
|
|
"""
|
|
def __init__(self, max_requests_per_second: int = 5):
|
|
"""
|
|
初始化速率限制器
|
|
|
|
Args:
|
|
max_requests_per_second: 每秒最大请求数
|
|
"""
|
|
self.max_requests = max_requests_per_second
|
|
self.requests = deque()
|
|
self.lock = threading.RLock()
|
|
|
|
def acquire(self):
|
|
"""
|
|
获取请求权限，如果需要则等待
|
|
"""
|
|
with self.lock:
|
|
now = time.time()
|
|
|
|
# 清理超过1秒的请求记录
|
|
while self.requests and now - self.requests[0] >= 1.0:
|
|
self.requests.popleft()
|
|
|
|
# 如果当前请求数已达上限，等待
|
|
if len(self.requests) >= self.max_requests:
|
|
# 计算需要等待的时间
|
|
wait_time = 1.0 - (now - self.requests[0])
|
|
if wait_time > 0:
|
|
time.sleep(wait_time)
|
|
return self.acquire() # 递归调用以重新检查
|
|
|
|
# 记录这次请求
|
|
self.requests.append(now)
|
|
|
|
def can_make_request(self) -> bool:
|
|
"""
|
|
检查是否可以立即发起请求（非阻塞）
|
|
"""
|
|
with self.lock:
|
|
now = time.time()
|
|
|
|
# 清理超过1秒的请求记录
|
|
while self.requests and now - self.requests[0] >= 1.0:
|
|
self.requests.popleft()
|
|
|
|
return len(self.requests) < self.max_requests
|
|
|
|
|
|
# 全局速率限制器实例
|
|
global_rate_limiter = RateLimiter(max_requests_per_second=5)
|
|
|
|
|
|
def wait_for_request():
|
|
"""
|
|
等待直到可以发起请求
|
|
"""
|
|
global_rate_limiter.acquire()
|
|
|
|
|
|
def can_request_now() -> bool:
|
|
"""
|
|
检查是否可以立即发起请求
|
|
"""
|
|
return global_rate_limiter.can_make_request()
|