仅用Python标准库实现多域名并发请求速率限制方案问询

核心设计思路 #

针对单队列阻塞问题，核心是按域名隔离限流逻辑，同时通过全局调度器统一管理可执行/延迟任务：

• 每个域名绑定独立的速率控制器，负责跟踪请求频率与退避状态
• 全局维护两个任务存储结构：
  • 待执行队列：存放当前可立即执行的任务（所属域名未触发限流）
  • 延迟任务堆：按可执行时间排序，存放处于退避期的域名任务
• 专用调度线程负责将到期的延迟任务移至待执行队列，工作线程仅需处理待执行队列，无需关心限流逻辑

数据结构设计 #

1. DomainRateLimiter：每个域名的限流控制器，包含：
   • 请求时间记录窗口、最大请求数、退避时长
   • 退避截止时间标记
   • 线程安全的锁，保证多线程下状态一致
2. TaskScheduler：全局任务调度器，包含：
   • 待执行队列（queue.Queue）：线程安全的阻塞队列，供工作线程获取任务
   • 延迟任务堆（heapq）：按可执行时间升序排列，快速取出最早到期的任务
   • 域名-限流控制器映射表
   • 条件变量与锁：实现调度线程的高效等待与唤醒，避免忙轮询

代码实现 #

1. 域名限流控制器 #

import threading
import heapq
import time
from queue import Queue
from typing import List, Tuple, Callable

class DomainRateLimiter:
    def __init__(self, max_requests: int, window_seconds: float, backoff_seconds: float = 5.0):
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        self.backoff_seconds = backoff_seconds
        self.request_timestamps: List[float] = []
        self.backoff_until: float = 0.0  # 0表示无退避
        self.lock = threading.Lock()

    def can_request(self) -> bool:
        """判断当前是否可向该域名发起请求"""
        with self.lock:
            now = time.time()
            # 检查退避期
            if now < self.backoff_until:
                return False
            # 清理时间窗口外的请求记录
            self.request_timestamps = [t for t in self.request_timestamps if now - t < self.window_seconds]
            return len(self.request_timestamps) < self.max_requests

    def record_request(self) -> None:
        """记录请求，触发限流则设置退避"""
        with self.lock:
            now = time.time()
            self.request_timestamps.append(now)
            # 超过速率限制时启动退避
            if len(self.request_timestamps) >= self.max_requests:
                self.backoff_until = now + self.backoff_seconds

    def get_backoff_end(self) -> float:
        """获取退避结束时间"""
        with self.lock:
            return self.backoff_until

2. 全局任务调度器 #

class TaskScheduler:
    def __init__(self):
        self.pending_queue = Queue()
        self.delayed_heap = []
        self.domain_limiters = {}
        self.lock = threading.Lock()
        self.condition = threading.Condition(self.lock)
        self.running = True
        # 启动后台调度线程
        threading.Thread(target=self._scheduler_loop, daemon=True).start()

    def add_domain_limiter(self, domain: str, max_requests: int, window_seconds: float, backoff_seconds: float = 5.0):
        """为域名配置限流规则"""
        with self.lock:
            self.domain_limiters[domain] = DomainRateLimiter(max_requests, window_seconds, backoff_seconds)

    def add_task(self, domain: str, task: Callable):
        """添加下载任务，自动判断立即执行或延迟"""
        with self.lock:
            limiter = self.domain_limiters.get(domain)
            if not limiter:
                raise ValueError(f"No rate limiter for domain: {domain}")
            
            now = time.time()
            if limiter.can_request():
                self.pending_queue.put((domain, task))
                self.condition.notify()
            else:
                # 计算延迟时间并加入堆
                execute_time = limiter.get_backoff_end()
                heapq.heappush(self.delayed_heap, (execute_time, domain, task))
                self.condition.notify()

    def _scheduler_loop(self):
        """调度线程：处理到期的延迟任务"""
        while self.running:
            with self.lock:
                # 计算下一个延迟任务的等待时长
                wait_time = None
                if self.delayed_heap:
                    earliest_time = self.delayed_heap[0][0]
                    wait_time = max(0.0, earliest_time - time.time())
                
                # 等待唤醒：新任务添加或延迟任务到期
                self.condition.wait(timeout=wait_time)

                # 处理所有到期的延迟任务
                now = time.time()
                while self.delayed_heap and self.delayed_heap[0][0] <= now:
                    execute_time, domain, task = heapq.heappop(self.delayed_heap)
                    limiter = self.domain_limiters[domain]
                    if limiter.can_request():
                        self.pending_queue.put((domain, task))
                    else:
                        # 仍在退避期，重新计算延迟时间放回堆
                        new_execute_time = limiter.get_backoff_end()
                        heapq.heappush(self.delayed_heap, (new_execute_time, domain, task))

    def get_task(self) -> Tuple[str, Callable]:
        """工作线程获取可执行任务"""
        return self.pending_queue.get()

    def task_done(self):
        """工作线程标记任务完成"""
        self.pending_queue.task_done()

3. 工作线程与示例用法 #

def worker(scheduler: TaskScheduler):
    """工作线程：执行任务并更新限流状态"""
    while True:
        domain, task = scheduler.get_task()
        try:
            task()
            # 记录请求，更新域名限流状态
            scheduler.domain_limiters[domain].record_request()
        except Exception as e:
            print(f"Task failed for {domain}: {str(e)}")
        finally:
            scheduler.task_done()

if __name__ == "__main__":
    # 初始化调度器
    scheduler = TaskScheduler()

    # 配置域名限流规则
    scheduler.add_domain_limiter("domain1.com", max_requests=2, window_seconds=1.0)
    scheduler.add_domain_limiter("domain2.com", max_requests=5, window_seconds=1.0)

    # 启动3个工作线程
    for _ in range(3):
        threading.Thread(target=worker, args=(scheduler,), daemon=True).start()

    # 模拟添加任务：先加domain1的任务，再加domain2的任务
    def download_task(domain: str, idx: int):
        print(f"[{time.strftime('%H:%M:%S')}] Downloading {domain} task {idx}")
        time.sleep(0.1)  # 模拟下载耗时

    # 添加domain1的任务
    for i in range(5):
        scheduler.add_task("domain1.com", lambda d="domain1.com", i=i: download_task(d, i))
        time.sleep(0.2)

    # 添加domain2的任务
    for i in range(5):
        scheduler.add_task("domain2.com", lambda d="domain2.com", i=i: download_task(d, i))
        time.sleep(0.1)

    # 等待所有任务完成
    scheduler.pending_queue.join()
    print("All tasks completed.")

方案优势 #

• 无跨域名阻塞：某域名触发退避时，其他域名任务正常执行，解决单队列的阻塞问题
• 高效调度：调度线程通过条件变量等待，仅在必要时唤醒，避免忙轮询的资源浪费
• 任务顺序保障：同一域名的任务按提交顺序处理，延迟任务按到期时间优先调度
• 动态扩展：支持随时添加新域名的限流规则与任务，无需重启调度器
• 资源可控：工作线程数量固定，不会因域名过多导致线程资源耗尽

内容的提问来源于stack exchange，提问作者Ed_