异步处理数据（异步队列消费者）

前言

通过一个异步队列消费者来处理数据

Snippets

• main.py

from contextlib import asynccontextmanager
from fastapi import FastAPI
import asyncio
import uvicorn
import logging
from typing import List
from consumer import ConsumerPool

logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    # 在应用状态中创建队列和消费者池
    async_queue = asyncio.Queue()
    app.state.queue = async_queue
    
    # 创建消费者池，默认包含3个消费者
    consumer_pool = ConsumerPool(async_queue, pool_size=3)
    app.state.consumer_pool = consumer_pool
    
    # 启动消费者池
    consumer_pool.start()
    
    yield
    
    # 在应用关闭时停止消费者池
    await consumer_pool.stop()


app = FastAPI(lifespan=lifespan)

# 导入并注册路由
from api.routes import router
app.include_router(router)

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000)

• api/routes.py

from fastapi import APIRouter
import asyncio
from typing import Annotated
from fastapi import Request
from consumer import ConsumerPool

# 创建API路由器
router = APIRouter(prefix="/api")

@router.post("/events")
async def add_event(request: Request, event: str):
    queue = request.app.state.queue
    await queue.put(event)
    return {"status": "event added"}

@router.get("/queue_size")
async def queue_size(request: Request):
    queue = request.app.state.queue
    return {"size": queue.qsize()}

@router.get("/consumer_pool_status")
async def consumer_pool_status(request: Request):
    pool = request.app.state.consumer_pool
    active_consumers = sum(1 for c in pool.consumers if c.running)
    return {
        "pool_size": pool.pool_size,
        "active_consumers": active_consumers,
        "queue_size": request.app.state.queue.qsize()
    }

@router.post("/resize_pool")
async def resize_pool(request: Request, new_size: int):
    pool = request.app.state.consumer_pool
    if new_size <= 0:
        return {"error": "Pool size must be greater than 0"}
    
    if new_size == pool.pool_size:
        return {"message": f"Pool size is already {new_size}"}
    
    # 停止当前池
    await pool.stop()
    
    # 创建新池
    new_pool = ConsumerPool(pool.queue, pool_size=new_size)
    request.app.state.consumer_pool = new_pool
    new_pool.start()
    
    return {"message": f"Pool resized from {pool.pool_size} to {new_size}"}

• consumers.py

import asyncio
import logging
from typing import List

logger = logging.getLogger(__name__)

class AsyncQueueConsumer:
    def __init__(self, queue: asyncio.Queue, consumer_id: int):
        self.queue = queue
        self.consumer_id = consumer_id
        self.running = True

    async def consume(self):
        logger.info(f"Consumer {self.consumer_id} started")
        while self.running:
            try:
                item = await self.queue.get()
                if not self.running:
                    # 如果不再运行，重新放回队列或丢弃
                    break
                await asyncio.sleep(1)  # 模拟异步处理
                print(f"Consumer {self.consumer_id} consumed: {item}")
                self.queue.task_done()
            except asyncio.CancelledError:
                logger.info(f"Consumer {self.consumer_id} was cancelled")
                break
            except Exception as e:
                logger.error(f"Error in consumer {self.consumer_id}: {e}")
                # 确保不会阻塞队列
                self.queue.task_done()  # 确保不会阻塞队列

class ConsumerPool:
    def __init__(self, queue: asyncio.Queue, pool_size: int = 3):
        self.queue = queue
        self.pool_size = pool_size
        self.consumers: List[AsyncQueueConsumer] = []
        self.tasks: List[asyncio.Task] = []

    def start(self):
        """启动消费者池中的所有消费者"""
        self.consumers = [
            AsyncQueueConsumer(self.queue, i) for i in range(self.pool_size)
        ]
        self.tasks = [
            asyncio.create_task(consumer.consume()) for consumer in self.consumers
        ]
        logger.info(f"Started consumer pool with {self.pool_size} consumers")

    async def stop(self):
        """停止消费者池中的所有消费者"""
        logger.info("Shutting down consumer pool...")
        
        # 停止所有消费者
        for consumer in self.consumers:
            consumer.running = False

        # 取消所有任务
        for task in self.tasks:
            task.cancel()

        # 等待所有任务完成
        for task in self.tasks:
            try:
                await task
            except asyncio.CancelledError:
                logger.info("Consumer task cancelled successfully")

        logger.info("Consumer pool shutdown completed")

测试

curl --request GET \
 --url http://127.0.0.1:8000/api/queue_size

curl --request POST \
  --url 'http://127.0.0.1:8000/api/events?event=asdf1234'