基于asyncio的批量队列处理器实现遇阻，求技术指导

问题概述 #

尝试用AsyncIO构建批量队列处理器，核心需求是：将单个请求推入队列，处理器按token阈值批量整合请求，处理完成后将结果返回给对应的调用任务，每个调用任务需等待自身结果可用。当前代码运行后出现卡住现象，无法完成所有请求的处理。

核心错误分析 #

• 类属性与实例属性混淆：lock、batch_processed、results、counter被定义为类属性，导致所有实例共享状态，引发并发逻辑混乱。
• 事件未及时重置：batch_processed事件触发后未清除，未获取到结果的调用任务会进入无限循环，持续占用事件循环，阻止处理器继续处理队列剩余请求。
• Token计算错误：count_tokens传入整个Event字典而非payload，导致token计数不符合预期的内容长度规则。
• 队列任务未标记完成：调用queue.get()后未执行queue.task_done()，违反AsyncIO队列的规范使用方式，可能导致队列状态异常。
• 返回类型标注错误：add_to_batch实际返回处理结果，但类型标注为-> None，不符合代码逻辑。

修复后的完整代码 #

import asyncio
from typing import TypedDict, List, Any, Dict

class Event(TypedDict):
    id: int
    payload: Any

# 按请求内容计算token数
count_tokens = lambda payload: len(str(payload))

class BatchProcessor:
    def __init__(self, tokens_per_batch=100) -> None:
        # 所有状态属性改为实例属性，保证每个实例独立
        self.queue: asyncio.Queue[Event] = asyncio.Queue()
        self.lock = asyncio.Lock()
        self.batch_processed = asyncio.Event()
        self.results: Dict[int, Any] = {}
        self.counter: int = 0
        self.tokens_per_batch = tokens_per_batch
        # 存储当前批次和累计token数，方便flush方法访问
        self._current_batch: List[Event] = []
        self._current_tokens: int = 0
        self.processor_task = asyncio.create_task(self.processor())
        
    async def next_counter(self):
        async with self.lock:
            self.counter += 1
            return self.counter

    async def processor(self) -> None:
        while True:
            event = await self.queue.get()
            try:
                self._current_batch.append(event)
                self._current_tokens += count_tokens(event['payload'])
                print(f"Tokens: {self._current_tokens}")
                
                # 达到阈值触发批量处理
                if self._current_tokens >= self.tokens_per_batch:
                    print(f"Triggering batch of {len(self._current_batch)} items")
                    batch_results = await self.process_batch(self._current_batch)
                    async with self.lock:
                        self.results.update(batch_results)
                    print(f"Notifying batch processed, results available: {list(self.results.keys())}")
                    # 触发事件后立即重置，避免无效循环
                    self.batch_processed.set()
                    self.batch_processed.clear()
                    self._current_batch = []
                    self._current_tokens = 0
            finally:
                # 标记队列任务完成，维护队列状态
                self.queue.task_done()
                
    async def process_batch(self, batch: List[Event]) -> Dict[int, Any]:
        print(f"Processing batch of {len(batch)} items")
        await asyncio.sleep(1)
        return {event['id']: f"{event['payload']} processed" for event in batch}
    
    async def add_to_batch(self, payload: Any ) -> Any:
        req_id = await self.next_counter()
        self.queue.put_nowait(Event(id=req_id, payload=payload))
        print(f"Added request {req_id} to batch")
        
        while True:
            await self.batch_processed.wait()
            async with self.lock:
                if req_id in self.results:
                    result = self.results.pop(req_id)
                    print(f"Result for request {req_id}: {result}")
                    return result

    async def flush(self):
        # 等待队列中所有请求被处理器接收
        await self.queue.join()
        # 处理剩余未达阈值的批次
        if self._current_batch:
            print(f"Flushing remaining batch of {len(self._current_batch)} items")
            batch_results = await self.process_batch(self._current_batch)
            async with self.lock:
                self.results.update(batch_results)
            self.batch_processed.set()
            self.batch_processed.clear()
            self._current_batch = []
            self._current_tokens = 0

async def main():
    batch_processor = BatchProcessor(tokens_per_batch=100)

    # 模拟10个客户端请求
    coros = [batch_processor.add_to_batch(f"data_{i}") for i in range(1, 11)]

    # 等待所有请求结果
    results = await asyncio.gather(*coros)
    print("Results:", results)
    
    # 可选：测试flush处理剩余请求
    # await batch_processor.flush()

# 运行主函数
asyncio.run(main())

关键修复点说明 #

1. 状态实例化：将所有状态属性改为实例属性，确保每个BatchProcessor实例拥有独立的队列、锁、事件等资源，避免多实例状态冲突。
2. 事件重置机制：每次触发batch_processed事件后立即调用clear()，防止未获取结果的调用任务进入无限循环，保证事件循环能正常调度处理器任务。
3. Token计算修正：传入event['payload']计算token数，符合按请求内容长度统计的预期逻辑。
4. 队列规范使用：用try/finally确保每次get()后都调用task_done()，维护队列的正确状态。
5. Flush方法实现：添加flush()方法，等待队列清空后处理剩余未达阈值的批次，确保所有请求都能被处理，避免请求残留。

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